PaperlayerView exact revision on arXiv

Source-linked AI summary

When Teacher Guidance Misleads: Reward-Aligned On-Policy Distillation

Siyuan Gan, Yuhan Li, Xiran Wang, Linjian Meng, Boyan Wang, Zhen Zhao, Jing Huo, Yang Gao

arXiv:2608.27960v1cs.AI

TL;DR

OPD can apply teacher guidance that conflicts with outcome reward, potentially steering students away from correct trajectories or toward incorrect ones. RA-OPD checks trajectory-level distillation returns against outcome rewards and filters misaligned trajectories, outperforming tested OPD methods across math and code benchmarks. It reuses existing rollouts and teacher probabilities, but comparisons with Uni-OPD highlight limitations of prior relative-consistency criteria.

  • Problem

    Teacher guidance on student-generated prefixes is not always aligned with outcome reward, which can mislead optimization and degrade performance.

  • Method

    RA-OPD independently checks each trajectory’s distillation return against its outcome reward and filters trajectories with inconsistent guidance.

  • Results

    Across seven math and three code benchmarks, RA-OPD significantly outperforms standard OPD and other tested OPD variants.

  • Takeaways & Limitations

    RA-OPD selects more reliable training signals without additional student rollouts or teacher evaluations, while remaining applicable to single-trajectory and reward-homogeneous prompts.

  • Takeaways & Limitations

    Uni-OPD’s relative-consistency criterion requires multiple rollouts and cannot be established when sampled trajectories share the same reward.

Abstract

from arXiv · show

On-policy distillation (OPD) has recently emerged as a popular post-training paradigm for large language models (LLMs), providing an efficient way to transfer the knowledge and capabilities of teacher models into student models. However, teacher guidance on student-generated prefixes is not always reliable. Training should optimize the model to generate responses that are more likely to be correct, or equivalently, to get higher outcome rewards. But during OPD, the teacher model may provide guidance that discourages the student from moving toward correct trajectories or moves the student toward incorrect ones, which is misaligned with outcome reward. Such misaligned guidance is unreliable, as it would mislead the optimization process and ultimately degrade model performance. To mitigate misaligned teacher guidance, we propose Reward-Aligned On-Policy Distillation (RA-OPD). The key insight is to keep only trajectories whose induced updates move the student toward correct trajectories or discourage the student from moving toward incorrect ones. Specifically, for each sampled trajectory, RA-OPD checks whether its trajectory-level distillation return is consistent with its outcome reward and then filters out the misaligned trajectories. RA-OPD selects more reliable trajectories to improve student model performance without requiring additional computational cost. We evaluate RA-OPD on math and code benchmarks using models from the Qwen3 family and the DeepSeek-R1 family. Across seven math benchmarks and three code benchmarks, RA-OPD significantly outperforms standard OPD and other tested OPD variants.

1 Introduction

OPD efficiently transfers teacher capabilities using student-generated prefixes, but teacher guidance can conflict with outcome reward and mislead optimization. RA-OPD filters such trajectories and improves performance across evaluated benchmarks without additional computational cost.

  • OPD transfers teacher knowledge and capabilities while training on student-generated prefixes under a reverse KL objective.
  • Misaligned teacher guidance can discourage correct trajectories or encourage incorrect ones, degrading OPD optimization and performance.
  • RA-OPD retains trajectories whose trajectory-level distillation returns agree with outcome rewards and filters reward-misaligned trajectories.It reuses existing student trajectories and teacher probabilities, so no additional student rollouts or teacher evaluations are required.
  • RA-OPD makes an absolute reliability decision for each trajectory, including when a prompt has one trajectory or lacks reward diversity.
  • Across seven math and three code benchmarks, RA-OPD significantly outperforms other tested OPD methods.The evaluation uses models from the Qwen3 and DeepSeek-R1 families.

2 Preliminaries

