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

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

Many widely used post-training methods for language models rely on off-policy supervision from fixed datasets, including supervised fine-tuning on demonstration data ( Ouyang et al., 2022 ; Guha et al., 2025 ) and conventional knowledge distillation on teacher-provided targets ( Hinton et al., 2015 ; Sanh et al., 2019 ) . In autoregressive generation, however, such fixed-prefix training creates a mismatch between the contexts seen during training and the prefixes induced by the student’s own test-time rollouts, motivating subsequent work on on-policy distillation (OPD) ( Agarwal et al., 2024 ) , where supervision is computed on student-generated trajectories. Agarwal et al. (2024) formulate generalized knowledge distillation for autoregressive models, while Gu et al. (2024) derive an effective on-policy optimization procedure. More recently, OPD and closely related variants have appeared in large-scale post-training pipelines, including Qwen3 ( Yang et al., 2025 ) , Thinking Machines Lab ( Lu and Lab, 2025 ) , MiMo-V2-Flash ( Xiao et al., 2026 ) , and GLM-5 ( Zeng et al., 2026 ) , underscoring the practical value of on-policy supervision for reasoning transfer, continual adaptation, and capability recovery across training stages. Subsequent work broadens this family in several directions: Ye et al. (2025) , Yang et al. (2026b) , and Ko et al. (2026) extend OPD to black-box teachers, more flexible reward formulations, and relaxed imitation objectives, respectively. Ye et al. (2026b) further connect OPD with context distillation, showing that on-policy supervision can also be used to internalize context- or prompt-conditioned behaviors. Related self- or hindsight-distillation approaches replace the external teacher with guidance derived from privileged information or feedback ( Zhao et al., 2026 ; Hübotter et al., 2026 ; Wang et al., 2026b ; Yang et al., 2026a ) . In contrast, our work does not introduce a new supervision regime; instead, we revisit the standard sampled-token estimator itself and improve its token-level learning signal through teacher top- K K local support matching.

##### KL computation for LLM post-training.

Sampled KL estimators in LLM post-training are commonly based on forms such as K 1 K_{1} , K 2 K_{2} , and K 3 K_{3} ( Schulman, 2020 ) . These estimators are lightweight, but they retain only the sampled token and therefore discard the richer next-token distribution available from teacher logits in white-box settings. When teacher logits are available, one can in principle match the full next-token distribution, but materializing full-vocabulary logits is often prohibitively memory-intensive for long responses and large models ( Hinton et al., 2015 ; Kim and Rush, 2016 ; Agarwal et al., 2024 ) . At one extreme, Zhao et al. (2026) use full-vocabulary logit distillation in a self-distillation setting. To reduce the cost of full-distribution matching, Zhang and Ba (2026) derive an unbiased top- K K KL estimator that combines exact computation on a top- K K set with a sampled correction for the tail. Related self-distillation methods such as SDPO and OpenClaw-RL also adopt top- K K approximations for efficiency, but without explicitly incorporating the sampled token into the truncated support ( Hübotter et al., 2026 ; Wang et al., 2026b ) . Our method instead computes a truncated reverse-KL objective on the teacher’s local top- K K support, with both teacher and student distributions renormalized within the selected set. This yields a simple distribution-level alternative to sampled-token OPD without incurring the cost of full-vocabulary matching.

##### KL divergence in RL for LLMs.

The choice of KL divergence can substantially affect distribution matching in on-policy distillation and related post-training methods for language models. Gu et al. (2024) argue that reverse KL is often better suited than the forward-KL objective commonly used in conventional distillation because it discourages the student from overestimating the teacher’s low-probability regions and tends to yield more precise generation. Agarwal et al. (2024) further show that different divergences induce a quality–diversity trade-off: moving from forward KL toward reverse KL through generalized Jensen–Shannon divergence (JSD) yields increasingly mode-seeking behavior and lower diversity. A related perspective is provided by Chen et al. (2025) , who connect the reverse-KL-like, mode-seeking character of on-policy RL to reduced forgetting during post-training. More recently, Jin et al. (2026a) argue that pure reverse-KL-style OPD can become brittle when the teacher distribution has high entropy, and propose augmenting reverse KL with forward KL on high-entropy tokens to better preserve diversity. At the same time, the preferred divergence can depend on the supervision setting. In self-distillation with privileged information, Zhao et al. (2026) report that forward KL outperforms reverse KL and JSD, whereas SDPO adopts JSD as a stability-oriented design choice ( Hübotter et al., 2026 ) . In this work, we adopt a reverse-KL formulation, but focus on a different question: how to obtain a stable local reverse-KL-style training signal under sampled-token OPD.

## Appendix D Bias and variance analysis of token-level versus sequence-level OPD

### D.1 Bias of the token-level estimator

