PaperlayerView exact revision on arXiv

Source-linked AI summary

ReNIO: Reweighting Negative Trajectory Importance for LLM On-Policy Distillation

Chen Lin, Kedi Chen, Wei Zhang

arXiv:2606.23104v1cs.LGcs.AI

TL;DR

On-policy distillation lacks a way to distinguish more informative negative trajectories while retaining its short-prefix training advantage. ReNIO uses student-to-teacher probability ratios to reweight likely negative trajectories, improving OPD and OPSD across mathematical reasoning and code generation tasks.

  • Problem

    Incorrect trajectories can preserve more exploratory and self-corrective reasoning than correct trajectories, motivating their greater use in on-policy distillation.

  • Method

    ReNIO uses prefix-conditioned student-to-teacher probability ratios to identify pivotal tokens and assign normalized sample weights emphasizing likely negative trajectories.

  • Results

    10.00% relative improvement on AIME25 was obtained by DS-R1-Qwen-7B with OPSD, while ReNIO consistently improved OPD or OPSD across mathematical reasoning and code generation tasks.

  • Takeaways & Limitations

    ReNIO emphasizes negative trajectories without final-answer correctness, preserving on-policy distillation’s short-prefix training advantage over full-rollout reinforcement learning.

  • Takeaways & Limitations

    Due to hardware constraints, the effectiveness of ReNIO has not been verified on larger-scale models.

Abstract

from arXiv · show

On-policy distillation (OPD) improves LLM reasoning by training a student model on its own generated outputs, but standard OPD treats all student-generated outputs (SGOs) equally regardless of their informativeness. We observe a consistent asymmetry in controlled filtering experiments: in both OPD and on-policy self distillation (OPSD), training only on incorrect SGOs outperforms training only on correct ones. Our further analysis suggests that models trained on correct-only SGOs tend to generate shorter reasoning traces and show weaker reflection behavior, while incorrect SGOs better preserve exploratory reasoning near the model's capability boundary. To exploit this signal without requiring full answer-containing rollouts, we introduce ReNIO, which Reweights Negative trajectory Importance for LLM On-policy distillation. By using the student-to-teacher probability ratio, ReNIO identifies pivotal tokens leading to wrong reasoning traces and aggregates their information into a normalized sample weight, inherently assigning larger weights to likely negative trajectories without observing the correctness of final-answer. Since Re-NIO only uses prefix-conditioned token probabilities, it preserves OPD's prefix training advantage over full-rollout reinforcement learning. Across both mathematical reasoning and code generation tasks, ReNIO improves both OPD and OPSD, with representative relative gains of up to 8.90% for Qwen3-1.7B and 10.00% for R1-Distill-Qwen-7B on mathematical reasoning benchmarks. Code repo: https://github.com/BDML-lab/ReNIO.

1 Introduction

The introduction argues that incorrect student-generated outputs provide more useful reasoning signals than correct ones, motivating ReNIO, a prefix-computable reweighting method that emphasizes likely negative trajectories without final-answer labels. ReNIO improves both OPD and OPSD across mathematical reasoning and code generation tasks.

  • Introduction: OPD aligns student and teacher token distributions on student-generated outputs, providing denser supervision than reward-based reinforcement learning.OPD can also learn from short prefixes, whereas reinforcement learning generally requires observing final answers before assigning rewards.
  • Introduction: ReNIO preserves on-policy distillation’s short-prefix training advantage without additional full rollouts or reward-labeling cost.Correctness-based weighting would require rolling out each trajectory until its answer appears, undermining prefix efficiency.
  • Introduction: Incorrect student-generated outputs can be more useful than correct outputs for improving reasoning in both OPD and OPSD.Correct-only training produces shorter responses with fewer reflection-style markers, while incorrect-only training better preserves exploratory and self-corrective reasoning.
  • Introduction: ReNIO reweights samples using prefix-conditioned student-to-teacher probability ratios to identify pivotal tokens in likely negative trajectories.High ratios indicate tokens the student favors but the teacher rejects; the method avoids requiring final-answer correctness labels.
  • Introduction: 8.90% relative gains are reported for Qwen3-1.7B and 10.00% for R1-Distill-Qwen-7B on mathematical reasoning benchmarks.The introduction reports improvements for both OPD and OPSD across mathematical reasoning and code generation tasks.

2 Preliminary Analysis of On-Policy Distillation

