PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Token-Local Imitation: Reward-Compatible Temporal Credit Assignment for On-Policy Distillation

Shiqi Liu, Zeyu He, Letian Tao, Guojian Zhan, Jiaxin Gao, Feihong Zhang, Jingliang Duan, Wei Xiong, Kehua Sheng, Bo Zhang, Yang Guan, Shengbo Eben Li

arXiv:2609.16937v1cs.LGcs.AIcs.PL

TL;DR

Existing OPD methods trade stable local supervision against faithful but high-variance sequence-level credit. The paper unifies them through temporal credit, proposes γOPD with discounted credit and RBM, and reports consistent improvements across vanilla, size-mismatched, and multi-teacher distillation settings.

  • Problem

    OPD faces a trade-off between computationally efficient token-level supervision that alters the objective and sequence-level supervision whose variance increases with trajectory length.

  • Method

    The paper introduces γOPD, a discounted temporal-credit surrogate with a horizon-independent variance bound, plus RBM for combining discounted teacher guidance with verifiable outcome rewards.

  • Results

    γOPD consistently improves over existing OPD methods in vanilla, size-mismatched, and multi-teacher distillation across mathematical and code reasoning settings.

  • Takeaways & Limitations

    Temporal credit assignment and reward-compatible mixing extend OPD beyond purely local or teacher-dependent supervision while retaining optimization stability.

  • Takeaways & Limitations

    The experiments are limited to models with fewer than 10B parameters, leaving larger-scale evaluation and more flexible discounting for future work.

Abstract

from arXiv · show

On-policy distillation (OPD) has emerged as an effective approach for large language model post-training, yet existing objectives face a trade-off between objective fidelity and optimization stability. Token-level OPD provides stable but local supervision, whereas sequence-level OPD captures future credit at the cost of horizon-dependent variance. We establish a unified temporal-credit view of these formulations, showing that practical token-level OPD can be interpreted as a temporal approximation to the sequence-level reverse-KL gradient. Building on this connection, we propose $γ$OPD, which uses discounted temporal credit assignment to balance long-horizon supervision and optimization stability, while admitting a horizon-independent variance bound. We further develop a reward-compatible bounded mixing (RBM) mechanism for $γ\mathrm{OPD}$ that balances verifiable outcome feedback with the discounted OPD advantage to move beyond purely teacher-dependent optimization. Experiments on mathematical and code reasoning demonstrate consistent improvements over existing OPD methods across vanilla, size-mismatched, and multi-teacher distillation settings.

1 INTRODUCTION

OPD provides dense on-policy teacher supervision for reasoning, but existing estimators trade objective fidelity against optimization stability. The paper unifies these views and introduces discounted temporal credit with reward-compatible outcome mixing.

  • OPD trains on student-sampled trajectories, providing dense teacher supervision that promotes reasoning quality while mitigating exposure bias.
  • Token-level OPD is computationally efficient but biased relative to the sequence-level objective, whereas sequence-level estimation preserves fidelity but becomes unstable on long trajectories.
  • γOPD uses discounted temporal credit to interpolate between token-level and sequence-level supervision while admitting a horizon-independent variance bound.
  • The paper establishes a temporal-credit view in which practical token-level OPD is a temporally truncated approximation to the sequence-level reverse-KL gradient.
  • RBM combines discounted teacher-derived credit with verifiable outcome rewards while preventing teacher supervision from dominating task-level feedback.

2 PRELIMINARIES

The paper studies student-policy optimization under a fixed teacher on autoregressive reasoning trajectories. It situates OPD’s dense teacher supervision alongside RLVR’s sparse verifier-grounded feedback.

  • 2 PRELIMINARIES: The setup optimizes a student policy πθ on complete autoregressive response trajectories generated for prompts from a reasoning dataset.Each trajectory consists of tokens y1 through yT from a discrete vocabulary.
  • 2 PRELIMINARIES: OPD minimizes reverse KL from the student policy to a fixed teacher policy using trajectories induced by the current student.On-policy trajectories reduce exposure bias and suit long chain-of-thought reasoning tasks.
  • 2 PRELIMINARIES: Practical OPD converts reverse-KL minimization into an RL-style surrogate using the token-level advantage log π∗(yt | ht) − log πθ(yt | ht).This advantage compares teacher and student probabilities for the sampled token at each decoding step.
  • 2 PRELIMINARIES: RLVR supplies sparse verifier-grounded sequence rewards, while OPD provides dense token supervision that can be constrained by suboptimal teacher behavior.Verifiers may include code compilers or mathematical rule checkers, producing rewards in {−1, 1}.

3 METHODOLOGY

