当上下文回归:面向策略内蒸馏的稳健内化

arXiv cs.LG 论文

摘要

论文发现,将特权上下文重新引入蒸馏后的学生模型会导致性能下降(上下文诱导退化),并提出了一种轻量级一致性正则化器,该正则化器锚定无上下文输出以缓解此问题,从而在12种配置中提高了稳健性。

arXiv:2606.11627v1 Announce Type: new\n摘要:最近的研究表明,策略内蒸馏可以将特权上下文(如系统提示或任务提示)内化到学生模型中,从而在推理时不再需要该上下文。尽管这种方法成功提高了学生的无上下文性能,但我们发现了一个有趣且此前未被研究过的现象:在许多设置中,将原始特权上下文重新引入蒸馏后的学生模型实际上会降低其性能,即使是在它没有上下文时已经正确解决的实例上也是如此。我们将此称为上下文诱导退化,并认为稳健内化不仅要求匹配教师的上下文条件行为,还要求在上下文重新引入时保持稳定,我们将这一属性称为上下文可移除性。基于这一观察,我们提出了一种轻量级一致性正则化器,首先通过停止梯度锚定学生的无上下文输出,然后通过前向KL散度惩罚与无上下文输出偏离的上下文条件输出。这一简单的添加在每个训练步骤中仅需一次额外的前向传播,却能有效缓解上下文诱导退化,并且在许多情况下甚至提高了无上下文性能。在涵盖不同领域和模型家族的12种配置中,我们的方法在大多数设置中提高了上下文条件准确率,在12种设置中的11种中减少了上下文诱导的损害,并有效消除了响应长度膨胀。一项机制性案例研究进一步证实,上下文可移除性是在表示层面实现的,无论上下文是否存在,隐藏状态几乎保持相同。
查看原文
查看缓存全文

缓存时间: 2026/06/11 13:50

# 面向在策略蒸馏中的鲁棒内化
来源:https://arxiv.org/html/2606.11627
## 当上下文回归:面向在策略蒸馏中的鲁棒内化
Xun Wang¹, Ruishuo Chen¹, Zhuoran Li¹, Yu Chen¹, and Longbo Huang¹\{\}^\{1\\,\\text\{\\faIcon\{envelope\}\}\} ¹清华大学交叉信息研究院
通讯作者:[email protected]

###### 摘要
近期研究表明,在策略蒸馏能够将特权上下文(如系统提示或任务提示)内化到学生模型中,使得推理时不再需要该上下文。虽然这种方法成功地提升了学生在无上下文情况下的表现,但我们发现了一个有趣且此前未被研究过的现象:在许多设置中,将原始特权上下文重新引入蒸馏后的学生模型,实际上会*降低*其性能,甚至对于已经在无上下文中正确解决的实例也是如此。我们将此称为*上下文引发的退化*,并认为鲁棒内化不仅要求学生匹配教师基于上下文的行为,还要求在上下文被重新引入时保持稳定,我们将这一属性称为*上下文可去除性*。受此观察启发,我们提出了一种轻量级一致性正则化方法,该方法首先通过停止梯度锚定学生无上下文的输出,然后通过前向KL散度惩罚基于上下文的输出偏离该锚定。这一简单的添加每个训练步骤仅需一次额外的前向传播,却能有效缓解上下文引发的退化,并且在许多情况下甚至能提升无上下文的表现。在涵盖不同领域和模型家族的12种配置中,我们的方法在大多数设置中提升了基于上下文的准确率,在12种设置中的11种中减少了上下文引发的损害,并有效消除了响应长度膨胀。一项机制性案例研究进一步证实,上下文可去除性是在表示层面实现的,无论上下文是否存在,隐藏状态几乎保持一致。

