Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes — part 1/4

Cleaned-text code points [0, 15897). Bibliography omitted here; source.md includes it.

# Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

###### Abstract

On-policy distillation (OPD) is increasingly used in LLM post-training because it can leverage a teacher model to provide dense supervision on student rollouts. The standard implementation, however, usually reduces distribution matching to a sampled-token log-ratio, which can make the learning signal fragile on long rollouts whose prefixes drift away from the teacher’s typical support. We revisit this formulation from both theoretical and implementation perspectives. Theoretically, token-level OPD is biased relative to sequence-level reverse-KL minimization, but admits a substantially tighter worst-case variance bound; a controlled synthetic study further shows that stronger future-reward coupling increases gradient variance and destabilizes training. Empirically, we identify three failure modes of sampled-token OPD: imbalanced token-level supervision, unreliable teacher guidance on student-generated prefixes, and tokenizer or special-token mismatch. These findings motivate teacher top- K K local support matching, a truncated reverse-KL objective that compares teacher and student distributions over a teacher-supported token set at each prefix, together with top- p p rollout sampling and special-token masking. Across single-task reasoning and multi-task benchmarks spanning agentic and reasoning settings, this objective improves optimization stability and yields a +19.8% performance gain over standard sampled-token OPD baselines, providing a practical recipe for more stable on-policy distillation.

## 1 Introduction

On-policy distillation (OPD) is becoming an increasingly common component of LLM post-training, especially for reasoning and agentic models. Recent public reports from Thinking Machines Lab ( Lu and Lab, 2025 ) , Qwen3 ( Yang et al., 2025 ) , MiMo-V2-Flash ( Xiao et al., 2026 ) , and GLM-5 ( Zeng et al., 2026 ) suggest a broader shift toward supervision on model-generated trajectories, or closely related on-policy variants, alongside off-policy distillation and reinforcement learning. By training on student rollouts while evaluating them with a stronger teacher, OPD combines on-policy data collection with dense token-level feedback at relatively low cost ( Agarwal et al., 2024 ; Gu et al., 2024 ) . This profile is attractive in practical post-training pipelines, where training efficiency matters and models often need to combine or recover capabilities across domains and training stages ( Xiao et al., 2026 ; Zeng et al., 2026 ; Wang et al., 2026a ; DeepSeek-AI, 2026 ) .

Most OPD implementations in current LLM pipelines use a token-level estimator, even though it is biased relative to sequence-level reverse-KL ( Lu and Lab, 2025 ) . A basic reason is that sequence-level objectives couple each token update to many future rewards, which can make optimization substantially noisier in long-horizon settings. We make this trade-off explicit: token-level OPD removes future-reward coupling and is therefore biased, yet it admits more favorable worst-case variance scaling. Our toy experiment shows the same pattern empirically, with stronger future coupling leading to higher gradient variance and less stable optimization. This suggests a practical design principle for long-horizon training: keep supervision token-level to control variance.

In current LLM pipelines, this token-level objective is usually instantiated by sampled-token comparison: at each training step, the update is driven by the teacher–student log-probability difference on the sampled token ( Lu and Lab, 2025 ; Xiao et al., 2026 ) . This implementation is simple and efficient, but its learning signal can become brittle once rollouts grow long. Gu et al. (2024) report degraded outputs such as repetition, which is consistent with our observations. Recent work also reports entropy collapse under sampled-token OPD in certain cases ( Ko et al., 2026 ; Jin et al., 2026b ) . More recently, full-vocabulary distillation has been reported to outperform sampled-token variants in some settings ( Zhao et al., 2026 ; DeepSeek-AI, 2026 ) , suggesting that the one-token formulation can leave useful teacher information unused.

Our empirical analysis shows that this brittleness stems from several recurring failure modes. In particular, we identify two objective-level issues: the one-token signal is often highly imbalanced, and teacher guidance can become unreliable on student-generated prefixes. We also observe an additional implementation-level issue from tokenizer or special-token mismatch, which can further distort one-token comparisons. Taken together, these results point to a practical question: how can we retain the variance advantage of token-level OPD while making its supervision signal less brittle in practice?

