PaperlayerView exact revision on arXiv

Source-linked AI summary

TV-Regulated OPD: Direction Matters in On-Policy Distillation

Han Xiao, Yifan Niu, Dongyi Liu, Chang Luo, Jia Li

arXiv:2609.08341v1cs.LG

TL;DR

OPD’s token-level supervision can be high-variance because log-ratio magnitudes are unbounded and may destabilize training. The paper separates directional information from global strength, derives sign-only supervision as conditional TV descent, and regulates global scale with TV. Across tested settings, TV-OPD remains competitive with raw OPD while improving training stability and late-stage retention, though its guarantees are limited by teacher quality, stopped occupancy, and evaluation scope.

  • Problem

    OPD relies on unbounded teacher–student log-ratio coefficients whose high variance and extreme values can destabilize training and limit performance.

  • Method

    TV-OPD keeps the sign of the sampled-token log-ratio for conditional TV descent and uses a discrepancy-responsive global scale without token-wise likelihood-gap weighting.

  • Results

    Across tested settings, TV-OPD remains competitive with raw OPD while exhibiting greater training stability and superior late-stage performance.

  • Takeaways & Limitations

    Separating local teacher-relative direction from global update intensity provides a simple basis for OPD objectives with explicit directional supervision and global control.

  • Takeaways & Limitations

    The study covers limited model pairs, benchmarks, and seeds, and does not establish a higher capability ceiling.

Abstract

from arXiv · show

On-Policy Distillation (OPD) facilitates the transfer of knowledge from domain expert to student in the post-training phase of Large Language Models (LLMs). However, the supervision signals in mainstream OPD methods suffer from high variance and noise which is generally instable during training. In this work, we systematically investigated what really matters to the performance and the fundamental mechanisms behind the instability during training. We found that retaining only the sign of token-level advantages is sufficient to achieve the performance comparable to standard OPD. Meanwhile, smoother and bounded advantages can stabilize the training process without sacrificing its performance. These motivated us to shape the advantages using the Total Variation (TV) and propose a robust TV regulated On-Policy Distillation (TV-OPD) method. Benefiting from the bounded and diminished advantages, TV-OPD exhibits stable training dynamics and steady late-stage performance. We conducted comprehensive experiments and found that, across various settings, TV-OPD consistently achieved better performance and lower variance in the late-stage of training.

1 INTRODUCTION

OPD uses teacher–student log-probability differences as token-level supervision, but their unbounded magnitude can create high-variance, extreme updates. Controlled analysis finds that sign and global scale matter more consistently than precise token-wise magnitude allocation, motivating TV-OPD.

  • Motivation: OPD uses teacher–student log-probability differences as token-wise supervision, with sign determining promotion or suppression and magnitude scaling the gradient.The coefficient has no uniform upper bound, and observed extreme values can destabilize training.
  • Motivation: Existing methods regulate token-wise weights through clipping, compression, power transformations, control variates, or divergence-corrected coefficients.These approaches assume token-level reweighting is highly informative, an assumption the authors find may not hold.
  • Key observation: Controlled interventions decompose supervision into sign, relative magnitude, and global magnitude, showing that sign and global scale are the key factors.Sign determines encouragement or suppression, relative magnitude controls token weighting, and global magnitude controls update intensity.
  • Key observation: Retaining only sign enables stable and effective learning, while a global scale that decays as teacher and student converge supports later-stage convergence.Fine-grained token-wise magnitude allocation yields no consistent benefit.
  • Proposed method: TV-OPD retains sign-based directional supervision and uses total variation to regulate global update intensity.The resulting bounded, low-variance signal is reported to match baseline performance while improving training stability and later-stage performance in 1.5B and 8B settings.

2 RELATED WORK

Related work develops OPD objectives, estimators, and advantage transformations for matching teacher and student distributions on student-generated states. This paper instead isolates the information carried by sampled-token coefficients and separates direction from fine-grained magnitude.

  • Distribution matching: Knowledge distillation transfers a teacher’s predictive distribution to a student, while OPD applies supervision on student-generated prefixes.This motivates matching conditional distributions on states the student visits rather than relying only on fixed teacher responses.
  • Distribution matching: Distribution-matching approaches include reverse-KL distillation and flexible divergence choices combined with student-generated data.MiniLLM uses reverse KL, while Generalized Knowledge Distillation supports flexible divergences.
  • Supervision estimation: Full-vocabulary OPD sums over token probabilities, whereas sampled-token OPD uses the generated token’s teacher–student log-ratio as a lightweight coefficient.Top-k approximations trade vocabulary coverage for computational cost, and variance-reduction methods modify sampled-token estimators.
  • Advantage shaping: Other methods reshape advantages or add information through power transformations, clipping, compression, corrected divergence coefficients, or teacher top-k injection.These methods address estimation, objective design, and missing supervision, whereas this work studies the coefficient’s directional and magnitude information.

