Technical Summary: Toward Robust Internalization in On-Policy Distillation

1. Problem Statement

Recent advances in on-policy distillation (OPD) have enabled student models to internalize privileged context (e.g., system prompts, task hints, or game-state descriptions) into their parameters, allowing them to perform well without the context at inference time. However, this paper identifies a previously unstudied phenomenon termed context-induced degradation.

Contrary to the intuition that internalized knowledge should make context redundant (an idempotent operation), the authors observe that reintroducing the privileged context to a distilled student often degrades its performance and inflates response lengths, even on instances the model solves correctly without context. This occurs because standard OPD optimizes exclusively for privileged fidelity (matching the teacher's context-conditioned output) but imposes no constraints on the student's behavior when the context is present. Consequently, the student fails to achieve context removability: the property where the model's output remains stable regardless of whether the privileged context is provided.

The authors categorize post-distillation context interactions into three regimes:

• Regime A: The student performs well without context, but reintroducing context causes significant accuracy drops and high "harm rates" (flipping correct predictions to incorrect). This is the most prevalent regime in their experiments.
• Regime B: Context still acts as a useful scaffold, indicating incomplete internalization.
• Regime C: The student has successfully internalized the context, and performance is stable with or without it.

2. Methodology: No-Context Anchoring (NCA)

To address context-induced degradation, the authors propose No-Context Anchoring (NCA), a lightweight consistency regularizer. The core idea is to enforce that the student's context-conditioned output ($q_c$) remains consistent with its no-context output ($q_x$), treating the latter as an anchor.

Objective Function

Standard OPD minimizes the KL divergence between the student's no-context output and the teacher's context-conditioned distribution:
$$\mathcal{L}_{OPD} = \mathbb{E}_{x \sim D} [\text{KL}(q_x \parallel f(y \mid x, c))]$$

NCA augments this with a consistency term:
$$\mathcal{L} = \mathcal{L}_{OPD} + \beta \cdot \mathcal{L}_{NCA}$$
where $\mathcal{L}_{NCA}$ is defined as:
$$\mathcal{L}_{NCA} \triangleq \text{KL}(\text{sg}[g_\theta(\cdot \mid x)] \parallel g_\theta(\cdot \mid x, c))$$

Key Design Choices

1. Stop-Gradient Anchoring: The no-context view $q_x$ is passed through a stop-gradient operator (sg). This ensures $q_x$ receives gradients only from the primary OPD loss, while $q_c$ receives gradients only from the NCA loss. This eliminates gradient competition between the two objectives.
2. Forward KL Divergence: The authors choose $\text{KL}(q_x \parallel q_c)$ (forward KL) rather than reverse KL. This choice is "mode-covering," preventing $q_c$ from collapsing modes of $q_x$. Furthermore, the expectation falls under $q_x$, allowing the reuse of on-policy samples generated for OPD without requiring a separate rollout from $q_c$.
3. Computational Efficiency: NCA requires only one additional forward pass per training step (to compute $q_c$). The logits for $q_x$ are detached from the computation graph, and the on-policy sequences are reused, making the overhead negligible.

3. Experimental Setup

The authors evaluated NCA across 12 configurations spanning:

• Domains: System-prompt QA (MedMCQA, Safety datasets) and Text Games (Sokoban, FrozenLake).
• Models: Various Llama-3 and Qwen models (ranging from 1.7B to 8B parameters), including both self-distillation and cross-model distillation scenarios.
• Metrics: Accuracy/Win Rate with and without context, Harm Rate (probability of context flipping a correct answer), Both-Correct/Won rates, and relative response length changes.

4. Key Results

Across the 12 settings, NCA demonstrated the following effects:

• Mitigation of Context-Induced Degradation: NCA reduced context-induced harm in 11 out of 12 settings. In Regime A settings (where degradation was severe), NCA significantly reduced the accuracy drop upon context reintroduction.
• Improved Stability: NCA consistently increased the "BothCorrect" or "BothWon" rates, indicating the model maintains correct behavior under both views.
• Elimination of Response-Length Inflation: Standard OPD often caused the model to generate significantly longer responses when context was reintroduced (up to +39.5% length increase). NCA effectively eliminated this inflation in most settings, often reducing the length change to near zero.
• Enhanced No-Context Performance: In 7 out of 12 settings, NCA improved the primary training objective (no-context accuracy) by 0.6 to 3.6 percentage points. The authors attribute this to the consistency constraint acting as a regularizer that prevents parameter updates from accommodating conflicting behaviors across views.
• Boundary Cases: In one Regime B setting (FrozenLake, cross-model), NCA slightly increased harm, suggesting that when context still carries unique, non-internalized information, forcing strict removability may trade off some context-conditioned performance for stability.

5. Mechanistic Analysis

The paper provides a mechanistic case study on the Sokoban task to explain how NCA achieves context removability:

• Representation Level: Hidden state analysis revealed that while standard OPD reduced the divergence between context and no-context views, a noticeable gap remained in middle and late layers. NCA, however, maintained near-perfect cosine similarity ($\ge 0.997$) across all layers, confirming that context removability is achieved at the representation level.
• Gradient Alignment: In standard OPD, the gradients for the distillation objective and the context-conditioned view were nearly orthogonal. NCA introduces a new gradient direction that actively balances the forces of leveraging context for distillation and remaining invariant to it.
• Parameter Updates: Both OPD and NCA concentrate parameter updates in shallow layers (particularly $k\_proj$ and $q\_proj$). However, the specific updates driven by NCA are distinct from those of OPD, suggesting that context removability is primarily achieved through modifications to shallow-layer knowledge representations that OPD alone does not produce.

6. Significance and Claims

The authors claim that their work shifts the evaluation paradigm for privileged-context distillation. They argue that robust internalization requires not just matching the teacher's behavior (privileged fidelity) but also ensuring context removability.

• Novelty: This is the first study to identify and address the phenomenon where reintroducing internalized context degrades model performance.
• Practicality: The proposed method (NCA) is simple, requiring only a single forward pass and a standard consistency regularizer, yet it effectively solves the degradation problem without sacrificing the primary distillation goals.
• Generalizability: The method improves performance across diverse model families, tasks, and distillation settings (self and cross-model).

The paper concludes by urging the community to evaluate distilled models not only by their no-context performance but also by their robustness to context reintroduction, ensuring that internalized knowledge does not become a liability when the context is accidentally or intentionally provided.