The analysis above suggests a targeted modification of the standard sampled-token objective. We replace one-token supervision with teacher top- K K local support matching, where teacher and student are compared over a teacher-supported token subset at each prefix rather than only on the sampled token. We implement this objective as truncated reverse-KL, together with top- p p rollout sampling and special-token masking. The resulting update remains local and inexpensive, while providing a less brittle training signal.

Overall, our main contributions are as follows: • We clarify the theoretical trade-off in OPD: token-level OPD is biased relative to sequence-level OPD, but has substantially better worst-case variance scaling with sequence length, making it attractive for long-horizon LLM post-training.

• We present an empirical analysis of why sampled-token OPD can be unstable in practice, highlighting two recurring objective-level issues—imbalanced one-token supervision and unreliable teacher guidance on student-generated prefixes—along with an additional implementation issue from tokenizer or special-token mismatch.

• We propose teacher top- K K local support matching as an analysis-driven revision of sampled-token OPD, implemented with truncated reverse-KL, top- p p rollouts, and special-token masking, and show that it yields more stable optimization and stronger empirical performance than sampled-token OPD in single-task math reasoning and multi-task agentic-plus-reasoning training.

## 2 Understanding Sampled-Token OPD: Trade-offs and Failure Modes

### 2.1 From reverse-KL to token-level OPD

On-policy distillation aims to transfer the capabilities of a stronger teacher model into a student model by minimizing the reverse-KL from the student to the teacher. For a prompt x x , the OPD objective is J OPD ( θ ) = 𝔼 x ∼ D [ D KL ( π θ ( ⋅ ∣ x ) ∥ q ( ⋅ ∣ x ) ) ] , J_{\mathrm{OPD}}(\theta)=\mathbb{E}_{x\sim D}\left[D_{\mathrm{KL}}\left(\pi_{\theta}(\cdot\mid x)\,\|\,q(\cdot\mid x)\right)\right], where π θ \pi_{\theta} and q q denote the student and teacher models, respectively. Its gradient can be written as ∇ θ J OPD ( θ ) = 𝔼 x , y ∼ π θ ( ⋅ ∣ x ) [ ( log π θ ( y ∣ x ) − log q ( y ∣ x ) ) ∇ θ log π θ ( y ∣ x ) ] . \nabla_{\theta}J_{\mathrm{OPD}}(\theta)=\mathbb{E}_{x,\,y\sim\pi_{\theta}(\cdot\mid x)}\left[\big(\log\pi_{\theta}(y\mid x)-\log q(y\mid x)\big)\,\nabla_{\theta}\log\pi_{\theta}(y\mid x)\right]. For each decoding step t t , let c t = ( x , y < t ) c_{t}=(x,y_{<t}) denote the prefix context, and define s t = ∇ θ ​ log ​ π θ ​ ( y t ∣ c t ) , r t = log ⁡ π θ ​ ( y t ∣ c t ) q ⁡ ( y t ∣ c t ) . s_{t}=\nabla_{\theta}\log\pi_{\theta}(y_{t}\mid c_{t}),\qquad r_{t}=\log\frac{\pi_{\theta}(y_{t}\mid c_{t})}{q(y_{t}\mid c_{t})}. Using the autoregressive factorization log ⁡ π θ ​ ( y ∣ x ) − log ⁡ q ⁡ ( y ∣ x ) = ∑ t ′ = 1 T r t ′ , ∇ θ ​ log ​ π θ ​ ( y ∣ x ) = ∑ t = 1 T s t , \log\pi_{\theta}(y\mid x)-\log q(y\mid x)=\sum_{t^{\prime}=1}^{T}r_{t^{\prime}},\qquad\nabla_{\theta}\log\pi_{\theta}(y\mid x)=\sum_{t=1}^{T}s_{t}, we obtain the sequence-level estimator g ^ seq = ∑ t = 1 T ( ∑ t ′ = 1 T r t ′ ) ​ s t . \hat{g}_{\mathrm{seq}}=\sum_{t=1}^{T}\left(\sum_{t^{\prime}=1}^{T}r_{t^{\prime}}\right)s_{t}. (1) For t ′ < t t^{\prime}<t , we have 𝔼 ⁡ [ r t ′ ​ s t ] = 0 \mathbb{E}[r_{t^{\prime}}s_{t}]=0 , because r t ′ r_{t^{\prime}} depends only on the prefix before step t t , while 𝔼 [ s t ∣ x , y < t ] = ∑ y t π θ ( y t ∣ c t ) ∇ θ log π θ ( y t ∣ c t ) = 0 . \mathbb{E}[s_{t}\mid x,y_{<t}]=\sum_{y_{t}}\pi_{\theta}(y_{t}\mid c_{t})\,\nabla_{\theta}\log\pi_{\theta}(y_{t}\mid c_{t})=0. The same gradient can therefore be expressed in causal reward-to-go form 𝔼 ⁡ [ g ^ seq ] = 𝔼 ⁡ [ ∑ t = 1 T ( ∑ t ′ = t T r t ′ ) ​ s t ] , \mathbb{E}[\hat{g}_{\mathrm{seq}}]=\mathbb{E}\left[\sum_{t=1}^{T}\left(\sum_{t^{\prime}=t}^{T}r_{t^{\prime}}\right)s_{t}\right], where each token update is coupled to all future rewards along the trajectory.

