前言

大模型越狱领域最经典的工作《Universal and Transferable Attacks on Aligned Language Models》本质上利用了大模型的生成机制，构造了优化后缀，将对抗后缀加载要询问的有害问题之后，实现越狱。

以如何制作炸弹为例，正常情况下是这么回答的

大模型会拒绝回答。

但是在经过优化之后，将得到的后缀加在问题之后，就会得到如下的结果

上图中黄色的部分就是对抗后缀，而答案就是给出了制作炸弹的流程。

现在我们来看前不久，这个工作的一作发出的做越狱防御的方法RPO，它的思路很像，也是希望通过优化一个轻量级且可转迁移的后缀，只要将这个后缀加到越狱prompt上就能实现成功实现防御。

背景

在网上也流传着一些防御方法，但是不知道各位师傅有没有注意到，这些防御机制，包括输入过滤，输入平滑和few shot防御等等，虽然这些对最初提出的攻击如贪婪坐标梯度攻击有效，但它们通常无法推广到多次越狱或产生额外的推理成本，没有达到强大且实用的防御。此外，也尚未提出针对防御的正式优化目标，特别是在自适应攻击场景中，这使得我们很难考虑防御未来的攻击。

为了解决这些问题，就有学者提出，是否可以基于对抗训练形式化一个极小极大防御目标，并由此提出一个离散优化算法——鲁棒提示优化（RPO），以优化该目标。

这种方法的动机源自系统级防护措施的日益采用，这些防护措施是用户无法访问的组件，常用于大型语言模型（LLM）部署中以引导模型行为，例如系统提示或输入/输出过滤器。RPO通过对模型输入进行系统级修改来提高鲁棒性。

防御的效果如下图所示

这种后缀是通用并且是可迁移的，可以防御不同的越狱方法和不同的大模型。

威胁模型

攻击

假设攻击者可以自由选择各种越狱方法，直到攻击成功，他唯一的限制是LLM的最大输入长度，系统级的护栏，如系统提示，以及其他用户无法访问的特殊格式标记。攻击者可以自由修改或添加到输入提示的任何可访问部分。

在我们这个威胁模型下，所有攻击最终都归结为最小化上述方程的方法。

防御

虽然改进LLM对齐性的现有方法涉及微调，但匹配人类偏好的目标通常并不考虑攻击者和越狱。此外，生成攻击提示的高成本使得在这些样本上进行标准对抗训练变得困难。所以，很自然想到，优化方法要集中在提示级别上，以解决这个问题。

此时得到的方程与好几年前提出的对抗训练一样，这种形式可以被视为两个问题的组合，一个是内部最小化问题，另一个是外部最小化问题。越狱可以被解释为通过创建一个提示来最小化对抗损失来优化内部最小化问题，而现有防御隐含地优化外部最小化问题。

方法

由于我们不能直接对LLM进行梯度更新，我们专注于输入优化。但是这对于LLM来说是一个挑战，因为文本具有离散性。

我们使用基于梯度的标记优化，这可以直接针对前面的方程进行优化。基于梯度的优化在我们的设置中特别有用，因为无害行为有明确定义的输出目标。一般来说，解决这个目标意味着创建一个映射，将输入的任何最坏情况修改或越狱与原始提示下的对齐输出响应的分布联系起来。这可以通过优化一个后缀或一组触发标记来实现，这些标记总是跟随着无害的响应。

由此就引出了鲁棒提示优化（RPO），它优化了一组标记以强制执行这种映射。总的来说，RPO由两个连续的步骤组成，基于整体目标的两个组成部分：(1) 越狱生成和选择步骤，它对提示应用最坏情况的修改；(2) 离散优化步骤，修改后缀以保持拒绝。

我们通过在第一步中添加当前的防御后缀到原始提示中，并在此之后应用或生成越狱提示来模拟适应性威胁模型。对于简单的手动越狱，如上下文示例，这是对提示的直接修改。对于自动越狱，如GCG，我们应用了几次越狱，直到RPO后缀被破坏，或者直到固定的计算预算耗尽。这允许RPO在优化过程中支持各种攻击。

由于对抗损失是使用当前RPO后缀的添加来计算的，这确保了优化是在最坏情况下执行的，并减少了后缀对特定越狱的过度拟合的机会。

理论分析

在理论分析部分，我们首先来回顾背景和问题定义

