Reducing the Safety Tax in LLM Safety Alignment with On-Policy Self-Distillation

📄 arXiv: 2605.15239v1 📥 PDF

作者: Yu Fu, Longxuan Yu, Haz Sameen Shahgir, Zhipeng Wei, Hui Liu, N. Benjamin Erichson, Yue Dong

分类: cs.LG

发布日期: 2026-05-14

备注: 20 pages, 5 figures


💡 一句话要点

提出OPSA:通过在线自蒸馏减少LLM安全对齐中的安全税。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 安全对齐 在线自蒸馏 大型语言模型 安全税 教师翻转率

📋 核心要点

  1. 现有安全对齐方法存在“安全税”问题,即提高安全性的同时降低了模型的推理能力。
  2. 提出在线自蒸馏方法OPSA,通过模型自身生成轨迹并进行token级别KL散度监督,缓解分布不匹配问题。
  3. 实验表明,OPSA在多个模型和规模上,相比离线蒸馏和外部教师蒸馏,实现了更好的安全性和推理能力权衡。

📝 摘要(中文)

安全对齐通常以牺牲推理能力为代价来提高对有害查询的鲁棒性,这种权衡被称为安全税。一个常见原因是分布不匹配:监督微调在人类、外部模型或固定自生成轨迹产生的安全演示上训练目标模型,而不是在其自身策略采样的轨迹上训练。本文将离线训练不匹配确定为安全税的第二个来源,并研究用于安全对齐的在线自蒸馏,称之为OPSA。该模型生成自己的rollout,并从自身的一个冻结教师副本接收密集的token级别KL散度监督,该教师副本以特权安全上下文为条件。因为这个教师必须比采样的学生轨迹更安全,所以引入了教师翻转率:一个衡量特权上下文将不安全响应转换为安全响应的频率的标准。使用这个信号来搜索激活潜在安全推理而不是仅仅引出看起来安全的演示的上下文。在两个推理模型系列和五个模型规模上,OPSA在匹配数据和全参数微调下,比离线自蒸馏和外部教师蒸馏实现了更强的安全-推理权衡,在较小模型上获得了最大的收益(R1-Distill-1.5B上+8.85点,Qwen3-0.6B上+5.49点)。这些收益在训练集大小和自适应越狱评估中持续存在。Token级别的分析进一步表明,OPSA将更新集中在早期的合规决策token附近,为在保持一般推理的同时提高安全性提供了一种机制。

🔬 方法详解

问题定义:大型语言模型(LLM)的安全对齐旨在使其对有害查询更具鲁棒性,但通常会降低其推理能力,即所谓的“安全税”。现有的监督微调方法通常使用人类标注、外部模型或固定自生成轨迹来训练模型,导致训练数据与模型自身策略生成的分布不匹配,从而影响模型的泛化能力和推理能力。

核心思路:本文的核心思路是通过在线自蒸馏来解决安全对齐中的分布不匹配问题。具体来说,模型通过与自身交互生成轨迹,并使用自身的“教师”模型(一个更安全的版本)来监督训练过程。这种在线学习的方式可以减少训练数据与模型自身策略之间的差异,从而提高模型的安全性和推理能力。

技术框架:OPSA(On-Policy Self-Distillation for Safety Alignment)的整体框架包括以下几个主要步骤:1) 学生模型生成文本序列;2) 教师模型(学生模型的冻结副本,并以安全上下文为条件)生成相应的文本序列;3) 计算学生模型和教师模型输出之间的token级别KL散度;4) 使用KL散度作为损失函数来更新学生模型。关键在于教师模型使用了“特权安全上下文”,引导其生成更安全的响应。

关键创新:OPSA的关键创新在于使用在线自蒸馏来解决安全对齐中的分布不匹配问题,并引入了“教师翻转率”这一指标来衡量安全上下文的有效性。教师翻转率用于评估在特定上下文中,教师模型将不安全响应转换为安全响应的频率,从而帮助选择更有效的安全上下文。

关键设计:OPSA的关键设计包括:1) 使用token级别的KL散度作为损失函数,鼓励学生模型模仿教师模型的行为;2) 引入“教师翻转率”来选择有效的安全上下文;3) 使用冻结的教师模型,避免教师模型与学生模型过度拟合;4) 通过实验确定合适的KL散度系数和训练迭代次数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,OPSA在多个模型和规模上都优于离线自蒸馏和外部教师蒸馏。例如,在R1-Distill-1.5B模型上,OPSA的性能提升了8.85个点,在Qwen3-0.6B模型上提升了5.49个点。此外,OPSA在自适应越狱评估中也表现出更强的鲁棒性,表明其具有更好的安全性。

🎯 应用场景

OPSA方法可应用于各种需要安全对齐的大型语言模型,例如聊天机器人、智能助手等。通过提高模型的安全性和推理能力,可以减少有害内容的生成,提升用户体验,并降低潜在的风险。该研究对于构建更安全、更可靠的人工智能系统具有重要意义。

📄 摘要(原文)

Safety alignment often improves robustness to harmful queries at the cost of reasoning ability, a tradeoff known as the safety tax. A common cause is distributional mismatch: supervised fine-tuning trains the target model on safety demonstrations produced by humans, external models, or fixed self-generated traces, rather than on trajectories sampled from its own policy. We identify off-policy training mismatch as a second source of this tax and study on-policy self-distillation for safety alignment, which we call OPSA. The model generates its own rollouts and receives dense per-token KL supervision from a frozen teacher copy of itself conditioned on a privileged safety context. Because this teacher must be safer than the sampled student trajectory, we introduce \emph{teacher flip rate}: a criterion that measures how often a privileged context converts unsafe responses into safe ones. We use this signal to search for contexts that activate latent safety reasoning rather than merely elicit safe-looking demonstrations. Across two reasoning-model families and five model scales, OPSA achieves a stronger safety--reasoning tradeoff than off-policy self-distillation and external-teacher distillation under matched data and full-parameter fine-tuning, with the largest gains on smaller models (+8.85 points on R1-Distill-1.5B and +5.49 points on Qwen3-0.6B). The gains persist across training-set sizes and adaptive jailbreak evaluations. Token-level analyses further show that OPSA concentrates updates near early compliance-decision tokens, providing a mechanism for improving safety while preserving general reasoning.