Another approximation in LLM training retains only the immediate term at each position g ^ tok = ∑ t = 1 T r t ​ s t . \hat{g}_{\mathrm{tok}}=\sum_{t=1}^{T}r_{t}s_{t}. (2) We refer to Eq. ( 2 ) as the token-level estimator. This approximation removes future-reward coupling, so the update for token y t y_{t} depends only on its immediate reward. Consequently, it is biased relative to the sequence-level reverse-KL estimator, but exhibits lower variance in long-horizon settings. Under bounded rewards and bounded gradients, the worst-case variance upper bound of token-level OPD scales as O ⁡ ( T 2 ) O(T^{2}) , whereas the sequence-level estimator scales as O ⁡ ( T 4 ) O(T^{4}) . We provide a detailed derivation in Appendix D . To interpolate between these extremes, we consider the discounted return-to-go estimator g ^ γ = ∑ t = 1 T ( ∑ t ′ = t T γ t ′ − t ​ r t ′ ) ​ s t . γ ∈ [ 0 , 1 ] \hat{g}_{\gamma}=\sum_{t=1}^{T}\left(\sum_{t^{\prime}=t}^{T}\gamma^{t^{\prime}-t}r_{t^{\prime}}\right)s_{t}.\qquad\gamma\in[0,1] (3) The case γ = 0 \gamma=0 recovers token-level OPD, while γ = 1 \gamma=1 recovers the causal sequence-level estimator. We further validate this trade-off in a two-task toy experiment (Figure 1 ): stronger future coupling leads to substantially higher gradient variance and less stable optimization. This motivates our focus on token-level supervision in the remainder of the paper, where the main question becomes how to improve its local training signal in practical LLM settings. Additional experimental details are provided in Appendix E .

### 2.2 Why sampled-token OPD is brittle in practice

Although token-level OPD is appealing from a bias–variance perspective, the standard sampled-token formulation can be brittle in practice. We isolate three failure modes: (1) a highly imbalanced token-level distillation signal, (2) unreliable teacher guidance on student-generated prefixes, and (3) distortions introduced by tokenizer or special-token mismatch. These observations come from sampled-token OPD experiments on math reasoning, using Qwen2.5-7B-Instruct ( Qwen et al., 2024 ) as the student and OpenThinker3-7B ( Guha et al., 2025 ) , an SFT model built on Qwen2.5-7B-Instruct, as the teacher.

##### A highly imbalanced sampled-token signal.

In sampled-token OPD, the update at step t t is driven by the log-ratio on a single sampled token: log ⁡ q ⁡ ( y t ∣ c t ) − log ⁡ π θ ​ ( y t ∣ c t ) . \log q(y_{t}\mid c_{t})-\log\pi_{\theta}(y_{t}\mid c_{t}). Negative rewards arise whenever the student assigns higher probability to a sampled token than the teacher. As shown in Figure 2 , most sampled tokens receive negative reward. This leaves optimization dominated by a small subset of locally positive tokens, so training becomes sensitive to high-frequency fillers and short continuations that can receive favorable local scores while contributing little to trajectory-level quality.