Recall the sequence-level estimator in causal return-to-go form g ^ seq = ∑ t = 1 T ( ∑ t ′ = t T r t ′ ) ​ s t . \hat{g}_{\mathrm{seq}}=\sum_{t=1}^{T}\left(\sum_{t^{\prime}=t}^{T}r_{t^{\prime}}\right)s_{t}. Expanding the inner sum gives g ^ seq = ∑ t = 1 T r t ​ s t + ∑ t = 1 T ∑ t ′ = t + 1 T r t ′ ​ s t . \hat{g}_{\mathrm{seq}}=\sum_{t=1}^{T}r_{t}s_{t}+\sum_{t=1}^{T}\sum_{t^{\prime}=t+1}^{T}r_{t^{\prime}}s_{t}. Since the token-level estimator keeps only the first term, g ^ tok = ∑ t = 1 T r t ​ s t , \hat{g}_{\mathrm{tok}}=\sum_{t=1}^{T}r_{t}s_{t}, their expectation gap is 𝔼 ⁡ [ g ^ seq ] − 𝔼 ⁡ [ g ^ tok ] = 𝔼 ⁡ [ ∑ t = 1 T ∑ t ′ = t + 1 T r t ′ ​ s t ] . \mathbb{E}[\hat{g}_{\mathrm{seq}}]-\mathbb{E}[\hat{g}_{\mathrm{tok}}]=\mathbb{E}\left[\sum_{t=1}^{T}\sum_{t^{\prime}=t+1}^{T}r_{t^{\prime}}s_{t}\right]. This makes explicit that token-level OPD removes the future-reward coupling terms and is therefore generally biased with respect to the sequence-level objective.

### D.2 Worst-case variance upper bounds

Assume there exist constants B r , B s > 0 B_{r},B_{s}>0 such that | r t | ≤ B r , ∥ s t ∥ ≤ B s for all t . |r_{t}|\leq B_{r},\qquad\|s_{t}\|\leq B_{s}\quad\text{for all }t. For the token-level estimator, ‖ g ^ tok ‖ ≤ ∑ t = 1 T | r t | ​ ‖ s t ‖ ≤ T ​ B r ​ B s , \|\hat{g}_{\mathrm{tok}}\|\leq\sum_{t=1}^{T}|r_{t}|\,\|s_{t}\|\leq TB_{r}B_{s}, which implies 𝔼 ​ ‖ g ^ tok ‖ 2 ≤ T 2 ​ B r 2 ​ B s 2 . \mathbb{E}\|\hat{g}_{\mathrm{tok}}\|^{2}\leq T^{2}B_{r}^{2}B_{s}^{2}. Using Var ⁡ ( X ) ≤ 𝔼 ​ ‖ X ‖ 2 \mathrm{Var}(X)\leq\mathbb{E}\|X\|^{2} , we obtain Var ⁡ ( g ^ tok ) = O ⁡ ( T 2 ) . \mathrm{Var}(\hat{g}_{\mathrm{tok}})=O(T^{2}).

For the sequence-level estimator, define R = ∑ t = 1 T r t , S = ∑ t = 1 T s t , g ^ seq = R ​ S . R=\sum_{t=1}^{T}r_{t},\qquad S=\sum_{t=1}^{T}s_{t},\qquad\hat{g}_{\mathrm{seq}}=RS. Then | R | ≤ T ​ B r , ‖ S ‖ ≤ T ​ B s , |R|\leq TB_{r},\qquad\|S\|\leq TB_{s}, so ‖ g ^ seq ‖ ≤ T 2 ​ B r ​ B s , 𝔼 ​ ‖ g ^ seq ‖ 2 ≤ T 4 ​ B r 2 ​ B s 2 . \|\hat{g}_{\mathrm{seq}}\|\leq T^{2}B_{r}B_{s},\qquad\mathbb{E}\|\hat{g}_{\mathrm{seq}}\|^{2}\leq T^{4}B_{r}^{2}B_{s}^{2}. Therefore, Var ⁡ ( g ^ seq ) = O ⁡ ( T 4 ) . \mathrm{Var}(\hat{g}_{\mathrm{seq}})=O(T^{4}).

### D.3 Discussion

The sequence-level estimator is closer to the exact trajectory-level objective, but it couples each score term with many future rewards. In worst-case scaling, this changes variance growth from quadratic to quartic in sequence length. The argument is conservative, but it captures why stronger reward coupling can become problematic in long-horizon training.

## Appendix E Toy experiment details

### E.1 Environment

We use a two-task one-dimensional continuous-control environment to visualize how stronger return coupling changes OPD optimization. The student policy is a three-layer MLP with roughly 4K parameters. Its input is a three-dimensional vector containing task identity, current position, and normalized time step. The policy outputs the mean and standard deviation of a Gaussian action distribution, and the state transition is s t + 1 = s t + δ , δ ∼ 𝒩 ⁡ ( μ , σ ) . s_{t+1}=s_{t}+\delta,\qquad\delta\sim\mathcal{N}(\mu,\sigma). The two tasks are mirror images of each other: the left task starts from + 2 +2 and targets − 3 -3 , while the right task starts from − 2 -2 and targets + 3 +3 . We first train separate teachers with REINFORCE and then distill them into a shared student with alternating-task OPD.

