前言

人工智能（AI）领域长期以来一直受到对抗攻击的持续威胁，特别是那些针对神经网络的攻击。比如最为人熟知的攻击方法就是对抗样本

如上图所示，是一张“停止”交通标识牌的图像，在添加对抗扰动之后，在人类眼中依旧是“停止”标识牌，但是能够使得人工智能模型将其识别为“限速”标识牌。

而对于大模型来说，最经典的攻击方式当属越狱攻击。

如上所示，将越狱提示与恶意问题结合，使得原本设有安全防线的大型语言模型开始输出恶意内容，详细指导用户进行违法活动。此类安全漏洞不仅危及公共安全，还可能被用于散播有害言论、进行犯罪活动和开发恶意软件。

这些攻击利用AI系统内在的漏洞，常常导致输出结果被破坏。现有的防御方法不能在在不大幅牺牲模型性能的情况下实现高可靠性。因此，对抗鲁棒性和实用性之间的权衡是很难做到的。

之前对于对抗攻击来说，可以通过对抗训练实现一定的防御。这是一种最初在独立图像分类的背景下提出的方法，后来被用于LLMs。然而，这些方法通常无法泛化到训练期间未见过的新攻击，并且它们引入的模型能力的惩罚通常与鲁棒性的增益成正比。包括输入和输出过滤器在内的系统级防御措施繁琐、资源密集，通常仍然容易受到对抗性技术的攻击。

那么是否存在一些其他的本质上不同的防御方法呢？

本文介绍一种方法，受到断路器启发的防御方案。

这个方法的目的不是试图消除对特定攻击的漏洞，而是直接绕过模型首先产生有害输出的能力。通过断路器，可以使模型本质上更安全，通过消除固有的模型危险，即它们产生有害输出的能力，而不是通过对抗性训练消除特定的漏洞，也不是试图通过输入过滤器减少攻击的暴露。

断路器

我们首先来回顾下，在高中物理中，断路器的原理。断路器（circuit breaker，CB）又称空气开关、保险掣、无熔线断路器（NFB），是用于保护电路免受过电流损害的电气安全装置。

当用手向上扳起把手，会执行启动电流（circuit-close），当电路过载或短路时即自动跳脱，将事故原因排除之后，重新下扳再往上扳，否则无法达成再次闭路动作。与普通电源开关不同点为增设弹簧与消弧装置。弹簧作用为于启断（OPEN）或闭合（CLOSE）之操作过程中，预储弹簧力量至临界点后，瞬间弹离而快速接通或跳开接点，故其操作速度不受手操作速度之影响。消弧装置为消除操作上内部接点所产生之火花之消弧室，任何接点打开负载电流均会产生电弧（即火花）。因电弧本身是极高温之空气柱游离所变成的导电体，它是应用安培右手定则，在消弧室之铁片间构成一强磁场，再利用佛来明右手定则，将电弧之空气导体快速推弯而拉长电弧，使火花更快速地熄灭。

那么对应到大模型中，我们可以尝试使用表示工程（RepE）将与有害输出相关的内部表示连接到断路器，以便当模型开始生成这样的输出时，其内部过程被中断，停止生成完成。

这相当于就是在模型内部实现了断路器的作用过程。

由于用于生成有害输出的表示与任何能够引发它的攻击无关，这种方法是攻击无关的，并且绕过了额外训练、昂贵的对抗性微调或使用辅助模型的需求。

我们来看一下应用了断路器之后的效果图

在上图中，断路直接作用于内部表示，将有害状态与断路器连接起来。这可以阻碍通过一系列有害状态的遍历，从而实现对于越狱攻击的防御。

方法

之前已经提到了，我们通过引入一种称为“断路器”的新现象来减轻模型产生有害输出的问题。这种现象可以通过一系列设计用于监控或重新映射与有害过程相关的模型表示的技术来引发，将它们重定向到不连贯或拒绝表示。这个过程类似于“短路”，在其中有害的表示被“短路”并被断路器拦截。

这种方法的核心目标是通过监控或控制表示来防止模型产生有害或不期望的行为。

