PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning from Think-Mode Advantage via On-Policy Distillation

Wanqi Ren, Jianxiang Wang, Danxuan Liu, Linyi Ding, Huaixiao Tou

arXiv:2609.37044v2cs.CL

TL;DR

A shared privileged think trace may provide mismatched supervision across student responses, even when their prefixes are on-policy. ThinkOPD routes supervision using reward gain and a trace–response compatibility proxy, improving over Uniform ThinkOPD across mathematics and code in same- and cross-model settings.

  • Problem

    Uniform think-mode OPD reuses one privileged trace across sibling responses, but on-policy prefixes do not ensure that the trace is compatible with each complete response.

    Source: §1 · p. 1, §3.3 · p. 3

  • Method

    ThinkOPD uses reward gain and TRD, a response-level proxy for trace–response compatibility, to route supervision and normalize response weights within each prompt group.

    Source: §1 · pp. 1–2

  • Results

    Across mathematics and code, ThinkOPD improves over Uniform ThinkOPD in both same- and cross-model settings; five-benchmark Overall avg@16 is 32.0 versus 29.0.

    Source: ThinkOPD Outperforms Matched Baselines · p. 5, §5 · p. 7

  • Takeaways & Limitations

    Think-enabled OPD provides a controlled setting for studying how teacher advantage becomes transferable along student responses.

    Source: §5 · p. 7

Abstract

from arXiv · show

Explicit intermediate reasoning gives large language models (LLMs) a stronger problem-solving mode. We study learning from this think-mode advantage via on-policy distillation (OPD). OPD preserves student-generated trajectories and provides dense token-level teacher targets at student-visited prefixes. Privileged reasoning is used during distillation rather than student inference. Uniform ThinkOPD, a natural think-enabled OPD baseline, conditions a fixed teacher on one shared think trace and uniformly distills every sibling student response. Although its prefixes are on-policy, the trace need not follow a route compatible with every complete response: the same privileged trace can induce different teacher-student discrepancies even when responses reach the same outcome. We summarize this interaction with trace-response divergence (TRD) and introduce ThinkOPD, which routes supervision at the response level by combining group-relative reward gain with a TRD-based compatibility proxy. Final response weights are normalized within each rollout group. Across mathematical reasoning and code generation, ThinkOPD outperforms Uniform ThinkOPD in both same-model settings and both cross-model teacher-student pairs, and it exceeds representative rationale and self-distillation baselines in a controlled comparison. Controlled interventions show that outcome benefit and the TRD-based proxy provide complementary routing signals in this setting. Think-enabled OPD provides a controlled setting for studying how teacher advantage becomes transferable along student responses.

1 Introduction

Think-enabled on-policy distillation pairs student exploration with privileged teacher reasoning, but a shared teacher trace may fit sibling responses differently. ThinkOPD routes supervision using outcome benefit and trace–response compatibility, improving over uniform routing in the reported settings.

1 Introduction

  • Think-enabled OPD pairs a no-think student’s own trajectories with privileged supervision from a think-mode teacher during distillation.Student inference does not condition on the teacher’s privileged think trace.Source: §1 · p. 1
  • A shared privileged trace can suit one student response yet conflict with another, even when both responses reach the same outcome.The student-scored prefixes are on-policy, but the trace follows the teacher’s solution route.Source: §1 · p. 1
  • The paper frames think-enabled OPD as a controlled setting for studying how outcome benefit and trace–response compatibility interact during transfer.Its contributions define TRD over complete student paths and introduce a group-relative router combining reward gain with a TRD-based proxy.Source: §1 · pp. 1–2
  • ThinkOPD combines group-relative reward gain with a TRD-based compatibility proxy to route response-level supervision, then normalizes weights within each prompt group.Uniform ThinkOPD instead weights sibling responses equally under the same teacher view.Source: §1 · p. 2, Figure 1 · p. 2
  • Across both same-model settings, ThinkOPD improves five-benchmark Overall avg@16 over Uniform ThinkOPD by 1.7–3.0 percentage points.Controlled interventions also find that TRD-only routing falls below uniform weighting, reward-only routing remains below the full method, and reversing the TRD preference is harmful.Source: §1 · p. 2