Controlled studies show that incorrect-only student-generated outputs outperform correct-only outputs in both OPD and OPSD, while producing longer, more reflective reasoning. These findings motivate weighting trajectories that reveal stronger student–teacher deviations.

  • Correct versus incorrect SGOs: Incorrect-only training outperforms correct-only training under both OPD and OPSD, with gains across AIME24, AIME25, and the benchmark average.Under OPD, gains are +3.60, +3.89, and +2.59 Avg@12 points; under OPSD, they are +1.94, +3.34, and +2.50 points, respectively.
  • Inference-time effects: Models trained on incorrect-only SGOs consistently generate longer responses than models trained on correct-only SGOs.The comparison evaluates average response length and reflection-style marker counts at inference time.
  • Inference-time effects: Incorrect-only training produces more reflection-style markers, indicating more frequent checking, revision, and alternative exploration.The marker set includes terms such as “but,” “wait,” “perhaps,” “actually,” “check,” and “alternatively.”
  • Motivation for ReNIO: Correct-only training may consolidate successful behavior and reduce exploration, whereas incorrect-only training preserves cautious reasoning from failed but structured attempts.Standard OPD and OPSD weight all trajectories equally, motivating greater emphasis on trajectories with stronger student–teacher deviations.

3 Methodology

ReNIO estimates the importance of student-generated outputs from prefix-level student–teacher disagreement, identifies rare pivotal tokens with a fixed threshold, and aggregates them into normalized sample weights. This answer-free weighting preserves truncated-prefix training and redistributes emphasis without changing the average batch update scale.

  • Prefix-computable importance signal: ReNIO uses the student-to-teacher probability ratio to identify pivotal local decisions where the student is confident but the teacher rejects the sampled token.These decisions can indicate informative failures without directly observing final-answer correctness.
  • Prefix-computable importance signal: The token log ratio ℓ_t is both a disagreement signal and a corrective weight under the prefix-level reverse-KL objective.ReNIO therefore uses ℓ_t as prefix-computable evidence for sample-level SGO importance.
  • Key-token selection: ReNIO retains a token as pivotal when its log ratio exceeds the fixed threshold τ, filtering out the many low-ratio tokens and preserving rare high-ratio tokens.Both OPD and OPSD exhibit long-tailed log-ratio distributions in which most tokens have values near zero and only a small subset has large values.
  • Weight aggregation and normalization: ReNIO aggregates selected token information with a geometric mean and sets the SGO weight to 1 when no key tokens are selected.Geometric-mean aggregation averages log student–teacher ratios before exponentiating.
  • Weight aggregation and normalization: Batch normalization gives the weights mean 1, redistributing emphasis across SGOs without changing the average update scale.The resulting sample weights are applied to the unweighted on-policy distillation objective, where standard OPD otherwise weights every SGO equally.
  • Integration into distillation: Because ReNIO uses only prefix-conditioned token probabilities and not final answers, it remains compatible with truncated-prefix training and retains OPD’s efficiency advantage over reward-based RL.The method is designed to estimate informative failures without requiring complete answer-containing rollouts.

4 Experiments

Experiments across mathematical reasoning and code generation show that ReNIO consistently improves both teacher-based OPD and teacher-free OPSD, while retaining prefix-based efficiency. Ablations further show that key-token selection, clipping, and batch-wise normalization each contribute to its effectiveness and stability.

  • Main results: ReNIO improves both OPD and OPSD across evaluated model families and task domains.Experiments cover Qwen3 and DeepSeek-R1-Distill-Qwen models on mathematical reasoning and code generation.
  • Main results: 15.44% relative improvement: OPD with DS-R1-Qwen-1.5B improves on HMMT25; 10.00% relative improvement: OPSD with DS-R1-Qwen-7B improves on AIME25.These results demonstrate benefits for both teacher-based distillation and teacher-free self-distillation.
  • Main results: OPSD+ReNIO improves over both standard OPSD and the base model on all three math benchmarks for DS-R1-Qwen-1.5B.Standard OPSD slightly decreases the math average relative to the base model, indicating that ReNIO stabilizes noisy self-distillation.
  • Prefix efficiency: 40.83 vs. 38.70: 1024-token OPSD outperforms 4096-token OPSD without ReNIO, while ReNIO adds 1.95 and 1.86 points to OPSD and 1.67 and 0.93 points to OPD at 1024 and 4096 tokens, respectively.Prefix-based distillation is substantially cheaper than GRPO, and ReNIO computes weights from prefixes.
  • Ablations: 40.37 to 42.04: the full ReNIO pipeline gives the best average score in the OPD ablation.Removing clipping reduces the average to 40.09, removing key-token selection reaches 40.83, and removing normalization drops the average to 39.54.

5 Related Work

