Source-linked AI summary
ReTaCo: Residual-Target Control for On-Policy Distillation
Zixiang Ni, Zhuo Hu, Renjie Cao, Weijie Ren, Binqin Shi, Weijia Zhang, Shuheng Cao, Zhicheng Shi, Zhenhao Zhang, Haomin Wen, Zhiyuan Hu
TL;DR
On-policy distillation uses teacher feedback on student-generated prefixes, but transmitting full-vocabulary probabilities is costly and top-k supervision leaves omitted probability mass unresolved. ReTaCo makes that residual mass explicit and controls how much is transferred to selected tokens. Its fixed-prefix theory and experiments support controlled selected-mass behavior and higher point estimates than EOPD on most mathematics and code benchmarks.
Problem
On-policy distillation uses teacher feedback on student-generated prefixes, but full-vocabulary probability transmission is costly and top-k supervision leaves the treatment of omitted mass open.
Source: §1 · pp. 1–2Method
Results
Across three teacher–student pairs, ReTaCo has higher point estimates on most mathematics and code benchmarks, with improvements across all three OOD benchmarks for Qwen3.5-2B.
Source: §6 · p. 10Takeaways & Limitations
At a fixed prefix, the population objective has a unique optimum whose selected mass increases monotonically with β, while β = 0 still gives selected tokens recovery gradients.
Source: §1 · p. 2, §6 · pp. 9–10The fixed-prefix population equilibrium theorem does not guarantee convergence with evolving prefixes, stale policies, or clipped updates.
Source: §5 · p. 9
Abstract
from arXiv · showhide
On-policy distillation (OPD) trains a student on its own generated prefixes with token-level teacher feedback, but transmitting or storing the teacher's full-vocabulary distribution at every token is costly. Entropy-aware OPD (EOPD) adds forward supervision to reverse KL to help the student recover plausible tokens it underestimates, using only the teacher's top-$k$ probabilities to limit cost. Because EOPD renormalizes these probabilities, its target assigns no mass to the omitted vocabulary. We prove that the resulting loss keeps pushing the student's top-$k$ mass toward one even after the student matches the teacher's relative probabilities within the top-$k$ set, so the teacher itself is not a stationary point whenever the omitted tokens have positive teacher probability. We propose ReTaCo (Residual-Target Control), which keeps the top-$k$ tokens individually and groups the remaining tokens into one residual symbol, and pairs this forward target with a single-sample estimator whose expectation equals the full-vocabulary reverse KL. With teacher top-$k$ mass $m$, the residual target is $(1-\beta)(1-m)$ for $\beta\in[0,1]$: $\beta=0$ preserves the teacher's mass, and larger $\beta$ moves more mass onto the top-$k$ tokens without changing their relative probabilities. At a fixed prefix, we prove that the population objective has a unique optimum whose top-$k$ mass lies between $m$ and $m+\beta(1-m)$ and increases monotonically with $\beta$; at $\beta=0$, underestimated top-$k$ tokens still receive non-vanishing recovery gradients. Numerical optimization confirms these predictions, and across three teacher-student pairs, ReTaCo outperforms EOPD on most mathematics and code benchmarks.
1 Introduction
On-policy distillation trains students on their own prefixes with teacher feedback, but full-vocabulary transmission is costly and top-k supervision leaves the treatment of omitted mass unresolved. The introduction shows how renormalized forward supervision pressures selected mass upward, then presents ReTaCo’s explicit residual target and its theoretical and empirical evaluation.
1 Introduction
- OPD uses student-generated prefixes and token-level teacher feedback, but transmitting full-vocabulary probabilities at every token is costly.Sending only teacher top-k probabilities reduces that cost but leaves open how the objective should treat omitted probability mass.Source: §1 · pp. 1–2
- Renormalized top-k forward KL targets selected mass one, so the teacher is not stationary when omitted tokens have positive probability.Even after conditional distributions match, the loss continues pushing selected mass upward; reverse KL moderates but does not remove this pressure.Source: §1 · p. 2
- ReTaCo groups omitted tokens into a residual symbol with target mass (1 −β)(1 −m), preserving selected-token relative probabilities while controlling transferred tail mass.At β = 0, the target preserves the teacher’s mass allocation; larger β transfers more tail mass to the selected set.Source: §1 · p. 2
- The introduction reports numerical verification, evaluation across three model pairs, and one-epoch ablations on Qwen3-1.7B.The passage describes evaluation scope, not benchmark outcomes.Source: §1 · p. 3
2 Preliminaries
The preliminaries frame OPD as token-level supervision at student-generated prefixes and separate probability-mass allocation from within-set distribution mismatch. They then describe EOPD’s entropy-gated top-k forward term and its full-vocabulary normalization context.
2.1 On-Policy Distillation and Selected Support
- OPD studies the per-prefix reverse-KL objective DKL(p∥q), with teacher and student next-token distributions q and p fixed at a student-generated prefix.This isolates token-level supervision from rollout sampling, stale policies, and clipping in the practical loop.Source: §2.1 · p. 3
- The analysis assumes full-support teacher and student distributions and 0 < m, P < 1, interpreting endpoint identities by continuity where defined.Here m and P are teacher and student mass on the selected top-k set.Source: §2.1 · p. 3
- Conditioning on selected and residual sets separates disagreement in mass allocation from disagreement within each set.This distinction separates recovering selected tokens from changing their total probability.Source: §2.1 · p. 3
2.2 Entropy-Aware On-Policy Distillation
- EOPD adds forward supervision because reverse-KL recovery signals vanish as a student token probability approaches zero, whereas full forward-KL gradients approach −q_j.EOPD applies its top-k forward term at high-teacher-entropy positions and combines it with a PPO-style reverse update.Source: §2.2 · p. 3
3 Method
ReTaCo makes omitted probability mass explicit while retaining selected-token recovery supervision, and combines this target with sampled reverse KL. Its fixed-prefix objective has a unique, controllable selected-mass optimum; at β = 0, underestimated selected tokens still receive recovery gradients.
3.1 The Renormalization Effect in EOPD
- EOPD’s renormalized forward loss targets unit selected mass, so even matched conditional shapes leave pressure to increase selected mass; the teacher is not stationary when residual mass is positive.The entropy gate determines where supervision applies, but does not remove the normalization effect at active prefixes.Source: §3.1 · p. 4
3.2 Residual Representation and Target Construction
- ReTaCo groups tokens outside the selected set into one loss-only residual symbol and assigns it probability (1 − β)(1 − m), preserving the teacher’s selected-token conditional shape.β = 0 preserves teacher mass; positive β transfers β(1 − m) to the selected set, with zero residual target at β = 1.Source: §3.2 · p. 5
3.3 Reverse-KL Estimation and the Combined Objective
- A single-sample reverse-KL estimator complements the aggregate residual target, and its expectation gives the full-vocabulary reverse KL in the unclipped, fixed-prefix objective.The estimator uses the teacher probability of the student-sampled token, which may be outside the selected set.Source: §3.3 · p. 5
3.4 Population Optimum and Token Recovery
- For α > 0, the fixed-prefix population objective has a unique full-support minimizer with teacher conditional shapes and selected mass between m and rβ; at β = 0, P* = m.As a selected token’s student probability approaches zero at β = 0, its forward-term gradient approaches −q_j.Source: §3.4 · p. 6
3.5 Implementation
- At each unmasked response position, the implementation transmits O(k) teacher values and uses O(k) additional loss computation, though the teacher still computes full-vocabulary normalization and top-k selection.The teacher supplies top-k IDs and raw log-probabilities, while the student rollout supplies the sampled token for the reverse term.Source: §3.5 · p. 6
4 Experiments
Across three teacher–student pairs, ReTaCo generally leads on mathematics and code benchmarks, while out-of-domain results vary. Qwen3-1.7B ablations show task-dependent effects from residual trimming and mass supervision, alongside numerical validation of the predicted population optimum.
4.1 Experimental Setup
- ReTaCo and Sampled OPD share a direct reverse estimator, whereas EOPD uses entropy-gated forward supervision and PPO-style reverse updates, making this a comparison of complete procedures.The evaluation uses mathematics, code, and out-of-domain benchmarks; OOD scores use mathematics-distilled checkpoints.Source: §4.1 · p. 7
- ReTaCo leads both baselines on all five mathematics tasks for Qwen3.5-2B and Gemma4-E2B, four for Qwen3-1.7B, and on HumanEval+ and MBPP+ across all three pairs.Sampled OPD remains slightly higher on Qwen3-1.7B AIME25.Source: §4.1 · p. 7
- Out-of-domain results are less uniform: ReTaCo leads all three tasks for Qwen3.5-2B, but leadership is split across tasks for Qwen3-1.7B and Gemma4-E2B.For Qwen3-1.7B, ReTaCo leads MMLU-Pro but trails Sampled OPD on ARC-C and GPQA-Diamond; for Gemma4-E2B, EOPD leads ARC-C and MMLU-Pro.Source: §4.1 · pp. 7–8
4.2 Residual Targets and Component Contributions
- On Qwen3-1.7B, the default forward term improves four of five mathematics scores over reverse-only training, but residual trimming’s effects vary by task.Preserving teacher mass outperforms partial or complete trimming on MATH500, OlympiadBench, and AMC; trimming lowers AIME24 by 4.58 and 5.00 points at β = 0.5 and β = 1, respectively.Source: §4.2 · p. 8, Table 2 · p. 8, Figure 3 · p. 9
- Conditional-only supervision has the highest scores on MATH500, OlympiadBench, AMC, and AIME24, while mass-only supervision leads on AIME25.Removing the mass term lowers MATH500 from 83.78% to 83.55% and AIME24 from 28.75% to 27.92%, but raises AMC from 52.56% to 54.37%.Source: §4.2 · p. 8, Table 2 · p. 8
- For a uniform 80-token teacher with k = 16, the renormalized target gives P* = 0.5821 versus m = 0.2; joint optimization’s maximum conditional-distribution ℓ1 error is 4.13 × 10^-8.The tested equilibria lie between m and rβ, with P* = m at β = 0.Source: §4.2 · p. 8, Figure 3 · p. 9
5 Discussion and Limitations
ReTaCo’s theoretical guarantees apply to a fixed prefix under the population objective, not to training with evolving prefixes, stale policies, or clipped updates. The discussion also distinguishes task-dependent mathematics ablations from a response-length experiment on a different objective variant.
5 Discussion and Limitations
- The fixed-prefix population theorem does not guarantee convergence with evolving prefixes, stale policies, or clipped updates.At β = 0, the forward objective also cannot resolve the conditional shape of the tail.Source: §5 · p. 9
- Mathematics ablations show task-dependent effects of residual trimming and stronger forward supervision.The response-length experiment uses an entropy-gated forward term and PPO-style reverse objective, not default ReTaCo.Source: §5 · p. 9
6 Conclusion
ReTaCo controls residual probability explicitly, pairs its forward target with an unbiased single-sample reverse-KL estimator, and has a unique fixed-prefix optimum whose selected mass increases with β. The conclusion reports improved point estimates on most mathematics and code benchmarks, while the appendix notes how residual aggregation changes the KL decomposition.
6 Conclusion
- ReTaCo groups omitted tokens into a residual symbol, and its single-sample reverse estimator equals full-vocabulary reverse KL in expectation.The unique fixed-prefix optimum has selected mass increasing with β; at β = 0, selected tokens still receive recovery gradients.Source: §6 · pp. 9–10
- Across three teacher–student pairs, ReTaCo has higher point estimates on most mathematics and code benchmarks, with improvements across all three OOD benchmarks for Qwen3.5-2B.One-epoch component ablations indicate conditional recovery explains most accuracy gains, while the mass term does not improve accuracy uniformly.Source: §6 · p. 10
AI use statement
- The authors report that AI tools assisted with language polishing, editing, and figure preparation, while they retained responsibility for technical content and the final manuscript.This statement also assigns responsibility for experimental results and references to the authors.Source: AI use statement · p. 11
A.1 KL chain rules
- Residual aggregation removes the conditional tail term from the KL decomposition while preserving the tail’s aggregate mass.The appendix states that this follows by adding the selected- and residual-set contributions.Source: §A.1 · p. 13
A.2 The sampled reverse-KL estimator
The appendix establishes that the ideal sampled reverse-KL estimator has the full-vocabulary reverse KL as its expectation. It also shows why renormalized top-k forward KL favors unit selected mass and characterizes the resulting equilibrium when combined with reverse KL.
A.2 The sampled reverse-KL estimator
- The ideal sampled reverse-KL estimator has the full-vocabulary reverse KL as its expected value.The proof uses a detached sampled token and applies to the ideal unclipped per-prefix estimator without differentiating through prefix sampling.Source: §A.2 · p. 13
A.3 The implicit target of renormalized forward KL
- After the selected conditional distributions match, the renormalized forward term still favors unit selected mass, so positive teacher residual mass makes teacher matching non-stationary.At teacher matching, the conditional KL vanishes while −log P decreases strictly until selected mass reaches one.Source: §A.3 · p. 13
A.4 Equilibrium under the implicit unit-mass target
- When combined with reverse KL, the implicit unit-mass target has a unique equilibrium with selected mass greater than the teacher’s mass.The student’s residual becomes smaller than the teacher’s; as teacher selected mass approaches one, the residual ratio obeys ϵ = e^−αδ + o(δ).Source: §A.4 · p. 14
B.1 Unique population optimum and monotonic control
The ReTaCo population objective has a unique optimum, with selected mass equal to the teacher’s at β = 0 and strictly between the teacher mass and the forward target for β > 0. The forward-gradient intervention scales with β(1−m), while selected tokens can still receive recovery gradients at β = 0.
B.1 Unique population optimum and monotonic control
- The fixed-prefix population objective has a unique interior optimum whose selected mass equals m at β = 0 and lies strictly between m and rβ when β > 0.The selected mass increases monotonically with β, and the optimum remains interior at the parameter endpoints.Source: §B.1 · pp. 14–15
B.2 Forward gradient and intervention size
- At β = 0, selected tokens receive non-vanishing recovery gradients as their student probabilities approach zero.At teacher matching, both selected and residual gradient branches are zero.Source: §B.2 · p. 15
- The intervention size is δ = β(1−m): each selected target probability changes by δqS, and the forward-logit-gradient change has ℓ1 size 2δ.The absolute changes in the selected and residual branches each sum to δ.Source: §B.2 · p. 15
C Reference Implementation
The per-token implementation builds a residual-aware forward target from teacher-selected tokens and combines its divergence with a sampled reverse term. Numerical stabilization preserves normalization when rounding makes selected mass appear to be one.
C.1 Per-token computation
- The implementation gathers teacher-selected student probabilities, recovers selected and residual log-probabilities, then forms the residual and selected target log-probabilities.The residual probabilities use a stable log(1 − exp(·)) routine.Source: §C.1 · p. 16
- The computation sums selected-token and residual contributions to the forward divergence, adds the sampled reverse term, and averages over unmasked response tokens.The rollout token, selected IDs, teacher top-k log-probabilities, and set construction are detached, while the full-vocabulary student softmax lets aggregate tail gradients reach individual residual logits.Source: §C.1 · p. 16
C.2 Numerical stability
- When rounding makes selected mass indistinguishable from one, selected probabilities are rescaled to 1 − ε and ε is assigned to the residual.The implementation uses ε = 10^-7 in mixed-precision training and ε = 0 in float64 checks; this preserves conditional shape and normalization while perturbing aggregate mass.Source: §C.2 · p. 16
D.1 Mass and Conditional-Shape Recovery
Synthetic optimization tests the predicted selected mass and conditional shape, while gradient checks compare automatic differentiation with closed-form gradients. These diagnostics use settings distinct from the main benchmark protocol.
D.1 Mass and Conditional-Shape Recovery
- The synthetic test jointly optimizes selected mass and conditional distributions for β ∈ {0, 0.5, 1}, comparing predicted and optimized student masses at teacher selected mass m = 0.6802847194.The comparison is reported in Table 4.Source: §D.1 · p. 16, Table 4 · p. 16
- The largest conditional-distribution ℓ1 error over the selected set is 4.13 × 10^-8 in the joint optimization test.The test includes the conditional tail term, so it evaluates both components of the population solution.Source: §D.1 · p. 16
D.2 Gradient Verification
- Automatic differentiation matches closed-form gradients in float64 for full-support and vanishing-residual distributions, reproducing the equilibrium curves in Figure 3.The checks use the configuration in a separate long-run diagnostic variant, not the main benchmark protocol.Source: §D.2 · p. 16, §E.1 · p. 16
E.2 Effect of Residual-Mass Trimming
Residual-mass trimming changes the target mass assigned outside the selected set, including the aggregate probability of an end-of-sequence token in that set. The supplied diagnostics also report response-length behavior and policy entropy under specific training and evaluation settings.
E.2 Effect of Residual-Mass Trimming
- Increasing β reduces residual target mass from 1 − m to (1 − β)(1 − m), and trimming may affect stopping when the end-of-sequence token belongs to the residual set.The end-of-sequence probability is included in the aggregate residual target.Source: §E.2 · p. 16
E.3 Long-Run Behavior
- The β = 0 run maintains a mean response length of approximately 478 tokens and a response truncation rate below 0.5%.The table reports training diagnostics, with truncation denoting responses that reach the generation limit.Source: §E.3 · p. 17, Table 6 · p. 17
F.1 Benchmarks and Evaluation
- Mathematics evaluation uses eight sampled responses per problem, while HumanEval+ and MBPP+ use one zero-shot greedy completion per task.The mathematics protocol uses temperature 0.7 and top-p 0.95; code completions are scored with the extended EvalPlus test suites.Source: §F.1 · p. 17, Table 7 · p. 18
- Over the final 20 steps, ReTaCo with β = 0 records 0.399 nats of actor entropy versus 0.305 nats for Sampled OPD.This Qwen3-1.7B run uses a Qwen3-30B-A3B-Instruct-2507 teacher and k = 16 over their shared 277-step interval.Source: Figure 4 · p. 18
F.2 Training-Dynamics Diagnostic
The training diagnostics report higher policy entropy for ReTaCo than Sampled OPD in one run, alongside declining conditional loss and mass error in a separate run. The supplied optimization, ablation, hyperparameter, and related-work passages qualify the scope and comparisons of these findings.
F.2 Training-Dynamics Diagnostic
- Over the final 20 steps, ReTaCo with β = 0 records 0.399 nats of actor entropy versus 0.305 nats for Sampled OPD.Both use a Qwen3-1.7B student, Qwen3-30B-A3B-Instruct-2507 teacher, and k = 16; the curves cover steps 1–277.Source: §F.2 · p. 18
- In a separate 300-batch run, the conditional forward component and absolute mass error both decline, but the diagnostic uses a 7,168-token response limit and differs from the benchmark protocol.The run uses k = 16, α = 1, β = 0; no run averaging or uncertainty intervals are shown, and numerical tail projection is active at 31.65% of response-token positions on average over the last 20 steps.Source: §F.2 · pp. 18–19, Figure 5 · p. 19
G.1 Component Contributions
- All three component-ablation variants reach the response limit on nearly half of their completions.The variants train for one epoch, and truncation means reaching the 8,192-token generation limit.Source: §G.1 · p. 20, Table 9 · p. 20
G.2 Hyperparameter Sensitivity
- The default α = 1 outperforms α = 0.5 on all five tasks, while α = 2 lowers MATH500 and AIME24 accuracy but raises AIME25 accuracy.For α = 2, MATH500 changes from 83.78% to 70.45%, AIME24 from 28.75% to 25.00%, and AIME25 from 21.67% to 22.50%; increasing k from 8 to 32 raises AIME24 from 23.33% to 27.50% but lowers AIME25 from 22.50% to 18.33%.Source: §G.2 · p. 20, Table 10 · p. 21
H Related Work
- ReTaCo pairs residual aggregation with a single-sample estimator whose expectation equals full-vocabulary reverse KL, and parameterizes the residual target to control mass preservation versus trimming.The related-work comparison distinguishes this combination from prior sparse-target and tail-aware distillation approaches.Source: H Related Work · p. 21