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

Not All Disagreement Is Learnable

When teaching a small AI model using a big one, not all mistakes are useful to learn from. Some mistakes are just "noise." These are instances where the large teacher model suggests something the small student model cannot even comprehend or reach. This paper introduces a way to pick only the "teachable" mistakes. This makes the distillation process significantly more efficient.

In the field of knowledge distillation, we often use On-Policy Distillation (OPD). This involves training a student model on its own generated text (rollouts) using the teacher's probability distributions as a guide. Previously, the state of the art involved "selective OPD." Engineers tried to save compute by only supervising certain tokens. They targeted tokens with high entropy (uncertainty) or high KL divergence (where the teacher and student disagree sharply). The open question was whether these "salient" tokens actually drive learning. This paper finds they aren't. Raw disagreement is a blunt instrument that often picks useless noise.

The Problem

The status quo assumes that significant teacher-student disagreement equals a high-value training signal. However, the authors argue that raw KL divergence is a coarse proxy. It conflates two fundamentally different scenarios.

As illustrated in [.A], there is "learnable disagreement." Here, the teacher's corrective signal lands on candidates the student already considers plausible.

There is also "incompatible disagreement." In this case, the teacher's preferred mass falls entirely outside the student's current predictive support (the set of tokens the student thinks are likely). Current selective methods often target the latter. This wastes the training budget on signals the student is structurally incapable of absorbing in a single update. Existing approaches like TIP miss the nuance of what is actually absorbable.

How It Works

The authors first develop a "fixed-context diagnostic" to decouple token-level learning from sampling noise. They freeze a bank of student-generated prefixes. They then measure how much the teacher-student KL reduction actually improves on those specific, static contexts.

To solve the selection problem, they propose Teachability-Aware OPD (TA-OPD). The mechanism follows these steps:

1. Identify Local Support: The system identifies the student's "local support." These are the top-$K$ tokens the student currently considers likely candidates [.A].
2. Decompose Disagreement: The algorithm calculates the teacher's probability mass on that student support ($C_t$). It then decomposes total disagreement into two components: learnable disagreement ($D_L$) and incompatible disagreement ($D_I$) .

3. Calculate Teachability: A teachability score is derived by combining these values ($s_t^{teach} = \tilde{D}_t \tilde{C}_t$). This prioritizes tokens where the disagreement is both large and compatible with the student's current state.
4. Apply Budgeted Mask: During training, the OPD loss is applied only to the top $n$ positions ranked by teachability. The value of $n$ is determined by a predefined token budget.

This method is lightweight. It relies on top-$K$ statistics already available during the forward pass. It requires no external reward models or expensive verifiers [.B].

Numbers

The headline result is a massive gain in sample efficiency. The authors report that TA-OPD can match or even surpass full-token OPD (supervising every single token). This occurs even while using only 5% of the token budget.

In a Qwen3-4B to Qwen3-1.7B setting, TA-OPD achieves an average benchmark score of 44.89 at a 10% budget. This outperforms the 42.37 achieved by full-token OPD. The efficiency is even more striking in the "gain per kept token" metric. [.B] shows that TA-OPD maintains a much higher efficiency curve than entropy-based or divergence-based baselines.

The robustness of this signal is also documented. Through regression analysis, the authors show that learnable disagreement ($D_L$) is a much stronger predictor of improvement. Specifically, $D_L$ has roughly twice the standardized coefficient of incompatible disagreement ($D_I$) [.A].

This confirms that the "teachability" metric is a statistically superior predictor of improvement than raw KL divergence.

What's Missing

There are a few gaps a practitioner should note. First, the evaluation is centered on the Qwen family and DeepSeek-R1-distilled models. It focuses specifically on math-heavy reasoning tasks. It is unclear if the "teachability" signal remains as clean in creative writing or open-ended dialogue.

Second, the paper notes that the token budget refers to the amount of supervision. It does not necessarily represent a proportional reduction in wall-clock training time. While you are computing fewer losses, you are still performing the full forward/backward passes. Unless you integrate this with strategies to prune actual transformer computations, you won't see a direct speedup in GPU utilization.

Finally, the "fixed-context" diagnostic is a controlled environment. In a real production run, the student's policy is constantly shifting. This means the "support" is a moving target. The interaction between shifting support and the teachability gate in a live environment remains partially unobserved.

Should You Prototype This

Yes, if you are running on-policy distillation and hitting a compute bottleneck. The method is mathematically elegant and computationally cheap. It only requires top-$K$ log-probabilities. The fact that it outperforms full-token distillation with only 5% of the tokens suggests it could drastically reduce costs.

Code is reportedly available at https://github.com/wyy-code/TA-OPD. If you are currently using entropy or simple KL-divergence masking, switching to a support-aware selector is a low-risk experiment for your next training run.

Figures from the paper