Technical Summary: Not All Disagreement Is Learnable: Token Teachability in On-Policy Distillation

Problem Statement

On-policy distillation (OPD) trains a student language model on its own rollouts using token-level supervision from a teacher policy. While recent selective OPD methods attempt to improve efficiency by prioritizing "salient" tokens—such as those with high entropy or high teacher-student disagreement (KL divergence)—these approaches rely on the assumption that high disagreement equates to high learning value.

The paper identifies a critical flaw in this assumption: raw KL disagreement is a coarse proxy for learning. It conflates two qualitatively different scenarios:

1. Learnable Disagreement: The teacher assigns corrective probability mass to tokens already within the student's current top-$K$ predictive support.
2. Incompatible Disagreement: The teacher assigns mass to tokens far outside the student's current support. While this creates a large KL gap, the student cannot easily absorb this correction, leading to weak or negative local learning effects.

The core problem is distinguishing which token-level teacher signals are actually learnable (i.e., capable of inducing local improvement) versus those that are merely salient or incompatible.

Methodology

1. Fixed-Context Diagnostic

To isolate the learning value of specific tokens from rollout noise and downstream context shifts, the authors introduce a fixed-context diagnostic.

• Protocol: Student-generated prefixes are collected and frozen as a context bank. Both the initial (pre-update) and trained student models are rescored against the same teacher distribution on these fixed contexts.
• Metric: The Fixed-Context Token Gain ($G^{\text{fix}}_{i,t}$) is defined as the reduction in KL divergence between the teacher and the student after training, measured on the same prefix.
• Goal: This allows the authors to correlate pre-update signals (like raw KL) with actual local improvement, controlling for sampling variance.

2. Formalizing Token Teachability

Using the diagnostic, the authors decompose disagreement into two components based on local support compatibility:

• Local Support ($S^S_t(K)$): The set of the student's top-$K$ tokens.
• Compatibility Mass ($C_t$): The total probability mass the teacher assigns to the student's local support $S^S_t(K)$.
• Decomposition:
  • Learnable Disagreement ($D^L_t$): High disagreement where the teacher's correction is aligned with the student's support (High $C_t$).
  • Incompatible Disagreement ($D^I_t$): High disagreement where the teacher's correction falls off the student's support (Low $C_t$).

The Teachability Score is defined as the product of normalized disagreement and compatibility:
$$s^{\text{teach}}_t = \tilde{D}_t \cdot \tilde{C}_t$$
where $\tilde{D}_t$ is the normalized disagreement and $\tilde{C}_t$ is the normalized compatibility mass.

3. Teachability-Aware OPD (TA-OPD)

Based on the finding that $D^L_t$ predicts fixed-context gain better than raw KL, the authors propose TA-OPD, a lightweight token-position selection method:

• Selection: For a given budget (e.g., 5% or 10% of tokens), TA-OPD selects the top-$N$ positions with the highest teachability scores ($s^{\text{teach}}_t$).
• Implementation: It requires only teacher and student token probabilities (specifically top-$K$ log-probs) available during standard OPD. It does not require reward models, verifiers, or additional labels.
• Contrast: Unlike previous methods that select based on entropy or raw divergence (TIP), TA-OPD explicitly filters for "support-aligned" corrections.

Key Contributions

1. Fixed-Context Diagnostic: A method to measure same-context KL reduction, linking pre-update token signals to local OPD improvement without confounding rollout noise.
2. Token Teachability Formalization: The demonstration that raw KL disagreement conflates learnable and incompatible signals. The paper proves that support-aligned disagreement is the reliable predictor of useful learning.
3. TA-OPD Algorithm: A practical, lightweight selection method that applies OPD loss only to high-teachability positions, improving efficiency without external tools.
4. Empirical Validation: Extensive experiments showing that selecting for teachability yields better downstream performance than selecting for salience, even with significantly reduced token budgets.

Experimental Results

The authors evaluated TA-OPD across four teacher-student settings involving Qwen3 (4B, 8B, 14B) and Qwen2.5 (3B) models, using benchmarks like AIME, GPQA-Diamond, HumanEval, and MATH-500.

• Performance vs. Full OPD: In many settings, TA-OPD matched or surpassed full-token OPD while using only 5% of the tokens. For example, in the Qwen3-4B $\to$ 1.7B setting, TA-OPD achieved an average score of 44.89 compared to Full OPD's 42.37, despite using a 10% budget (and outperforming full OPD in specific metrics).
• Comparison to Baselines: TA-OPD consistently outperformed entropy-based, divergence-based, and TIP (Token Importance in OPD) baselines.
  • In the cross-backbone setting (DeepSeek-R1-Distill-Qwen-14B $\to$ Qwen2.5-3B), TA-OPD improved the average score to 30.62, significantly outperforming the Base model (29.98) and Full OPD (28.76).
• Budget Sensitivity: The results indicate that OPD performance is token-quality limited rather than purely budget-limited. Increasing the token budget beyond 10% often did not yield further gains and sometimes degraded performance, reinforcing that not all additional supervision is beneficial.
• Robustness: The positive "high-support gain gap" (the difference in performance between high-teachability and low-teachability tokens) remained consistent across different contexts (math, GSM8K-COT), teacher scales (8B, 14B), and support sizes ($K=8, 16, 32$).

Significance and Claims

The paper reframes selective on-policy distillation not as a problem of selecting "salient" or "uncertain" tokens, but as a problem of identifying locally learnable teacher signals.

• Core Claim: The gains from OPD depend on the compatibility between the teacher's correction and the student's current predictive support. High KL divergence alone is insufficient; the teacher must assign mass to tokens the student is already considering plausible.
• Practical Impact: TA-OPD offers a way to drastically reduce the computational cost of distillation (by filtering out incompatible tokens) while maintaining or improving model performance. It achieves this without the overhead of reward models or verifiers, making it applicable to standard training pipelines.
• Limitations: The authors note that their analysis focuses on math-heavy reasoning and Qwen-family models. While the diagnostic is model-agnostic, broader validation on dialogue, code, and non-Qwen backbones is needed to confirm generalizability. Additionally, the "token budget" refers to supervision density, not necessarily wall-clock speedup, as the method selects positions rather than pruning transformer computation.

In conclusion, the work suggests that effective distillation requires a compatibility-aware supervision strategy, where the student learns best from teacher signals that are both corrective and locally reachable.