3 PRELIMINARY

The preliminary formulation defines OPD on student-visited states using teacher and student conditional token probabilities. Raw OPD estimates a reverse-KL descent signal, while the paper distinguishes conditional TV from sequence-level distribution distances and states the scope of its gradient identities.

  • Student rollouts: Student rollouts define visited states from prompts and generated prefixes, while the fixed teacher evaluates those same prefixes rather than generating training states.The teacher and student share a vocabulary, and the student samples the next token at each visited state.
  • Sampled-token coefficient: The sampled-token OPD coefficient is the teacher–student log-probability difference, Δ(s,a) = log p(a | s) − log q̄θ(a | s).A positive coefficient promotes a token favored more by the teacher, while a negative coefficient suppresses it.
  • Reverse-KL signal: Holding student-visited state occupancy fixed, the conditional reverse-KL objective yields the expected sampled-token log-ratio multiplied by the student score gradient.Raw OPD therefore estimates a reverse-KL descent signal at visited states.
  • Gradient scope: Only the student log-probability receives gradients; teacher outputs, sampled prefixes, and coefficients remain fixed in the stated surrogate.The exact identities apply to the on-policy, pre-optimizer signal with stopped state occupancy, not the full sequence-level derivative.
  • Conditional TV: Conditional total variation measures next-token distribution mismatch averaged over student-visited states, distinct from TV between complete response distributions.This distinction frames the paper’s later TV-based objective and regulator.

4 OBSERVATION: DIRECTION AND MAGNITUDE IN OPD

The observation study holds teacher-relative direction fixed while intervening on coefficient magnitude. Its results indicate that precise token-wise magnitude allocation is not consistently useful, whereas direction preserves effective supervision and motivates separate global-strength control.

  • Direction and magnitude: The coefficient decomposes into direction and magnitude, with direction controlling promotion or suppression and magnitude controlling token allocation and overall scale.The interventions hold direction fixed to test whether precise magnitude allocation is necessary.
  • Controlled interventions: Raw, Sign, Group-Constant, and Permuted variants preserve direction while respectively retaining raw, unit, sign-group, or reassigned magnitude patterns.The comparison isolates fine-grained magnitude under matched experimental conditions.
  • Diagnostic results: 36.29% is the Sign variant’s highest two-benchmark trajectory average, ahead of Raw at 35.40%, Group-Constant at 34.67%, and Permuted at 34.35%.These comparisons use the common evaluation horizon in the two-seed diagnostic.
  • Diagnostic results: 50.00 ± 1.18 on AIME 2024 and 37.50 ± 0.00 on AIME 2025 are Sign’s per-seed best-over-training accuracies, versus 47.50 ± 0.00 and 35.83 ± 1.18 for Raw.Sign therefore leads Raw on both reported per-seed best-over-training accuracies.
  • Interpretation: Token-wise magnitude reweighting provides no consistent benefit in this diagnostic, supporting teacher-relative direction without precise magnitude allocation.The evidence is specific to the tested model pair and does not directly measure supervision noise.
  • Implication: Sign-TV removes raw log-ratio magnitude to produce a conditional TV descent signal, while TV-OPD restores global attenuation through a shared TV-responsive regulator.This separates local directional supervision from global update strength without restoring token-wise likelihood-gap weights.

5 METHOD: TV-REGULATED OPD

TV-OPD retains sign-based token supervision for a conditional-TV descent direction and uses estimated TV to regulate global update strength. The resulting estimator is bounded, sampled-token compatible, and subject to explicit pooling and variable-length qualifications.

  • From sign supervision to TV optimization: Sign-based coefficients encourage student-underrepresented tokens and suppress overrepresented tokens, yielding a conditional TV descent direction.The sign of the teacher–student log-probability difference equals the sign of the probability difference.
  • From sign supervision to TV optimization: The sign coefficient is bounded by one, limiting extreme teacher–student log-ratios from assigning arbitrarily large token weights.The bound applies to the teacher-induced coefficient; score norms and the optimizer still affect the actual update.
  • Estimating TV: Sampled-token log-probabilities provide an unbiased per-state estimate of TV without requiring full-vocabulary teacher scoring.The one-sided estimator lies in [0, 1] and has conditional variance at most 1/4.
  • Global regulation: TV-OPD pools token estimates and smooths their ratio with an EMA to obtain a shared, discrepancy-responsive supervision scale.The shared coefficient attenuates updates below a reference discrepancy while remaining at one above it.
  • Global regulation: The detached coefficient is fixed before sampling and uniformly scales the policy-gradient loss without changing relative token weights.For variable-length responses, active-token pooling is a consistent ratio estimator under the stated rollout and masking conditions but need not be unbiased at finite batch size.

