大模型的水印方法《A Watermark for Large Language Models》解读
1. 引言
随着 ChatGPT、Llama 等大模型的广泛应用,如何区分人类写作与AI生成文本成为一个重要问题。马里兰大学团队在论文A Watermark for Large Language Models中提出了一种水印机制,可以在不影响文本可读性的情况下,让AI输出的文本带上“隐形标记”,从而在检测阶段以统计方法验证其来源。
2. 大模型水印的目的与应用场景
大语言模型水印机制的根本目标是在不影响文本可读性和流畅性的前提下,通过在生成过程中嵌入隐形特征,使 AI生成的内容能够被可靠识别和追溯。这不仅能区分人类写作与机器生成文本,还能够有效应对由大模型带来的社会风险。 具体而言,水印技术具有以下应用价值:
(1)内容来源验证与风险防范
- 在社交媒体平台,可以用于识别和标记自动化生成的大规模虚假信息,防止谣言传播、选举操纵等问题。
- 在新闻与出版行业,可以为稿件提供溯源机制,确保新闻报道、出版物的真实性与可信度。
(2)学术与教育诚信保障
- 教育机构和教师能够利用水印检测技术判别学生作业、论文是否由AI生成,从而维护学术诚信,避免作弊和学术不端行为。
(3)训练数据质量维护
- 在大规模数据收集过程中,水印检测能够帮助研究人员识别并过滤掉AI合成内容,保证训练语料仍以真实人类创作数据为主,防止模型“自我污染”。
(4)法律、监管与版权保护
- 水印技术可作为一种数字签名或溯源工具,在法律纠纷或版权争议中提供证据支持。
- 监管部门也可利用水印检测手段,实现对AI内容的合规审查与责任追溯。
3. 论文核心方法
该论文的核心思想是基于红/绿词表机制通过在生成过程中嵌入隐形特征,使 AI 生成的内容能够被可靠识别。
(1)词表划分
语言模型的词表记为VVV,大小为∣V∣|V|∣V∣。 在生成第ttt个token时,首先利用前一个token的s(t−1)s(t-1)s(t−1)的哈希值作为随机数种子,来决定词表的划分:
- 绿色集合G⊂VG \subset VG⊂V:推荐生成的token集合
- 红色集合R=V∖GR = V \setminus GR=V∖G:非推荐生成的token集合
其中绿色集合的大小为:∣G∣=γ∣V∣,0<γ<1|G| = \gamma |V|, \quad 0 < \gamma < 1∣G∣=γ∣V∣,0<γ<1红色集合大小为:∣R∣=(1−γ)∣V∣|R| = (1 - \gamma)|V|∣R∣=(1−γ)∣V∣这里γ\gammaγ控制绿名单比例,一般取γ=0.5\gamma = 0.5γ=0.5。哈希函数的作用可以分为以下三个作用:
- 根据前一个 token 生成随机种子,保证每一步划分不同。
- 可复现:检测方也能用相同哈希函数重建红/绿集合。
- 提供隐蔽性:划分看似随机,人类难以察觉,但统计特征可检测。
(2)水印嵌入规则
在得到大语言模型输出的logits向量 l(t)∈R∣V∣l^{(t)} \in \mathbb{R}^{|V|}l(t)∈R∣V∣后,使用以下两种方式进行干预:
- 硬红名单(Hard Red List)
直接屏蔽红色集合RRR中的所有 token,使其概率为零:
pk(t)={elk(t)∑i∈Geli(t),k∈G0,k∈Rp^{(t)}_k = \begin{cases}\frac{e^{l^{(t)}_k}}{\sum_{i \in G} e^{l^{(t)}_i}}, & k \in G \\0, & k \in R\end{cases} pk(t)=⎩⎨⎧∑i∈Geli(t)elk(t),0,k∈Gk∈R
即生成过程中只允许选择绿色集合的token。 - 软红名单(Soft Red List)
给绿色集合中的 logits 增加一个偏置δ>0\delta > 0δ>0,从而在 softmax 中提升绿色集合 token 的概率:p^k(t)={elk(t)+δ∑i∈Reli(t)+∑i∈Geli(t)+δ,k∈Gelk(t)∑i∈Reli(t)+∑i∈Geli(t)+δ,k∈R\hat{p}^{(t)}_k =\begin{cases}\dfrac{e^{l^{(t)}_k + \delta}}{\sum_{i \in R} e^{l^{(t)}_i} + \sum_{i \in G} e^{l^{(t)}_i + \delta}}, & k \in G \\\dfrac{e^{l^{(t)}_k}}{\sum_{i \in R} e^{l^{(t)}_i} + \sum_{i \in G} e^{l^{(t)}_i + \delta}}, & k \in R\end{cases}p^k(t)=⎩⎨⎧∑i∈Reli(t)+∑i∈Geli(t)+δelk(t)+δ,∑i∈Reli(t)+∑i∈Geli(t)+δelk(t),k∈Gk∈R
其中,δ\deltaδ控制偏置强度,δ\deltaδ越大,越倾向于绿色集合。这种方式在高熵文本(多种词都有可能出现)时影响较大,而在低熵文本(确定性强的文本)时影响几乎为零,因此不破坏语义流畅性。
(3)检测方法
检测方无需访问模型,只需知道哈希函数(用于划分红/绿集合)和参数γ\gammaγ(绿名单比例)即可对文本进行检测,检测过程为给定文本序列s=(s(1),…,s(T))s = (s(1), \dots, s(T))s=(s(1),…,s(T)),统计其中落入绿色集合的token数量:∣s∣G=∑t=1T1{s(t)∈Gt}|s|_G = \sum_{t=1}^T \mathbf{1}\{ s(t) \in G_t \}∣s∣G=t=1∑T1{s(t)∈Gt}其中GtG_tGt是由s(t−1)s(t-1)s(t−1)的哈希值生成的绿色集合。
(4)统计检验(zzz检验)
在无水印的假设下(H0H_0H0:文本非 AI 生成),绿色 token 出现概率期望为γ\gammaγ。 因此:
E[∣s∣G]=γT,Var(∣s∣G)=Tγ(1−γ)\mathbb{E}[|s|_G] = \gamma T, \quad \mathrm{Var}(|s|_G) = T\gamma(1-\gamma)E[∣s∣G]=γT,Var(∣s∣G)=Tγ(1−γ)实际观测的绿色 token 数量为∣s∣G|s|_G∣s∣G,构造zzz统计量:
z=∣s∣G−γTTγ(1−γ)
z = \frac{|s|_G - \gamma T}{\sqrt{T \gamma (1-\gamma)}}
z=Tγ(1−γ)∣s∣G−γT若zzz显著大于零(例如 z>4z > 4z>4,对应的ppp值约为 3×10−53 \times 10^{-5}3×10−5),则拒绝H0H_0H0,说明文本中存在水印。即使攻击者替换部分token,也难以完全破坏水印,因为每个替换可能会影响后续多个token的红/绿划分。 数学分析表明,若文本长度为TTT,则攻击者需要修改约T/4T/4T/4的 token 才可能有效去除水印。
4.示例介绍
为了更直观地理解论文中Algorithm1的算法细节,通过一个简单的实例展示模型如何在3步迭代中嵌入水印。 假设词表为: V={cat,dog,apple,banana,car,bus}V = \{\mathrm{cat, dog, apple, banana, car, bus}\}V={cat,dog,apple,banana,car,bus},绿名单比例取γ=0.5\gamma = 0.5γ=0.5,即词表在每一步都会被随机划分为一半绿色集合(允许生成),一半红色集合(禁止生成)。
Step 1:基于Prompt生成第一个词
- 输入的prompt为
The
,模型在没有水印干预的情况下,输出logits对应的概率分布为:cat (0.3), dog (0.25), apple (0.2), banana (0.15), car (0.1)
- 在正常情况下,模型最有可能生成
cat
,但由于使用了水印机制,我们需要先根据哈希函数进行划分。取前一个token=cat
,计算哈希值为H(H(H(cat
)=17)= 17)=17,用这个种子随机划分词表:- Green = {dog, apple, car}
- Red = {cat, banana, bus}
- 因为是硬红名单,红色集合中的 token 完全禁止采样。因此虽然
cat
概率最高,但它在红色集合中,被屏蔽掉。最终模型只能从Green列表中选择概率最大的结果是dog
。
Step 2:继续生成第二个词
- 当前序列为
The dog
。上一个生成的token是dog
,所以再次利用哈希函数计算为H(H(H(dog
)=42) = 42)=42,划分结果为:- Green = {cat, banana, bus}
- Red = {dog, apple, car}
- 模型在这一位置的概率分布是:
apple (0.35), bus (0.25), car (0.2), banana (0.1), cat (0.1)
- 从分布中可以看到
apple
的概率最高,但是由于它属于Red集合,因此被禁止选择;同样car
也被排除。 在Green集合中,bus
的概率最大,因此最终生成结果为bus
。
Step 3:生成第三个词
-
当前序列为
The dog bus
,上一个生成的 token 是bus
,计算哈希值为H(H(H(bus
)=93) = 93)=93,划分结果为:- Green = {apple, car, dog}
- Red = {banana, cat, bus}
-
模型在这一位置的概率分布为:
car (0.4), cat (0.3), apple (0.2), banana (0.1)
-
观察分布可以发现,
car
概率最高,cat
次之。但由于cat
属于红色集合,因此被禁止选择;banana
也在Red中。因此在Green集合中概率最高的car
被选中。
🔹 最终生成结果
经过以上3步迭代,最终的生成序列为:
The dog bus car
水印机制并不会直接篡改文本语义,而是通过 限制采样范围,在概率层面微妙地调整生成结果。
5.代码实现
下面给出一个基于Hugging Face的最小实现。支持硬红名单和软红名单两种水印生成方式,并提供检测方法。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
import hashlib
import argparse
from scipy.stats import norm
# --------------------------
# 工具函数
# --------------------------
def hash_token(token, seed=0):"""将 token 转换成伪随机种子"""h = hashlib.sha256((str(token) + str(seed)).encode()).hexdigest()return int(h, 16)def partition_vocab(vocab_size, seed, gamma=0.5):"""将词表分成 green list 和 red list"""rng = np.random.default_rng(seed)perm = rng.permutation(vocab_size)split = int(gamma * vocab_size)green = set(perm[:split])red = set(perm[split:])return green, red# --------------------------
# 水印采样器
# --------------------------def watermark_sampling(logits, prev_token, gamma=0.5, delta=2.0, hard=False):"""根据水印规则修改 logits"""vocab_size = logits.shape[-1]seed = hash_token(prev_token) % (2**32)green, red = partition_vocab(vocab_size, seed, gamma)if hard:# 硬红名单:直接屏蔽红色集合mask = torch.full_like(logits, float("-inf"))mask[list(green)] = 0logits = logits + maskelse:# 软红名单:给绿色集合加 δ 偏置bias = torch.zeros_like(logits)bias[list(green)] = deltalogits = logits + biasreturn logits# --------------------------
# 文本生成
# --------------------------def generate_with_watermark(model, tokenizer, prompt, max_new_tokens=50, gamma=0.5, delta=2.0, hard=False):input_ids = tokenizer(prompt, return_tensors="pt").input_idsgenerated = input_ids.clone()for _ in range(max_new_tokens):outputs = model(generated)logits = outputs.logits[:, -1, :].squeeze(0)prev_token = int(generated[0, -1])logits = watermark_sampling(logits, prev_token, gamma, delta, hard)probs = torch.softmax(logits, dim=-1)next_token = torch.multinomial(probs, num_samples=1).unsqueeze(0) # shape [1,1]generated = torch.cat([generated, next_token], dim=1)return tokenizer.decode(generated[0], skip_special_tokens=True)# --------------------------
# 水印检测
# --------------------------def detect_watermark(text, tokenizer, gamma=0.5):tokens = tokenizer(text, return_tensors="pt").input_ids[0].tolist()T = len(tokens) - 1green_count = 0for i in range(1, len(tokens)):seed = hash_token(tokens[i - 1]) % (2**32)green, red = partition_vocab(tokenizer.vocab_size, seed, gamma)if tokens[i] in green:green_count += 1# z 检验expected = gamma * Tvar = T * gamma * (1 - gamma)z = (green_count - expected) / np.sqrt(var)p_value = 1 - norm.cdf(z) # 单尾检验return z, p_value# --------------------------
# 主程序
# --------------------------def main():parser = argparse.ArgumentParser()parser.add_argument("--model_path", type=str, default="./models/Phi-3-mini-128k-instruct")parser.add_argument("--prompt", type=str, default="The future of AI is")parser.add_argument("--max_new_tokens", type=int, default=50)parser.add_argument("--hard", action="store_true", help="使用硬红名单")args = parser.parse_args()print(f"Loading model from {args.model_path}...")tokenizer = AutoTokenizer.from_pretrained(args.model_path)model = AutoModelForCausalLM.from_pretrained(args.model_path, torch_dtype=torch.float16, device_map="auto")print("Generating text with watermark...")text = generate_with_watermark(model, tokenizer, args.prompt, max_new_tokens=args.max_new_tokens, hard=args.hard)print("\nGenerated text:\n", text)print("\nDetecting watermark...")z, p = detect_watermark(text, tokenizer)print(f"z-score: {z:.2f}, p-value: {p:.2e}")if p < 0.01:print("水印检测结果:文本可能是 AI 生成的")else:print("水印检测结果:无法确认文本为 AI 生成")if __name__ == "__main__":main()
🔹 运行示例
python watermark_demo.py --model_path ./models/Phi-3-mini-128k-instruct --prompt "The future of AI is" --max_new_tokens 30
🔹 输出示例:
Generated text:
The future of AI is shaping industries rapidly...z-score: 6.42, p-value: 1.6e-10
✅ 水印检测结果:文本可能是AI生成的