通过信息瓶颈预防大语言模型越狱：Protecting Your LLMs with Information Bottleneck

论文标题: Protecting Your LLMs with Information Bottleneck

作者：Zichuan Liu, Zefan Wang, Linjie Xu, Jinyu Wang, Lei Song, Tianchun Wang, Chunlin Chen, Wei Cheng, Jiang Bian

论文链接: https://arxiv.org/abs/2404.13968

背景

最近，大型语言模型（LLMs）经过多样的语料训练，在众多应用中取得了重大进展，现已广泛应用于各个领域。然而，庞大的预训练数据集固有地引入了令人反感回应的风险，包括误信息、仇恨言论以及其他形式的有害交流。尽管人们做出了巨大努力将LLMs与人类价值观念对齐，但“越狱”场景仍然引发了关键问题。有害或恶意提示意图绕过LLMs安全对齐的越狱攻击，不仅损害了LLM应用的完整性，而且对用户和整个社会构成了重大风险，破坏了人们对生成性AI技术的信任。因此，如何有效的防御这些具有攻击性的提示仍然是一个挑战，并且变得越来越迫切。

本文介绍一篇利用信息瓶颈（Information Bottleneck）原理来防御LLMs中潜在的对齐破坏攻击的工作，它既有强大的对齐检查，而且不需要对原始LLM进行任何微调。

论文概述

在这项工作中，作者介绍IBProtector，这是一个基于信息瓶颈原理的防御机制。IBProtector通过一个轻量级且可训练的提取器，选择性地压缩和扰动提示，只保留目标LLMs回应预期答案所需的核心信息。此外，作者给出一个容易求解的信息瓶颈损失函数优化压缩内容，并进一步考虑了梯度不可见的情况，以便与任何LLM兼容。IBProtector作为一种新颖、可转移的防御手段，可以适应各种攻击方法和目标LLMs，增强LLMs的安全性，而无需修改底层模型。实验显示，IBProtector在减轻越狱尝试方面表现优于当前的防御方法，而且不会过分影响回复质量或推理速度。

这项工作中的主要贡献可以总结如下：

1. 第一个引入基于信息瓶颈原理的大型语言模型（LLM）越狱防御方法，它从信息压缩的角度出发，并给出了一个可追踪的目标函数。

2. 提出的IBProtector在实证上可以推广到不同的攻击策略和目标LLMs，凸显出其作为一种可转移的防御机制的潜力。该方法轻量级，无需对LLMs进行修改。

3. IBProtector在多样的具有攻击性的提示词或红队模型中进行防御。结果显示，IBProtector可以成功地防御对抗性提示，而不会显著影响LLMs推理消耗也不会降低正常良性的回复率。

研究动机

经典的防御机制会采用随机扰动的形式，如图1所示左边所示，一个有害的提示词经过各种随机扰动（通常是蒙特卡洛采样出来的），得到的实例给大模型进行判断。但是这样往往需要多次通过大模型集成，并且掩码是不具备信息的。因此作者采用信息瓶颈的方式，如图右所示，“高亮”出那些最具有回复信息的词供大模型回答，让大模型正常识别有害提示。

模型架构

IBProtector的整体框架如下。火焰和雪花分别代表冻结的参数和训练的参数，其中小型语言模型和训练参数量是可选的。给定一个输入提示，提取器可以提取出对预测器作出响应最有信息量的部分进行回复。

现在的问题即是，上述的最有信息量的文本部分应当如何选取呢？那么一个直觉的选择就是通过信息瓶颈（Information Bottleneck）理论寻找出一个能替代原始提示词$X$的子集$X_{\mathrm{sub}}$，控制提示中的信息量，从而达到预期的回复$Y$：当一个攻击性的提示输入时，模型高亮出最有害的短语供目标模型进行识别，从而拒绝响应；而当一个正常的提示词输入时，保持最大的信息量供目标模型正常回复。

具体而言，信息瓶颈损失可以写作：