##### The teacher signal can become unreliable on student-generated prefixes.

Sampled-token OPD assumes that teacher probability on a student-generated token is a useful proxy for trajectory quality. This proxy becomes unreliable on prefixes that are common under the student but uncommon for the teacher. In such regions, tokens with high teacher probability can remain rewarded even after the trajectory has drifted into repetition, self-resetting reasoning, or other meaningless continuations (Figure 3 ; Appendix H ). This creates an objective mismatch between token-level teacher agreement and trajectory-level quality.

We hypothesize that two factors amplify this issue: sharp teacher distributions, where modest teacher–student mismatch can produce large log-ratio values, and growing teacher–student divergence along long rollouts. Consistent with this view, Figure 4 shows that the distribution of teacher–student gaps becomes wider later in the sequence.

##### Tokenizer and special-token mismatch.

Sampled-token OPD compares the exact token generated by the student using the teacher distribution. When the two models use different tokenizations, the same raw text can be segmented differently, so a student-generated token may not correspond to a natural token under the teacher ( Boizard et al., 2025 ; Patiño et al., 2025 ; Minixhofer et al., 2025 ) . For example, the student may generate ‘<think>’ as ‘<’, ‘think’, ‘>’ , while the teacher expects ‘<th’, ‘ink’, ‘>’ . Then token ‘<’ receives low probability from the teacher, even though both models produce the same semantic content. Similar mismatches arise for special tokens such as end-of-sequence markers. In this setting, a one-token comparison confuses semantic disagreement with tokenizer mismatch. Since supervision is applied on a single token, such mismatch can distort the reward signal.

These observations motivate moving beyond one-token supervision: instead of comparing only the sampled token, we compare teacher and student over a set of teacher-supported next-token continuations at each prefix, while retaining token-level updates for stability.

## 3 Method

Our method is designed as a direct response to the failure modes above: it retains token-level OPD, but replaces one-token supervision with a distribution-level comparison over a teacher-selected support set at each prefix. This yields a truncated reverse-KL objective that preserves the efficiency of local updates while reducing dependence on any single sampled token. Section 3.1 introduces the objective, and Section 3.2 describes the practical choices that ensure stable training.

### 3.1 Teacher top- K K local support matching

Instead of comparing teacher and student on a single sampled token, we compare their next-token distributions over a teacher-defined local support set. A natural starting point is the full-vocabulary reverse-KL at prefix c t c_{t} : ℒ full ​ ( c t ) = ∑ v ∈ 𝒱 π θ ​ ( v ∣ c t ) ​ log ⁡ π θ ​ ( v ∣ c t ) q ⁡ ( v ∣ c t ) . \mathcal{L}_{\mathrm{full}}(c_{t})=\sum_{v\in\mathcal{V}}\pi_{\theta}(v\mid c_{t})\log\frac{\pi_{\theta}(v\mid c_{t})}{q(v\mid c_{t})}. (4) Sampled-token OPD can be viewed as a Monte Carlo approximation to this quantity: ℒ sample ( c t , y t ) = log π θ ​ ( y t ∣ c t ) q ⁡ ( y t ∣ c t ) , y t ∼ π θ ( ⋅ ∣ c t ) . \mathcal{L}_{\mathrm{sample}}(c_{t},y_{t})=\log\frac{\pi_{\theta}(y_{t}\mid c_{t})}{q(y_{t}\mid c_{t})},\qquad y_{t}\sim\pi_{\theta}(\cdot\mid c_{t}). (5) This approximation is computationally attractive, but it concentrates the entire update on a single sampled token. We instead compare teacher and student over a teacher-supported candidate set at each prefix.

For each prompt x x , we sample a group of outputs { o i } i = 1 G \{o_{i}\}_{i=1}^{G} using the student inference policy. Let c i , t = ( x , y i , < t ) c_{i,t}=(x,y_{i,<t}) be the prefix at position t t of output o i o_{i} , and define the teacher support set S ⁡ ( c i , t ) = TopK q ​ ( c i , t ) , S(c_{i,t})=\mathrm{TopK}_{q}(c_{i,t}), (6) which contains the K K highest-probability tokens under the teacher at that prefix.