Classical knowledge distillation transfers teacher behavior using fixed data or teacher-generated targets, whereas on-policy distillation samples student trajectories and supervises visited prefixes to reduce exposure bias. Recent OPD research includes self-distillation.

  • Classical and on-policy distillation: Classical knowledge distillation transfers teacher behavior to students on fixed data or teacher-generated targets.This offline setting can suffer from exposure bias in LLM post-training.
  • Classical and on-policy distillation: On-policy distillation samples trajectories from the student policy and applies teacher supervision on visited prefixes.This approach reduces the mismatch associated with offline training.
  • Recent OPD studies: Recent on-policy distillation studies include self-distillation.The passage identifies self-distillation as one area covered by recent OPD research.

6 Conclusion

The paper finds that incorrect student-generated outputs can provide especially useful supervision for on-policy distillation. ReNIO emphasizes negative trajectories using prefix-computable student-to-teacher log-ratio weights without observing final-answer correctness.

  • Incorrect SGOs can provide especially useful supervision for on-policy distillation, motivating emphasis on negative trajectories without observing final-answer correctness.
  • ReNIO identifies pivotal tokens through student-to-teacher log ratios and aggregates them into normalized sample weights.The reweighting method is computable from prefixes and does not require an SGO’s answer for weighting.

Limitations · A Reverse-KL Gradient Interpretation of the Student–Teacher Ratio · B ReNIO Algorithm

ReNIO’s student-to-teacher log ratio has a reverse-KL gradient interpretation: it provides token-level corrective weights that are aggregated into normalized sample weights. The algorithm computes these ratios, selects key-token information, aggregates it, and the method remains unverified on larger-scale models because of hardware constraints.

  • Limitations: ReNIO shows consistent gains across multiple model families and task domains, but its effectiveness on larger-scale models remains unverified.The limitation is attributed to hardware constraints.
  • A Reverse-KL Gradient Interpretation of the Student–Teacher Ratio: Under reverse-KL distillation, the token-specific gradient component is controlled by log pS(u) − log pT(u) + 1, with a shared distribution-level baseline.The derivation considers a fixed prefix and student and teacher next-token distributions.
  • A Reverse-KL Gradient Interpretation of the Student–Teacher Ratio: Because the constant +1 is a removable baseline, the effective token-level weight is exactly the student-to-teacher log ratio.The score-function form uses the zero-mean property of the softmax-gradient term.
  • A Reverse-KL Gradient Interpretation of the Student–Teacher Ratio: For an on-policy SGO, tokens assigned much higher probability by the student than the teacher receive larger reverse-KL corrective emphasis.ReNIO uses this log-ratio term as token-level evidence for identifying pivotal decisions.
  • B ReNIO Algorithm: The ReNIO algorithm takes a student, teacher, and on-policy batch as input and outputs normalized sample weights.Its stated inputs include the threshold τ and clip bounds ϵmin and ϵmax.
  • B ReNIO Algorithm: In Phase I, ReNIO computes each generated token’s log ratio as ℓt = log pS(yt | xi, y<t) − log pT(yt | xi, y<t).The computation is performed for every token position in each trajectory.
  • B ReNIO Algorithm: The remaining stages select key-token information and aggregate token information before producing normalized sample weights.The algorithm explicitly labels these as Phase II and Phase III, followed by batch-level weight processing.

C Detailed Experimental Setup · D Additional Analyses and Discussions · D.1 Teacher Confidence on High-Weight Trajectories

The experiments specify standardized hardware, optimization, precision, distillation, mode, and generation settings across GRPO, OPD, and OPSD. Additional analysis shows ReNIO assigns high weights to student–teacher disagreements that remain confidently correctable, while highly errant trajectories receive low weights.

  • C Detailed Experimental Setup: Qwen3 OPD and OPSD training uses Thinking-Mode-off students and Thinking-Mode-on teachers, while evaluation uses Thinking Mode on.Section 2 experiments otherwise follow the Tables 5 and 7 settings, except for maximum completion length.
  • C Detailed Experimental Setup: All experiments use four H200 GPUs with gradient checkpointing, Flash Attention 2, AdamW, and bfloat16 precision.Full-vocabulary logit distillation is applied in every OPD and OPSD experiment.
  • C Detailed Experimental Setup: Training and evaluation configurations for GRPO, OPD, and OPSD are reported separately for Qwen3 models, DS-Distill-Qwen models, and evaluation.These configurations appear in Tables 5, 6, and 7.
  • C Detailed Experimental Setup: For Qwen3-1.7B with Thinking Mode off, Section 2 experiments set maximum generation length to 4096 so rollouts include final answers.Identifying correctness requires rolling out the whole trajectory.
  • D.1 Teacher Confidence on High-Weight Trajectories: ReNIO’s high weights correlate with lower teacher entropy in both OPD and OPSD, indicating sharp correction signals rather than teacher uncertainty.Teacher entropy is used as a confidence proxy, with lower entropy indicating sharper and more reliable token-level guidance.
  • D.1 Teacher Confidence on High-Weight Trajectories: A structured but incorrect trajectory can receive a large weight when the student favors a pivotal token that the teacher assigns low probability.This creates a large student-to-teacher probability ratio while preserving teacher confidence in the local correction.
  • D.1 Teacher Confidence on High-Weight Trajectories: Extremely erroneous trajectories receive small ReNIO weights because flatter, high-entropy teacher distributions do not sharply suppress a particular wrong token.Thus, ReNIO does not strongly amplify trajectories outside the teacher’s reliable reasoning distribution.