2 Related Work

Prior transfer approaches differ in whether they use fixed teacher targets, task feedback, or teacher distributions at student-visited states. ThinkOPD fixes the teacher view and routes a shared trace across sibling responses, separating response routing from teacher evolution and rationale generation.

2 Related Work

  • Sequence and rationale distillation use fixed teacher targets, reinforcement learning uses task feedback, and OPD aligns teacher distributions to student trajectories (Hinton, Vinyals, and Dean 2015; Kim and Rush 2016; Hsieh et al. 2023; Ho, Schmid, and Yun 2023; Yu et al. 2024; Xu et al. 2025a; Yu et al. 2025; Agarwal et al. 2024; Gu et al. 2024).Speculative KD interleaves student proposals and teacher corrections, while OPCD adds context-conditioned teachers (Agarwal et al. 2024; Gu et al. 2024; Xu et al. 2025b; Ye et al. 2026).Source: §2 · p. 2
  • The matched controls span demonstration-conditioned on-policy SDFT, fixed privileged OPSD, and feedback-conditioned SDPO (Shenfeld et al. 2026; Zhao et al. 2026; Hübotter et al. 2026).ThinkOPD holds the teacher view fixed and routes one trace across sibling responses, separating response routing from teacher evolution and rationale generation.Source: §2 · p. 2
  • Selective supervision adapts local targets using entropy, teachability, position, or near-future guidance, while BRTS searches teacher trajectories (Jin et al. 2026; Wang et al. 2026; Liu et al. 2026; Jiang and Ferraro 2026; Zhang et al. 2026).ThinkOPD instead aggregates discrepancy across each response and routes a shared trace by outcome benefit and compatibility.Source: §2 · pp. 2–3

3 Method

ThinkOPD routes dense teacher supervision across student responses using both outcome benefit and trace–response compatibility. It retains student-generated trajectories while making the routing decision at the complete-response level.

3 Method

  • ThinkOPD treats each shared-trace/student-response pair as the routing unit, combining response-level trace compatibility with outcome benefit to allocate dense OPD targets.Source: §3 · p. 3

3.1 Problem Formulation

  • A no-think student samples sibling responses while a fixed think-mode teacher scores their visited prefixes under a shared trace; sparse forward KL compares their token distributions.Scoring covers valid generated-token positions across each complete response, excluding prompt and padding tokens.Source: §3.1 · p. 3

3.2 Uniform ThinkOPD

  • Uniform ThinkOPD gives every trace–response pair equal weight, although on-policy coverage applies to student prefixes, not necessarily to the teacher’s shared solution route.Source: §3.2 · p. 3, §3.3 · p. 3

3.3 Trace–Response Divergence

  • Figure 2 shows lower TRD for the route-aligned response than the alternative correct route, despite the same answer and reward; TRD aggregates persistent disagreement across a complete response.Its discounted continuation distinguishes sustained route-level departures from temporary token-level discrepancies.Source: Figure 2 · p. 3, §3.3 · pp. 3–4

3.4 Group-Relative Response Routing

  • ThinkOPD standardizes reward gain and TRD among siblings, increasing response weights with relative benefit and decreasing them with relative TRD before within-group normalization.The smooth TRD factor has no hard threshold, and its contrast is controlled by η.Source: §3.4 · p. 4
  • Figure 3’s weight map assigns the largest coefficients where positive reward advantage coincides with low TRD; at comparable benefit, weights fall as TRD rises.The figure bins 768 responses from 192 four-response Qwen3-1.7B prompt groups by reward advantage and TRD.Source: §3.4 · p. 4, Figure 3 · p. 4

3.5 Training Objective and Procedure

  • Each response coefficient is shared across its valid generated-token positions and held fixed during the update; setting every coefficient to one recovers Uniform ThinkOPD.Source: §3.5 · p. 4
  • Response-level routing changes each path’s aggregate contribution while retaining dense token-level supervision, with current student paths receiving recomputed response-conditioned targets.Source: §3.5 · p. 4