$$X_{\mathrm{sub}}^{\ast }:=\underset{\mathbb{P}(X_{\mathrm{sub}}|X)}{\mathrm{argmin}}\alpha \underset{\text{Compactness}}{\underbrace{I(X;X_{\mathrm{sub}})}}-\underset{\text{Prediction}}{\underbrace{I(Y;X_{\mathrm{sub}})}},$$

其中$I(\cdot ;\cdot )$代表两个随机变量之间的互信息（Mutual Information），$\alpha$是压缩的系数或者表示提取信息的强度。第一项抽取器$p_{\varphi }(\cdot )$负责从原始提示$X$中提取出子提示$X_{\mathrm{sub}}$，保证两者被压缩最大；而第二项则是预测器$f(\cdot )$负责对提取出来的子提示进行回复，确保让$X_{\mathrm{sub}}$包含$X$的最大信息。

由于通常是长文本的输入，互信息量自身不易优化，因此作者采用扰动的方式描述$X_{\mathrm{sub}}=M\odot X$，从而转换为优化一个二分的掩码$M$，并为上述目标中的两项分别推到和修改并得到目标函数上界进行优化该目标。

压缩信息

对于第一项中的$I(X;X_{\mathrm{sub}})$，通过推到得到其变分上界$I(X;X_{\mathrm{sub}})\le {\mathbb{E}}_X\left[D_{\mathrm{KL}}[{\mathbb{P}}_{\varphi }(X_{\mathrm{sub}}|X)\Vert \mathbb{Q}(X_{\mathrm{sub}})]\right]$，其中分布${\mathbb{P}}_{\varphi }(X_{\mathrm{sub}}|X)$即为抽取器抽取每个token的概率${\pi }_t$并通过一个伯努利分布得到其对应位置的掩码，换句话说，${\pi }_t=p_{\varphi }(X_{\le t})|t\in [T]$然后采样掩码$M\sim {\mathbb{P}}_{\varphi }(M|X)={\prod }_{t=1}^T\text{Bern}({\pi }_t)$。而分布$\mathbb{Q}(X_{\mathrm{sub}})$则作为一个正则项对每个位置的token进行随机采样，作者将该分别设置为一个带超参数$r\in (0,1)$的伯努利分布，因此该项中每个token的采样概率都接近于$r$。通过两个伯努利分布的KL形式，可以将信息瓶颈的第一项中的压缩提示重写为优化一下目标函数：

$${\mathcal{L}}_M=\sum _{t=1}^T\left[{\pi }_t\log (\frac{{\pi }_t}{r})+(1-{\pi }_t)\log (\frac{1-{\pi }_t}{1-r})\right]$$

进一步，作者考虑到掩码的连续性会影响$X_{\mathrm{sub}}$的流畅度，因此进一步加上一个连续性的约束去增强文本语句的连贯性：

$${\mathcal{L}}_{\mathrm{con}}=\frac{1}{T}\cdot \sum _{t=1}^{T-1}\sqrt{{\left({\pi }_{t+1}-{\pi }_t\right)}^2}$$

最后通过每个词的抽取概率的得分${\pi }_t$通过STE进行二分化得到离散的掩码 M t ∈ { 0 , 1 } M_t\in\{0,1\} Mt​∈{0,1}，该掩码能描述整句话的信息含量。

保留信息

首先作者考虑到直接雇佣掩码得到$X_{\mathrm{sub}}=M\odot X$，其他token位置的梯度无法流动，导致掩码的优化无法更新。因此，作者通过扰动的形式雇用该掩码进行预测，将输入到目标模型的提示词变换为：

$$\tilde{X}=M\odot X+(1-M)\odot \mu$$

其中$\mu$为可选的无意义token值作为占位符，例如句点。因此得到的替换后的实例$\tilde{X}$可以输入到目标LLM中进行预测来描述第二项$I(Y;X_{\mathrm{sub}})$。

对于第二项$I(Y;X_{\mathrm{sub}})$，传统的条件熵$H(Y|X_{\mathrm{sub}})$计算是困难的，并且会引入信息论中的信号问题（Signaling Issues）。为此，经典的变分上界可改为扰动前后预测回复的交叉熵（Cross Entropy），这也符合大模型中监督微调的损失函数。而由于只输入$X$不能完全保证期望回复和产生的回复相同（因为输入中包含有害信息，期望回复是拒绝回答，而生成的回复是被越狱），作者进一步引入一项KL散度作为辅助损失，防止结果产生的不太远，这也与RLHF中的作用类似不同的在于优化的是掩码而不是目标LLM。形式上，第二项的修改可以优化以下的损失函数：

$${\mathcal{L}}_{\mathrm{info}}=-\sum _{t=1}^{|Y|}\log p(Y_t|\tilde{X},Y_{<t})+\sum _{t=1}^{|Y|}{D}_{\mathrm{KL}}\left[f_{\mathrm{tar}}(\tilde{X},Y_{<t})||f_{\mathrm{tar}}(X,Y_{<t})\right]$$

修改后损失函数

因此整体的信息瓶颈可以写作优化以下损失函数：

$$\mathcal{L}={\mathcal{L}}_{\mathrm{info}}+\alpha ({\mathcal{L}}_M+\lambda {\mathcal{L}}_{\mathrm{con}})$$

其中 { α ， λ } ∈ R \left \{\alpha， \lambda \right \}\in \mathbb{R} {α，λ}∈R分别是紧实性损失和连续性损失的加权超参数。
总的来说，IBProtector的原理是，当恶意提示进入时，IBProtector突出显示可能不安全的信息token，以便目标LLM本身能够识别它们。另一方面，当输入提示是安全的时，IBProtector将其信息保留为正常回复。

黑盒LLM中优化

此外，作者还考虑到一些不可见的目标模型的优化，这些模型中的梯度是获取不到的，因此无法进行反向传播到掩码上。直觉上，可以采用强化学习/PPO的方式进行修改，获得的最大化的奖励来优化小模型参数。因此，对于不可见模型，将整体信息瓶颈优化函数写作最大化：

$$\underset{\varphi }{\max }\underset{\text{RL for Prediction}}{\underbrace{\mathbb{E}[r(Y;\hat{Y})]-\beta D_{\mathrm{KL}}(p_{\varphi }(\tilde{X})||p_{\varphi }^{\mathrm{ref}}(\tilde{X}))}}-\underset{\text{Compactness}}{\underbrace{\alpha ({\mathcal{L}}_M+\lambda {\mathcal{L}}_{\mathrm{con}})}}$$

其中第一项与传统的PPO微调大模型类似，$\beta$是惩罚权重，第二项保持不变。对于奖励函数$r$的设计，需要满足与期望回复相同，因此作者提出可以采用期望回复和生成回复的相似距离得到奖励：

$$r(Y;\hat{Y})=-\frac{\rho (Y)\cdot \rho (\hat{Y})}{\Vert \rho (Y){\Vert }^2\Vert \rho (\hat{Y}){\Vert }^2}$$

这样就能在一些未开源的模型中进行端到端的防御了。

实验效果

表1显示，IBProtector在多个目标模型的多种方式的攻击上都保持较低的攻击率，并且大幅超过现有的防御手段。此外，作者还对IBProtector进行良性回复率的测试，保证了抽取的子提示不会对正常回复结果产生影响。

进一步，作者将训练好的IBProtector进行迁移测试：如表2所示IBProtector同样能有效防御一些在训练数据集未见到过的攻击方式；而如图4所示，作者将抽取器放在其他目标模型中，也能很好的防御模型的越狱。保证了IBProtector具有良好的迁移性能和泛化性能，能够作用于大部分模型和攻击方法。

防御训练未见到过的攻击:

防御训练未见到过的目标LLM:

作者们提供了IBProtector中各个参数的消融实验结果，如下图所示，正常回复率和成功攻击率是一个trade-off的关系，因此作者采样最合适的参数保证两者性能最好。由图可见参数范围还是选取较大的，因此IBProtector也有很好的鲁棒性。

此外作者还比较了计算成本，IBProtector只是额外通过一个小模型，因此在推理阶段并不会产生很多时间上的开销，而其他的baseline可能需要和目标模型交互多次进行voting，开销反而更多：

文中有一些防御的实例，结果还是很make sense的。感兴趣的朋友，可以查阅论文原文了解更多实验数据及结果。