STAIR:Improving Safety Alignment with Introspective Reasoning
| 原文 | STAIR: Improving Safety Alignment with Introspective Reasoning |
|---|---|
| 会议 | ICML |
| 时间 | 2025 |
| 更新 | 2025-08-12 10:04:47 |
| 状态 | Done |
作者想将System 1 thinking(快思考)变为System 2 thinking(慢思维)，使得LLM真正理解风险指令
安全对齐已成为增强大语言模型安全性和无害性的关键解决方案，典型方法包括SFT、基于偏好的优化如RLHF和DPO等。
尽管现有的安全对齐方法使LLMs能够直接拒绝具有明显风险的查询，但它们往往无法抵抗能够避免初始拒绝tokens的越狱攻击。
作者的三个贡献:
-
提出了 STAIR 框架，将安全对齐与内省推理相结合，使大语言模型能够通过具有安全意识的思维链分析来检查潜在风险，并通过安全意识的系统2思维确保输出无害。
-
引入SI-MCTS，并设计了满足特定属性的安全感知奖励函数，通过迭代自我改进机制，在逐步推理数据上进行偏好优化，增强了模型的安全意识推理能力。
-
通过实验证明了 STAIR 在多个基准测试中有效增强了对各种有害查询的抵抗力，减轻了安全-性能权衡。
方法¶
结构图
-
在由 GPT-4o 生成的结构化思维链数据上对模型进行初始训练。
-
通过自我生成和自我奖励构建SI-MCTS。此过程中的安全感知奖励函数将安全性信息与实用性信息整合到内部搜索节点中。从构建的搜索树中，通过阈值采样收集逐步偏好数据集，用于通过步骤级 DPO 优化模型。这个自我改进过程可以重复K=3次迭代。
-
可以基于相同的搜索树进一步训练一个过程奖励模型（PRM），并通过测试时搜索算法引导最后一次迭代的模型生成更好、更安全的响应
2.1 结构化思维链格式对齐¶
为使模型能以系统2思维分析风险，首先通过在少量安全性和实用性数据上进行微调，赋予模型推理能力。
2.2 安全蒙特卡洛自我提升¶
蒙特卡洛搜索树: 选择→扩展→模拟→反向传播
初始化: 建立根节点，用LLM生成的初始答案构造，来最小化模型过拟合的趋势。
选择：为了进一步探索，利用一个函数对未探索的节点的答案进行价值排序，随后选择价值最高的节点进行进一步的探索。
扩展: 如果该节点不是最终节点，则创建新的节点，表示下一步可能的动作。
模拟: 从新扩展的节点开始，随机选择其中一个节点，进行模拟，直到得到结果。
反向传播: 模拟结果按照路径返回，并获得奖励分数。
形式上，设对\(f\)的评估包含两个奖励函数：有用性奖励函数\(H(f) \in [-1,1]\)和安全性奖励函数\(S(f) \in [-1,1]\)。我们假设当答案\(f\)安全时\(S(f) >0\)；否则 \(S(f) <0\)。由\(H\)和\(S\)参数化的安全感知奖励函数 \(R:[-1,1]\times[-1,1]\rightarrow\mathbb{R}\)定义为：对于任意最终答案\(f\)，\(R(f) := R(H(f),S(f))\)。该函数需满足以下三个性质：
-
安全优先原则：安全响应始终获得比不安全响应更高的奖励，即 \(\forall f_1, f_2, S(f_1) >0 > S(f_2) \Rightarrow R(f_1) > R(f_2)\)；
-
帮助性双重单调性：预期帮助性取决于响应安全性，即 \(\forall S >0, \frac{\partial R}{\partial H} >0\)且 \(\forall S <0, \frac{\partial R}{\partial H} <0\)；
-
单目标退化特性：当仅关注单一维度时，可将奖励设为与该维度奖励保持恒定差值，即 \(\exists C_1 \in [-1,1], \text{令} S \equiv C_1, \forall f_1, f_2, R(f_1) - R(f_2) = H(f_1) - H(f_2)\)；\(\exists C_2 \in [-1,1], \text{令} H \equiv C_2, \forall f_1, f_2, R(f_1) - R(f_2) = S(f_1) - S(f_2)\)。
为了设计这样一个函数满足上述性质，作者通过证明得到了一个定理
设常数 \(C_1, C_2 \in [-1,1]\)且 \(C_1 \neq0\)。若函数 \(R: [-1,1] \times [-1,1] \rightarrow \mathbb{R}\)二阶可微，并满足偏微分关系\(\frac{\partial R}{\partial H} = F(S)\)，其中 \(F: [-1,1] \rightarrow \mathbb{R}\)为连续函数。则后两个性质成立的充要条件为：
且满足边界条件：\(F(0) =0, \quad F(C_1) =1, \quad \forall S >0, F(S) >0, \quad \forall S <0, F(S) <0\)其中\(c\)为常数。
我们注意到，在方程(1)中取 \(C_1 =1\)、\(C_2 = -1\)、\(F(S) = S\)、\(c =0\)时，\(R(H, S) = S \cdot H +2S\)是满足"安全优先"第一特性的最简形式。
逐步偏好优化¶
当 MCTS 搜索预算耗尽时，从搜索树中提取\(s_i=(z_1,...,z_i)\)共享相同历史路径的节点对。\((z^w_{i+1}, z^l_{i+1})\)根据估值函数对节点排序，并且通过节点之间的函数差值和正样本绝对值作为约束条件来保证偏好样本的质量。对于一个配对样本\((x, \mathbf{s}_i, z^w_{i+1}, z^l_{i+1}) \sim \mathcal{D}_1\)，训练目标为：
迭代式自我改进¶
通过基于高质量偏好数据的多轮迭代，我们可以持续强化安全对齐效果。具体而言，我们采用step-wise DPO方法对模型 \(\pi^k_\theta\)（\(k =1,...,K\)进行迭代优化，其偏好数据集\(D^k\)由上一轮经SI-MCTS训练的模型\(\pi^{k-1}_\theta\)生成。
2.3 Test-time scaling¶
训练一个奖励模型，并指导语言模型进行测试时搜索。
首先，通过采样相同深度的部分解对\((x, s_w^i, s_l^i)\)，我们构建了用于奖励建模的偏好数据集\(\mathcal{D}_R\)。然后训练得到过程奖励模型(PRM)，在其上评估部分解\((x,s_i)\)。
后续基于该奖励模型给出的信号进行Best-of-N搜索或束搜索。
-Best-of-N: 首先独立采样\(N\)个答案，然后使用PRM对每个答案进行评分，对所有的得到相同最终答案的解决方案，将他们的分数累加，累加总分最高的答案就作为最终输出。
- Beam search: 这是一种逐步解决方案的方法，保留束宽(beam width)个最高概率候选项，每个保留的候选项生成下一个可能的步骤，直到得到答案后，从所有候选序列中选取奖励最高的那个。