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

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

Within this support, we renormalize both teacher and student distributions: π ^ θ ​ ( v ∣ c i , t ) = π θ ​ ( v ∣ c i , t ) ∑ u ∈ S ⁡ ( c i , t ) π θ ​ ( u ∣ c i , t ) , q ^ ​ ( v ∣ c i , t ) = q ⁡ ( v ∣ c i , t ) ∑ u ∈ S ⁡ ( c i , t ) q ⁡ ( u ∣ c i , t ) . \hat{\pi}_{\theta}(v\mid c_{i,t})=\frac{\pi_{\theta}(v\mid c_{i,t})}{\sum_{u\in S(c_{i,t})}\pi_{\theta}(u\mid c_{i,t})},\qquad\hat{q}(v\mid c_{i,t})=\frac{q(v\mid c_{i,t})}{\sum_{u\in S(c_{i,t})}q(u\mid c_{i,t})}. (7) Our local support matching (LSM) objective averages the truncated reverse-KL over all rollout positions: ℒ LSM = 𝔼 x , { o i } ∼ π θ , infer ​ [ 1 ∑ i = 1 G | o i | ​ ∑ i = 1 G ∑ t = 1 | o i | ∑ v ∈ S ⁡ ( c i , t ) π ^ θ ​ ( v ∣ c i , t ) ​ log ⁡ π ^ θ ​ ( v ∣ c i , t ) q ^ ​ ( v ∣ c i , t ) ] . \mathcal{L}_{\mathrm{LSM}}=\mathbb{E}_{x,\,\{o_{i}\}\sim\pi_{\theta,\mathrm{infer}}}\left[\frac{1}{\sum_{i=1}^{G}|o_{i}|}\sum_{i=1}^{G}\sum_{t=1}^{|o_{i}|}\sum_{v\in S(c_{i,t})}\hat{\pi}_{\theta}(v\mid c_{i,t})\log\frac{\hat{\pi}_{\theta}(v\mid c_{i,t})}{\hat{q}(v\mid c_{i,t})}\right]. (8) Relative to sampled-token OPD, this replaces a one-token point estimate with a distribution-level comparison over teacher-supported candidates at the same prefix. The update is therefore no longer determined by the sign and magnitude of a single sampled-token log-ratio, while remaining much cheaper than full-vocabulary KL.

### 3.2 Practical stabilization choices

##### Support-set renormalization.

Renormalization is necessary because the objective is evaluated on a truncated support rather than the full vocabulary. Specifically, we apply a separate softmax over the logits inside the support set, so gradients do not directly propagate to tokens outside this set. Without this step, optimization can become unstable because the teacher and student probability masses inside the support are not directly comparable.

##### Top- p p rollout sampling.

We generate rollouts with top- p p sampling. Unconstrained sampling occasionally produces very low-probability tokens, creating prefixes on which the teacher signal becomes less informative and optimization less stable. Top- p p sampling keeps trajectories closer to typical continuations and makes the teacher signal more reliable.

##### Special-token masking.

We mask problematic special tokens to reduce false negatives caused by incompatible tokenization conventions. This is an orthogonal practical fix: it materially helps sampled-token OPD in our experiments, while our local support objective is much less sensitive to it. In principle, one could also merge multi-token marker variants or average over equivalent tokenizations, but we do not pursue those tokenizer-specific remedies here because masking is the simplest model-agnostic correction.

## 4 Experiments

We evaluate local support matching in three settings: single-task math reasoning (Section 4.2 ), alternating multi-task training over math and agentic tasks (Section 4.3 ), and an additional single-task agentic setting on a smaller student model (Appendix G.1 ). We also present ablations in Section 4.4 and provide training-dynamics analysis in Appendix G.2 .

### 4.1 Setup

We implement local support matching on top of the verl-agent framework ( Feng et al., 2025 ) , using Qwen2.5-Instruct models as students. We consider two main settings. The first is single-task math reasoning, where OpenThinker3-7B ( Guha et al., 2025 ) serves as the teacher and training uses the English portion of DAPO-Math-17K ( Yu et al., 2025 ) with a maximum context length of 16K. The second is a multi-task setting that alternates batches between math reasoning and a multi-turn agentic task based on ALFWorld ( Shridhar et al., 2021 ) . In this setting, math uses OpenThinker3-7B ( Guha et al., 2025 ) as the teacher, while the agentic side uses the released GiGPO-Qwen2.5-7B-Instruct-ALFWorld checkpoint ( Feng et al., 2025 ) .