D.2 Additional Discussion on ReNIO Weighting

The weighting analysis finds that ReNIO’s student-to-teacher signal is most effective when aggregated into a sample-level weight. Token-level application is ineffective, while reversing the ratio improves over OPSD but remains 1.21 points below ReNIO.

  • Alternative weighting strategies: Under Qwen3-1.7B OPSD, the study compares student-to-teacher token weighting with teacher-to-student sample weighting against ReNIO’s sample-level strategy.The alternatives preserve either ReNIO’s ratio signal or its sample-level form while changing the other component.
  • Token-level weighting: Applying the student-to-teacher signal directly to token losses is ineffective, producing an average score below the OPSD baseline.The result indicates that the ratio should not independently rescale every token loss.
  • Reversed-ratio weighting: 1.21 points: teacher-to-student sample weighting improves over OPSD but still underperforms ReNIO.This variant retains sample-level aggregation while reversing the information signal.

D.3 Additional Ablation Studies

Additional ablations identify robust ReNIO hyperparameters for OPD on Qwen3-1.7B: a clipping bound of 3.0 and key-token threshold of 0.8 yield the best average across three mathematical reasoning benchmarks. The results indicate that ReNIO benefits from emphasizing salient student–teacher disagreements without over-weighting extreme ratios or routine tokens.

  • Experimental Setup: The ablation varies clipping bound and key-token threshold for OPD on Qwen3-1.7B, with per-benchmark results summarized across three mathematical reasoning benchmarks.These experiments test ReNIO’s two main hyperparameters.
  • Hyperparameter Ablations: A clipping bound of 3.0 achieves the best average across three mathematical reasoning benchmarks, balancing useful disagreement against extreme-ratio influence.Small bounds suppress useful disagreement, while loose bounds give extreme ratios too much influence.
  • Hyperparameter Ablations: A key-token threshold of 0.8 achieves the best average, avoiding dilution from routine tokens while retaining informative disagreements.Lower thresholds admit routine tokens; higher thresholds discard informative disagreements.

E Additional Experimental Results · F Qualitative Comparison

Additional experiments show that incorrect on-policy trajectories become increasingly longer and contain more epistemic markers than correct ones, while a qualitative counting example shows ReNIO preserving complementary reasoning steps to reach the correct answer.

  • E Additional Experimental Results: Section 2.2’s full experimental results are provided in Figures 7 and 8.These results cover the additional experimental analysis.
  • E Additional Experimental Results: ReNIO’s per-benchmark hyperparameter ablation results under OPD are reported for Qwen3-1.7B mathematical reasoning.The results are presented in Figure 6.
  • E Additional Experimental Results: Incorrect trajectories are consistently longer than correct trajectories across AIME24, AIME25, and HMMT25.Figure 7 compares average response lengths during training across all three benchmarks.
  • E Additional Experimental Results: The response-length gap widens over training steps, indicating that incorrect trajectories accumulate more deviation tokens.This pattern is shown for AIME24, AIME25, and HMMT25.
  • E Additional Experimental Results: Figure 8 compares the average count of Epistemic Markers in correct and incorrect on-policy trajectories during training.The comparison covers AIME24, AIME25, and HMMT25.
  • F Qualitative Comparison: On a mathematical counting problem, GRPO finds a valid family but misses symmetric placements, whereas OPD handles symmetry but misses duplicate removal.The qualitative comparison contrasts the complementary strengths and omissions of the two methods.
  • F Qualitative Comparison: OPD+ReNIO gives the correct answer by preserving both symmetric-placement handling and duplicate removal.The example shows ReNIO combining the two reasoning steps missed separately by GRPO and OPD.
Loading 2606.23104v1…