Efficient Refusal Ablation in LLM through Optimal Transport
(ICLR 2026 Trustworthy AI)
motivitaion
此前的表征级别越狱方法(如 RFA,拒绝特征消除)通过计算有害和无害提示词激活状态之间的“均值差”,提取出一个一维的“拒绝方向”,并在所有模型层中将其消除 。
要么是沿着$$d=\mu_2-\mu_1$$方向平移激活值,要么是通过正交投影$T _ {proj}(x;d)=(I-dd^T)/||d||^2x$。
这些方法完全忽略了模型激活空间中丰富的多维分布和协方差几何结构 。
方法
作者提出了一种全新的视角,将“越狱”视为一个分布匹配问题,并通过以下三个核心创新点来实现 :
- 引入高斯最优传输 (Gaussian OT):目标是寻找一个映射函数,将有害激活的分布完全转化为无害激活的分布 。在假设分布为高斯分布的情况下,最优传输可以通过一个闭式解的仿射变换 $T(x)=Ax+b$ 来实现 。矩阵 $A$ 负责对齐协方差结构,而向量 $b$ 负责对齐均值 。
- 结合 PCA 降维以提升效率:由于 LLM 的表征空间通常高达数千维(如 4096 到 8192 维),直接计算最优传输会导致协方差矩阵病态且计算成本极高 。因此,作者先使用主成分分析(PCA)将数据投影到一个低维子空间(仅保留 $k$ 个主成分),在低维空间计算出映射关系后,再将其映射回高维的原空间 。
- 特定层干预 (Layer-selective intervention):与传统方法在网络的所有层进行干预不同,该方法仅在网络深度的 40% 到 60% 处选择 1 到 2 个特定层进行干预 。
目标是找到一个映射函数 $T^{(l)}: \mathbb{R}^d \rightarrow \mathbb{R}^d$,将有害激活的经验分布(记为 $\mu$)转换为无害激活的经验分布(记为 $\nu$),同时最小化传输成本 $\mathbb{E} _ {x\sim\mu}[||x - T(x)||^2]$ 。
为了简便,假设两个分布均为多维高斯分布(均值分别为 $\mu_1, \mu_2$,协方差矩阵分别为 $\Sigma_1, \Sigma_2$),最优传输映射具有闭式解(仿射变换)形式 :
$$T(x) = Ax + b$$
其中,传输矩阵 $A$ 和平移向量 $b$ 的计算公式为:
$$A = \Sigma_1^{-1/2}(\Sigma_1^{1/2}\Sigma_2\Sigma_1^{1/2})^{1/2}\Sigma_1^{-1/2}$$
$$b = \mu_2 - A\mu_1$$
但在 LLM 中,激活维度 $d$ 极高(数千维),而用于提取特征的样本量($n_h, n_s$)通常只有几百个 。这就导致:
- 经验协方差矩阵极度病态(ill-conditioned),容易过拟合噪声 。
- 对数千维的矩阵计算平方根和逆矩阵,计算成本高昂且数值不稳定 。
作者使用了PCA,将传输过程限制在一个低维子空间($k \ll d$)内 。
首先计算有害和无害数据的混合池化均值:
$$\mu _ {pool} = \frac{n_h\mu _ {\mathcal{H}} + n_s\mu _ {\mathcal{S}}}{n _ {\mathcal{H}} + n _ {\mathcal{S}}}$$
有点类似target encoding
然后,将两组数据用该池化均值进行中心化,拼接成矩阵$Z = [X _ {\mathcal{H}} - \mu _ {pool};X _ {\mathcal{S}} - \mu _ {pool}]$,并提取前 $k$ 个主成分,构成投影矩阵 $P \in \mathbb{R}^{d \times k}$ 。
将原始激活值投影到 $k$ 维子空间:
$$Y _ {\mathcal{H}} = (X _ {\mathcal{H}} - \mu _ {pool})P, \quad Y _ {\mathcal{S}} = (X _ {\mathcal{S}} - \mu _ {pool})P$$
在这个稳定的低维空间中,利用经验协方差计算出 $k$ 维的最优传输矩阵 $A_k \in \mathbb{R}^{k \times k}$ 和平移向量 $b_k \in \mathbb{R}^k$ 。
将低维映射关系“提升”(lift)回原始的 $d$ 维空间:
$$A _ {full} = PA_kP^\top$$
$$b _ {full} = \mu_2 - A _ {full}\mu_1$$
最终在模型推理时,应用于目标层的变换即为 $T^{(l)}(x) = A _ {full}x + b _ {full}$ 。
理论
Fisher线性判别分析(Linear Discriminant Analysis, LDA)是机器学习和统计学中极其经典且优雅的算法,主要用于降维和分类。
$$
w^\ast=argmax_w\frac{w^TS_Bw}{w^TS_ww}
$$
其中$S_B=(\mu_\mathcal{H}-\mu_\mathcal{S})(\mu_\mathcal{H}-\mu_\mathcal{S})^T$为类间散度矩阵,$S_w=\Sigma_\mathcal{H}+\Sigma _ {\mathcal{S}}$为合并的类内散度矩阵。
采用前面的均值中心化的PCA方法自然地能体现出这种判别结构。定义合并均值 $\boldsymbol{\mu} _ {pool} = \frac{n_H\boldsymbol{\mu}_H + n_S\boldsymbol{\mu}_S}{n_H+n_S}$ 以及中心化后的合并数据矩阵:
$$
\mathbf{Z} = \begin{bmatrix} \mathbf{X}_H - \boldsymbol{\mu} _ {pool} \\ \mathbf{X}_S - \boldsymbol{\mu} _ {pool} \end{bmatrix} \in \mathbb{R}^{(n_H + n_S) \times d}
$$
$\mathbf{Z}$ 的经验协方差分解为:
$$
\begin{aligned}
\frac{1}{n_H + n_S} \mathbf{Z}^\top \mathbf{Z} &= \frac{n_h}{n_H + n_S} (\boldsymbol{\Sigma}_H + (\boldsymbol{\mu}_H - \boldsymbol{\mu} _ {pool})(\boldsymbol{\mu}_H - \boldsymbol{\mu} _ {pool})^\top) \\
&\quad + \frac{n_s}{n_H + n_S} (\boldsymbol{\Sigma}_S + (\boldsymbol{\mu}_S - \boldsymbol{\mu} _ {pool})(\boldsymbol{\mu}_S - \boldsymbol{\mu} _ {pool})^\top) \\
&= \text{(within-class variance)} + \text{(between-class variance)}.
\end{aligned}
$$