For math reasoning, we report pass@1 on five benchmarks: Math500 ( Hendrycks et al., 2021 ) , AIME24 ( Zhang and Math-AI, 2024 ) , AIME25 ( Zhang and Math-AI, 2025 ) , Minerva ( Lewkowycz et al., 2022 ) , and OlympiadBench ( He et al., 2024 ) . For ALFWorld ( Shridhar et al., 2021 ) , we report success rate by default. In a small number of cases, we additionally report avg@32 on the math benchmarks. More details on experimental setups can be found in Appendix F .

### 4.2 Single-task math reasoning

Table 1 shows that local support matching improves over sampled-token OPD in single-task math reasoning. Sampled-token OPD already improves the student from 28.2 to 36.4 average score, but still remains substantially below the teacher. Applying special-token masking to sampled-token OPD further improves the average to 40.7, indicating that tokenizer-related mismatch is a meaningful part of the failure.

Both variants of our method outperform sampled-token OPD and its masked variant in average score. This shows that the gain is not solely attributable to mismatch handling, and instead supports the role of a stronger distribution-level distillation signal. In addition, masking changes our method only modestly (41.7 vs. 41.5), consistent with the view that local support matching is less sensitive to tokenizer mismatch than one-token supervision. Additional evidence on WebShop ( Yao et al., 2022 ) is shown in Appendix G.1 .

### 4.3 Multi-task agentic-plus-reasoning training

Table 2 shows that the effect of local support matching differs across the two task families in alternating multi-task training. The unmasked version raises the average math score from 34.8 to 41.7 (+19.8%), while maintaining competitive ALFWorld performance. The masked version achieves the best ALFWorld score at 97.7, but gives up part of the math improvement. This pattern suggests that local support matching is especially helpful for the reasoning side of the mixture, where sampled-token signals are more exposed to prefix drift; masking, in contrast, mainly shifts the trade-off toward the agentic task in this run. Beyond evaluation performance, our objective also yields consistently better optimization dynamics. We defer the full learning curves and diagnostic plots to Appendix G.2 .

### 4.4 Ablations

Table 3 and Figure 6 suggest that the gains arise from several design choices rather than any single modification. Teacher top- K K comparison alone is not sufficient: the rollout policy must also remain in a stable region, and adding top- p p sampling turns an initially weaker top- K K variant into a stronger configuration. Under the same top- p p rollout conditions, teacher top- K K local support matching still improves AIME24 avg@32 from 21.6 to 23.6. Renormalization inside the truncated support is essential, as removing it leads to rapid collapse. Performance is not especially sensitive to the exact support size once K K is large enough, but training becomes unstable when the support is too small or rollouts are fully unconstrained.

##### Top- K K support variants.

Our main experiments define the truncated expectation on the teacher’s top- K K support. Alternative objectives can compute the expectation over teacher top- K K (used in the main results), student top- K K , or teacher top- K K augmented with the student-sampled token. We provide a preliminary comparison in Table 4 under both the single-task and multi-task settings.

In the single-task setting, the three variants achieve broadly comparable results. The teacher top- K K + sampled-token variant attains the highest average pass@1 score, while student top- K K performs competitively on several individual benchmarks, and the original teacher top- K K gives the best AIME24 avg@32 among the three. The multi-task results are less uniform. In that setting, the original teacher top- K K variant performs substantially better than the other two alternatives, while both student top- K K and teacher top- K K + sampled-token variants degrade considerably on the math benchmarks. We therefore treat these results as evidence that support construction matters, but do not over-interpret the ranking among variants.

Overall, these results suggest that the choice of where the KL expectation is computed can matter in nontrivial ways, especially in the multi-task setting. We treat this as a partial ablation and a preliminary exploration, and discuss possible causes, including support-set construction and remaining off-policy effects, in Appendix A .

## 5 Conclusion

This work revisits on-policy distillation (OPD) for LLM post-training from both theoretical and implementation perspectives. Our theoretical analysis shows why token-level OPD is an attractive approximation for long-horizon training: it is biased relative to sequence-level reverse-KL, but avoids future-reward coupling and has substantially better worst-case variance scaling. At the same time, our empirical study shows that the standard sampled-token implementation can provide brittle supervision because its signal is imbalanced, can remain misleading on student-drifted prefixes, and is sensitive to tokenizer or special-token mismatch. Teacher top- K K local support matching addresses these issues by preserving local token-level updates while replacing one-token supervision with a truncated distribution-level comparison. Across single-task math reasoning and alternating agentic-plus-reasoning training, this simple modification improves optimization stability and downstream performance over sampled-token OPD, while also clarifying where teacher matching remains an imperfect proxy for task success.