### E.2 Gradient variance estimation

At each training step, we split a batch of B = 64 B=64 trajectories into M = 8 M=8 micro-batches. For each micro-batch m m , we compute a loss ℒ m \mathcal{L}_{m} and the corresponding gradient vector 𝐠 m \mathbf{g}_{m} on the output layer parameters. We then estimate gradient variance by Var ⁡ ( 𝐠 ) = 1 M ​ ∑ m = 1 M ‖ 𝐠 m − 𝐠 ¯ ‖ 2 , 𝐠 ¯ = 1 M ​ ∑ m = 1 M 𝐠 m . \mathrm{Var}(\mathbf{g})=\frac{1}{M}\sum_{m=1}^{M}\left\|\mathbf{g}_{m}-\bar{\mathbf{g}}\right\|^{2},\qquad\bar{\mathbf{g}}=\frac{1}{M}\sum_{m=1}^{M}\mathbf{g}_{m}. We use this quantity only as a qualitative proxy, but it is sufficient for comparing relative variance across different γ \gamma settings.

### E.3 Additional Results of Toy Experiments

Figures A1 , A2 , and A3 report gradient-variance curves and corresponding state-visitation heatmaps for different OPD estimators ( γ ∈ { 0.0 , 0.25 , 0.5 , 0.75 , 1.0 } \gamma\in\{0.0,0.25,0.5,0.75,1.0\} ) across three random seeds. Although the exact magnitudes vary by seed, the qualitative pattern is consistent. All settings exhibit large variance spikes during early optimization, and larger γ \gamma typically remains at a higher variance level later in training. In several runs, the variance under γ = 0.75 \gamma=0.75 or γ = 1.0 \gamma=1.0 stays one to several orders of magnitude above that of smaller γ \gamma values. Across runs, token-level OPD ( γ = 0 \gamma=0 ) consistently learns trajectories that move toward the target states for both tasks. Intermediate values of γ \gamma remain qualitatively similar but become more diffuse. When γ \gamma approaches the sequence-level case ( γ = 1.0 \gamma=1.0 ), the learned trajectories often deviate from the desired direction and stabilize around sub-optimal regions.

## Appendix F Experiment Setups

We build our training pipeline on verl 1 1 1 https://github.com/verl-project/verl ( Sheng et al., 2025 ) and the verl-agent framework 2 2 2 https://github.com/langfengQ/verl-agent ( Feng et al., 2025 ) , and conduct all experiments on a node with 8 NVIDIA H100 GPUs. We report the settings and hyperparameters used in our experiments in Table A1 .

## Appendix G Additional Results

### G.1 Additional Evidence on Single-task WebShop

To test whether the same trend transfers beyond the main math and ALFWorld settings, we additionally evaluate single-task WebShop ( Yao et al., 2022 ) using Qwen2.5-1.5B-Instruct as the student. As shown in Table A2 , local support matching improves over sampled-token OPD on both task score and success rate, increasing success rate from 50.0 to 57.8. Although the gap to the teacher remains substantial, the result suggests that the method transfers beyond the main math/agentic settings. We do not include a masking variant in this setting because the teacher is obtained by RL training from the same base model, so tokenizer or special-token mismatch is less likely to be a dominant factor here.

### G.2 Training Dynamics and Alignment

Figures A4 , A5 , and A6 provide a more detailed view of the optimization dynamics.

##### Better learning curves.

On math reasoning, our method improves both training reward and evaluation performance across most of training rather than only at the final checkpoint. This pattern holds in both the single-task setting and the alternating multi-task setting.

##### More stable optimization.

Our method yields smaller gradient norms and lower clipping-boundary fractions while maintaining sufficient policy entropy, indicating more stable optimization. We also observe that special-token masking substantially reduces the clipping-boundary fraction of sampled-token OPD during early and middle training, while having only minor effects on our method.

##### Improved teacher–student alignment.

The teacher–student log-probability gap on sampled tokens also moves closer to zero, suggesting that the truncated local support objective improves alignment even under the sampled-token diagnostic used by the baseline.

### G.3 Additional Results of Top- K K Variants

We primarily evaluate five top- K K variants in our experiments. Below, we present their formulations under unified notation and evaluate them in both single-task math and multi-task settings. All experimental setups are kept the same as in the main experiments.

#### G.3.1 Formulations of Five Top- K K Variants

##### Shared notation.