The methodology unifies token-level and sequence-level OPD as different temporal-credit assignments, then introduces discounted γOPD and reward-compatible bounded mixing to balance long-horizon credit, variance, and task verification.

  • 3.1 TOKEN-LEVEL OPD AS A TEMPORAL APPROXIMATION: Practical token-level OPD omits future-token effects, whereas sequence-level OPD accounts for each token’s indirect influence on future prefixes.The difference arises because token-level OPD applies stop-gradient to the student-induced prefix distribution, while sequence-level OPD differentiates through future autoregressive prefixes.
  • 3.2 TEMPORAL-CREDIT OPD: γOPD introduces γ ∈ [0, 1] to control future credit, recovering sequence-level OPD at γ = 1 and token-level OPD at γ = 0.Intermediate γ values define a biased surrogate that interpolates between the endpoint gradients and controls the temporal horizon of credit assignment.
  • 3.2 TEMPORAL-CREDIT OPD: Sequence-level OPD credit can have horizon-dependent variance, whereas γOPD has a horizon-independent variance bound for any fixed γ < 1.This yields a principled trade-off between training stability and long-horizon credit propagation.
  • 3.3 REWARD-COMPATIBLE BOUNDED MIXING: γOPD remains fundamentally teacher-derived, motivating reward-compatible bounded mixing that incorporates verifiable task rewards without allowing teacher advantages to dominate.RBM calibrates the γOPD advantage by response-wise mean absolute magnitude, applies a softsign transformation, and combines the bounded signal with the task reward.
  • 3.3 REWARD-COMPATIBLE BOUNDED MIXING: RBM combines sparse verifiable rewards for task-level learning with bounded γOPD for dense token-level credit assignment.The verifiable reward determines whether a sampled response is reinforced or suppressed, while bounded γOPD modulates token-level update strength.

4 EXPERIMENTS

Experiments evaluate γOPD across vanilla, size-mismatch, and multi-teacher distillation, with ablations and analyses of training dynamics, discount sensitivity, and token-level credit. γOPD consistently improves benchmark performance while providing more stable optimization and more outcome-aligned supervision.

  • Vanilla Distillation: γOPD improves over the strongest baseline by 2.30% in AvgAcc and 1.80% in AvgPass under vanilla distillation, surpassing the teacher on both averaged metrics.The results suggest RBM leverages dense teacher supervision while enabling improvement beyond direct imitation.
  • Size-Mismatch Distillation: γOPD achieves the best size-mismatch performance, improving AvgAcc by 4.92% and AvgPass by 2.06% over the strongest baselines while narrowing the teacher gap.All student methods remain below the teacher because the student has substantially lower capacity.
  • Multi-Teacher Distillation: γOPD improves TotalAvg by 1.87% over AOPD in multi-teacher distillation, leads on HumanEval+ and LCB, and outperforms the teacher on average across mathematics and code.The method remains effective when jointly training a shared student with domain-specialized teachers.
  • Training Dynamics and Token Advantages: γOPD maintains the highest verifiable reward, stable absolute advantage, smallest gradient-norm fluctuations, and a shorter, more stable response-length range during training.Token visualizations show discounted credit spreading positive signal toward key derivations and concentrating negative credit on erroneous formula derivations.
  • Ablation Study: Temporal discounting alone improves AIME average accuracy by 2.04 points, while removing it from the full method reduces the gain from 3.91 to 2.00 points.Naive reward mixing adds only 0.05 points, whereas bounded normalization with temporal discounting and mixing yields the strongest overall performance.
  • Sensitivity to γ: γ = 0.99 achieves the best discount-sensitivity performance, whereas γ = 1 produces larger gradient norms and rapid entropy collapse under long response horizons.Local OPD and discounted variants steadily improve, indicating that a large sub-unity discount balances long-range credit and optimization stability.

5 RELATED WORK

Related work addresses unreliable student trajectories and reward integration in OPD through truncation, entropy-aware objectives, and verifier-guided supervision. These approaches can reduce unreliable continuation influence or add outcome feedback, but may involve additional estimation or computational complexity.

  • Truncation-based OPD methods terminate rollouts early or mask later-token signals, while entropy-aware OPD augments reverse-KL distillation.
  • Reward-guided methods combine verifiable rewards with teacher- or student-weighted objectives, using mechanisms such as clipping, entropy-aware sampling, routing, and scheduling.
  • Existing reward-guided approaches may add group-based advantage estimation, log-ratio correction, or verifier-weighted teacher KL objectives, potentially increasing computational overhead.

6 CONCLUSION

The paper unifies token-level and sequence-level OPD through temporal credit assignment, proposes γOPD for balancing long-horizon supervision with optimization stability, and adds RBM for verifiable outcome rewards. Experiments show effectiveness across teacher–student configurations and distillation scenarios, while evaluation remains limited to models below 10B parameters.

  • γOPD unifies token-level and sequence-level OPD from a temporal-credit perspective to balance long-horizon supervision and optimization stability.
  • RBM combines discounted teacher guidance with verifiable outcome rewards.
  • Experiments on mathematical and code reasoning show γOPD remains effective across different teacher–student configurations and distillation scenarios.
  • The experiments are limited by computational constraints to models with fewer than 10B parameters, leaving larger-scale and adaptive discounting evaluations for future work.

A PROOF OF PROPOSITION 3.1