6 EXPERIMENTS

Experiments compare TV-OPD with Raw and Sign-TV across two teacher–student pairs and multiple training stages. TV-OPD achieves strong aggregate benchmark results, with its clearest advantage in late-stage retention, while excessive attenuation harms performance.

  • Main results: TV-OPD reaches 51.67 ± 1.18 on AIME 2024 and 38.75 ± 2.08 on AIME 2025 on the JustRL pair, averaging 45.21%.The selected checkpoints retain the high AIME 2024 result while recovering AIME 2025 performance relative to the unrestricted choice.
  • Main results: TV-OPD reaches 70.00 ± 1.17 on AIME 2024 and 58.33 ± 1.67 on AIME 2025 on the Qwen pair, averaging 64.17%.This exceeds the strongest baseline by 1.67 points on AIME 2024 and 1.25 points on AIME 2025.
  • Evaluation scope: The available validation exports do not include the TV-OPD scale, conditional TV, or realized update norm, so mechanism diagnostics are analyzed separately.The limitation concerns observed retention versus the proposed mechanism’s direct variables.
  • Stage-wise dynamics: TV-OPD is higher by 3.47 points on AIME 2024 and 0.90 points on AIME 2025 in the late stage, although it does not improve every stage.Sign-TV performs best early, while Raw leads on AIME 2025 in the middle stage.
  • Late-stage retention: TV-OPD improves LateMean from 40.87 ± 0.83 to 43.06 ± 0.10 and reduces PeakDrop from 3.51 ± 1.13 to 2.36 ± 0.49.These statistics summarize sustained late-stage performance and the gap from peak performance over steps 500–625.
  • Regulator sensitivity: Increasing α from 0.5 to 4.0 lowers the common-horizon mean by 4.80 points across the two benchmarks.The reported decreases are 5.68 points on AIME 2024 and 3.92 points on AIME 2025, consistent with stronger attenuation.

7 DISCUSSION

The discussion separates token update direction from training intensity and positions TV-OPD as a bounded, globally regulated alternative. Its evidence supports improved optimization dynamics and late-stage retention, but not a higher capability ceiling.

  • Sign information determines the teacher-relative update direction, while magnitude controls token allocation and global training strength.
  • TV-OPD restores one global scale without reintroducing likelihood-gap weighting across tokens.
  • The study reports improved optimization dynamics and late-stage retention, while not establishing a higher capability ceiling.

8 CONCLUSION

The work separates token update direction from training strength in sampled-token OPD. Sign-TV and TV-OPD remain competitive with raw OPD while improving late-stage retention in the completed regulated comparison.

  • Sign-TV provides a principled conditional total-variation objective for sampled-token OPD.
  • TV-OPD adds a discrepancy-responsive global scale without restoring token-wise magnitude.
  • TV-OPD remains competitive with raw OPD and improves late-stage retention across the tested settings.

A THEORETICAL PROPERTIES AND PROOFS

The theoretical analysis derives TV-based OPD signals with bounded teacher influence, establishes behavioral preservation under TV closeness, and characterizes bounded on-policy TV estimation. It also states scope conditions involving teacher quality, support, occupancy, finite-step optimization, and fixed-horizon sampling.

  • Sign-based token signals preserve every nonzero token direction while bounding the teacher coefficient, unlike unbounded log-ratio coefficients.
  • At fixed global scale, changing the teacher can flip a token direction but cannot induce arbitrarily large relative weight.
  • 2Amax(s)DTV(˜π, p) bounds the loss in expected task advantage, so DTV(˜π, p) < ∆T(s)/(2Amax(s)) preserves positivity.
  • Positive advantage at one state is not a global return guarantee, and improvement requires conditions across visited states, occupancy control, and teacher quality.
  • The fixed-horizon analysis assumes positive support and specified sampling conditions; truncation, changing occupancy, stochastic updates, nonsmooth points, and random EOS lengths require separate treatment.
  • The on-policy TV estimator is bounded in [0, 1], unbiased for d(s), and has variance at most d(s)(1 − d(s)) ≤ 1.

