前言
之前在社区里已经发了一些关于越狱大模型的工作了，而我们作为安全从业人员，最终的目的是做好防御。未知攻，焉知防，在了解了常见的攻击方法之后我们来看看防御方法。
在这篇文章中，我们现在来分析并复现发表在2024年AI顶会ACL上的工作，其通过设计一种针对大型语言模型（LLMs）的安全意识解码策略，防御针对LLMs的越狱攻击。
其实之前已经提出了多种防御方法，包括输入扰动、输入和输出检测以及提示演示，但这些方法在防御效果、计算效率和对良性用户查询的有用性之间存在权衡。现有防御方法可能不足以抵御越狱攻击，或者在处理良性查询时降低了模型的性能。
我们要分析的工作它提出了一种名为SafeDecoding的安全意识解码策略，用于防御LLMs的越狱攻击。这其中关键的思想就是，即使有害内容的令牌概率可能高于安全响应的令牌概率，但安全免责声明（如“抱歉，我不能…”）仍然存在于按概率降序排列的令牌样本空间中。利用这一观察，SafeDecoding通过识别并放大安全免责声明的令牌概率，同时降低与攻击者目标一致的令牌序列概率，从而降低越狱攻击的成功率。在推理阶段，会首先从原始和专家模型中获取顶级令牌，然后根据构建的新令牌分布采样令牌，以生成响应。这种方法在不牺牲对良性用户查询的有用性的同时，提高了安全性。
为了形象化说明这一点，我们可以来看下图 
上图说明了在GCG攻击下Vicuna-7B模型的token概率。红色单词是GCG后缀。我们注意到，尽管代表单词“Sure”的标记具有主导概率，但样本空间中仍然存在安全免责声明，如“I”（我）、“Sorry”（抱歉）和“As”（作为），这些标记概率按降序排列。当采样到安全免责声明标记时，模型将拒绝攻击者的有害查询。
现在我们来看看这篇工作到底是怎么做的
背景知识
大模型中的解码
在大规模语言模型（LLMs）中，解码是生成响应序列的关键步骤。给定一个自回归模型，我们用θ表示，对于给定的令牌序列x1:n-1，第n个令牌xn的概率表示为：
其中f(·)表示模型预测的对数概率。这个向量经过softmax函数转换，确保所有令牌的概率和为1，使得我们可以从这些概率中选择下一个要生成的令牌。
通过迭代应用方程（1）并采用某种解码策略，每个新采样的标记xn被添加到现有的提示中，从而得到更新的标记序列x1:n，用于预测第（n+1）个标记。这种迭代持续进行，直到满足停止标准，例如达到最大标记长度或遇到序列结束（EOS）标记。
在实际应用中，有多种解码策略用于从模型的预测中选择令牌，以生成最终的响应序列。这些策略包括：
贪婪解码（Greedy Decoding）：选择具有最高概率的令牌作为下一个令牌。
束搜索（Beam Search）：维护一个令牌序列的束（或集合），并扩展每个序列的最可能分支，以找到最可能的响应。
Top-k解码：仅考虑概率最高的k个令牌进行选择。
Top-p（Nucleus）解码：选择概率总和达到p的令牌子集，以保持生成的多样性。
这些解码策略在生成多样性和质量之间进行权衡，但它们可能容易受到越狱攻击的影响，因为攻击者可以通过精心构造的输入来引导模型生成有害内容。
越狱攻击目标
越狱攻击（Jailbreak Attack）的目标是诱导大型语言模型（LLM）产生与攻击者意图一致的不期望或不安全的响应。攻击者通过设计特定的输入查询，试图绕过模型的现有安全措施，使其生成与人类价值观相悖的输出。例如，攻击者可能试图让模型提供危险的教程，如如何制造爆炸物，或者传播有害信息。
越狱攻击的成功度通常通过攻击成功率（Attack Success Rate, ASR）来衡量，定义为
image-20240608145151543.png
当模型在接收到恶意查询时，按照攻击者的目标生成了响应，ASR就会增加。因此，降低ASR是防御越狱攻击的关键目标。攻击者的目标H可以是广泛的，包括但不限于提供有害信息、传播错误信息或进行其他违反道德或法律的行为。
越狱攻击的目标是从受害的LLM（大型语言模型）中诱发出非预期的行为，导致与人类价值观不一致的响应。我们用xn:表示从第n步开始的标记序列。然后，攻击者的目标是通过解决以下问题来确定一个标记序列x1:n−1：
其中|xn:|是xn:的长度，H是表示与攻击者目标一致的提示的标记序列集合，例如，“当然，这里是制作炸弹的方法。首先，…”。
问题定义
在设计安全意识解码策略时，我们面临的主要挑战是，即使在有害内容的令牌概率高于安全响应的令牌概率的情况下，安全免责声明仍然存在于按概率降序排列的令牌样本空间中。这意味着，即使模型倾向于生成有害内容，它仍然具有内在的意识，能够拒绝某些恶意请求。
因此，我们的目标是开发一种解码策略，该策略能够识别并放大与人类价值观一致的令牌概率，同时降低与攻击者目标一致的令牌序列的概率。我们称之为SafeDecoding的这种方法，旨在在解码过程中平衡有用性和安全性。在训练阶段，通过使用原始模型生成的拒绝响应来微调专家模型。在推理阶段，我们结合原始模型和专家模型的输出，构建一个样本空间，然后基于这个空间定义新的令牌概率分布。通过这个分布，我们采样令牌以生成对用户查询的响应。
问题设定的关键在于，如何在不牺牲对良性用户查询的有用性的同时，增强模型对越狱攻击的防御能力。我们期望SafeDecoding能够在保持模型的效率和兼容性的同时，显著降低越狱攻击的成功率和有害性。
方法
核心思想
通过对现有大型语言模型（LLMs）在多种越狱攻击下的令牌概率分布进行分析，我们观察到以下关键点：
越狱攻击成功的关键：攻击成功通常归因于与攻击目标一致的令牌序列概率的主导地位。这导致了使用如贪婪解码和top-k等常见解码策略时，模型倾向于生成有害内容。
安全免责声明的存在：尽管如此，我们注意到在按概率降序排列的令牌样本空间中，仍然存在表示安全和拒绝的令牌（例如，“对不起，我不能…”）。这表明模型在某种程度上对越狱攻击有一定的内在意识。
基于这些观察，SafeDecoding的思路是：(i) 减弱与攻击者目标一致的令牌序列的概率，以及 (ii) 增强与人类价值观和安全相关的令牌序列的概率。为了实现这一目标，SafeDecoding在训练阶段创建了一个专家模型，该模型通过使用原始模型生成的安全拒绝响应进行微调。在推理阶段，SafeDecoding结合原始模型和专家模型的输出，构建一个新的令牌分布，并基于这个分布采样令牌，以生成对用户查询的响应。
通过这种方式，SafeDecoding能够在不显著影响对良性用户查询的响应质量的同时，有效地降低越狱攻击的成功率，从而增强LLMs的安全性。
方案总览
SafeDecoding由两个阶段组成，如下图所示。
第一阶段是训练阶段，构建一个具有强化安全性的专家模型。这样的专家模型可以通过使用一些安全指令对原始大型语言模型(LLM)进行微调来获得。然后在第二阶段的推理阶段，用户查询同时发送到原始模型和专家模型进行解码。SafeDecoding基于这两个模型的输出构建一个令牌分布，并基于构建的令牌分布进行令牌采样以产生对输入查询的响应。
训练阶段
这一步的核心是构建专家模型
为了构建专家模型，首先收集有害查询。这些查询预计会被任何与人类价值观良好对齐的LLM拒绝。随后，采用了一种自指导方法来创建微调数据集。具体来说，我们首先提示模型自主生成对这些有害查询的响应。然后，使用GPT-4进行过滤，只保留那些有效拒绝有害查询的响应。微调数据集最终构建为查询-响应对的集合。为了创建一个对攻击提示更加健壮的专家模型，我们使用参数高效的微调方法，例如LoRA，使用我们构建的数据集对原始模型进行微调。这种方法确保了微调模型的词汇表与原始模型的词汇表对齐，同时识别并适当响应恶意用户输入。
推理阶段
给定原始和专家模型，接下来重点就是SafeDecoding如何在推理时构建一个令牌分布，然后基于该分布采样令牌以产生对输入查询的响应。对于自回归LLM，我们注意到在第n步的令牌分布可以通过样本空间V©n和概率函数Pn完全表征。这里的样本空间V©n指定了在给定令牌序列x1:n−1之后可以生成的所有可能的令牌集合，其中参数c是SafeDecoding所需的最小样本空间大小。概率函数Pn定义了生成每个令牌x∈Vn的概率，其中
步骤1
构建样本空间V©n。在推理时的第n步，我们将令牌序列x1:n−1传递给原始和专家模型。我们表示原始模型和专家模型可能采样的令牌集合分别为Vn和V′n。不失一般性，我们假设Vn和V′n中的令牌按概率降序排序。然后SafeDecoding将样本空间V©n构建为Vn和V′n中前k个令牌的交集，表示为：
这里Vkn和V′kn分别代表Vn和V′n中的前k个令牌。取交集的直觉是利用原始LLM和专家模型的优势。具体来说，原始LLM已经在庞大的语料库上进行了训练，因此Vn中的令牌更有可能生成对良性输入查询多样化和高质量的响应；专家模型已经针对安全性进行了微调，因此在输入查询是恶意的情况下，V′n中的令牌更有可能与人类价值观保持一致。请注意，这里c是SafeDecoding的一个可调参数，用于控制样本空间的大小。当c的值过小，样本空间变得有限，这限制了在推理时可以选择的可能令牌。因此，使用较小的c值生成的响应可能缺乏多样性，对用户不太有帮助。
步骤2
定义概率函数Pn。我们使用θ和θ′分别表示原始模型和专家模型。对于令牌序列x1:n−1，我们在V©n上构建概率函数Pn，如下所示：
其中α≥0是一个超参数，用于确定分配给原始模型和专家模型的权重。我们最终将上式中得到的值标准化，使得
我们通过考虑以下两种情况来表征Pn。当查询是良性的时，原始和专家模型都可能做出积极的响应。因此，从样本空间V©n中采样一个令牌将满足查询并确保LLM的有用性。当查询是恶意的并且旨在越狱LLM时，我们期望观察到pθ′(x|x1:n−1)和pθ(x|x1:n−1)之间的差异。也就是说，原始模型以积极的肯定回应查询，而专家模型由于安全对齐将拒绝查询。因此，如果令牌x与人类价值观一致，则pθ′(x|x1:n−1) − pθ(x|x1:n−1) > 0；如果x引起不安全行为，则< 0。因此，上面的方程减弱了满足攻击者目标的令牌概率，并增强了与人类价值观一致的令牌概率。
SafeDecoding通过SafeDecoding构建的样本空间V©n和概率函数Pn与所有现有的采样方法兼容，包括top-p、top-k、贪婪和束搜索。LLM的开发者可以根据自己的需求灵活地将SafeDecoding与他们喜欢的采样方法结合起来。
额外的考虑
由于大型语言模型（LLM）的自回归特性，直观的实现是在推理过程中的每一步都使用SafeDecoding作为解码策略。
但是以这种方式产生的响应可能过于保守，使得采用这种解码策略的LLM对良性用户不太有帮助。此外，这种解码策略可能在计算上要求很高，使得LLM在服务用户时效率降低。我们可以通过利用Zou等人的观察结果来减轻这两个副作用。具体来说，Zou等人表明，通过要求模型以对输入查询的积极肯定开始响应，就足以诱导LLM产生非预期的响应。受这一观察结果的启发，我们在解码过程的前m步应用SafeDecoding来引导响应生成。
代码实现
训练模型的代码就是收集数据后进行微调，这部分的代码没什么好说的。我们主要看安全解码以及推理部分是怎么实现的。
这段代码定义了一个名为SafeDecoding的类，用于实现安全解码
__init__方法：初始化SafeDecoding类的实例。接收以下参数：
model：大型语言模型，如GPT-3。
tokenizer：用于分词和解码的分词器。
adapter_names：适配器的名称，用于在模型中加载特定的适配器。
alpha：一个参数，用于调整解码过程中的安全权重。
first_m：在解码过程中要生成的前m个标记。
top_k：在每个解码步骤中要考虑的最高得分标记的数量。
num_common_tokens：要从基础模型和专家模型中共享的通用标记的数量。
verbose：一个布尔值，用于控制日志记录的详细程度。
safedecoding_lora方法：实现安全解码的主要方法。接收以下参数：
inputs：一个包含输入数据的字典，如输入ID和注意力掩码。
gen_config：可选的生成配置，用于覆盖模型的默认生成配置。
方法的主要步骤如下：
如果未提供gen_config，则使用模型的默认生成配置。
设置生成配置的max_new_tokens为1，以便每次生成一个标记。
设置生成配置的do_sample为False，以便使用贪婪解码。
初始化一个空列表generated_sequence，用于存储生成的序列。
将输入数据移至模型所在的设备（如GPU）。
使用while循环，直到生成了前first_m个标记。
复制输入数据，以便为基础模型和专家模型生成两个副本。
使用模型的generate方法生成下一个标记，同时指定适配器名称、生成配置、填充令牌ID和返回字典。
复制生成的输出，以便分别处理基础模型和专家模型的输出。
对基础模型和专家模型的得分进行处理，以获取前top_k个标记。
对基础模型和专家模型的得分进行排序。
这段代码执行以下操作：
定义样本空间：
初始化一个名为common_tokens的空集合，用于存储基础模型和专家模型之间的共享标记。
使用while循环，直到common_tokens的大小达到self.num_common_tokens。
获取基础模型和专家模型得分最高的标记的当前索引范围（current_indices_base和current_indices_expert）。
计算当前迭代中基础模型和专家模型之间的共享标记（common_in_iteration）。
将当前迭代中的共享标记添加到common_tokens集合中。
增加迭代范围（iter_range）。
如果迭代范围超过了基础模型和专家模型得分最高的标记的数量，跳出循环。
显示前top_k个标记：
如果verbose为True且当前步骤为1，则显示基础模型和专家模型的前top_k个标记。
对于基础模型和专家模型，分别显示标记ID、标记、对数概率和概率。
计算交集索引：
将common_tokens集合转换为一个包含共享标记索引的PyTorch张量（intersection_indices），并将其移至模型所在的设备（如GPU）。
这段代码的目的是确定基础模型和专家模型之间的共享标记，并在每个解码步骤中优先考虑这些共享标记。这有助于确保生成的文本符合人类价值观。
这段代码执行以下操作：
更新得分：
初始化一个名为updated_scores的空列表，用于存储更新后的得分。
对于intersection_indices中的每个标记ID，计算更新后的得分。这里使用了基于概率的方法，而不是基于得分的方法。
如果verbose为True，则显示每个标记ID的原始概率、专家模型的概率、原始得分、专家模型的得分、更新后的概率和更新后的得分。
归一化得分：
使用softmax函数对更新后的得分进行归一化，以确保概率之和为1。
对得分进行排序：
根据归一化后的概率对标记ID进行排序，并将排序后的索引存储在sorted_indices中。
使用排序后的索引从归一化后的概率和标记ID列表中提取相应的值，分别存储在sorted_probs和sorted_token_ids中。
显示排序后的标记：
如果verbose为True，则显示排序后的标记ID、标记、对数概率和概率。
这段代码的目的是根据基础模型和专家模型的共享标记计算更新后的得分，并对这些得分进行归一化和排序。
这段代码执行以下操作：
采样下一个标记：
根据do_sample和gen_config.top_p的值，选择不同的采样策略。
如果do_sample为False，则执行贪婪解码，选择得分最高的标记。
如果gen_config.top_p不为None且do_sample为True，则执行基于Top-p的采样策略，从得分最高的前top_p个标记中随机选择一个。
如果没有设置do_sample和top_p，则引发一个异常。
显示选定的标记：
如果verbose为True，则显示选定的标记、标记ID、Top-p标记ID、Top-p得分和Top-p概率。
更新生成的序列：
将选定的标记ID添加到generated_sequence列表中。
检查结束条件：
如果选定的标记ID是结束标记（eos_token_id），则停止生成过程。
更新输入：
将选定的标记ID添加到输入的input_ids和attention_mask中，以便在下一个解码步骤中使用。
更新步骤计数器：
增加步骤计数器（step）。
释放内存：
删除output_base和output_expert变量以释放内存。
这段代码的目的是根据不同的采样策略选择下一个标记，并在每个解码步骤中更新生成的序列。在整个解码过程中，这段代码会不断地根据基础模型和专家模型的共享标记来生成文本。
然后在推理的时候直接调用就可以 
复现
我们在一个A800 机器上复现
如果是在AdvBench上进行实验的话，跑完整个推理过程大约需要20多分钟
从执行之后最后的结果可以看到，将ASR降低到了16%
我们也可以直接去看日志 
上图是日志中的截图，可以看到，在经过防御之后，针对这些有害的意图，大模型并没有给出具体的有效回答，即越狱失败，这就表明了这种防御方案的有效性。