## Appendix A Discussion and Limitations

##### Sampled-token augmentation of the teacher top- K K support.

In our current implementation, the support is S ⁡ ( c t ) = TopK q ​ ( c t ) S(c_{t})=\mathrm{TopK}_{q}(c_{t}) : teacher and student distributions are renormalized on this teacher-selected set before computing the KL term. However, this is not the only possible design. One can instead use an augmented support S + ​ ( c t ) = S ⁡ ( c t ) ∪ { y t } S^{+}(c_{t})=S(c_{t})\cup\{y_{t}\} , or include the sampled token through an importance-weighted correction that more explicitly preserves unbiasedness with respect to the full-vocabulary reverse-KL ( Zhang and Ba, 2026 ) . We provide formal definitions of several support-augmentation variants, including alternatives based on teacher- or student-centered supports and EMA-PG-style corrections, and evaluate them in Appendix G.3 . Overall, our default renormalized formulation remains competitive across both single-task and multi-task settings, which is why we adopt it in the main experiments.

##### Support-set KL is a truncated objective.

Our formulation computes the KL term only over a restricted token subset, i.e., the teacher top- K K support, rather than over the full vocabulary. This means that the expected gradient is formed within that subset: tokens outside the support do not receive gradient contributions from this objective. Relative to full-vocabulary reverse-KL, this introduces bias; more generally, it changes where the expectation is taken and which parts of the vocabulary participate in the update. We describe this here as a property of the estimator, rather than as a settled benefit or drawback.

##### Training–inference mismatch.

Prefixes are generated under a rollout policy, e.g., top- p p sampling in the vLLM engine, while the training engine updates the model without correcting for this sampling process. This may introduce a training–inference mismatch ( Liu et al., 2025 ; Yao et al., 2025 ) , which we leave to ongoing work.

##### Teacher matching remains an imperfect proxy for task success.

Even when OPD is well defined as a teacher-matching objective, the resulting reward can still diverge from the underlying notion of successful behavior. Our reward-hacking cases make this gap concrete: locally teacher-preferred continuations can remain rewardable even when the overall trajectory is already unhelpful or harmful. A noticeable gap to the teacher also remains in our experiments, which suggests that better local supervision is only one part of the distillation problem, especially when teacher and student differ substantially. Closing that gap may require stronger rollout control, better handling of distribution shift, better use of teacher uncertainty, and combinations with outcome-verifiable rewards.

## Appendix B Future Directions

##### Top- p p truncation as an adaptive support.

Another route toward a more compute-efficient reverse-KL objective is top- p p truncation instead of top- K K truncation, where the KL divergence is computed on the subset of tokens whose cumulative probability mass reaches a prescribed top- p p threshold.

##### OPD versus RL in multi-task settings.

Our multi-task results motivate a more direct comparison between OPD and RL as transfer mechanisms. In RL, positive or negative transfer can be read directly from environment reward across tasks. In OPD, the optimization target remains teacher-derived, so transfer is filtered through what the teacher regards as locally preferable behavior. This distinction may help explain why our multi-task gains are strongest on the math side and why nearby support-set definitions become less uniform in that setting. A matched-task, matched-compute comparison between OPD and RL would help clarify when teacher-guided transfer tracks environment-level generalization and when the teacher–reward gap becomes the bottleneck.

##### Continual learning as a testbed.

Continual learning is another natural setting for OPD. A teacher-guided on-policy objective could act as a retention mechanism while the student adapts to new tasks, but that regime would also stress exactly the issues surfaced in this paper: distribution shift, teacher staleness, and the accumulation of approximation error over long adaptation horizons. Testing OPD there would therefore probe not only whether local support matching mitigates forgetting, but also whether teacher-based objectives remain useful once the student keeps moving away from the teacher’s original domain.

##### Relation to other stabilization directions.

This work is complementary to directions such as reward-hacking mitigation, EMA-anchor stabilization with top- K K KL ( Zhang and Ba, 2026 ) , perturbation-based off-policy correction ( Ye et al., 2026a ) , and logit-level fusion between teacher and student rollouts ( Zhang et al., 2026 ) . These methods address different parts of the same broader problem: how to keep teacher-derived signals useful once teacher and student policies begin to diverge ( Li et al., 2026 ) . We view our method as one component in that larger toolbox, rather than as a replacement for those stabilization strategies.

## Appendix C Related Work

##### On-policy distillation.