STAIR: Improving Safety Alignment with Introspective Reasoning

📄 arXiv: 2502.02384v2 📥 PDF

作者: Yichi Zhang, Siyuan Zhang, Yao Huang, Zeyu Xia, Zhengwei Fang, Xiao Yang, Ranjie Duan, Dong Yan, Yinpeng Dong, Jun Zhu

分类: cs.CL

发布日期: 2025-02-04 (更新: 2025-06-27)

备注: 22 pages, 8 figures, ICML2025 Oral

🔗 代码/项目: GITHUB

💡 一句话要点

STAIR：通过内省推理提升大语言模型的安全对齐

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 安全对齐 内省推理 大语言模型 蒙特卡洛树搜索 思维链 越狱攻击 奖励模型

📋 核心要点

  1. 现有安全对齐方法在安全性和性能之间存在权衡，且易受越狱攻击，主要原因是依赖直接拒绝恶意查询。
  2. STAIR框架通过内省推理，使LLM能够逐步分析识别安全风险，并利用安全意识的CoT推理进行自我改进。
  3. 实验表明，STAIR在减轻有害输出的同时，更好地保留了有用性，并在越狱攻击防御上达到与Claude-3.5相当的水平。

📝 摘要（中文）

确保大型语言模型（LLM）的安全性和无害性，与它们在应用中的性能同等重要。然而，现有的安全对齐方法通常面临安全性和性能之间的权衡，并且容易受到越狱攻击，这主要是因为它们依赖于对恶意查询的直接拒绝。本文提出了STAIR，一个将安全对齐与内省推理相结合的新框架。STAIR使LLM能够通过逐步分析来识别安全风险，通过具有安全意识的思维链（CoT）推理进行自我改进。STAIR首先使模型具备结构化的推理能力，然后通过在我们新提出的安全信息蒙特卡洛树搜索（SI-MCTS）生成的步级推理数据上进行迭代偏好优化来推进安全对齐。我们进一步在此数据上训练过程奖励模型，以指导测试时搜索，从而改进响应。大量实验表明，与本能对齐策略相比，STAIR有效地减轻了有害输出，同时更好地保留了有用性。通过测试时扩展，STAIR在针对流行的越狱攻击时，实现了与Claude-3.5相当的安全性能。本文的相关资源可在https://github.com/thu-ml/STAIR获取。

🔬 方法详解

问题定义：现有的大语言模型安全对齐方法，如直接拒绝有害请求，往往导致模型在安全性和实用性之间做出妥协。此外，这些方法容易受到对抗性攻击（越狱攻击），攻击者可以通过精心设计的提示绕过安全机制，诱导模型产生有害输出。因此，如何提升LLM在复杂场景下的安全性和鲁棒性，同时保持其生成有用信息的能力，是一个亟待解决的问题。

核心思路：STAIR的核心思路是赋予LLM内省推理的能力，使其能够像人类专家一样，逐步分析问题，识别潜在的安全风险。通过这种方式，模型不仅可以识别显而易见的有害请求，还可以识别隐藏在复杂问题中的潜在风险。此外，STAIR还利用安全意识的思维链（CoT）推理，引导模型生成更安全、更合理的响应。

技术框架：STAIR框架包含以下几个主要模块：1) 结构化推理能力构建：通过训练使模型具备逐步分析问题的能力。2) 安全信息蒙特卡洛树搜索（SI-MCTS）：用于生成带有安全信息的步级推理数据，该数据用于后续的偏好优化。3) 迭代偏好优化：利用SI-MCTS生成的数据，通过迭代优化模型，使其更倾向于生成安全的响应。4) 过程奖励模型：训练一个奖励模型，用于评估模型在推理过程中的每一步的安全性，并在测试时指导搜索，从而改进响应。

关键创新：STAIR的关键创新在于将内省推理和安全对齐相结合。与传统的直接拒绝方法不同，STAIR通过逐步分析问题，识别潜在的安全风险，从而更有效地防御越狱攻击。此外，SI-MCTS的引入，使得模型能够在探索潜在响应的同时，考虑到安全性，从而生成更安全、更合理的响应。

关键设计：SI-MCTS在传统的蒙特卡洛树搜索基础上，引入了安全奖励函数，用于评估每个节点的安全性。在搜索过程中，模型不仅考虑生成有用信息的可能性，还考虑生成有害信息的风险。过程奖励模型的设计至关重要，它需要准确评估模型在推理过程中的每一步的安全性。具体实现中，可以使用人工标注的数据或者基于规则的方法来训练过程奖励模型。此外，迭代偏好优化的具体算法（如Direct Preference Optimization, DPO）以及超参数的选择也会影响最终的安全对齐效果。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，STAIR在减轻有害输出的同时，更好地保留了LLM的有用性。在针对流行的越狱攻击的测试中，STAIR实现了与Claude-3.5相当的安全性能，显著优于传统的安全对齐方法。此外，通过测试时扩展，STAIR的安全性可以进一步提升，表明其具有良好的可扩展性。

🎯 应用场景

STAIR框架可应用于各种需要安全保障的大语言模型应用场景，例如智能客服、内容生成、代码生成等。通过提升LLM的安全性和鲁棒性，STAIR可以减少有害信息的传播，提高用户信任度，并降低潜在的法律风险。未来，STAIR可以进一步扩展到其他类型的AI系统，例如机器人和自动驾驶汽车，以确保这些系统的安全可靠运行。

📄 摘要（原文）

Ensuring the safety and harmlessness of Large Language Models (LLMs) has become equally critical as their performance in applications. However, existing safety alignment methods typically suffer from safety-performance trade-offs and the susceptibility to jailbreak attacks, primarily due to their reliance on direct refusals for malicious queries. In this paper, we propose STAIR, a novel framework that integrates SafeTy Alignment with Itrospective Reasoning. We enable LLMs to identify safety risks through step-by-step analysis by self-improving chain-of-thought (CoT) reasoning with safety awareness. STAIR first equips the model with a structured reasoning capability and then advances safety alignment via iterative preference optimization on step-level reasoning data generated using our newly proposed Safety-Informed Monte Carlo Tree Search (SI-MCTS). We further train a process reward model on this data to guide test-time searches for improved responses. Extensive experiments show that STAIR effectively mitigates harmful outputs while better preserving helpfulness, compared to instinctive alignment strategies. With test-time scaling, STAIR achieves a safety performance comparable to Claude-3.5 against popular jailbreak attacks. Relevant resources in this work are available at https://github.com/thu-ml/STAIR.