Standard OPD applies teacher guidance to student-generated rollout prefixes, while Uni-OPD aggregates guidance across trajectories and enforces reward-based relative consistency. Uni-OPD’s criterion adds rollout cost and is undefined for reward-homogeneous groups.

  • 2.1 On-Policy Distillation: Standard OPD samples trajectories from the student policy and applies teacher guidance to the resulting student-generated prefixes.The rollout distribution consists of response prefixes visited by current student rollouts.
  • 2.1 On-Policy Distillation: OPD’s guidance is computed from sampled trajectories as a fixed training signal: positive values increase sampled-token probability, while negative values decrease it.
  • 2.2 Uni-OPD: Uni-OPD aggregates teacher guidance over complete trajectories into a trajectory-level distillation return representing the teacher-to-student probability log-ratio.Its sign indicates whether the teacher assigns higher or lower probability to the complete trajectory than the student.
  • 2.2 Uni-OPD: Uni-OPD partitions trajectories by binary outcome reward and compares returns within rollout groups containing both correct and incorrect outcomes.
  • 2.2 Uni-OPD: Uni-OPD enforces a positive safety margin through margin masking, which discards violating groups, or margin shifting, which adjusts their returns.
  • 2.2 Uni-OPD: Uni-OPD requires multiple rollouts per prompt and falls back to original OPD when sampled trajectories share the same reward.These constraints increase rollout overhead and prevent reliable assessment for single-rollout or reward-homogeneous prompt groups.

3 Method

RA-OPD aggregates token-level teacher guidance into a trajectory-level return, checks its consistency with verified outcome rewards, and filters misaligned student trajectories before optimization. It preserves standard OPD updates on kept trajectories while avoiding extra rollouts or teacher evaluations.

  • 3.1 Motivation: RA-OPD aggregates token-level OPD guidance into a trajectory-level distillation return and compares it with each trajectory’s verified outcome reward.The return summarizes teacher guidance over the trajectory, while the outcome verifier supplies the binary reward.
  • 3.2 Details of RA-OPD: Correct trajectories are kept when their distillation return is non-negative, while incorrect trajectories are kept when it is non-positive; conflicting trajectories are filtered out.A zero return is treated as neutral and retained.
  • 3.2 Details of RA-OPD: The reward-alignment mask determines whether a complete trajectory contributes to optimization, without changing token-level OPD rewards within kept trajectories.Positive token-level OPD rewards increase sampled-token log-probability under gradient descent, and negative rewards decrease it.
  • 3.3 Discussion: RA-OPD reuses student trajectories and teacher probabilities already computed for OPD, requiring no additional student rollouts or teacher evaluations.Its absolute per-trajectory criterion remains defined for single-rollout prompts and reward-homogeneous prompt groups.
  • 3.3 Discussion: RA-OPD’s binary consistency mask provides an interpretable decision: filtering reflects either discouraging a correct response or encouraging an incorrect response.The method can also be applied after aggregating effective token-level signals from other OPD-reliability methods.

4 Experiments

RA-OPD is evaluated against standard OPD and other variants across math and code benchmarks, model families, and training settings. It achieves the strongest reported performance while retaining standard-OPD-like efficiency and filtering reward-misaligned trajectories.

  • Qwen3 math results: RA-OPD achieves the strongest average math reasoning performance for both Qwen3-4B-Base and Qwen3-8B-Base students.For Qwen3-4B-Base, average avg@k is 45.88, 5.20 points above standard OPD and 3.52 points above ExOPD; for Qwen3-8B-Base, it is 49.43, 5.09 and 2.48 points higher, respectively.
  • DeepSeek-R1 math results: RA-OPD achieves the strongest overall performance on both avg@k and pass@k for the DeepSeek-R1 family.With Skywork-OR1-Math-7B as teacher and DeepSeek-R1-Distill-Qwen-7B as student, average avg@k is 69.34 and average pass@k is 83.48.
  • Training efficiency: RA-OPD maintains training efficiency comparable to standard OPD and uses less than one third of Uni-OPD’s total training time.OPD, ExOPD, and RA-OPD remain in a similar per-step timing regime, while Uni-OPD is slower because of greater rollout and optimization cost.
  • Filtering dynamics: Reward-misaligned trajectories remain prevalent throughout training in both model families.P(G_i < 0) averages 94.25% in the Qwen3 setting and 92.46% in the DeepSeek-R1 setting, while overall dropped fractions average 48.63% and 68.23%, respectively.
  • Code results: RA-OPD achieves the strongest average performance among the tested OPD methods on code generation benchmarks.Its average score is 58.97, exceeding standard OPD by 3.93 points and ExOPD by 2.87 points; mean pass@4 is 65.17.
  • Ablation study: Masking both correct-negative and incorrect-positive conflicts yields the strongest ablation performance.Training only on conflict trajectories improves over the untrained base model but underperforms standard OPD and RA-OPD, indicating useful information with lower signal-to-noise ratio.