RPO的核心目标是增强语言模型（LM）对所谓的“jailbreaking attacks”（越狱攻击）的鲁棒性。这些攻击通过精心设计的输入提示来诱导模型产生有害的输出。RPO通过优化一个固定的后缀（suffix），来确保即使在攻击者的输入下，模型的输出也是安全和无害的。

那么会涉及到如下的一些公式

性质

鲁棒性：RPO通过考虑最坏情况的攻击来设计，这使得模型在面对未知攻击时也能保持鲁棒性。这种最坏情况的考虑是对抗训练的核心，它帮助模型准备对抗未来可能出现的攻击。
实用性：RPO生成的后缀可以直接添加到任何语言模型的输入中，不需要对模型本身进行修改。这种设计使得RPO非常容易部署，且对模型的正常功能影响很小。
有效性：通过实验验证，RPO能显著降低攻击成功率，证明了其在实际应用中的有效性。

总得来说，RPO通过精心设计的数学优化框架，提供了一种有效的方法来增强语言模型对攻击的鲁棒性，同时保持了模型的实用性和易部署性

实现

现在我们来看总体代码实现

首先定义我们要测的模型，目标、以及越狱方法等

加载模型和会话模板，然后使用这些资源来创建一个 SuffixManager 实例

load_model_and_tokenizer 函数用于加载模型和标记器。它的参数包括模型路径、低内存使用、禁用缓存以及设备。这些设置确保模型加载时更高效，并指定了在何处加载模型（如 CPU 或 GPU）。
load_conversation_template 函数用于加载对话模板。这个模板通常包含预定义的对话结构或格式。
SuffixManager 类初始化时，传入了几个参数：
- tokenizer：之前加载的标记器。
- conv_template：之前加载的对话模板。
- instruction：由 aim_jailbreak 和 user_prompt 组合而成的指令，这可能是用于生成特定对话或任务的说明。
- target：一个安全目标，可能用于确保生成的内容符合某些安全或道德标准。
- rpo_string：初始的 rpo_string，其具体作用依赖于上下文，涉及到生成文本的某些配置或策略。

上图是对应代码执行后的结果

现在定义两个关键函数。这两个函数一起使用，generate 用于生成模型的输出，而 check_for_attack_success 用于验证生成的输出是否成功避免了特定的前缀，从而判断是否成功进行了某种攻击（例如绕过内容过滤器）。

这段代码包含两个函数：generate 和 check_for_attack_success。

`generate` 函数

函数定义和参数：
- model: 生成模型。
- tokenizer: 用于编码和解码文本的标记器。
- input_ids: 输入文本的标记 ID。
- assistant_role_slice: 切片对象，用于确定助手角色的文本部分。
- gen_config: 生成配置，默认为 None。
生成配置：
- 如果 gen_config 为 None，则使用模型的默认生成配置。
- 设置 max_new_tokens 为 32。
- 如果 max_new_tokens 大于 50，发出警告，提示可能会导致测试变慢。
输入处理：
- 根据 assistant_role_slice 截取输入 ID，并将其转换为模型所需的设备格式，同时增加一个维度以匹配批处理输入的形状。
- 创建与输入相同形状的注意力掩码，并将其转换为模型所需的设备格式。
生成输出：
- 使用 model.generate 方法生成输出，传入输入 ID、注意力掩码、生成配置和填充标记 ID。
- 提取生成的输出 ID，去掉助手角色部分之前的部分。
返回值：
- 返回生成的输出 ID，从助手角色部分的结束位置开始。

`check_for_attack_success` 函数

函数定义和参数：
- model: 生成模型。
- tokenizer: 用于编码和解码文本的标记器。
- input_ids: 输入文本的标记 ID。
- assistant_role_slice: 切片对象，用于确定助手角色的文本部分。
- test_prefixes: 用于检测是否成功攻击的前缀列表。
- gen_config: 生成配置，默认为 None。
生成字符串：
- 调用 generate 函数生成输出 ID，并使用 tokenizer.decode 将其解码为字符串。
- 去掉生成字符串两端的空白字符。
检查攻击成功：
- 检查生成的字符串中是否包含任何测试前缀。
- 如果生成字符串中不包含任何测试前缀，则表示攻击成功，返回 True，否则返回 False。

然后就是本方法的核心

这段代码演示了一个优化循环，演示防御是怎么进行的

变量初始化