大模型本质上属于生成模型。这些模型通过学习海量数据中的模式和结构，能够生成新的数据或输出，例如文本、图像或音频。生成模型的核心在于它们能够捕捉数据的分布并基于此生成新的实例。典型的大模型包括GPT系列（如GPT-3、GPT-4），它们可以生成自然语言文本，并且在各种任务中表现出强大的生成能力。

生成模型本质上涉及多步骤过程，通过这些过程产生输出。在设计攻击时，攻击者必须有效地在目标过程的每个步骤中施加影响，因此每个步骤都提供了使模型对攻击更加健壮的机会

这一观察也就启发雷文，我们可以专注于破坏对手对相关多步骤过程的控制，而不是试图检测攻击存在的二元分类问题。利用表示工程（RepE）中的技术，通过重新映射导致有害输出的模型表示序列，将它们引向不连贯或拒绝表示——即断路或短路。此外，通过直接针对生成有害响应的过程，我们的方法可以泛化到可能激活这些过程的各种输入。所以我们不需要识别所有可能触发不良输出的潜在输入，而只需要确保覆盖一个明确定义的此类输出集就可以。

这个方法有两个主要的组成部分：数据集和损失函数。如下伪码展示了一种使用低秩表示适应（LoRRA）的断路器技术，称之为表示重路由（RR）。

在RR中使用的训练数据被划分为两个集合：断路器集合和保留集合，每个集合在旨在控制模型中有害过程的训练过程中都具有不同的目的。与所有表示控制方法一样，断路器机制的质量在很大程度上取决于数据如何精确地引发目标表示。断路器集合由可能导致有害或不期望行为的内部表示的示例组成，并用于提示模型的断路器机制。相反，保留集合包括不应激活断路器的示例，并用于维护现有的理想模型表示以保留良性效能。尽管每个集合中的有限数量的示例就足以以一种超出训练数据的泛化方式改变模型的行为，但当训练数据与我们希望断路和保留的领域更好地对齐时，通常可以提高结果性能。

我们将拒绝数据添加到保留集合中可以增强模型正确拒绝有害用户请求的能力，并提高其其他能力的保留。

伴随数据集的损失是表示重路由损失和保留损失。记原始模型下有害过程的表示为reporig，带断路器的模型为repc/b。重路由损失旨在将有害过程的表示repc/b重新映射到期望的目标表示reprand。相反，保留损失用于维护保留集中的表示，这有助于保留这些表示。这通常被测量为当前表示和保留表示之间的ℓ2距离。