5 Conclusion

The paper identifies reward-misaligned teacher guidance as a reliability problem in OPD and proposes RA-OPD to filter trajectory-level updates inconsistent with outcome accuracy. RA-OPD reuses existing trajectories and improves performance across the evaluated math and code benchmarks.

  • Reward-misaligned teacher guidance can discourage correct trajectories or encourage incorrect ones, misleading OPD optimization and degrading performance.
  • RA-OPD reuses student trajectories and corresponding teacher probabilities, requiring no additional student rollouts or teacher evaluations.
  • RA-OPD significantly outperforms other tested OPD methods across seven math benchmarks and three code benchmarks using Qwen3- and DeepSeek-R1-family models.
  • RA-OPD checks each trajectory’s distillation return against its final outcome accuracy, retaining updates aligned with correct responses or avoiding incorrect ones.
  • Unlike relative-comparison approaches, RA-OPD makes an absolute reliability decision for each trajectory and remains applicable with one trajectory per prompt.

B Additional Experiments

Additional experiments show that misaligned trajectories contain some useful local information but reduce performance when used alone, while masking both conflict types yields the full improvement. RA-OPD also preserves standard OPD’s entropy profile despite filtering many trajectories.

  • Figure 5 compares mean avg@k across seven math benchmarks for RA-C, RA-I, and RA-Inv ablations defined by which conflict types they mask or retain.
  • Misaligned trajectories retain useful token-level information, but using them alone performs below standard OPD and RA-OPD.RA-Inv reaches 43.41 average avg@k, compared with 44.34 for standard OPD and 49.43 for RA-OPD.
  • Masking either correct-negative or incorrect-positive conflicts improves average avg@k, while masking both is necessary for the complete RA-OPD improvement.RA-C reaches 45.31 and RA-I reaches 45.83, compared with 49.43 for complete RA-OPD.
  • Figure 6 presents actor entropy for RA-OPD and standard OPD, with separate panels for Qwen3-8B-Base and DeepSeek-R1-Distill-Qwen-7B.
  • RA-OPD closely tracks standard OPD’s actor-entropy profile across Qwen3 and DeepSeek-family settings despite filtering a large fraction of trajectories.

C Training Parameters

This section distinguishes model-training settings from settings used only for final evaluation.

  • The paper separates settings used for model training from those used only for final evaluation.

C.1 Training-Time Parameters

Standard OPD and RA-OPD use shared on-policy training procedures, while ExOPD and Uni-OPD retain their prescribed baseline-specific settings.

  • For standard OPD and RA-OPD, students sample one on-policy response per prompt, receive teacher scores, and update with the sampled-token k1 OPD objective.
  • Every student is trained for two epochs under the reported training procedure.
  • Baseline-specific exceptions include ExOPD’s reward-scaling factor λ = 1.25 and Uni-OPD’s use of four responses per prompt.

C.2 Evaluation-Time Parameters

