1 引言
该研究聚焦于大语言模型中的无失真水印（详情见博客斯坦福大学推出无失真大模型水印技术），这一机制通过隐藏密钥在不改变模型分布的前提下嵌入水印，被认为比传统修改分布的方法更安全。然而，论文《Toward Breaking Watermarks in Distortion-free Large Language Models》指出这种方法并不牢靠：攻击者能够通过少量样本逆推出隐藏密钥，并利用其实施伪造攻击 ，使检测器错误判断文本已被水印。作者提出了一种基于混合整数线性规划 (MILP) 的攻击框架，不仅能在无扰动情况下恢复密钥，还能应对密钥偏移和样本损坏等复杂场景。该工作创新性地展示了即便是先进的无失真水印技术，也存在被破解的风险，从而对现有水印方法的稳健性提出了严峻挑战。
2 论文方法解析
2.1 符号与设置
词表记为 V V V，其中划分出两个不相交的子集 G ⊂ V G \subset V G⊂V与 R = V ∖ G R = V \setminus G R=V∖G，满足 G ∩ R = ∅ G \cap R = \varnothing G∩R=∅且 G ∪ R = V G \cup R = V G∪R=V。设自回归语言模型为 p ( ⋅ ∣ y : j − 1 ) p(\cdot \mid y_{:j-1}) p(⋅∣y:j−1)，它在给定前缀 y : j − 1 y_{:j-1} y:j−1的条件下输出下一个 token 的分布；水印密钥为 ξ = ( ξ 1 , ... , ξ m ) \xi = (\xi_1,\dots,\xi_m) ξ=(ξ1,...,ξm)，其中每个 ξ j ∈ [ 0 , 1 ] \xi_j \in [0,1] ξj∈[0,1]。对于样本 i i i的第 j j j个位置，定义 q i j ≡ P j ( i ) ( G ) = ∑ x ∈ G p ( x ∣ y : j − 1 ( i ) ) q_{ij} \equiv P_j^{(i)}(G) = \sum_{x\in G} p(x\mid y^{(i)}_{:j-1}) qij≡Pj(i)(G)=∑x∈Gp(x∣y:j−1(i))，即在该条件下下一个token落在 G G G中的概率质量。
2.2 二进制水印生成
无失真水印算法的核心思想是：在不改变语言模型原始分布的前提下，通过一个随机密钥序列来决定每一步的采样集合。具体来说，在每一步生成时，语言模型会计算集合 G G G的概率： q i j = P j ( G ) = ∑ x ∈ G p ( x ∣ y : j − 1 ( i ) ) , q_{ij} = P_j(G) = \sum_{x \in G} p(x \mid y^{(i)}{:j-1}), qij=Pj(G)=x∈G∑p(x∣y:j−1(i)),其中 y : j − 1 ( i ) y^{(i)}{:j-1} y:j−1(i) 表示样本 i i i 的前 j − 1 j-1 j−1 个已生成 token。 然后，算法将该概率与密钥值 ξ j ∈ [ 0 , 1 ] \xi_j \in [0,1] ξj∈[0,1]进行比较，从而生成一个二进制序列： y j ( i ) = 1 { ξ j > q i j } . y^{(i)}j \;=\; \mathbf{1}\{\xi_j > q{ij}\}. yj(i)=1{ξj>qij}.这一规则可以解释为：
- 若 ξ j > q i j \xi_j > q_{ij} ξj>qij，则在第 j j j 步生成时，强制从集合 G G G 中采样，记为 y j ( i ) = 1 y^{(i)}_j=1 yj(i)=1；
- 若 ξ j ≤ q i j \xi_j \le q_{ij} ξj≤qij，则在第 j j j 步生成时，强制从集合 R R R 中采样，记为 y j ( i ) = 0 y^{(i)}_j=0 yj(i)=0。
因此，整个水印过程把连续的概率分布映射为一个二进制序列 y ( i ) = ( y 1 ( i ) , y 2 ( i ) , ... , y m ( i ) ) y^{(i)}=(y^{(i)}_1,y^{(i)}_2,\dots,y^{(i)}_m) y(i)=(y1(i),y2(i),...,ym(i))，而该序列与隐藏密钥 ξ \xi ξ 一一对应。
2.3 攻击者获取的信息
攻击者可以直接获得生成的文本序列 Y ( i ) Y^{(i)} Y(i)，因此能够把每个输出token根据是否属于已知的集合 G G G或 R R R映射成二进制序列 y ( i ) y^{(i)} y(i)。
- 当攻击者可以从API获得logits或概率分布时，计算每个位置的 q i j = P j ( i ) ( G ) = ∑ x ∈ G p ( x ∣ y ( i ) : j − 1 ) q_{ij}=P_j^{(i)}(G)=\sum_{x\in G} p(x\mid y^{(i)}{:j-1}) qij=Pj(i)(G)=x∈G∑p(x∣y(i):j−1)是直接且精确的；在该情形下，观测对 ( q i j , y ( i ) j ) (q{ij},y^{(i)}j) (qij,y(i)j)可立即构成对密钥 ξ j \xi_j ξj 的不等式约束。
- 当攻击者只能从API中获得输出的token，仍可通过对相同前缀 y : j − 1 y_{:j-1} y:j−1进行重复采样或对多组相似前缀做蒙特卡洛估计来近似 q i j q_{ij} qij：即统计在大量采样中落在 G G G的次数除以总采样次数得到经验概率 P ^ j ( i ) ( G ) ≈ q i j \hat P_j^{(i)}(G)\approx q_{ij} P^j(i)(G)≈qij。
需要注意的是，在无logits的黑盒场景下若攻击者只能获得每个前缀的一次采样且无法重现或重复调用模型，则无法直接用频率估计 q i j q_{ij} qij；此外，估计精度受采样次数、采样预算、以及模型随机性的影响（采样次数越多、估计越精确），因此在该场景下恢复 ξ \xi ξ仍可行但通常需要更多查询与更高的计算/查询成本。
2.4 构建不等式约束
每个观测对 ( q i j , y j ( i ) ) (q_{ij}, y^{(i)}_j) (qij,yj(i)) 都会为未知的密钥分量 ξ j \xi_j ξj提供一个单边不等式约束。
- 若 y j ( i ) = 1 y^{(i)}j = 1 yj(i)=1，说明第 j j j步生成的token落入集合 G G G，则一定有 ξ j > q i j \xi_j > q{ij} ξj>qij；
- 若 y j ( i ) = 0 y^{(i)}j = 0 yj(i)=0，说明第 j j j步生成的token落入集合 R R R，则一定有 ξ j ≤ q i j \xi_j \le q{ij} ξj≤qij。
也就是说，每次采样都为攻击者缩小了 ξ j \xi_j ξj的可能范围。随着样本数量 n n n的增加，约束条件逐步积累并收紧，逼近真实的水印密钥。 收集 n n n个样本后，对每个索引 j j j可以得到如下的约束集合： { ξ j > q i j ∀ i : y j ( i ) = 1 } ∪ { ξ j ≤ q i j ∀ i : y j ( i ) = 0 } . \{\, \xi_j > q_{ij} \;\; \forall i: y^{(i)}j=1 \,\} \;\;\cup\;\; \{\, \xi_j \le q{ij} \;\; \forall i: y^{(i)}_j=0 \,\}. {ξj>qij∀i:yj(i)=1}∪{ξj≤qij∀i:yj(i)=0}.
2.5 整数规划估计密钥
论文将水印密钥的恢复问题形式化为一个在线性约束下的优化问题，即在所有位置 j j j上同时求上下界的 LP/MILP 模型，从而给出对密钥 ξ \xi ξ的估计。该方法的核心思想是：每个观测到的 ( q i j , y j ( i ) ) (q_{ij}, y^{(i)}j) (qij,yj(i))都为 ξ j \xi_j ξj提供了一个不等式约束，收集多个样本后，攻击者可以通过优化问题求解一组最紧的上下界，并据此逼近真实密钥。
在最理想的情况下，观测到的二进制序列 y ( i ) y^{(i)} y(i) 没有发生任何偏移或损坏，攻击者可以直接利用所有样本建立线性规划模型来求解上下界。
（1）下界估计：
minimize ∑ j = 1 m ξ ‾ j s.t. ξ ‾ j ≥ 1 { y j ( i ) = 1 } q i j , ∀ i , j , ξ ‾ j ∈ R . \begin{aligned}\text{minimize}\quad & \sum{j=1}^m \underline{\xi}_j \\\text{s.t.}\quad & \underline{\xi}_j \ge \mathbf{1}\{y^{(i)}j=1\}\,q{ij}, \quad \forall i,j, \\& \underline{\xi}j \in\mathbb{R}.\end{aligned} minimizes.t.j=1∑mξjξj≥1{yj(i)=1}qij,∀i,j,ξj∈R.求解以上不等式组即可以得到所有下界的最大值。
（2）上界估计：
maximize ∑ j = 1 m ξ ‾ j s.t. ξ ‾ j ≤ 1 { y j ( i ) = 0 } ( q i j − 1 ) + 1 , ∀ i , j . \begin{aligned} \text{maximize}\quad & \sum{j=1}^m \overline{\xi}_j \\ \text{s.t.}\quad & \overline{\xi}_j \le \mathbf{1}\{y^{(i)}j=0\}(q{ij}-1)+1, \quad \forall i,j. \end{aligned} maximizes.t.j=1∑mξjξj≤1{yj(i)=0}(qij−1)+1,∀i,j.求解以上不等式组即可以得到所有上界的最小值。最终估计取上下界的中点： ξ ^ j = ξ ‾ j + ξ ‾ j 2 . \hat{\xi}_j = \frac{\underline{\xi}_j + \overline{\xi}_j}{2}. ξ^j=2ξj+ξj.这种方法相当于在保证所有样本约束成立的前提下，找到 ξ j \xi_j ξj 的最大可能区间，并以中点作为估计值。
2.6 检测与假设检验
给定候选密钥 ξ ^ \hat{\xi} ξ^和待测输入文本 Y ~ \tilde{Y} Y~，检测器首先将 Y ~ \tilde{Y} Y~映射为相应的二进制子序列，然后在文本中枚举或滑动取若干长度为 k k k 的子块，将这些子块分别与 ξ ^ \hat{\xi} ξ^的对应片段对齐并计算对齐代价，最终基于这些代价构造检验统计量 ϕ ( Y ~ , ξ ^ ) \phi(\tilde{Y},\hat{\xi}) ϕ(Y~,ξ^)，并输出相应的p-value。
- H 0 H_0 H0： Y ~ \tilde{Y} Y~与候选密钥 ξ ^ \hat{\xi} ξ^无关（即 Y ~ \tilde{Y} Y~未被水印）；
- H 1 H_1 H1： Y ~ \tilde{Y} Y~与候选密钥 ξ ^ \hat{\xi} ξ^有统计相关性（即 Y ~ \tilde{Y} Y~被水印）。
通常的检测流程包括：1) 将 Y ~ \tilde{Y} Y~二进制化并提取所有长度为 k k k的子序列；2) 对每个子序列计算与 ξ ^ \hat{\xi} ξ^对应片段的对齐代价（例如汉明距离或加权差异）；3) 以这些代价的最小值或某种汇总统计作为 ϕ ( Y ~ , ξ ^ ) \phi(\tilde{Y},\hat{\xi}) ϕ(Y~,ξ^)；4) 在零假设 H 0 H_0 H0下通过模拟或理论分布得到该统计量的参考分布并据此计算 p-value。若 p-value 小于预设显著性阈值（例如 0.05 0.05 0.05），则拒绝 H 0 H_0 H0，认为文本被水印。
3 破解无失真水印示例介绍
3.1 构造词表与集合划分
- 大模型词表： V = { good , excellent , bad , terrible } V = \{\text{good}, \text{excellent}, \text{bad}, \text{terrible}\} V={good,excellent,bad,terrible}
- 攻击者词表划分： G = { good , excellent } , R = { bad , terrible } G = \{\text{good}, \text{excellent}\}, \quad R = \{\text{bad}, \text{terrible}\} G={good,excellent},R={bad,terrible}
- 集合关系： G ∩ R = ∅ , G ∪ R = V G \cap R = \varnothing, \; G \cup R = V G∩R=∅,G∪R=V
3.2 计算集合概率
假设语言模型在每一步生成时，下一个token属于集合 G G G 的概率为条件概率： P j ( G ) = ∑ x ∈ G p ( x ∣ s : j − 1 ) P_j(G) = \sum_{x \in G} p(x \mid s_{:j-1}) Pj(G)=x∈G∑p(x∣s:j−1)其中 s : j − 1 s_{:j-1} s:j−1表示前 j − 1 j-1 j−1个已生成的token，迭代3次则有： P 1 ( G ) = 0.35 = ∑ x ∈ G p ( x ) P 2 ( G ) = 0.65 = ∑ x ∈ G p ( x ∣ s 1 ) P 3 ( G ) = 0.20 = ∑ x ∈ G p ( x ∣ s 1 , s 2 ) \begin{aligned}P_1(G) &= 0.35 = \sum_{x \in G} p(x ) \\ P_2(G) &= 0.65 = \sum_{x \in G} p(x \mid s_1) \\ P_3(G) &= 0.20 = \sum_{x \in G} p(x \mid s_1, s_2)\end{aligned} P1(G)P2(G)P3(G)=0.35=x∈G∑p(x)=0.65=x∈G∑p(x∣s1)=0.20=x∈G∑p(x∣s1,s2)
3.3 大模型生成水印序列
假设密钥序列长度为3，即 ξ = [ 0.5 , 0.3 , 0.7 ] \xi = [0.5, 0.3, 0.7] ξ=[0.5,0.3,0.7], 且每个位置 ξ j ∈ [ 0 , 1 ] \xi_j \in [0,1] ξj∈[0,1]。生成水印二进制序列规则： y j = I { ξ j > P j ( G ) } y_j = \mathbb{I}\{\xi_j > P_j(G)\} yj=I{ξj>Pj(G)} 其中 I { ⋅ } \mathbb{I}\{\cdot\} I{⋅}为指示函数，水印序列的计算过程为： 0.5 > 0.35 ⟹ y 1 = 1 0.3 < 0.65 ⟹ y 2 = 0 0.7 > 0.20 ⟹ y 3 = 1 \begin{aligned}0.5 > 0.35 &\Longrightarrow y_1 = 1 \\ 0.3 < 0.65 &\Longrightarrow y_2 = 0 \\ 0.7 > 0.20 &\Longrightarrow y_3 = 1 \end{aligned} 0.5>0.350.3<0.650.7>0.20⟹y1=1⟹y2=0⟹y3=1其中对应的 y j y_j yj是由水印密钥 ξ j \xi_j ξj与 q j q_j qj 的比较关系决定的二进制值，最终生成水印序列表示为：
y = [ 1 , 0 , 1 ] y = [1, 0, 1] y=[1,0,1]
3.4 攻击者获取水印序列
攻击者可以通过访问大模型的API可以获得生成的文本序列 和集合划分信息 （其中攻击者已知 G , R G, R G,R的定义），则攻击者可以立刻把输出序列转化为二进制序 y y y。当大模型API只返回token信息，不返回logits信息时，但攻击者依然可以通过模拟方式估计概率 P j ( G ) P_j(G) Pj(G)，对同一个上下文 y : j − 1 y_{:j-1} y:j−1，多次调用大模型的API，统计 G G G和 R R R中token的出现频率：
P ^ j ( G ) = Number of tokens in G Total number of samples \hat{P}_j(G) = \frac{\text{Number of tokens in G}}{\text{Total number of samples}} P^j(G)=Total number of samplesNumber of tokens in G当采样次数 n n n 较大时， P ^ j ( G ) ≈ P j ( G ) \hat{P}_j(G) \approx P_j(G) P^j(G)≈Pj(G)，攻击者就能获得较精确的概率估计。假设攻击者收集到 n = 3 n = 3 n=3个语言模型生成的样本，这些3个样本经过二进制化即可得到3个序列：
| 样本 | q 1 q_1 q1 | q 2 q_2 q2 | q 3 q_3 q3 | y y y |
|---|---|---|---|---|
| 1 1 1 | 0.35 0.35 0.35 | 0.65 0.65 0.65 | 0.20 0.20 0.20 | [ 1 , 0 , 1 ] [1,0,1] [1,0,1] |
| 2 2 2 | 0.60 0.60 0.60 | 0.40 0.40 0.40 | 0.10 0.10 0.10 | [ 0 , 0 , 1 ] [0,0,1] [0,0,1] |
| 3 3 3 | 0.45 0.45 0.45 | 0.20 0.20 0.20 | 0.80 0.80 0.80 | [ 1 , 1 , 0 ] [1,1,0] [1,1,0] |
当攻击者收集到足够多的样本时，能够更加逼近真实的水印密钥序列。
3.5 攻击者构造约束条件
根据水印序列生成规则：
- 若 y j = 1 y_{j} = 1 yj=1，则说明在第 j j j步生成时，水印密钥满足 ξ j > q j \xi_j > q_{j} ξj>qj；
- 若 y j = 0 y_{j} = 0 yj=0，则说明在第 j j j步生成时，水印密钥满足 ξ j ≤ q j \xi_j \leq q_{j} ξj≤qj。
因此，每一个观测到的 ( q i j , y i j ) (q_{ij}, y_{ij}) (qij,yij) 对都会给水印密钥 ξ j \xi_j ξj提供一个不等式约束 。随着样本数量的增加，这些约束能够逐步缩小 ξ j \xi_j ξj 的可能取值范围，从而帮助攻击者逼近真实的水印密钥。逐个样本写出的约束如下：
样本1对应水印序列为 [ 1 , 0 , 1 ] [1,0,1] [1,0,1]：
y 1 = 1 ⇒ ξ 1 > 0.35 y_1=1 \;\Rightarrow\; \xi_1 > 0.35 y1=1⇒ξ1>0.35 （在集合 G G G中，即 ξ 1 \xi_1 ξ1大于 P 1 ( G ) = 0.35 P_1(G)=0.35 P1(G)=0.35）
y 2 = 0 ⇒ ξ 2 ≤ 0.65 y_2=0 \;\Rightarrow\; \xi_2 \le 0.65 y2=0⇒ξ2≤0.65 （在集合 R R R中，即 ξ 2 \xi_2 ξ2小于等于 P 2 ( G ) = 0.65 P_2(G)=0.65 P2(G)=0.65）
y 3 = 1 ⇒ ξ 3 > 0.20 y_3=1 \;\Rightarrow\; \xi_3 > 0.20 y3=1⇒ξ3>0.20 （在集合 G G G中，即 ξ 3 \xi_3 ξ3大于 P 3 ( G ) = 0.20 P_3(G)=0.20 P3(G)=0.20）
样本2对应水印序列为 [ 0 , 0 , 1 ] [0,0,1] [0,0,1]：y 1 = 0 ⇒ ξ 1 ≤ 0.60 y_1=0 \;\Rightarrow\; \xi_1 \le 0.60 y1=0⇒ξ1≤0.60 （在集合 R R R中，即 ξ 1 \xi_1 ξ1小于等于 P 1 ( G ) = 0.60 P_1(G)=0.60 P1(G)=0.60）
y 2 = 0 ⇒ ξ 2 ≤ 0.40 y_2=0 \;\Rightarrow\; \xi_2 \le 0.40 y2=0⇒ξ2≤0.40 （在集合 R R R中，即 ξ 2 \xi_2 ξ2小于等于 P 2 ( G ) = 0.40 P_2(G)=0.40 P2(G)=0.40）
y 3 = 1 ⇒ ξ 3 > 0.10 y_3=1 \;\Rightarrow\; \xi_3 > 0.10 y3=1⇒ξ3>0.10 （在集合 G G G中，即 ξ 3 \xi_3 ξ3大于 P 3 ( G ) = 0.10 P_3(G)=0.10 P3(G)=0.10）
样本3对应水印序列为 [ 1 , 1 , 0 ] [1,1,0] [1,1,0]：y 1 = 0 ⇒ ξ 1 > 0.45 y_1=0 \;\Rightarrow\; \xi_1 > 0.45 y1=0⇒ξ1>0.45 （在集合 R R R中，即 ξ 1 \xi_1 ξ1大于 P 1 ( G ) = 0.45 P_1(G)=0.45 P1(G)=0.45）
y 2 = 0 ⇒ ξ 2 > 0.20 y_2=0 \;\Rightarrow\; \xi_2 > 0.20 y2=0⇒ξ2>0.20 （在集合 G G G中，即 ξ 2 \xi_2 ξ2大于 P 2 ( G ) = 0.20 P_2(G)=0.20 P2(G)=0.20）
y 3 = 1 ⇒ ξ 3 ≤ 0.80 y_3=1 \;\Rightarrow\; \xi_3 \le 0.80 y3=1⇒ξ3≤0.80 （在集合 R R R中，即 ξ 3 \xi_3 ξ3小于等于 P 3 ( G ) = 0.80 P_3(G)=0.80 P3(G)=0.80）
根据以上三组不等式，整合约束区间综合得到：0.45 < ξ 1 ≤ 0.60 0.45 < \xi_1 \le 0.60 0.45<ξ1≤0.60
0.20 < ξ 2 ≤ 0.40 0.20 < \xi_2 \le 0.40 0.20<ξ2≤0.40
0.20 < ξ 3 ≤ 0.80 0.20 < \xi_3 \le 0.80 0.20<ξ3≤0.80
由此可见，每个样本都对 ξ j \xi_j ξj提供了不同方向的约束：有些是下界（ > > >），有些是上界（ ≤ \le ≤）。当攻击者收集到多个样本后，就能综合这些不等式，得到一个更窄的区间，从而逐步逼近真实的水印密钥序列。
3.6 线性规划估计密钥
攻击者在整合多个样本时，会将所有给出的不等式约束统一起来：每个 y i j = 1 y_{ij}=1 yij=1的情况为 ξ j \xi_j ξj提供一个下界，而每个 y i j = 0 y_{ij}=0 yij=0的情况为 ξ j \xi_j ξj提供一个上界，最终攻击者取所有下界中的最大值作为 ξ j ‾ \underline{\xi_j} ξj、所有上界中的最小值作为 ξ j ‾ \overline{\xi_j} ξj，并用二者的中点来作为对真实密钥 ξ j \xi_j ξj 的估计值： ξ j ^ = ξ j ‾ + ξ j ‾ 2 \hat{\xi_j} = \frac{\underline{\xi_j} + \overline{\xi_j}}{2} ξj^=2ξj+ξj
在示例则有：
- ξ 1 ^ = ( 0.45 + 0.60 ) / 2 = 0.525 \hat{\xi_1} = (0.45+0.60)/2 = 0.525 ξ1^=(0.45+0.60)/2=0.525
- ξ 2 ^ = ( 0.20 + 0.40 ) / 2 = 0.30 \hat{\xi_2} = (0.20+0.40)/2 = 0.30 ξ2^=(0.20+0.40)/2=0.30
- ξ 3 ^ = ( 0.20 + 0.80 ) / 2 = 0.50 \hat{\xi_3} = (0.20+0.80)/2 = 0.50 ξ3^=(0.20+0.80)/2=0.50
最终可得估计密钥为： ξ ^ = [ 0.525 , 0.30 , 0.50 ] \hat{\xi} = [0.525, 0.30, 0.50] ξ^=[0.525,0.30,0.50]与真实密钥 [ 0.5 , 0.3 , 0.7 ] [0.5, 0.3, 0.7] [0.5,0.3,0.7] 非常接近，即攻击者攻击成功。
4. 代码实现
这段代码实现了论文中无偏移、无损坏场景下的水印密钥恢复攻击，并展示了攻击者如何在理想条件下恢复大模型中的无失真水印密钥。
python
import argparse
from typing import List, Tuple, Optional
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# ---------- 辅助函数 ----------
def tokens_to_ids_single_token(tokenizer, token_str: str) -> Optional[int]:
"""
将一个字符串映射到单个 token id；若映射结果为多个 token，则返回 None 并警告。
建议 G 中的元素尽量选单 token。
"""
enc = tokenizer(token_str, add_special_tokens=False)
ids = enc["input_ids"]
if len(ids) == 1:
return ids[0]
else:
print(f"警告：'{token_str}' 被编码为多个 token {ids}. 建议选择单 token 字符。")
return None
# ---------- 生成并同时获取 q_j（白盒：从 logits） ----------
@torch.no_grad()
def generate_and_collect_q(model, tokenizer, prompt: str, m: int, G_token_ids: List[int],
device: torch.device, temperature: float = 1.0, top_k: Optional[int] = 50
) -> Tuple[List[int], List[float], List[int]]:
"""
在白盒模式下生成长度 m 的序列，并在每一步计算 q_j。
返回：
generated_ids: 生成的 token id 列表
q_list: 每步对应的 q_j
y_list: 每步对应的二进制值（token 属于 G -> 1，否则 0）
"""
model.eval()
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(device)
input_ids = inputs["input_ids"]
generated = input_ids.clone()
q_list, y_list, generated_ids = [], [], []
for j in range(m):
outputs = model(generated)
logits = outputs.logits[0, -1, :]
probs = torch.softmax(logits / max(temperature, 1e-8), dim=-1)
# 计算 q_j
qj = probs[torch.tensor(G_token_ids, device=probs.device)].sum().item()
q_list.append(qj)
# 采样下一个 token
if top_k is not None:
topk = torch.topk(logits, top_k)
topk_ids, topk_logits = topk.indices, topk.values
topk_probs = torch.softmax(topk_logits / max(temperature, 1e-8), dim=-1)
chosen_idx = torch.multinomial(topk_probs, num_samples=1)
next_id = topk_ids[chosen_idx]
else:
next_id = torch.multinomial(probs, num_samples=1)
# 保证 next_id 形状为 (1,1)
next_id = next_id.view(1, 1).to(device)
next_id_scalar = next_id.item()
inG = 1 if next_id_scalar in G_token_ids else 0
y_list.append(inG)
generated_ids.append(next_id_scalar)
# 拼接 token 到序列
generated = torch.cat([generated, next_id], dim=1)
return generated_ids, q_list, y_list
# ---------- 估计 xi（no-alteration 简化方法） ----------
def estimate_xi_from_samples(q_matrix: List[List[float]], y_matrix: List[List[int]]) -> Tuple[List[float], List[float], List[float]]:
"""
根据样本估计 xi 的上下界及中点。
"""
n = len(q_matrix)
if n == 0:
return [], [], []
m = len(q_matrix[0])
lower_bounds = [0.0] * m
upper_bounds = [1.0] * m
for j in range(m):
lowers = [q_matrix[i][j] for i in range(n) if y_matrix[i][j] == 1]
uppers = [q_matrix[i][j] for i in range(n) if y_matrix[i][j] == 0]
lower_bounds[j] = max(lowers) if lowers else 0.0
upper_bounds[j] = min(uppers) if uppers else 1.0
mid = [(lower_bounds[j] + upper_bounds[j]) / 2.0 for j in range(m)]
return lower_bounds, upper_bounds, mid
# ---------- 主流程 ----------
def main():
parser = argparse.ArgumentParser(description="Distortion-free watermark attack (no-alteration case)")
parser.add_argument("--model_path", type=str, default="./models/Phi-3-mini-128k-instruct",
help="本地模型路径，默认: ./models/Phi-3-mini-128k-instruct")
parser.add_argument("--prompt", type=str, default="Please recommend some classic movies",
help="输入提示语，默认: 'Please recommend some classic movies'")
parser.add_argument("--m", type=int, default=10, help="生成二进制序列长度")
parser.add_argument("--n_samples", type=int, default=10, help="收集的样本数量")
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# G 集合（示例：单 token）
G_tokens = ["good", "excellent"]
print("加载 tokenizer 与模型：", args.model_path)
tokenizer = AutoTokenizer.from_pretrained(args.model_path, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
args.model_path,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to(device)
# G 集合转为 token ids
G_ids = []
for t in G_tokens:
tid = tokens_to_ids_single_token(tokenizer, t)
if tid is not None:
G_ids.append(tid)
if not G_ids:
raise ValueError("G 集合在 tokenizer 下没有有效的单 token 映射。")
print("G token ids:", G_ids)
# 收集样本
q_matrix, y_matrix, generated_texts = [], [], []
for i in range(args.n_samples):
gen_ids, q_list, y_list = generate_and_collect_q(model, tokenizer, args.prompt, args.m, G_ids, device)
q_matrix.append(q_list)
y_matrix.append(y_list)
generated_texts.append(tokenizer.decode(gen_ids))
# 估计 xi
lower_bounds, upper_bounds, xi_hat = estimate_xi_from_samples(q_matrix, y_matrix)
print("\n=== 攻击结果 ===")
print("下界 lower_bounds:", lower_bounds)
print("上界 upper_bounds:", upper_bounds)
print("中点估计 xi_hat:", xi_hat)
# 简单判断攻击是否成功：如果区间长度普遍很窄，则认为成功
intervals = [upper_bounds[j] - lower_bounds[j] for j in range(len(lower_bounds))]
avg_interval = sum(intervals) / len(intervals)
print(f"\n平均区间宽度: {avg_interval:.4f}")
if avg_interval < 0.2: # 阈值可调整
print("✅ 攻击成功：估计的密钥区间足够精确。")
else:
print("❌ 攻击不成功：估计的密钥区间过宽。")
if __name__ == "__main__":
main()🔹 运行示例
bash
python attack_distortion_free.py --model_path ./models/Phi-3-mini-128k-instruct --prompt "Please recommend some classic movies."🔹 输出结果：
=== 攻击结果 ===
下界 lower_bounds: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
上界 upper_bounds: [2.288818359375e-05, 4.172325134277344e-07, 1.7881393432617188e-07, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
中点估计 xi_hat: [1.1444091796875e-05, 2.086162567138672e-07, 8.940696716308594e-08, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
平均区间宽度: 0.0000
✅ 攻击成功：估计的密钥区间足够精确。