4 Experiments

The experiments evaluate ThinkOPD against matched routing and external baselines across mathematical reasoning and code generation. ThinkOPD improves benchmark results, with controlled interventions supporting the complementary roles of outcome benefit and TRD-based compatibility.

4.1 Settings

  • Evaluation covers mathematical reasoning and code generation with disjoint training and evaluation data; external baselines use matched data, response budgets, checkpoint selection, and evaluation.The Qwen3-1.7B math comparison averages three independent runs; other model-task settings are matched validation-selected point estimates.Source: §4.1 · p. 5

4.2 Results and Analysis

  • 3.0 points: ThinkOPD raises Overall avg@16 from 29.0 to 32.0, above BRTS at 29.7 on Qwen3-1.7B.The comparison uses matched teacher, trace bank, student rollouts, and optimization protocol.Source: ThinkOPD Outperforms Matched Baselines · p. 5, Table 2 · p. 5
  • ThinkOPD gains on AIME24, AMC23, and Skywork while remaining competitive on AIME25 and LCBv5, rather than concentrating improvements on individual benchmarks.The Qwen3-1.7B result raises LCBv5 best@16 by 13.3 points.Source: ThinkOPD Outperforms Matched Baselines · p. 6, Gains Across Models and Tasks · p. 6
  • Across both same-model scales, ThinkOPD improves Overall by 1.7–3.0 points; across both cross-model pairs, pooled math avg@16 improves by 1.6–1.9 points.On Qwen3-1.7B, LCBv5 best@16 also rises by 13.3 points.Source: Gains Across Models and Tasks · p. 6, Table 4 · p. 6
  • For Qwen3-4B→1.7B, routing reaches 53.7 versus 51.8 for uniform; for Qwen3-1.7B→0.6B, it reaches 22.6 versus 21.0.The first pair has equal no-think OPD and uniform scores at 51.8; in the second pair, uniform rises from 15.0 to 21.0.Source: Gains Across Models and Tasks · p. 6, Table 4 · p. 6
  • In the controlled ablation, full ThinkOPD scores 48.0, while reversed TRD scores 42.5 below uniform routing at 44.8.Reward-only routing improves over uniform but remains below full ThinkOPD; removing reward gain also falls below uniform.Source: Benefit and Compatibility · p. 6, Table 5 · p. 7
  • Across tested future horizons, performance stays above uniform routing; routing strength η has an interior optimum at 0.5.Stronger η can suppress responses that combine high TRD with substantial outcome benefit.Source: Figure 4 · p. 7, Benefit and Compatibility · p. 7
  • ThinkOPD’s recurring cost is 1.31 versus 1.00 for standard OPD and 1.28 for Uniform ThinkOPD, while scoring 2.3 points higher than BRTS-4 with 37.6% less recurring time.BRTS-4 has recurring cost 2.10; ThinkOPD reuses rewards and token-level divergences already computed for distillation.Source: Figure 5 · p. 7, Response Routing Improves Supervision… · p. 7

5 Conclusion

Think-enabled OPD preserves student exploration but does not ensure a shared privileged trace suits every sibling response. ThinkOPD combines reward gain and TRD-based compatibility in group-normalized routing, improving over Uniform ThinkOPD across same- and cross-model settings.

5 Conclusion

  • ThinkOPD improves over Uniform ThinkOPD across mathematics and code in both same-model and both cross-model settings.The controlled ablation favors combining reward gain and the TRD-based proxy; reusing existing teacher signals keeps the router practical.Source: §5 · p. 7
  • Group-normalized routing combines reward gain for outcome benefit with TRD as a proxy for trace–response compatibility.ThinkOPD routes supervision across sibling responses while preserving student exploration.Source: §5 · p. 7
  • Think-enabled OPD provides a controlled setting for studying route-dependent compatibility with a fixed teacher across student paths.Source: §5 · p. 7

Related papers

Loading 2609.37044v2…