The proof rewrites the token-level OPD objective using a stop-gradient prefix distribution, differentiates the next-token KL term, and applies the score-function identity. The resulting on-policy expression is exactly the token-level OPD gradient, establishing Proposition 3.1.

  • The stop-gradient prefix distribution is fixed during differentiation, allowing the gradient to move inside the expectation over prefixes.
  • For a fixed prefix, the full-vocabulary log-ratio is introduced as the counterpart of the sampled token-level log-ratio Δ_t.
  • The score-function identity is used after the vanishing term is removed to derive the policy-gradient form.
  • Evaluating the fixed prefix distribution under the current policy permits on-policy trajectory sampling without propagating gradients through prefix sampling.
  • The resulting expression exactly matches the token-level OPD gradient, proving Proposition 3.1.

B PROOF OF PROPOSITION 3.2

The proof expands the sequence-level reverse-KL objective over response trajectories and derives its policy-gradient form. Conditioning and the score-function identity eliminate earlier-token terms, leaving only future and current credit.

  • B PROOF OF PROPOSITION 3.2: The proof expands reverse KL over complete response trajectories and differentiates the resulting sequence-level objective.This establishes the trajectory-level starting point for Proposition 3.2.
  • B PROOF OF PROPOSITION 3.2: Terms with t′ < t vanish in expectation because their credit is determined by earlier tokens and is measurable from the current history.Conditioning on h_t and applying the token-level score identity yields the cancellation.
  • B PROOF OF PROPOSITION 3.2: Keeping only t′ ≥ t terms produces exactly the sequence-level gradient stated in equation 6.This completes the proof of Proposition 3.2.

C PROOF OF THEOREM 3.3

The theorem bounds sequence-level OPD credit variance and shows that discounted γOPD avoids its horizon-dependent growth. The resulting discount factor controls the trade-off between long-horizon credit propagation and stability.

  • C PROOF OF THEOREM 3.3: The proof applies the L2 triangle inequality to sequence-level credit and constructs a shared-noise case attaining the quadratic dependence.The assumption sets the relevant future credit terms equal to a zero-mean random variable with variance σ².
  • C PROOF OF THEOREM 3.3: Sequence-level OPD variance can grow quadratically with the remaining horizon, whereas γOPD has a horizon-independent bound for every fixed γ < 1.The bound distinguishes undiscounted sequence-level credit from discounted temporal credit.
  • C PROOF OF THEOREM 3.3: The complete reward-enhanced γOPD procedure samples prompts and responses, computes discounted advantages, and updates the student policy.The algorithm uses the dataset, student and teacher policies, discount factor, batch size, and verifier rewards.

E EXPERIMENT DETAILS

The experiments use standardized implementations, reasoning datasets, and repeated-sampling evaluations for mathematics and code. The additional γOPD operations impose negligible measured training overhead.

  • E EXPERIMENT DETAILS: Mathematical reasoning uses filtered DeepMath-103K examples, while code reasoning uses Eurus-RL-Code with verifier outcomes mapped to rewards.Both domains are implemented with veRL, vLLM rollouts, and FSDP actor training.
  • E EXPERIMENT DETAILS: The OPD-based methods use reverse-KL token advantages without an additional KL reward penalty and share RADAR optimization settings unless otherwise specified.Training runs use a constant learning rate on 32 NVIDIA H20 GPUs.
  • E EXPERIMENT DETAILS: Mathematical evaluation covers AIME24, AIME25, AMC23, and MATH500 with repeated sampling and reports accuracy-based metrics.Sampling uses 32 responses for AIME24, AIME25, and AMC23, and 4 for MATH500, MinervaMath, and OlympiadBench.
  • E EXPERIMENT DETAILS: 0.1064% computation time and 0.00046% memory are added by γOPD relative to a full OPD update.The profile includes discounted credit computation, bounded advantage shaping, and verifier-reward mixing on a 4B student/teacher setup.

F.2 COMPLETE TOKEN-LEVEL ADVANTAGE VISUALIZATION

The visualizations compare local and temporally propagated token advantages across correct and incorrect responses. Discounted propagation strengthens credit near key derivations and errors while reducing penalties on unrelated tokens.

  • F.2 COMPLETE TOKEN-LEVEL ADVANTAGE VISUALIZATION: The visualizations display responses in token order across consecutive panels, with local advantages above and normalized temporally propagated advantages below.Blue and red encode negative and positive advantages, respectively.
  • F.2 COMPLETE TOKEN-LEVEL ADVANTAGE VISUALIZATION: For correct responses, A(0)_t is sparse, whereas A(0.99)_t propagates later-step information and gives stronger positive credit to the key mathematical derivation.The propagated signal also mildly penalizes redundant text.
  • F.2 COMPLETE TOKEN-LEVEL ADVANTAGE VISUALIZATION: For incorrect responses, A(0.99)_t assigns stronger negative credit to the erroneous formula derivation without excessively penalizing the final token.Local A(0)_t instead penalizes tokens largely unrelated to the actual error.
Loading 2609.16937v1…