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

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

##### Variant 2: student top- K K + sampled token + renormalization.

For this variant, the support is obtained by augmenting the student top- K K set with the sampled token when needed: 𝒮 π + y ​ ( c i , t ) = { 𝒯 K π ​ ( c i , t ) , y i , t ∈ 𝒯 K π ​ ( c i , t ) , 𝒯 K π ​ ( c i , t ) ∪ { y i , t } , y i , t ∉ 𝒯 K π ​ ( c i , t ) . \mathcal{S}_{\pi+y}(c_{i,t})=\begin{cases}\mathcal{T}_{K}^{\pi}(c_{i,t}),&y_{i,t}\in\mathcal{T}_{K}^{\pi}(c_{i,t}),\\[4.0pt] \mathcal{T}_{K}^{\pi}(c_{i,t})\cup\{y_{i,t}\},&y_{i,t}\notin\mathcal{T}_{K}^{\pi}(c_{i,t}).\end{cases} We then define ℓ π + y ​ ( c i , t , y i , t ) = ℓ renorm ​ ( c i , t , 𝒮 π + y ) , \ell_{\pi+y}(c_{i,t},y_{i,t})=\ell_{\mathrm{renorm}}(c_{i,t};\mathcal{S}_{\pi+y}), and ℒ π + y = ℒ ⁡ [ ℓ π + y ] . \mathcal{L}_{\pi+y}=\mathcal{L}[\ell_{\pi+y}].

##### Variant 3: teacher top- K K + sampled token + renormalization.

Similarly, this variant augments the teacher top- K K set with the sampled token: 𝒮 q + y ​ ( c i , t ) = { 𝒯 K q ​ ( c i , t ) , y i , t ∈ 𝒯 K q ​ ( c i , t ) , 𝒯 K q ​ ( c i , t ) ∪ { y i , t } , y i , t ∉ 𝒯 K q ​ ( c i , t ) . \mathcal{S}_{q+y}(c_{i,t})=\begin{cases}\mathcal{T}_{K}^{q}(c_{i,t}),&y_{i,t}\in\mathcal{T}_{K}^{q}(c_{i,t}),\\[4.0pt] \mathcal{T}_{K}^{q}(c_{i,t})\cup\{y_{i,t}\},&y_{i,t}\notin\mathcal{T}_{K}^{q}(c_{i,t}).\end{cases} Its local and rollout-level losses are ℓ q + y ​ ( c i , t , y i , t ) = ℓ renorm ​ ( c i , t , 𝒮 q + y ) , ℒ q + y = ℒ ⁡ [ ℓ q + y ] . \ell_{q+y}(c_{i,t},y_{i,t})=\ell_{\mathrm{renorm}}(c_{i,t};\mathcal{S}_{q+y}),\qquad\mathcal{L}_{q+y}=\mathcal{L}[\ell_{q+y}].

##### EMA-PG-style local loss.

The remaining two variants follow the EMA-PG decomposition ( Zhang and Ba, 2026 ) , where a top- K K truncated head term is combined with an importance-weighted sampled-token tail correction to obtain an unbiased estimator of the full-vocabulary reverse-KL. For any support set 𝒮 ⁡ ( c i , t ) ⊆ 𝒱 \mathcal{S}(c_{i,t})\subseteq\mathcal{V} , define the truncated head term ℓ head ​ ( c i , t , 𝒮 ) = ∑ v ∈ 𝒮 ⁡ ( c i , t ) π θ ​ ( v ∣ c i , t ) ​ log ⁡ π θ ​ ( v ∣ c i , t ) q ⁡ ( v ∣ c i , t ) . \ell_{\mathrm{head}}(c_{i,t};\mathcal{S})=\sum_{v\in\mathcal{S}(c_{i,t})}\pi_{\theta}(v\mid c_{i,t})\log\frac{\pi_{\theta}(v\mid c_{i,t})}{q(v\mid c_{i,t})}. Let sg ⁡ ( ⋅ ) \operatorname{sg}(\cdot) denote stop-gradient, and define r i , t = π θ ​ ( y i , t ∣ c i , t ) sg ⁡ ( π θ ​ ( y i , t ∣ c i , t ) ) , s i , t = sg ⁡ ( π θ ​ ( y i , t ∣ c i , t ) π infer ​ ( y i , t ∣ c i , t ) ) , r_{i,t}=\frac{\pi_{\theta}(y_{i,t}\mid c_{i,t})}{\operatorname{sg}\!\left(\pi_{\theta}(y_{i,t}\mid c_{i,t})\right)},\qquad s_{i,t}=\operatorname{sg}\!\left(\frac{\pi_{\theta}(y_{i,t}\mid c_{i,t})}{\pi_{\mathrm{infer}}(y_{i,t}\mid c_{i,t})}\right), where s i , t s_{i,t} is the sampling-policy correction (optionally clipped, as in EMA-PG). The sampled-token tail correction is ℓ tail ( c i , t , y i , t ; 𝒮 ) = s i , t [ y i , t ∉ 𝒮 ( c i , t ) ] r i , t sg ( log π θ ​ ( y i , t ∣ c i , t ) q ⁡ ( y i , t ∣ c i , t ) ) . \ell_{\mathrm{tail}}(c_{i,t},y_{i,t};\mathcal{S})=s_{i,t}\,\mathbf{1}\!\left[y_{i,t}\notin\mathcal{S}(c_{i,t})\right]\,r_{i,t}\,\operatorname{sg}\!\left(\log\frac{\pi_{\theta}(y_{i,t}\mid c_{i,t})}{q(y_{i,t}\mid c_{i,t})}\right). The corresponding local EMA-PG loss is ℓ EMA ​ ( c i , t , y i , t , 𝒮 ) = ℓ head ​ ( c i , t , 𝒮 ) + ℓ tail ​ ( c i , t , y i , t , 𝒮 ) . \ell_{\mathrm{EMA}}(c_{i,t},y_{i,t};\mathcal{S})=\ell_{\mathrm{head}}(c_{i,t};\mathcal{S})+\ell_{\mathrm{tail}}(c_{i,t},y_{i,t};\mathcal{S}).

##### Variant 4: student top- K K + EMA-PG correction.

This variant uses the student top- K K support 𝒮 π ​ -EMA ​ ( c i , t ) = 𝒯 K π ​ ( c i , t ) , \mathcal{S}_{\pi\text{-EMA}}(c_{i,t})=\mathcal{T}_{K}^{\pi}(c_{i,t}), and defines ℓ π ​ -EMA ​ ( c i , t , y i , t ) = ℓ EMA ​ ( c i , t , y i , t , 𝒮 π ​ -EMA ) , ℒ π ​ -EMA = ℒ ⁡ [ ℓ π ​ -EMA ] . \ell_{\pi\text{-EMA}}(c_{i,t},y_{i,t})=\ell_{\mathrm{EMA}}(c_{i,t},y_{i,t};\mathcal{S}_{\pi\text{-EMA}}),\qquad\mathcal{L}_{\pi\text{-EMA}}=\mathcal{L}[\ell_{\pi\text{-EMA}}].

##### Variant 5: teacher top- K K + EMA-PG correction.

This variant instead uses the teacher top- K K support 𝒮 q ​ -EMA ​ ( c i , t ) = 𝒯 K q ​ ( c i , t ) , \mathcal{S}_{q\text{-EMA}}(c_{i,t})=\mathcal{T}_{K}^{q}(c_{i,t}), and defines ℓ q ​ -EMA ​ ( c i , t , y i , t ) = ℓ EMA ​ ( c i , t , y i , t , 𝒮 q ​ -EMA ) , ℒ q ​ -EMA = ℒ ⁡ [ ℓ q ​ -EMA ] . \ell_{q\text{-EMA}}(c_{i,t},y_{i,t})=\ell_{\mathrm{EMA}}(c_{i,t},y_{i,t};\mathcal{S}_{q\text{-EMA}}),\qquad\mathcal{L}_{q\text{-EMA}}=\mathcal{L}[\ell_{q\text{-EMA}}].

##### Summary.

The five formulations differ only in how the local support is constructed and whether the sampled token is incorporated through support augmentation or through an EMA-PG-style tail correction. Variants 1–3 compute a renormalized reverse-KL on the resulting local support, while Variants 4–5 replace support renormalization by a truncated head term plus a sampled-token correction ( Zhang and Ba, 2026 ) .

#### G.3.2 Results under Single-task and Multi-task Settings

As shown in Table A3 , teacher top- K K performs reasonably well in the single-task setting and remains strong in the multi-task setting. By contrast, the EMA-PG variants, despite their unbiasedness motivation, do not translate into better empirical performance here. We therefore treat the variant comparison as a diagnostic result rather than a definitive ranking of support-set objectives.

## Appendix H Qualitative OPD reward-hacking case study

To complement the representative failures in the main text, we summarize a longer trajectory from multi-task training under sampled-token OPD. Read chronologically, the case exhibits the same pattern in several forms: the model continues reasoning after a valid answer has already been reached, falls into repetition loops such as wait , drifts into malformed continuations, and still receives high local teacher probability on those tokens. In these visualizations, uncolored tokens indicate positions where the displayed teacher and student probabilities are approximately matched.

The failure first appears as over-continuation . Even after the answer is effectively available, the local signal continues to place substantial mass on generic reasoning fillers and connective tokens, encouraging the model to keep going instead of stopping cleanly. The same pattern later appears on prefixes such as confirm , where the local signal still favors additional verification rather than termination. Some of this behavior may also reflect the teacher’s own output habits. Figure A7 illustrates several representative cases.

The trajectory then develops into hesitation loops and low-information continuations . Repeated wait tokens, punctuation-heavy continuations, and other semantically weak fillers can remain locally rewardable even after the overall trajectory has become unproductive. This is consistent with the repetition-loop discussion in Section 2.2 . We provide two similar cases in Figure A8 .

Finally, once the student drifts further off-distribution, the local signal can remain misleadingly positive rather than self-correcting. In the case study, this appears as degenerate continuations and malformed non-English text, yet many tokens still receive high teacher probability. An example is shown in Figure A9 .

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .