PaperlayerView exact revision on arXiv

Source-linked AI summary

ResOPD: Tail Residualization for Sparse On-Policy Distillation

Penghui Yang, Long Xing, Xuanlang Dai, Ziyu Liu, Kai Chen, Yuhang Zang

arXiv:2610.04882v1cs.LGcs.AI

TL;DR

Sparse on-policy distillation must work around costly dense teacher feedback, yet sparse Top-k objectives are biased and sampled-token estimators are noisy. ResOPD analytically aggregates the coarse tail and samples only the unresolved residual; across evaluated settings it reduces variance and improves downstream scores without increasing the payload.

  • Problem

    Dense teacher distributions are prohibitively costly, while sparse Top-k objectives omit tail information and introduce systematic gradient bias, and sampled-token estimators have severe variance.

    Source: §1 · pp. 1–2

  • Method

    ResOPD computes the exact gradient of a (k + 1)-way coarse distribution and samples only the within-tail residual, preserving unbiasedness with the same sparse payload.

    Source: §1 · pp. 2–3

  • Results

    Across evaluated settings, ResOPD reduces gradient variance and improves downstream scores; covariance trace falls 52–75% versus sampled-token OPD under Top-4 feedback.

    Source: §1 · p. 2, §5 · p. 15

  • Takeaways & Limitations

    ResOPD is a variance-reduction primitive for sparse OPD that adds no teacher queries or forward passes and does not expand the teacher payload.

    Source: §5 · p. 15

    Under the default recipe, ResOPD does not shift the ultimate divergence boundary at 1 × 10^-5, though it provides a wider stability margin in moderately aggressive regimes.

    Source: §4.7 · p. 14

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student model on self-generated trajectories, but transmitting dense teacher distributions across long reasoning traces creates prohibitive communication and memory bottlenecks. Practical systems therefore rely on sparse teacher interfaces, typically transmitting either the sampled-token score or a small Top-$k$ distribution. However, this sparse setting faces a fundamental dilemma: sampled-token estimators are unbiased but suffer from severe gradient variance, whereas directly optimizing Top-$k$ objectives introduces systematic bias. To improve this trade-off, we propose ResOPD (On-Policy Distillation with Tail Residualization), which provides unbiased full-vocabulary reverse KL gradient estimation under on-policy sampling, with substantial variance reduction under the same sparse payload in the evaluated settings. ResOPD aggregates the unobserved vocabulary into an observable coarse tail event, computes its exact aggregate gradient, and samples only the fine-grained within-tail residual, which requires no additional teacher queries or forward passes. Extensive experiments demonstrate that ResOPD substantially reduces gradient variance, stabilizes online training dynamics, and improves downstream performance across the evaluated settings. These results establish ResOPD as an efficient, plug-and-play variance reduction primitive for sparse on-policy distillation.

1. Introduction

Sparse OPD reduces teacher communication by exposing only Top-k information, but omitting the tail creates a choice between biased objectives and noisy unbiased estimators. ResOPD analytically handles the coarse tail event and samples only its fine-grained residual, reducing variance while preserving unbiasedness.

1. Introduction

  • Full-distribution OPD is prohibitively costly for modern LLM vocabularies, while sparse Top-k objectives carry systematic gradient bias because the tail is omitted.Payload grows with vocabulary size per token and rollout length per sequence.Source: §1 · p. 1
  • ResOPD scored best on average across four math benchmarks, 4.40 points above HETS at the same payload; its overall ranking varied across math, coding, and science.Its average advantage over HETS also held across tested support sizes k∈{1, 4, 16, 32}, with small margins.Source: Figure 1 · p. 2, §1 · p. 2
  • Unbiased alternatives also face high variance: sampled-token OPD is noisy, while HETS leaves the known aggregate tail mass unused and produces large spikes on tail hits.At 98.4% head coverage, HETS produces 460–470 spikes in a 30K-token reasoning trace.Source: §1 · pp. 1–2
  • ResOPD computes the exact gradient of a (k + 1)-way coarse distribution and adds a fine-tail residual correction only on tail hits, retaining an unbiased estimator with the same sparse payload.The estimator requires no additional teacher queries or forward passes.Source: §1 · pp. 2–3
  • Exact frozen-prefix audits found that ResOPD lowered gradient covariance trace by 52–75% versus sampled-token OPD under Top-4 feedback, below HETS and payload-matched vOPD.End-to-end distillation also showed lower online residual moments and improved downstream accuracy over evaluated sparse OPD baselines.Source: §1 · pp. 2–3

2. Related Work

Related work addresses OPD's dense-feedback cost with sparse teacher payloads and addresses gradient estimation with stochastic estimators and variance-reduction baselines. Sparse truncation and coarse-support methods remain biased, while existing unbiased estimators have variance or stability costs.

On-Policy Language-Model Distillation

  • OPD trains students on self-generated rollouts under dense teacher supervision, with subsequent work exploring skew divergences, adaptive targets, and reasoning diagnostics.Full-vocabulary OPD incurs prohibitive communication and memory footprints during distributed rollout generation.Source: §2 · p. 3

Sparse Divergences and Truncation Biases

  • Sparse Top-k payloads reduce communication, but renormalization and truncation optimize biased surrogates, while ghost-token methods discard fine-grained within-tail geometry.The passage cites sparse payloads [32], renormalization or truncation [32], and ghost-token aggregation [12].Source: §2 · p. 3

Stochastic Estimators and Variance Reduction

  • Sampled-token estimators recover the full-vocabulary rKL gradient unbiasedly but have extreme Monte Carlo variance; vOPD baselines and HETS address variance, yet HETS still has destabilizing tail-hit bursts.The passage cites sampled-token estimators, vOPD [23], and HETS in EMA-PG [30].Source: §2 · p. 3

3. Methodology

ResOPD uses the observable aggregate tail mass to construct an unbiased sparse estimator, replacing HETS’s zero tail control with a computable tail-bucket prediction. The methodology also separates between-event from within-tail variance and states the sampling conditions and limits of its conclusions.

3.1. Preliminaries: From sampled gradients to an exact sparse head

  • Sampled-token OPD is unbiased, while renormalized and truncated Top-k objectives are biased; HETS adds the entire tail correction only on tail draws.Large head coverage makes tail corrections rare per draw, even though long minibatches accumulate many corrections.Source: §3.1 · pp. 3–5

3.2. Tail residualization

  • The cited passage states that ĝ(B) is unbiased: its expectation over A drawn from p equals g, and identifies HETS as the special case B = 0.Source: §3.2 · p. 5
  • ResOPD has strictly lower variance than HETS when the tail-bucket vector predicts the true tail mean better than zero.The condition is ||μtail − μB||2 < ||μtail||2.Source: §3.2 · p. 5

3.3. Sampling assumptions and evidence boundaries

  • The unbiasedness identities assume on-policy sampling; rollout filtering preserves the variance decomposition but does not make filtered draws unbiased samples from pθ.For realized rollout distributions, ResOPD has lower variance than HETS when its tail control is closer to the realized tail mean than zero.Source: §3.3 · pp. 5–6
  • Online residual second moments M2 are mechanism-aligned diagnostics, not estimates of full gradient covariance, and lower covariance does not imply monotonic downstream accuracy gains.Online policies and visited prefixes change while sparse teacher feedback hides the fine tail.Source: §3.3 · p. 6

3.4.1. Separating between-event and within-tail controls

  • The Unified family separates controls: B changes between-event variance and γ changes within-tail variance; ResOPD uses (μB, 0) because the oracle controls are unavailable online.The oracle B requires the full Teacher tail, while per-prefix γ requires conditional covariance.Source: §3.4.1 · pp. 6–7

3.4.2. ResOPD and SC-TA as endpoints

  • SC-TA has lower trace covariance than ResOPD if and only if γ⋆ > 1/2, a distribution-dependent certificate rather than a universal variance ordering.The comparison is made within the two-control family.Source: §3.4.2 · pp. 7–8

4. Experiments

The experiments test whether residualization reduces gradient covariance, improves online training behavior, and benefits downstream performance through frozen-prefix audits and end-to-end distillation across mathematical, science, and coding tasks.

4.2. Exact frozen-prefix covariance audits

  • 52.5%–75.0%: ResOPD reduces aggregate covariance trace versus sampled-token OPD, with lower covariance trace and NSR than HETS across all five frozen-prefix panels.The sparse estimators share a Top-4 teacher budget.Source: Table 2 · p. 9, §4.2 · p. 9

4.3. End-to-end mathematical reasoning

  • 50.12% terminal average score and 47.96% learning-trajectory AUC: ResOPD leads the evaluated OPD baselines on Qwen3.5-27B-to-4B mathematical reasoning.It also avoids severe late-stage degradation on external mathematical competitions.Source: §4.3 · pp. 9–10, Table 3 · p. 10

4.3.1. Training dynamics

  • 32.3% support-mass-gap and 55.5% tail-log-mass-mismatch reductions versus HETS accompany ResOPD’s redistribution of predictable between-event corrections.The coarse tail-bucket correction preserves the full-rKL expectation under the stated on-policy condition.Source: §4.3.1 · p. 10, Table 4 · p. 10, Figure 3 · p. 11

4.3.2. Behavioral fidelity

  • 0.36 nats/token at step 192: ResOPD’s Teacher NLL declines monotonically, while HETS plateaus around 0.385 nats/token.The lower Teacher NLL indicates higher likelihood of student-generated trajectories under the teacher.Source: §4.3.2 · p. 11, Figure 4 · p. 12

4.4. Cross-domain transfer to science and coding

  • 44.15% terminal average score and 42.65% trajectory AUC: ResOPD leads aggregate cross-domain results, although sampled-token OPD performs best on science tasks.ResOPD and HETS nominally exceed the teacher on code tasks, but reported confidence intervals cross zero.Source: §4.4 · pp. 11–12, Table 5 · p. 12

4.5. Online support-size ablation

  • +0.11 to +1.56 points: ResOPD outperforms HETS across tested Top-k support sizes, while its expected variance-reduction benefit diminishes as k grows.ResOPD also lowers tail log-mass mismatch across supports, with the largest margin in tight-support regimes.Source: §4.5 · p. 12, Table 6 · p. 13, Figure 5 · p. 13

4.6. Rollout sampling and 4.7. Learning-rate sensitivity

  • Full-support rollouts lower AUC for all evaluated methods but preserve ResOPD’s advantage; at 5 × 10^-6, ResOPD retains 46.35% final score and 47.77% AUC as HETS degrades.At 1 × 10^-5, both methods diverge rapidly, so ResOPD does not shift the ultimate divergence boundary under the default recipe.Source: §4.6 · p. 13, §4.7 · pp. 13–14, Table 7 · p. 14

4.8. Comparing the two-control endpoints

  • Across three prompt domains, SC-TA has lower aggregate trace variance than ResOPD, but prefix-level ordering crosses over and SC-TA shows no consistent downstream advantage.ResOPD has the higher terminal score, while trajectory-level AUC remains similar.Source: §4.8.1 · p. 14, Table 9 · p. 15, §4.8.2 · p. 15

5. Conclusion

ResOPD estimates the full-vocabulary reverse-KL gradient from sparse teacher feedback, preserving unbiasedness under ancestral on-policy sampling without expanding the teacher payload or adding forward passes. In the evaluated settings, it reduces variance and achieves higher aggregate downstream scores.

