When Context Returns: Toward Robust Internalization in On-Policy Distillation
作者: Xun Wang, Ruishuo Chen, Zhuoran Li, Yu Chen, Longbo Huang
分类: cs.LG, cs.AI
发布日期: 2026-06-10
💡 一句话要点
提出轻量级一致性正则化以解决上下文引起的性能下降问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 上下文引起的退化 一致性正则化 蒸馏训练 模型稳定性 自然语言处理
📋 核心要点
- 现有的基于策略蒸馏方法在引入上下文时，可能导致模型性能下降，尤其是在已经能够正确解决的问题上。
- 论文提出了一种轻量级的一致性正则化方法，通过停止梯度锚定无上下文输出，并使用前向KL散度惩罚上下文条件输出的偏差。
- 在12个不同配置下，实验结果显示该方法在大多数情况下提高了上下文条件的准确性，并在11个设置中减少了上下文引起的负面影响。
📝 摘要（中文）
近期研究表明，基于策略的蒸馏可以将特权上下文（如系统提示或任务提示）内化到学生模型中，从而在推理时不再需要上下文。尽管这种方法成功提高了学生模型在无上下文条件下的表现，但我们发现重新引入原始特权上下文反而会降低其性能，这一现象被称为上下文引起的退化。为了解决这一问题，论文提出了一种轻量级的一致性正则化方法，通过停止梯度锚定学生的无上下文输出，并对上下文条件输出的偏差进行惩罚。实验结果表明，该方法在大多数设置中提高了上下文条件的准确性，并有效减少了上下文引起的负面影响。
🔬 方法详解
问题定义：论文要解决的问题是上下文引起的性能退化现象，即在蒸馏学生模型时，重新引入特权上下文反而降低了模型的表现。现有方法未能考虑到这一现象的影响，导致模型在实际应用中可能出现不稳定性。
核心思路：论文的核心解决思路是引入一种轻量级的一致性正则化方法，通过锚定无上下文输出并惩罚上下文条件输出的偏差，从而实现上下文的可移除性，确保模型在有无上下文时的表现稳定。
技术框架：整体架构包括两个主要阶段：第一阶段是通过停止梯度锚定学生模型的无上下文输出，第二阶段是计算上下文条件输出与锚定输出之间的KL散度，并对其进行惩罚。
关键创新：最重要的技术创新点在于提出了上下文可移除性这一概念，并通过一致性正则化有效地解决了上下文引起的退化问题。这与现有方法的本质区别在于，现有方法通常只关注模型在有上下文条件下的表现，而忽略了在无上下文条件下的稳定性。
关键设计：关键设计包括使用停止梯度技术来锚定无上下文输出，以及通过前向KL散度计算上下文条件输出的偏差。该方法仅需在每个训练步骤中增加一次前向传递，且在大多数情况下能够有效提高无上下文性能。
🖼️ 关键图片
📊 实验亮点
实验结果显示，提出的方法在12个不同配置中，大多数情况下提高了上下文条件的准确性，并在11个设置中有效减少了上下文引起的负面影响。此外，该方法成功消除了响应长度膨胀现象，展示了其在实际应用中的有效性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、对话系统和智能助手等场景，能够提升模型在实际应用中的稳定性和可靠性。未来，该方法有望在多种任务中推广应用，进一步推动智能系统的性能提升。
📄 摘要（原文）
Recent work has shown that on-policy distillation can internalize privileged context, such as system prompts or task hints, into a student model so that the context is no longer needed at inference time. Although this approach successfully improves the student's no-context performance, we identify an interesting and previously unstudied phenomenon: in many settings, reintroducing the original privileged context to the distilled student actually degrades its performance, even on instances it already solves correctly without context. We term this context-induced degradation and argue that robust internalization demands not only matching the teacher's context-conditioned behavior, but also remaining stable when the context is reintroduced, a property we call context removability. Motivated by this observation, we propose a lightweight consistency regularizer that first anchors the student's no-context output via stop-gradient, then penalizes the context-conditioned output for deviating from it via forward KL divergence. This simple addition requires only one extra forward pass per training step, yet it effectively mitigates context-induced degradation and, in many cases, even improves no-context performance. Across 12 configurations spanning diverse domains and model families, our method improves context-conditioned accuracy in the majority of settings, reduces context-induced harm in 11 out of 12 settings, and effectively eliminates response-length inflation. A mechanistic case study further confirms that context removability is achieved at the representation level, with hidden states remaining nearly identical regardless of whether the context is present.