## 1. 引言
语言模型常常通过特权上下文进行增强,例如编码行为约束的系统提示、提供专家推理过程的思维链轨迹(Wei等人,2022 (https://arxiv.org/html/2606.11627#bib.bib1)),或总结决策相关信息的游戏状态描述。这些上下文通过上下文学习(Brown等人,2020 (https://arxiv.org/html/2606.11627#bib.bib2);Dong等人,2024 (https://arxiv.org/html/2606.11627#bib.bib3))提升性能,但同时也增加了延迟、提高了服务成本,并存在向最终用户泄露敏感指令的风险。为了消除这些缺点,在策略自蒸馏(Zhao等人,2026a (https://arxiv.org/html/2606.11627#bib.bib4);Shenfeld等人,2026 (https://arxiv.org/html/2606.11627#bib.bib5);Hübotter等人,2026 (https://arxiv.org/html/2606.11627#bib.bib6))通过将模型自身基于上下文的输出蒸馏到无上下文通路中,训练模型内化上下文。Ye等人(2026 (https://arxiv.org/html/2606.11627#bib.bib7))将该思想推广到在策略上下文蒸馏,允许教师和学生模型不同,从而使学生能够在部署时*无需*上下文也能表现良好。

一个自然的问题随之而来:如果学生确实内化了特权信息,那么当上下文被*重新引入*时会发生什么?直观上,重新呈现已经内化的信息应该是一个幂等操作(Liu等人,2026 (https://arxiv.org/html/2606.11627#bib.bib8)):一个已经记住课本的学生,在被允许查阅课本时不应表现更差。然而,我们发现实际情况往往并非如此:在大多数设置中,重新引入蒸馏后的上下文不仅会膨胀模型输出长度,还会*降低*其准确率(图 1 (https://arxiv.org/html/2606.11627#S1.F1),左图)。我们将此现象称为*上下文引发的退化*。

参见图注
图1:左图:OPD训练学生模型在无法访问c的情况下匹配教师基于上下文的输出。在测试时,学生模型在没有上下文时表现良好,但当c被重新引入时,其性能可能下降。中图:在12种设置中平均的训练曲线。OPD的加上下文准确率(灰色虚线)持续落后于不加上下文准确率(灰色实线),而NCA(彩色线)缩小了这一差距。右图:最终准确率比较。NCA同时提升了无上下文(+2.1个百分点)和基于上下文(+3.0个百分点)的准确率。

我们将此归因于OPD仅针对我们所谓的*特权保真度*进行优化,即无上下文的学生应匹配有上下文的教师,但并未约束学生*带上下文*的行为。我们认为,鲁棒内化需要一个额外的属性,即*上下文可去除性*:如果模型确实将特权信息吸收到其参数中,那么上下文就变得冗余,学生的输出应无论上下文是否存在都保持稳定。有趣的是,我们发现无需对训练流程进行根本性更改,只需添加一个简单的一致性KL正则化项,就能有效缓解这种退化,且几乎不损失特权保真度。此外,在许多设置中,这个正则化项甚至*提升*了学生的无上下文性能,而这是主要的训练目标。具体而言,我们提出了*无上下文锚定(No-Context Anchoring, NCA)*,该方法将学生的无上下文输出视为停止梯度锚定,并惩罚基于上下文的输出偏离该锚定,每个训练步骤仅需一次额外的前向传播。在涵盖不同领域和模型家族的12种配置中,NCA在大多数设置中提升了基于上下文的准确率,在12种设置中的11种中减少了上下文引发的损害,并有效消除了响应长度膨胀(图 1 (https://arxiv.org/html/2606.11627#S1.F1),中图和右图)。进一步的机制分析证实,NCA不仅在输出层面实现了上下文可去除性,还在表示层面实现了这一点,无论上下文是否存在,都产生几乎相同的隐藏状态。总结起来,我们的主要贡献如下:

- •我们研究了蒸馏后学生模型在特权上下文被重新引入时性能下降的现象,称之为*上下文引发的退化*,并提出了*上下文可去除性*作为鲁棒内化的理想属性。
- •我们提出了无上下文锚定(NCA),一种简单而有效的一致性正则化方法,强制实现上下文可去除性,同时提升无上下文性能,每个训练步骤仅需一次额外的前向传播。
- •在涵盖不同领域和模型家族的12种配置中,NCA在大多数设置中提升了基于上下文的准确率,在12种设置中的11种中减少了上下文引发的损害,并有效消除了响应长度膨胀。一项机制性案例研究进一步证实,NCA在表示层面实现了上下文可去除性。

## 2. 相关工作
##### 知识蒸馏与在策略蒸馏。
知识蒸馏(Hinton等人,2015 (https://arxiv.org/html/2606.11627#bib.bib9))通过匹配输出分布将知识从教师传递给学生,已广泛应用于大语言模型,用于压缩推理(Guo等人,2025 (https://arxiv.org/html/2606.11627#bib.bib10))、指令遵循(Ouyang等人,2022 (https://arxiv.org/html/2606.11627#bib.bib11))和通用能力(Gu等人,2024 (https://arxiv.org/html/2606.11627#bib.bib12);Ko等人,2024 (https://arxiv.org/html/2606.11627#bib.bib14);Xu等人,2024 (https://arxiv.org/html/2606.11627#bib.bib13))。最近,在策略蒸馏(OPD)(Agarwal等人,2024 (https://arxiv.org/html/2606.11627#bib.bib15);DeepSeek-AI,2026 (https://arxiv.org/html/2606.11627#bib.bib16);Xiao等人,2026 (https://arxiv.org/html/2606.11627#bib.bib17);Fu等人,2026 (https://arxiv.org/html/2606.11627#bib.bib20);Li等人,2026b (https://arxiv.org/html/2606.11627#bib.bib19))从学生自身的分布中生成训练数据,减轻了离线策略方法中的训练-推理分布不匹配问题。当教师是模型本身基于特权上下文的条件版本时,这就成为在策略自蒸馏(Zhao等人,2026a (https://arxiv.org/html/2606.11627#bib.bib4);Shenfeld等人,2026 (https://arxiv.org/html/2606.11627#bib.bib5);Hübotter等人,2026 (https://arxiv.org/html/2606.11627#bib.bib6))。Ye等人(2026 (https://arxiv.org/html/2606.11627#bib.bib7))进一步放宽了这一设置,允许特权上下文教师和学生由不同模型实例化。Zhao等人(2026b (https://arxiv.org/html/2606.11627#bib.bib18))通过将推理工具本身视为待蒸馏的特权信息来扩展这一思想。关于OPD及其变体的更广泛讨论,请参阅Song and Zheng (2026 (https://arxiv.org/html/2606.11627#bib.bib21))。

##### 在策略蒸馏的扩展与分析。
越来越多的文献分析了在策略蒸馏背后的机制。例如,Li等人(2026b (https://arxiv.org/html/2606.11627#bib.bib19))识别了决定其成功与失败的条件,而Fu等人(2026 (https://arxiv.org/html/2606.11627#bib.bib20))则列举了采样令牌OPD的三种失败模式。作为这些分析的补充,最近的一些工作将在策略自蒸馏与基于可验证奖励的强化学习相结合(Yang等人,2026 (https://arxiv.org/html/2606.11627#bib.bib22);Li等人,2026a (https://arxiv.org/html/2606.11627#bib.bib23)),旨在进一步提升训练效率和推理性能。我们的关注点与这两条工作线正交:我们并非解释OPD何时成功、诊断其训练失败或改进其优化信号,而是研究成功蒸馏后当特权上下文被重新引入时会发生什么。我们表明,仅优化特权保真度而不引入显式的上下文可去除性目标是不够的。

##### 一致性正则化。
一致性正则化鼓励模型对语义等价的输入产生相似的输出(Berthelot等人,2019 (https://arxiv.org/html/2606.11627#bib.bib24);Sohn等人,2020 (https://arxiv.org/html/2606.11627#bib.bib25))。最近的应用包括大语言模型中对改写指令的鲁棒性(Zhao等人,2024 (https://arxiv.org/html/2606.11627#bib.bib26))和自奖励对齐(Wang等人,2025 (https://arxiv.org/html/2606.11627#bib.bib27)),以及持续学习中的幂等知识保持(Liu等人,2026 (https://arxiv.org/html/2606.11627#bib.bib8))。这些工作在模型行为的多个方面强制一致性。相比之下,我们将一致性原则应用于上下文蒸馏,确保蒸馏后的模型在推理时重新引入训练中使用的特权上下文时保持稳定。

## 3. 问题形式化
考虑一个由θ参数化的学生模型gθ,一个输入x,以及一个固定的特权上下文c(例如,系统提示、任务提示或游戏状态框架)。¹¹由于c对于每个任务是给定且固定的,方程1中的KL不包含不可约熵项(Yang等人,2026 (https://arxiv.org/html/2606.11627#bib.bib22))。然而即使在这种有利的设置下,仅优化特权保真度也无法保证上下文可去除性。我们将两种学生视图记为qx ≜ gθ(y∣x) 和 qc ≜ gθ(y∣x,c)。基于特权上下文的在策略蒸馏(Ye等人,2026 (https://arxiv.org/html/2606.11627#bib.bib7))通过最小化以下损失来训练学生:
LOPD = Ex∼D [ KL( qx ∥ f(y∣x,c) ) ], (1)
其中f(y∣x,c)是教师分布。实践中,KL是在令牌级别计算并沿序列求和。注意方程1仅约束qx。基于上下文的视图qc没有收到学习信号,也没有被直接约束。我们认为,成功内化的一个必要属性是*上下文可去除性*:如果模型确实将c吸收到其参数中,那么基于c的条件化就变得冗余,两个视图应该一致,即qx ≈ qc。尽管人们可能期望一个足够表达能力的模型在训练后满足这一属性,但我们发现在大多数设置中情况并非如此。

参见图注
图2:OPD后的上下文交互。x轴:重新引入上下文时的准确率下降;y轴:Harm 率。

##### 上下文交互的三种模式。
图2根据重新引入上下文时的准确率下降(Accx - Accx,c)与 Harm = P(qc 错误 | qx 正确)(即,一个正确的无上下文预测在引入上下文后变得错误的频率)绘制了我们的12种设置。出现了三种模式。在模式A(7/12设置)中,上下文充当扰动:学生已经能在没有上下文的情况下表现良好,但重新引入它会导致准确率下降和高伤害率。在模式B中,上下文仍然作为一个有用的框架(Accx,c > Accx),表明内化不完全。在模式C中,两个视图已经一致,伤害最小。模式A的普遍性表明了引入显式上下文可去除性目标的必要性。

## 4. 方法:无上下文锚定
### 4.1. 从约束目标到正则化目标
第3节的分析表明,标准训练目标不强制上下文可去除性。一个自然的修复是在优化特权保真度的同时显式约束两个视图之间的散度:
minθ LOPD(θ) s.t. D(qx, qc) ≤ ε, (2)
其中D是两个视图之间的散度。在实践中,我们将其松弛为带有固定权重β > 0的惩罚形式:
L = LOPD + β Ex∼D [ D(qx, qc) ]. (3)

### 4.2. 散度的选择
方程3中的散度D可以有多种形式。我们选择前向KL散度 D = KL(qx ∥ qc),它是模式覆盖的,因此防止qc坍缩qx的模式。另一个好处是期望落在qx之下,允许我们复用OPD已经生成的在策略样本,而无需从qc进行单独 rollout。为了进一步防止这一项通过qx干扰OPD目标的梯度,我们应用停止梯度来锚定无上下文视图。我们将由此产生的正则化器称为*无上下文锚定(No-Context Anchoring, NCA)*:
LNCA ≜ KL( sg[ gθ(·∣x) ] ∥ gθ(·∣x,c) ), (4)
其中sg[·]表示

相似文章

上下文蒸馏作为潜在记忆管理

arXiv cs.LG

本文将上下文蒸馏表述为一个潜在记忆管理问题,提出一个框架,将蒸馏后的上下文存储为独立的LoRA适配器,并通过检索、路由和自门控机制来提高鲁棒性和效率。

同策略蒸馏(5分钟阅读)

TLDR AI

本文引入同策略蒸馏,通过在教师提供的token级KL正则化下,在学生自身轨迹上训练学生模型,解决训练-推理分布不匹配问题,统一了前向KL、反向KL和JSD损失,其中反向KL更适用于较小的学生模型。

揭秘同策略蒸馏:其益处、危害及原因

Hugging Face Daily Papers

本文介绍了一种无需训练的框架,用于分析推理模型在逐token级别上的蒸馏信号。研究揭示,蒸馏引导在错误推理路径上更为有效,且其效果取决于学生模型的能力及任务上下文。