5. Conclusion

  • ResOPD integrates the observable coarse distribution analytically and samples only the unresolved fine-tail residual to estimate the full-vocabulary reverse-KL gradient.The estimator preserves exact unbiasedness under ancestral on-policy sampling without expanding the teacher payload or adding forward passes.Source: §5 · p. 15
  • Across the evaluated settings, ResOPD reduces variance while preserving exact unbiasedness under ancestral on-policy sampling.The conclusion is supported by theoretical analysis and exact covariance audits.Source: §5 · p. 15
  • ResOPD demonstrates higher aggregate downstream scores in the evaluated settings across math, science, and coding tasks.Source: §5 · p. 15

A. Detailed estimator derivations

The appendix details the bias of three Teacher Top-k reverse-KL forms relative to the full-vocabulary target gradient. It also supplies derivations for the single-control oracle calculation.

A. Detailed estimator derivations

  • The appendix provides detailed Top-k bias derivations and a single-control oracle calculation, including boundary-term cancellation for the complete coarse distribution.Source: A. Detailed estimator derivations · p. 18

A.1. Three biased forms of Teacher Top-k reverse KL

  • Renormalized Top-k reverse KL learns only the conditional head distribution, identifying neither head mass nor the omitted tail distribution.Its global optimum permits arbitrary head mass and tail shape.Source: §A.1 · p. 18
  • Scalar-truncated Top-k reverse KL leaves a tail-gradient bias that can push omitted logits upward; when Q is near one, its objective prefers roughly 1 − 1/e student mass outside the teacher head.The tail-push is an objective-level property, not merely a transient gradient sign.Source: §A.1 · p. 18
  • The stopped-coefficient Top-k gradient leaves the tail expectation uncorrected and the conditional tail distribution unconstrained, remaining biased when the tail vector is nonzero.HETS and ResOPD add a sampled tail correction to restore the missing expectation.Source: §A.1 · p. 19

A.2. Single-control residualization and its variance optimum

The single-control estimator is unbiased for any finite prefix-measurable control chosen before sampling. Its variance separates into an invariant within-tail term and a control-dependent between-event term.

A.2. Single-control residualization and its variance optimum

  • For any finite prefix-measurable vector B fixed before sampling the rollout action, the one-control estimator is unbiased.Source: §A.2 · p. 19
  • The variance decomposition has an invariant within-tail term and a between-event term determined by the control's distance from the true tail mean; the oracle minimizes trace covariance but requires the full conditional fine-tail mean.Source: §A.2 · p. 19
  • ResOPD has lower trace covariance than HETS exactly when its control is closer to the true tail mean than HETS's zero control; no unconditional ordering is asserted.HETS and ResOPD are the two payload-compatible choices.Source: §A.2 · p. 19
  • The complete coarse distribution over the head tokens and one tail bucket supplies an exact coarse gradient, even when the control differs from the true fine-tail mean.Boundary terms cancel because the tail pseudo-token completes the coarse support.Source: §A.2 · p. 19

B. Detailed training and evaluation configurations

This appendix section points readers to the detailed configurations for the primary mathematics and cross-domain runs and identifies where sampling assumptions and the matched full-support mathematics ablation are discussed.

B. Detailed training and evaluation configurations

  • Tables 11 and 12 collect primary mathematics and cross-domain configurations; sampling assumptions are discussed in Sections 3.3 and 4.1, and the matched full-support mathematics ablation appears in Section 4.6.Source: B. Detailed training and evaluation… · p. 20

C.1. Math OPD experiments

ResOPD finishes with the highest accuracy on both DeepMath held-out and external math competitions, while sampled-token OPD and HETS show late-stage deterioration. These trajectories are consistent with reduced stochastic gradient noise and improved late-stage optimization stability.

C.1. Math OPD experiments

  • ResOPD achieves the highest terminal accuracy on DeepMath held-out and decisively surpasses alternatives at the final step on external math competitions.Sampled-token OPD deteriorates after an early peak, and HETS has a late-stage drop on both suites.Source: §C.1 · p. 20
  • The trajectories are consistent with coarse-tail residualization reducing stochastic gradient noise and improving late-stage optimization stability.The passage relates these trajectories to lower fixed-prefix covariance trace and improved online training diagnostics reported elsewhere.Source: §C.1 · p. 20