B.7 IMPLEMENTATION SCOPE: PREPROCESSING AND SELECTION

The implementation defines when preprocessing preserves the TV estimator and clarifies how token selection changes the estimand. These conditions distinguish exact estimator equality from broader claims about bias.

  • Preprocessing and masking: The initial advantage is the teacher-minus-student log-probability, and the effective mask excludes sentinel positions while normalizing by max{N_k, 1}.An empty batch therefore contributes zero rather than a defined active-token TV estimate.
  • Preprocessing and masking: Exact equality with the TV estimator requires preserving negative log-ratios and mapping nonnegative ratios to nonnegative values under T(∆).Capping positive log-ratios at a nonnegative threshold preserves the statistic, whereas changing the negative branch generally does not.
  • Token selection: Pooling tokens selected by a current-token mask targets expected selected positive mass divided by expected selected-token count, not necessarily active-state TV or TV of renormalized policies.Current-token sentinel exclusions therefore require a selection-aware interpretation unless additional properties establish equality.
  • Token selection: If generation uses a sampling law different from the denominator distribution, the one-sided proof must be rederived using that law or appropriate weighting.

C FULL EXPERIMENTAL SETUP

The experiments use on-policy rollouts, fixed training controls, and matched method comparisons that isolate token-coefficient choices or a shared TV-OPD scale. The setup also specifies the optimizer, distributed training configuration, and evaluation protocol.

  • Training setup: Training uses batches of 64 prompts, one on-policy response per prompt, AdamW with learning rate 10−6, and a constant base schedule with warmup.Runs use bfloat16, gradient checkpointing, fully sharded data parallelism, and sequence parallelism on eight GPUs.
  • Evaluation setup: Evaluation samples four responses per problem on AIME 2024 and AIME 2025 and reports mean@4 across two paired seeds.The formal JustRL comparison includes Raw, Sign-TV, and TV-OPD.
  • Method controls: Raw and TV differ only by replacing ∆_i with sign(∆_i), while TV-OPD retains the TV signs and changes one shared pre-optimizer scale c_k.The completed Qwen and JustRL runs use α = 0.5, without mixing task rewards or verifiers into OPD coefficients.
  • Method controls: The advantage-scaling scheduler uses EMA retention β = 0.95, c_min = 0.1, and a reference frozen at the first valid optimizer step.The coefficient is fixed at step entry, refreshed for the next step, and multiplies the policy-gradient loss uniformly across microbatches before AdamW.
  • Diagnostic interventions: The diagnostic interventions preserve the teacher-relative sign while varying magnitudes: Raw keeps |∆_i|, Sign sets nonzero magnitudes to one, and Group-Constant averages within sign groups.

D.1 JUSTRL DIAGNOSTIC

The JustRL diagnostic compares magnitude manipulations over a jointly complete evaluation horizon, then examines a complementary allocation control and TV-OPD scale sensitivity. Sign leads the main trajectory diagnostic, while the allocation control only slightly improves the average trajectory.

  • Main diagnostic: The main diagnostic compares a JustRL-DeepSeek-1.5B teacher with a DeepSeek-R1-Distill-Qwen-1.5B student on AIME 2024 and 2025 using two seeds.Comparative summaries use the jointly complete horizon because run lengths differ; later reduced-coverage observations are dashed.
  • Main diagnostic: Sign is highest on both AIME 2024 and AIME 2025 in per-seed best-over-training accuracy within the common evaluation horizon.
  • Main diagnostic: 36.29 ± 0.15 is the trajectory-level two-benchmark mean for Sign, compared with 35.40 ± 0.11 for Raw, 34.67 ± 1.47 for Group-Constant, and 34.35 ± 0.13 for Permuted.These values average benchmarks at each checkpoint, then checkpoints within each seed, and finally the two seeds.
  • Allocation control: Sign-Mass/Raw-Allocation preserves each sign group’s total magnitude while retaining normalized raw magnitudes for within-group allocation.
  • Allocation control: 37.00 ± 0.28 is the trajectory-level mean for Sign-Mass/Raw-Allocation versus 36.94 ± 0.00 for Raw over the common 0–425 horizon.The control slightly improves the average trajectory but not the per-benchmark best-over-training values.
  • TV-OPD sensitivity: The α sweep changes only global feedback sensitivity in Eq. (11), with common-horizon results reported for the tested TV-OPD settings.
Loading 2609.08341v1…