LLM 水印技术指南：AI 生成内容的溯源

无水印:  P(token) → 采样 → 文本
水印:    P(token) + 绿名单偏置 → 采样 → 文本(含隐藏信号)
检测:    统计绿名单 token 频率 → 判断是否含水印

import torch

class KirchenbauerWatermark:
    def __init__(self, vocab_size=50272, greenlist_ratio=0.25,
                 strength=2.0, hash_key=15485863):
        self.vocab_size = vocab_size
        self.ratio = greenlist_ratio
        self.strength = strength
        self.hash_key = hash_key
    
    def _get_greenlist(self, prev_token: int) -> torch.Tensor:
        rng = torch.Generator()
        rng.manual_seed((self.hash_key * prev_token) % (2**32))
        perm = torch.randperm(self.vocab_size, generator=rng)
        size = int(self.vocab_size * self.ratio)
        return perm[:size]
    
    def watermark_logits(self, input_ids, logits):
        prev_token = input_ids[:, -1].item()
        greenlist = self._get_greenlist(prev_token)
        logits[:, greenlist] += self.strength
        return logits

class WatermarkDetector:
    def detect(self, text_tokens, config):
        green_count = 0
        total = 0
        for i in range(1, len(text_tokens)):
            greenlist = config._get_greenlist(text_tokens[i-1].item())
            if text_tokens[i].item() in greenlist:
                green_count += 1
            total += 1
        
        gamma = config.ratio
        z = (green_count - gamma * total) / (total**0.5 * gamma * (1-gamma)**0.5)
        return {'z_score': z, 'is_watermarked': z > 4.0, 'tokens': total}