C.2. Cross-Domain OPD Experiments

Cross-domain trajectories favor ResOPD on code and SuperGPQA, while Sampled-token OPD retains an advantage on GPQA-D. The appendix also interprets online residual moments as lightweight diagnostics and reports a finite enumeration check of estimator unbiasedness.

C.2. Cross-Domain OPD Experiments

  • On LiveCodeBench v5 and v6, ResOPD and HETS outperform Sampled-token OPD, with ResOPD having the highest overall trajectory AUC; science results are mixed.Sampled-token OPD has superior GPQA-D retention and final accuracy, while ResOPD has the best terminal SuperGPQA performance.Source: §C.2 · pp. 20–21

D.1. What the online residual moments measure

  • Under an exact on-policy draw, the support residual is conditionally zero-mean, with conditional variance P(1 − P).P denotes Student probability mass on the transmitted Teacher Top-k support, and H indicates whether the sampled Student token lands there.Source: §D.1 · p. 22
  • The two M2 metrics track the magnitude and sequence-level persistence of the random scalar modulating ResOPD’s bucket control.Token-level M2 captures instantaneous head-versus-tail event magnitude; trajectory-level M2 measures persistent signed imbalance along rollouts.Source: §D.1 · pp. 22–23
  • The M2 diagnostics are not full covariance estimators or replacements for fixed-prefix NSR, and direct Top-k rKL remains biased toward the true rKL target.Direct Top-k rKL lacks a sampled tail correction, so its zero action-sampling variance is not comparable to unbiased estimators.Source: §D.1 · p. 23

D.2. Finite gradient-enumeration check

  • Exact four-action enumeration makes HETS’s expected gradient match the dense full-rKL gradient to maximum absolute error 5.6 × 10^-17.ResOPD’s per-action correction relative to HETS sums to zero in this enumeration, recovering the same dense gradient and validating implementation semantics, not covariance or downstream accuracy.Source: §D.2 · p. 23

D.3. Complete frozen-prefix support diagnostics

Across the reported support-size settings, ResOPD outperforms HETS on covariance-trace elimination and NSR, including when head coverage is high or HETS fails to reduce variance.

D.3. Complete frozen-prefix support diagnostics

  • At k = 4 on HMMT26, HETS versus ResOPD variance reduction is 55.8% versus 75.3%, despite 98.8% head mass.At k = 1, HETS has −1.0% reduction while ResOPD eliminates 32.0% of gradient variance; Figure 8 reports lower ResOPD NSR across evaluated support sizes on GSM8K and HMMT.Source: §D.3 · pp. 23–24, Figure 8 · p. 24
  • Increasing k from 4 to 16 and 64 raises payload from 40 to 136 to 520 B/token, with diminishing marginal variance gains.ResOPD consistently outperforms HETS at matched support sizes across both suites.Source: §D.3 · p. 23

D.4. Training dynamics across learning rates

Across learning-rate settings, ResOPD maintains more stable training dynamics than HETS; at 5 × 10^-6 it has lower pre-clipping gradient norms and no gradient-clipping steps.

D.4. Training dynamics across learning rates

  • At 5 × 10^-6, ResOPD’s pre-clipping gradient norm mean is 0.43 versus 0.48 for HETS, and p99 is 0.97 versus 1.84.ResOPD has 0/64 clipping steps versus HETS’s 2/64; Figure 9 also shows more stable dynamics across the three tested learning rates.Source: §D.4 · p. 24, Figure 9 · p. 25
  • ResOPD curbs the support-mass gap and suppresses token- and trajectory-level residual second moments relative to HETS.Figure 9 shows the contrast as learning rates scale up to 5 × 10^-6.Source: §D.4 · p. 24, Figure 9 · p. 25

Related papers

Loading 2610.04882v1…