For each prompt x x , we sample a group of outputs { o i } i = 1 G \{o_{i}\}_{i=1}^{G} from the student inference policy π infer \pi_{\mathrm{infer}} . Let y i , t y_{i,t} denote the token at position t t in output o i o_{i} , and let c i , t = ( x , y i , < t ) c_{i,t}=(x,y_{i,<t}) be the corresponding prefix. We write 𝒯 K π ​ ( c i , t ) = TopK π ⁡ ( c i , t ) , 𝒯 K q ​ ( c i , t ) = TopK q ⁡ ( c i , t ) , \mathcal{T}_{K}^{\pi}(c_{i,t})=\operatorname{TopK}_{\pi}(c_{i,t}),\qquad\mathcal{T}_{K}^{q}(c_{i,t})=\operatorname{TopK}_{q}(c_{i,t}), for the student and teacher top- K K token sets at prefix c i , t c_{i,t} .

All variants below share the same rollout-level aggregation: ℒ ⁡ [ ℓ ] = 𝔼 x , { o i } ∼ π infer ​ [ 1 ∑ i = 1 G | o i | ​ ∑ i = 1 G ∑ t = 1 | o i | ℓ ⁡ ( c i , t , y i , t ) ] . \mathcal{L}[\ell]=\mathbb{E}_{x,\{o_{i}\}\sim\pi_{\mathrm{infer}}}\left[\frac{1}{\sum_{i=1}^{G}|o_{i}|}\sum_{i=1}^{G}\sum_{t=1}^{|o_{i}|}\ell(c_{i,t},y_{i,t})\right]. They differ only in how the local loss ℓ ⁡ ( c i , t , y i , t ) \ell(c_{i,t},y_{i,t}) is defined.

##### Renormalized local reverse-KL.

For any support set 𝒮 ⁡ ( c i , t ) ⊆ 𝒱 \mathcal{S}(c_{i,t})\subseteq\mathcal{V} , define the renormalized student and teacher distributions on 𝒮 ⁡ ( c i , t ) \mathcal{S}(c_{i,t}) by π ¯ θ ​ ( v ∣ c i , t ; 𝒮 ) = π θ ​ ( v ∣ c i , t ) ∑ u ∈ 𝒮 ⁡ ( c i , t ) π θ ​ ( u ∣ c i , t ) , q ¯ ​ ( v ∣ c i , t ; 𝒮 ) = q ⁡ ( v ∣ c i , t ) ∑ u ∈ 𝒮 ⁡ ( c i , t ) q ⁡ ( u ∣ c i , t ) , v ∈ 𝒮 ⁡ ( c i , t ) . \bar{\pi}_{\theta}(v\mid c_{i,t};\mathcal{S})=\frac{\pi_{\theta}(v\mid c_{i,t})}{\sum_{u\in\mathcal{S}(c_{i,t})}\pi_{\theta}(u\mid c_{i,t})},\qquad\bar{q}(v\mid c_{i,t};\mathcal{S})=\frac{q(v\mid c_{i,t})}{\sum_{u\in\mathcal{S}(c_{i,t})}q(u\mid c_{i,t})},\quad v\in\mathcal{S}(c_{i,t}). The corresponding local reverse-KL is ℓ renorm ​ ( c i , t , 𝒮 ) = ∑ v ∈ 𝒮 ⁡ ( c i , t ) π ¯ θ ​ ( v ∣ c i , t ; 𝒮 ) ​ log ⁡ π ¯ θ ​ ( v ∣ c i , t ; 𝒮 ) q ¯ ​ ( v ∣ c i , t ; 𝒮 ) . \ell_{\mathrm{renorm}}(c_{i,t};\mathcal{S})=\sum_{v\in\mathcal{S}(c_{i,t})}\bar{\pi}_{\theta}(v\mid c_{i,t};\mathcal{S})\log\frac{\bar{\pi}_{\theta}(v\mid c_{i,t};\mathcal{S})}{\bar{q}(v\mid c_{i,t};\mathcal{S})}.

##### Variant 1: teacher top- K K renormalization (our method).

Our default local support matching (LSM) method uses the teacher top- K K support 𝒮 LSM ​ ( c i , t ) = 𝒯 K q ​ ( c i , t ) , \mathcal{S}_{\mathrm{LSM}}(c_{i,t})=\mathcal{T}_{K}^{q}(c_{i,t}), and defines ℓ LSM ​ ( c i , t , y i , t ) = ℓ renorm ​ ( c i , t , 𝒮 LSM ) . \ell_{\mathrm{LSM}}(c_{i,t},y_{i,t})=\ell_{\mathrm{renorm}}(c_{i,t};\mathcal{S}_{\mathrm{LSM}}). The overall objective is ℒ LSM = ℒ ⁡ [ ℓ LSM ] . \mathcal{L}_{\mathrm{LSM}}=\mathcal{L}[\ell_{\mathrm{LSM}}].