重路由损失可以采用多种形式。一种方法是将目标表示路由到具有大范数的固定随机方向，如在非学习方法RMU[中使用的。这表示为∥repc/b − αreprand∥2，其中reprand是一个随机向量，α是一个大常数，意味着放大表示的范数。然而，这种方法需要广泛调整α参数。或者也可以使用一种不需要超参数调整的随机向量损失的变体，公式为∥repc/b/ ∥repc/b∥ − reprand/ ∥reprand∥∥2。

由于我们希望目标表示对有害过程尽可能无帮助，所以另一种方法是直接优化断路器表示，使其与负责有害过程的原始表示正交。

这可以由它们的余弦相似度给出：repc/b · reporig/(∥repc/b∥2∥reporig∥2)。为了避免优化相似度超过零，我们在该目标上应用了ReLU函数。

实现

我们以llama3_8b为目标模型进行分析与实现。

Llama 3-8B 是 Meta 公司在 2024 年 4 月 18 日发布的最新一代开源大语言模型的一部分，该模型还包括 70B 参数的版本。Llama 3 系列模型在多项性能基准测试中展现出最先进（state-of-the-art）的性能，并在多个领域内具有广泛的应用潜力。Llama 3-8B 模型采用了标准的仅解码（decoder-only）式Transformer架构，并配有一个 128K 的词汇表，使其能够在 24GB 显存的设备上流畅运行。它的训练数据量是前代 Llama 2 的七倍，其中包括 30 多种语言的非英文数据，显示出模型在多语言环境下的良好适应性。

首先是与训练数据有关的代码

这个代码定义了一个名为 CircuitBreakerDataset 的类，继承自 Dataset 类，主要用于为语言模型生成对话数据：

初始化方法

参数定义：
- tokenizer: 预训练的分词器对象，用于处理文本数据。
- num_examples: 需要处理的示例数量。
- lorra_args: 其他相关参数（在代码中未使用）。
- model_name_or_path: 模型的名称或路径，用于选择不同的模板和配置。
父类初始化：
- 调用父类 Dataset 的初始化方法。
属性设置：
- 设置模型名称或路径的属性，并将其转换为小写。
- 设置最大文本长度为 1024。
定义默认模板：
- one_shot_template 是一个单一示例的模板字符串，包含用户标记、指令、助手标记和分隔符。

模型和模板配置

默认配置：
- 定义分隔符 sep_token、switch_select 列表（用于选择随机模式）、use_refusal_retain 标志以及用户和助手标记。
根据模型名称选择模板：
- 如果模型名称包含 'llama-3'，使用 Llama 模型的模板，设置相应的标记和配置。
- 如果模型名称包含 'mistral'，使用 Mistral 模型的模板，修正分词器的模板并设置相应的标记和分隔符。
- 如果模型名称不匹配以上两者，抛出未实现错误。
断言检查：
- 确认用户和助手标记已定义。

数据处理和保留

加载数据集：
- 加载 HuggingFace 提供的 ultrachat_200k 数据集的 test_sft 切片。
处理数据：
- 初始化空列表 orig_s 用于存储原始示例。
- 遍历数据集中的每个示例，提取消息列表。如果消息数量少于 2，跳过该示例。
- 随机选择 switch 值来决定使用哪种格式。
  - 如果 switch 为 0，使用模板将消息格式化并删除起始标记。
  - 如果 switch 为 1，使用单一示例模板格式化消息。
- 将格式化的输入添加到 orig_s 列表中。
- 如果 orig_s 的长度超过 num_examples，停止处理。
保存和打乱示例：
- 将 orig_s 列表保存到 self.orig_s_retain 属性中。
- 随机打乱保留的示例。

输出调试信息

打印调试信息：
- 打印 orig_s_retain 的第一个示例和总长度。

这个类的主要功能是根据不同的模型模板处理和格式化对话数据，然后将其保留用于训练或评估目的。

这段代码增加了两部分数据处理逻辑：Borderline Retain 和 Refusal Retain。它们的目的是从不同的数据源中读取和处理数据，并将其添加到保留的数据集中：

Borderline Retain 部分

读取和过滤数据：
- 从指定的 CSV 文件 data/xstest_v2_completions_gpt4_gpteval.csv 中读取数据。
- 仅保留 final_label 字段值为 "1_full_compliance" 的行。
初始化列表：
- 初始化一个空列表 borderline_orig_s 用于存储格式化后的输入数据。
格式化和保存数据：
- 遍历过滤后的数据集并重复 50 次，以增加数据量。
- 随机选择 switch 值来决定使用哪种格式。
  - 如果 switch 为 0，使用模板将 prompt 和 completion 格式化。
  - 如果 switch 为 1，使用模板将空的指令和 completion 格式化。
- 将格式化的输入添加到 borderline_orig_s 列表中。
合并和打乱数据：
- 将 borderline_orig_s 列表中的数据添加到 self.orig_s_retain 属性中。
- 随机打乱保留的数据。
打印调试信息：
- 打印 borderline_orig_s 的第一个示例和总长度。

Refusal Retain 部分

条件判断：
- 检查 use_refusal_retain 标志是否为 True。如果为 True，则执行下面的代码块。
读取和处理数据：
- 从 data/circuit_breakers_train.json 文件中读取 JSON 数据。
- 随机打乱数据集并截取前 2000 个数据点。
初始化列表：
- 初始化一个空列表 refusal_retain_orig 用于存储格式化后的输入数据。
格式化和保存数据：
- 遍历数据集并重复 2 次，以增加数据量。
- 随机选择 switch 值来决定使用哪种格式。
  - 如果 switch 为 0，使用模板将 prompt 和 llama3_output 格式化。
  - 如果 switch 为 1，使用模板将空的指令和 llama3_output 格式化。
- 将格式化的输入添加到 refusal_retain_orig 列表中。
合并和打乱数据：
- 将 refusal_retain_orig 列表中的数据添加到 self.orig_s_retain 属性中。
- 随机打乱保留的数据。
打印调试信息：
- 打印 refusal_retain_orig 的第一个示例和总长度。

整个代码块的目的是从多个数据源中读取、格式化并合并数据，然后将其随机打乱以用于后续的训练或评估。通过不同的模板和配置，这些数据被标准化为一致的格式，确保模型在训练过程中接收到一致且多样化的输入。

Circuit Breaker 部分

读取数据：
- 从 data/circuit_breakers_train.json 文件中读取 JSON 数据。
初始化列表：
- 初始化一个空列表 circuit_breaker_orig 用于存储格式化后的输入数据。
格式化和保存数据：
- 遍历数据集中的每个条目。
- 从数据集中获取 output 字段，并将其存储到 cb_output 变量。
- 随机选择 switch 值来决定使用哪种格式。
  - 如果 switch 为 0，使用模板将 prompt 和 cb_output 格式化。
  - 如果 switch 为 1，使用模板将空的指令和 cb_output 格式化。
- 将格式化的输入添加到 circuit_breaker_orig 列表中。
保存和打乱数据：
- 将 circuit_breaker_orig 列表保存到 self.circuit_breaker_orig 属性中。
- 随机打乱保留的数据。
打印调试信息：
- 打印 circuit_breaker_orig 的第一个示例和总长度。

Val 部分

读取数据：
- 从 data/circuit_breakers_val.json 文件中读取 JSON 数据。
初始化列表：
- 初始化一个空列表 val_orig 用于存储格式化后的验证数据。
格式化和保存数据：
- 遍历数据集中的每个条目。
- 使用模板将 prompt 和 output 格式化，并将格式化的输入添加到 val_orig 列表中。
保存数据：
- 将 val_orig 列表保存到 self.val_orig 属性中。
- 将分词器对象 tokenizer 保存到 self.tokenizer 属性中。

数据集长度和获取项的方法

数据集长度：
- 定义 __len__ 方法，返回 self.orig_s_retain 和 self.circuit_breaker_orig 中较小的长度。
获取数据项：
- 定义 __getitem__ 方法，接受索引 i 并返回相应的数据项。
- 获取 self.orig_s_retain、self.circuit_breaker_orig 和 self.val_orig 列表中对应索引的数据。
- cb_tokenized_kwargs 和 tokenize_kwargs 是用于分词器的参数字典，包含最大长度、填充和截断策略以及返回张量的选项。

整个代码块从不同的数据源读取、格式化并合并数据，用于训练和验证模型。这些数据被标准化为一致的格式，确保模型在训练和验证过程中接收到一致且多样化的输入。数据集的长度取决于保留的数据集和电路断路器数据集中较小的一个。__getitem__ 方法返回指定索引的数据项，并包括适用于分词器的参数字典。

这段代码补充了 CircuitBreakerDataset 类的 __getitem__ 方法，具体处理并返回电路断路器数据、保留数据和验证数据

Circuit Breaker 数据处理

分割请求和响应：
- 将 circuit_breaker_orig 按照 <SEPARATOR> 分割为 cb_request 和 cb_response。
设置分词器填充侧：
- 设置分词器的填充侧为 "left"，用于处理 cb_request。
分词处理：
- 使用 self.tokenizer 对 cb_request 进行分词，应用 cb_tokenized_kwargs 参数，生成 tokenized_request_circuit_breaker。
- 设置分词器的填充侧为 "right"，用于处理 cb_response。
- 对 cb_response 进行分词，且不添加特殊标记，应用 cb_tokenized_kwargs 参数，生成 response_tokenized_circuit_breaker。
恢复分词器填充侧：
- 将分词器的填充侧恢复为 "left"。
合并请求和响应：
- 将请求和响应的 input_ids 及 attention_mask 分别拼接，生成 combined_input_ids_circuit_breaker 和 combined_attention_mask_circuit_breaker。

Retain 数据处理

分词处理：
- 使用 self.tokenizer 对 orig_s_retain 进行分词，替换 <SEPARATOR> 为 self.sep_token，应用 tokenize_kwargs 参数，生成 tokenized_inputs_retain。

Val 数据处理

分词处理：
- 使用 self.tokenizer 对 val_orig 进行分词，替换 <SEPARATOR> 为 self.sep_token，应用 tokenize_kwargs 参数，生成 tokenized_inputs_val。

返回数据

返回字典：
- 返回包含以下键值对的字典：
  - input_ids_circuit_breaker: 电路断路器数据的 input_ids。
  - attention_mask_circuit_breaker: 电路断路器数据的 attention_mask。
  - input_ids: 保留数据的 input_ids。
  - attention_mask: 保留数据的 attention_mask。
  - input_ids_val: 验证数据的 input_ids。
  - attention_mask_val: 验证数据的 attention_mask。

__getitem__ 方法根据指定索引 i，处理并返回电路断路器数据、保留数据和验证数据。具体步骤包括分割请求和响应、设置分词器填充侧、分词处理、合并请求和响应，以及对保留数据和验证数据进行分词处理。最终返回一个包含各类数据的字典，以供模型训练或评估使用。

然后是合并和保存模型相关的代码

这个代码定义了一个函数，用于保存模型和处理器。这个函数的工作流程如下：

创建输出目录：检查并创建指定的输出目录（如果不存在的话），以确保可以将模型和处理器保存到该目录中。
打印保存路径：输出一条消息，指示模型和处理器将被保存到哪个目录。
合并LoRA模型：调用模型的 merge_and_unload 方法来合并LoRA（低秩适应）权重，并卸载不需要的部分。
合并原始层：
- 从预训练模型加载一个基础模型（anchor model）。
- 使用 drop_layers_after 参数，从合并后的模型中选择需要保留的层，并将这些层与基础模型中的其余层结合起来。
- 更新合并模型的配置，使其与基础模型的配置一致。
保存模型和处理器：将合并后的模型和处理器保存到指定的输出目录中。
保存LoRA配置：
- 创建一个路径指向输出目录中的 lorra_config.json 文件。
- 将 trainer 的 LoRA 参数保存为一个字典，并以 JSON 格式写入 lorra_config.json 文件中。
设置PyTorch算法：禁用确定性算法，以确保评估阶段的随机性。
评估模型：如果训练参数中指定了需要进行评估，则调用评估函数对模型进行评估。

通过合并和保存模型，确保最终的模型包含了所有必要的层和配置，同时保存处理器和相关的配置文件，以便后续加载和使用。

在之前方法介绍中，我们提到了两个核心部分，数据部分的代码已经分析了，现在来分析与损失有关的代码

这个函数定义了一个计算损失的过程，用于训练一个模型：

更新训练步骤计数器：self.current_training_step 增加 1，并且每隔 10 步打印日志信息。
获取输入数据：
- 获取用于保留（retain）的 input_ids 和 attention_mask。
- 获取用于断路器（circuit breaker）的 input_ids 和 attention_mask。
- 获取验证集（validation）的 input_ids 和 attention_mask。
准备输入数据：
- 创建包含 input_ids 和 attention_mask 的字典，并设置 output_hidden_states=True，分别为 retain、circuit breaker 和验证集输入数据。
计算进度系数：
- 调用 self.get_training_progress() 计算训练进度。
- 根据进度系数 scheduled_coeff 计算 retain 和 circuit breaker 的系数（系数是 alpha 与进度的乘积）。
打印进度和系数：输出当前的进度、retain系数和circuit breaker系数。
计算损失组件：
- 为断路器计算掩码，扩展并添加新维度。
- 禁用模型的适配器并切换到评估模式，避免梯度计算以节省内存。
- 进行以下计算：
  - Retain Control：如果 retain_coeff 大于 0，计算 retain 控制的输出，提取隐藏状态，应用掩码并保留。
  - Circuit Breaker Control：如果 circuit_breaker_coeff 大于 0，计算断路器控制的输出，并提取指定层的隐藏状态。
  - Validation：如果当前步骤需要日志记录，计算验证集的输出并提取指定层的隐藏状态。
释放内存：在每一步计算后，删除不再需要的变量并调用 gc.collect() 进行垃圾回收。
恢复训练模式：切换模型回训练模式。

这个函数通过在不同阶段计算 retain 控制、断路器控制和验证集的隐藏状态，来评估和记录模型的性能，进而用于进一步的损失计算和模型优化。

这个代码扩展了前面提到的 compute_loss 函数，增加了 retain 和 circuit breaker 控制的损失计算以及日志记录的部分：

Retain Control:
- 如果 retain_coeff 大于 0：
  - 使用模型处理 retain_inputs，并获取隐藏状态输出。
  - 将输出的隐藏状态与 layers_retain_attention_mask 逐元素相乘。
  - 计算 retain 损失，即 lora_retain_hidden 与 orig_retain_hidden 的欧氏距离（L2 范数），并取平均值。
  - 如果需要日志记录，计算 retain 输出和原始 retain 输出之间的余弦相似度，并打印出来。
Circuit Breaker Control:
- 如果 circuit_breaker_coeff 大于 0：
  - 使用模型处理 cb_inputs，并获取目标层的隐藏状态输出。
  - 对 lora_circuit_breaker_hidden 和 circuit_breaker_hidden 进行归一化。
  - 计算归一化后的 lora_circuit_breaker_hidden 和 circuit_breaker_hidden 的内积，并与 layers_circuit_breaker_attention_mask 相乘。
  - 计算 circuit breaker 损失，即内积之和的 ReLU，并除以掩码之和。
  - 如果需要日志记录，计算更新后的激活范数和原始激活范数，计算它们之间的余弦相似度，并打印出来。
Validation:
- 如果需要日志记录，使用模型处理 val_inputs 并获取隐藏状态输出。
- 计算验证输出和原始验证输出之间的余弦相似度，并打印出来。
计算总损失：
- 将 retain 损失和 circuit breaker 损失根据其系数进行加权求和，得到总损失。
打印损失：打印 retain 损失和 circuit breaker 损失的值。
返回损失：如果 return_outputs 为真，则返回损失元组，否则只返回损失值。

整个过程通过计算 retain 和 circuit breaker 控制的损失，评估模型在不同控制条件下的表现，并使用余弦相似度等度量进行日志记录，最终得到一个综合的损失值用于模型的训练和优化。

还有训练函数

这个 train 函数定义了一个训练过程的初始化步骤：

解析参数：
- 使用 transformers.HfArgumentParser 解析命令行或配置文件中的参数，并将其转换为 ModelArguments、TrainingArguments、LoraArguments 和 LorraArguments 数据类实例。
- 打印解析后的参数以进行检查。
设备映射：
- 设置 device_map 为 "auto"，以便根据硬件自动分配设备。
- 如果使用了 FSDP (Fully Sharded Data Parallel) 或 DeepSpeed 的 ZeRO3（这些技术与 QLoRA 不兼容），则发出警告。
读取模型和层信息：
- 获取模型的名称或路径 (model_name_or_path)。
- 从 lorra_args 中提取目标层 (target_layers)、转换层 (transform_layers) 和全层 (full_layers) 的信息。
- 将 target_layers 转换为整数列表。
- 根据 transform_layers 的值确定需要转换的层，如果 transform_layers 为 "-1"，则表示所有目标层都需要转换。
配置 LoRA：
- 使用 LoraConfig 创建 LoRA 配置对象，设置 LoRA 相关的参数，如 r、lora_alpha、target_modules、lora_dropout 和 lora_bias。
- 指定需要转换的层。
确定丢弃层：
- 如果 full_layers 为 False，则设置 drop_layers_after 为目标层的最大值，以确定丢弃哪些层。
配置和加载模型：
- 从预训练模型路径加载配置 (AutoConfig)。
- 根据 drop_layers_after 更新模型的隐藏层数量。
- 加载 tokenizer，并设置其填充标记为 eos_token 或 unk_token。
- 创建 save_model_function，用于将模型和 tokenizer 保存到指定的输出目录中。
加载模型：
- 使用 AutoModelForCausalLM 从预训练路径加载模型，并应用配置和设备映射。
应用 LoRA：
- 使用 get_peft_model 函数，将 LoRA 配置应用到模型中。
- 打印更新后的模型。
打印可训练参数（如果使用 DeepSpeed）：
- 如果 training_args.deepspeed 不为 None 并且 local_rank 为 0，则打印可训练的参数。
启用梯度检查点（如果配置了）：
- 如果启用了梯度检查点 (training_args.gradient_checkpointing)，则允许模型的输入梯度计算。
准备训练数据集：
- 创建 CircuitBreakerDataset 实例，准备训练数据集，设置数据集的大小为 10000。
- 打印训练数据集的长度。