LLM 水印技术：AI 生成内容的溯源方案

水印嵌入过程：
1. 将词表 V 分为绿名单 G 和红名单 R（基于哈希）
2. 在生成时，对绿名单词的 logit 加 δ
3. 从调整后的分布中采样

水印检测过程：
1. 统计文本中绿名单词的比例
2. 使用 z 检验判断是否超过统计阈值

import torch
import torch.nn.functional as F
from hashlib import sha256

class KGWWatermark:
    """
    Kirchenbauer et al. 水印方案实现
    """

    def __init__(
        self,
        vocab_size: int,
        green_list_ratio: float = 0.25,
        delta: float = 2.0,
        seeding_scheme: str = "simple_1",
        hash_key: int = 15485863,  # 第100万个质数
    ):
        self.vocab_size = vocab_size
        self.green_list_ratio = green_list_ratio
        self.green_list_size = int(vocab_size * green_list_ratio)
        self.delta = delta
        self.seeding_scheme = seeding_scheme
        self.hash_key = hash_key

    def _get_green_list(self, prev_token: int) -> set:
        """
        基于前一个 token 生成绿名单
        使用哈希确保伪随机但确定性
        """
        h = sha256(f"{self.hash_key}{prev_token}".encode()).hexdigest()
        # 用哈希值生成伪随机序列
        torch.manual_seed(int(h[:8], 16))
        perm = torch.randperm(self.vocab_size)
        return set(perm[:self.green_list_size].tolist())

    def watermark_logits(self, input_ids: torch.Tensor, logits: torch.Tensor) -> torch.Tensor:
        """
        在 logits 上添加水印
        input_ids: [batch, seq_len] 已生成的 token
        logits: [batch, seq_len, vocab_size] 模型输出
        """
        batch_size, seq_len, _ = logits.shape
        watermarked_logits = logits.clone()

        for b in range(batch_size):
            for t in range(seq_len):
                if t == 0:
                    continue  # 第一个 token 不加水印

                prev_token = input_ids[b, t - 1].item()
                green_list = self._get_green_list(prev_token)

                # 对绿名单 token 的 logit 加 delta
                for token_id in green_list:
                    watermarked_logits[b, t, token_id] += self.delta

        return watermarked_logits

    def detect(self, text_tokens: torch.Tensor, z_threshold: float = 4.0) -> dict:
        """
        检测文本中是否包含水印
        """
        n = len(text_tokens)
        if n < 10:
            return {"detected": False, "reason": "文本过短"}

        green_count = 0
        for i in range(1, n):
            prev_token = text_tokens[i - 1].item()
            curr_token = text_tokens[i].item()
            green_list = self._get_green_list(prev_token)
            if curr_token in green_list:
                green_count += 1

        # z 检验
        expected = self.green_list_ratio * (n - 1)
        std = (self.green_list_ratio * (1 - self.green_list_ratio) * (n - 1)) ** 0.5
        z_score = (green_count - expected) / std if std > 0 else 0

        return {
            "detected": z_score > z_threshold,
            "z_score": z_score,
            "green_count": green_count,
            "total_tokens": n - 1,
            "green_ratio": green_count / (n - 1),
            "threshold": z_threshold,
        }


# 使用示例
watermarker = KGWWatermark(vocab_size=50257, green_list_ratio=0.25, delta=2.0)

# 生成时加水印
# watermarked_logits = watermarker.watermark_logits(input_ids, logits)
# next_token = sample_from(watermarked_logits)

# 检测时
# result = watermarker.detect(text_token_ids, z_threshold=4.0)
# print(f"水印检测: {result['detected']}, z-score: {result['z_score']:.2f}")

def synthid_watermark_sampling(logits, prev_token, watermark_key, gamma=0.5):
    """
    SynthID-Text 风格的水印采样
    使用 T 函数调整概率分布
    """
    # 1. 计算原始概率
    probs = F.softmax(logits, dim=-1)

    # 2. 生成基于前一个 token 的水印掩码
    torch.manual_seed(hash((watermark_key, prev_token)) % (2**32))
    mask = torch.rand(len(logits)) < gamma  # 选中 gamma 比例的 token

    # 3. 对选中 token 进行概率提升
    T = 0.5  # 偏移强度
    adjusted_probs = probs.clone()
    adjusted_probs[mask] = probs[mask] * (1 + T)

    # 4. 重新归一化
    adjusted_probs = adjusted_probs / adjusted_probs.sum()

    return adjusted_probs

def simulate_paraphrase_attack(original_text, paraphrase_model, n_iterations=3):
    """
    模拟多轮改写攻击，测试水印残留
    """
    current_text = original_text
    detection_results = []

    for i in range(n_iterations):
        # 检测水印
        tokens = tokenize(current_text)
        result = watermarker.detect(tokens)
        detection_results.append({
            "iteration": i,
            "text_length": len(current_text),
            "z_score": result["z_score"],
            "detected": result["detected"],
        })

        # 改写
        current_text = paraphrase_model.rewrite(current_text)

    return detection_results

# AI 生成图像的水印方案对比
image_watermark_methods = {
    "DCT域水印": {
        "原理": "在离散余弦变换域中嵌入信息",
        "容量": "中等",
        "鲁棒性": "高（抗压缩）",
        "代表": "Stable Signature (Meta, 2023)",
    },
    "扩散模型水印": {
        "原理": "在扩散过程中微调模型权重嵌入水印",
        "容量": "低",
        "鲁棒性": "非常高",
        "代表": "Tree-Ring Watermark (2023)",
    },
    "频域水印": {
        "原理": "在频域中嵌入伪随机模式",
        "容量": "中等",
        "鲁棒性": "中-高",
        "代表": "SynthID-Image (Google, 2024)",
    },
}

# AudioSeal (Meta, 2024) 原理
audio_watermark_pipeline = """
AI 生成音频
    │
    ├─→ 水印编码器（神经网络）
    │       │
    │       ▼
    │   水印音频（不可感知的差异）
    │
    └─→ 水印检测器
            │
            ▼
        检测结果 + 定位（时间戳级别）
"""

class WatermarkDetectionSystem:
    """
    多层水印检测系统
    """

    def __init__(self):
        self.detectors = {
            "kgw": KGWDetector(),
            "synthid": SynthIDDetector(),
            "statistical": StatisticalDetector(),
            "classifier": ClassifierDetector(),
        }

    def detect(self, text: str) -> dict:
        """
        多检测器融合
        """
        results = {}

        for name, detector in self.detectors.items():
            try:
                result = detector.detect(text)
                results[name] = result
            except Exception as e:
                results[name] = {"error": str(e)}

        # 融合决策
        scores = [r.get("confidence", 0) for r in results.values() if "error" not in r]
        avg_confidence = sum(scores) / len(scores) if scores else 0

        return {
            "is_ai_generated": avg_confidence > 0.7,
            "confidence": avg_confidence,
            "individual_results": results,
            "text_length": len(text),
            "recommendation": self._get_recommendation(avg_confidence),
        }

    def _get_recommendation(self, confidence: float) -> str:
        if confidence > 0.9:
            return "极高置信度 AI 生成，建议标记"
        elif confidence > 0.7:
            return "可能为 AI 生成，建议人工复核"
        elif confidence > 0.4:
            return "不确定，建议进一步分析"
        else:
            return "可能为人类撰写"