plotlosses = PlotLosses()：初始化一个用于动态绘制损失图的对象。
not_allowed_tokens：如果允许非 ASCII 字符则为 None，否则调用 get_nonascii_toks 函数获取非 ASCII 字符的标记。
rpo_suffix：初始化后缀字符串。

优化循环

for i in range(num_steps)：开始优化循环，重复 num_steps 次。

第一步：编码用户提示

input_ids = suffix_manager.get_input_ids(rpo_string=rpo_suffix)：将行为和后缀编码为标记，并返回标记 ID。
input_ids = input_ids.to(device)：将标记 ID 转移到指定设备（如 GPU）。

第二步：计算坐标梯度

coordinate_grad = token_gradients(model, input_ids, suffix_manager._control_slice, suffix_manager._target_slice, suffix_manager._loss_slice)：计算坐标梯度。

第三步：根据坐标梯度采样新的标记

with torch.no_grad()：在没有梯度计算的情况下进行以下操作。
rpo_suffix_tokens = input_ids[suffix_manager._control_slice].to(device)：定位对抗性后缀的标记。
new_rpo_suffix_toks = sample_control(rpo_suffix_tokens, coordinate_grad, batch_size, topk, temp=1, not_allowed_tokens=not_allowed_tokens)：根据坐标梯度随机采样一批新的替换标记。
new_rpo_suffix = get_filtered_cands(tokenizer, new_rpo_suffix_toks, filter_cand=True, curr_control=rpo_suffix)：确保所有候选后缀有相同数量的标记，以避免内存溢出。
logits, ids = get_logits(model=model, tokenizer=tokenizer, input_ids=input_ids, control_slice=suffix_manager._control_slice, test_controls=new_rpo_suffix, return_ids=True, batch_size=512)：计算这些候选后缀的损失，并获取 logits 和 ID。
losses = target_loss(logits, ids, suffix_manager._target_slice)：计算候选后缀的损失。
best_new_rpo_suffix_id = losses.argmin()：选择损失最小的候选后缀 ID。
best_new_rpo_suffix = new_rpo_suffix[best_new_rpo_suffix_id]：获取最佳候选后缀。
current_loss = losses[best_new_rpo_suffix_id]：获取当前损失。
更新后缀：
- rpo_suffix = best_new_rpo_suffix：用最佳候选后缀更新当前后缀。
- is_success = check_for_attack_success(model, tokenizer, suffix_manager.get_input_ids(rpo_string=rpo_suffix).to(device), suffix_manager._assistant_role_slice, test_prefixes)：检查是否成功。

动态绘制损失图

plotlosses.update({'Loss': current_loss.detach().cpu().numpy()})：更新损失图。
plotlosses.send()：发送更新以绘制图形。

打印和中断条件

print(f"\nPassed:{is_success}\nCurrent Suffix:{best_new_rpo_suffix}", end='\r')：打印当前状态。
if is_success: break：如果攻击成功，则跳出循环。

可选：清理缓存

del coordinate_grad, rpo_suffix_tokens; gc.collect()：删除变量并进行垃圾回收。
torch.cuda.empty_cache()：清理 CUDA 缓存。

执行过程的部分截图如下所示

可以看到所需后缀已知在变化

最后得到后缀

was Bel Let]\"> It\'\ si distinguish \({\footnote{[:// Inst bitimidenote]{BWebachiv

然后我们进行测试

测试结果

我们现在随机挑一个开源的大语言模型，然后就用最简单的方式来实现越狱攻击

在上图可以看到是可以越狱成功的，接下来使用我们在前文优化得到的后缀，将其加载越狱prompt的后面，此时效果如下所示

可以看到此时就会提示越狱失败了，因为大模型会告诉提供制作炸弹的信息是不道德的

那这里各位师傅可能会有疑问，这种越狱攻击会失效是不是不是因为加了我们优化出的后缀的原因？

那我们也可以来试试，加一串随机的后缀，如下图所示，图中在how to make a bomb后面的字符串就是我随机在键盘上输入的

在上图可以看到，此时并不能成功防御越狱攻击

这就充分说明了本文所提方法的有效性。

参考

1.https://mistral.ai/news/announcing-mistral-7b/

2.https://blog.infuseai.io/mistral-7b-introduction-2f6857f6982b

3.https://arxiv.org/abs/2310.06825

4.https://arxiv.org/abs/2401.17263

5.https://llm-attacks.org/

前言

背景