Final evaluation freezes each checkpoint and compares methods under identical decoding, without teacher inference or gradient computation. Math and code outputs are judged with task-specific evaluators, using benchmark-specific avg@k and pass@k settings.

  • Checkpoints are evaluated without teacher inference or gradient computation after training.
  • All methods use the same decoding configuration so comparisons reflect learned checkpoints rather than test-time sampling differences.
  • Math answers use math_verify, while code generations use EvalPlus tests or the LiveCodeBench judge.
  • Evaluation reports both avg@k and pass@k, with k = 32 for AIME and AMC, k = 8 for other math benchmarks, and k = 4 for code benchmarks.
  • Table 6 summarizes the sampling parameters used for final evaluation.
  • Qualitative analysis presents four prompts where OPD, ExOPD, and Uni-OPD fail while RA-OPD succeeds.

D.1 The Case of Qwen3-4B-Base Student Model

The Qwen3-4B-Base examples contrast baseline rollouts that abandon valid ingredients or guess with RA-OPD rollouts that connect the relevant geometric identities and reach the reference answer.

  • Baseline rollout: The geometry problem asks for AB · AC given circumradius 13, inradius 6, and IA ⊥ OI; the reference answer is 468.
  • Baseline rollout: The baseline computes OI^2 = 13 and abc = 312s but never combines these with the perpendicularity, instead cycling through incomplete formulas.
  • Baseline rollout: One baseline finishes by selecting AB · AC = 169 without derivation, identifying it only with R^2.
  • Baseline rollout: Another baseline incorrectly asserts A = 90° and obtains 338 from an already incorrect formula.
  • RA-OPD rollout: RA-OPD uses x = bc = AB · AC and the relation ax = 312s to reduce the problem to a one-variable equation.
  • Comparison: RA-OPD reaches x = 468 by connecting perpendicularity, the vertex–incenter distance identity, and the area formula.

D.2 The Case of Qwen3-8B-Base Student Model

The Qwen3-8B-Base examples show baselines guessing or making modular-arithmetic errors, whereas RA-OPD performs a valid congruence lift and reaches the reference answer 110.

  • Problem: The number-theory problem asks for the least m such that p^2 divides m^4 + 1, with reference answer 110.
  • Baseline rollout: Another baseline rejects its own modulo-17 contradiction and assumes m = 130 because the modulo-289 calculation seems tedious.
  • Baseline rollout: A separate baseline makes the false congruence 2401 = 25 · 96 + 1 and consequently selects p = 5 and m = 7.
  • Baseline rollout: One baseline assumes m = 44 without checking the congruence, and the candidate already fails modulo 17.
  • RA-OPD rollout: RA-OPD lifts a valid root modulo 17 by writing m = 8 + 17k and reducing the fourth power modulo 289.
  • Comparison: RA-OPD obtains m = 110 through the congruence lift, unlike baselines that guess unsupported candidates or make direct modular errors.

D.3 The Case of DeepSeek-R1-Distill-Qwen-7B Student Model

The DeepSeek-R1-Distill-Qwen-7B cases compare baseline failures with RA-OPD’s corrected geometric reasoning and code handling. RA-OPD preserves full tangent spans in the math case and handles exclusions, singleton values, and both circular directions in code.

  • D.3 The Case of DeepSeek-R1-Distill-Qwen-7B Student Model: The circle problem asks for m + n for two tangent-circle arrangements, with reference answer 197.
  • D.3 The Case of DeepSeek-R1-Distill-Qwen-7B Student Model: RA-OPD counts tangent-center gaps as 7 · 68 and 2023 · 2, while retaining endpoint offsets in the full-span relations.
  • D.3 The Case of DeepSeek-R1-Distill-Qwen-7B Student Model: RA-OPD derives m + n = 197, correcting the baselines’ factor-of-two or unsupported-scaling errors.
  • D.4 The Case of Qwen3-4B-Base Code Student Model: In the circular-array task, one baseline returns inf for a singleton because it lacks a post-loop infinity check, yielding [2, inf, 3] and a wrong answer.
  • D.4 The Case of Qwen3-4B-Base Code Student Model: RA-OPD skips the queried index, compares both circular directions, maps infinity to −1, and produces the required public behavior.
  • D.4 The Case of Qwen3-4B-Base Code Student Model: On this code prompt, 0 of 16 baseline responses pass versus 2 of 16 RA-OPD responses.
Loading 2608.27960v1…