Source-linked AI summary

TeachBack: Learning to Teach via Bidirectional Preference Optimization for On-Policy Self-Distillation

Linfei Ji, Jiayun Li, Haozhao Wang, Siyu Jiang, Zehao Wang, Longtao Xiao, Yu-Jie Yuan, Ruixuan Li

arXiv:2610.29167v1cs.AIcs.CLcs.LG

TL;DR

OPSD can yield accurate privileged Teacher trajectories that are not necessarily suitable for a problem-only Student. TeachBack shapes Teacher generation with partial reference reasoning and bidirectional preferences, and improves Student performance across mathematical and logical reasoning tasks over existing OPSD and related baselines.

  • Problem

    Privileged conditioning in OPSD can shorten responses and degrade native reasoning behaviors, while accurate Teacher trajectories are not necessarily suitable for Student learning.

    Source: §1 · pp. 1–2
  • Method

    TeachBack varies reference-suffix length and prompt style, then uses paired Teacher and Student trajectories in bidirectional preference learning to shape both Student transfer and future Teacher supervision.

    Source: §1 · pp. 2–3
  • Results

    TeachBack consistently improves Student reasoning performance over existing OPSD and related baselines on mathematical and logical reasoning tasks.

    Source: §1 · p. 3
  • Takeaways & Limitations

    A useful privileged Teacher trajectory should be reliable and rich enough for a problem-only Student to learn from.

    Source: §1 · p. 2

Abstract

from arXiv · show

On-policy self-distillation (OPSD) enables a language model to learn from a Teacher instantiated from the same model but given access to privileged information, such as a reference solution. We find that access to sufficient privileged information allows the Teacher to generate highly accurate trajectories; however, accuracy alone does not guarantee that these trajectories are suitable for Student learning. Direct access to a complete solution may encourage shortened, reference-conditioned derivations that omit the exploration, verification, and self-correction required to solve problems without privileged information. We introduce TeachBack, a bidirectional preference-learning framework for on-policy self-distillation that learns not only from the Teacher but also how the Teacher should teach. TeachBack shapes the Teacher's generation behavior by controlling the length of the reference-reasoning suffix injected into the prompt and by varying the prompt style. While maintaining comparable Teacher accuracy, it encourages the Teacher to explore and re-solve problems, thereby reducing its reliance on the reference solution. At each rollout refresh, the shared policy generates both a privileged Teacher trajectory and a problem-only Student trajectory. A pairwise preference objective uses privileged Teacher trajectories to improve Student learning, while the reversed preference uses the Student's behavior to shape future Teacher supervision under the privileged condition. Experiments on mathematical and logical reasoning tasks show that TeachBack consistently improves Student performance over existing OPSD methods and related baselines.

1 INTRODUCTION

OPSD can produce shorter reasoning that suppresses exploration and verification, and existing approaches leave open how to construct a PI-conditioned Teacher whose reasoning remains useful to a problem-only Student. TeachBack shapes Teacher trajectories and uses Student feedback to improve their suitability as supervision.

1 INTRODUCTION

  • Existing OPSD remedies regulate how the Student uses an already generated Teacher signal, leaving Teacher construction as an upstream question.The paper frames the open problem as using privileged information to guide reasoning without making it an authoritative solution to copy.Source: §1 · p. 2
  • A useful PI-conditioned Teacher should remain accurate while relying less on the reference and producing exploratory reasoning reproducible without PI.It should re-solve the problem, verify its path, and produce a complete trajectory that a Student can learn from.Source: §1 · p. 2
  • TeachBack presents partial reference reasoning as optional working notes and varies prompt style so the Teacher re-solves rather than simply follows the reference.The approach targets Teacher construction and transfer of privileged trajectories.Source: §1 · p. 2
  • Teacher-origin trajectories guide Student learning, while Student-origin trajectories provide feedback that shapes future Teacher supervision.Experiments on mathematical and logical reasoning tasks show consistent Student-performance improvements over existing OPSD and related baselines.Source: §1 · p. 3

2 RELATED WORK

Prior work addresses thinking collapse by modifying optimization, the Teacher signal, or how the Student learns from it. TeachBack instead treats construction of a Student-compatible Teacher as an upstream problem.

2 RELATED WORK

  • Existing OPSD approaches anchor updates, correct reference-induced style shifts, remove reference-only signal components, or use preferences for Student learning.These methods address degradation through optimization, supervision-signal changes, or Student-side learning choices.Source: §2 · p. 3
  • Prior methods filter or transfer an already formed Teacher signal, leaving construction of a Student-compatible Teacher unresolved.The related work identifies this as an upstream problem.Source: §2 · p. 3
  • TeachBack uses incomplete working notes, prompt framing, and cross-condition preferences to shape Teacher trajectories and transfer useful ones to the Student.The Student-side objective transfers Teacher trajectories under the problem-only condition; the reversed objective shapes future privileged Teacher trajectories.Source: §2 · p. 3

3 METHOD

TeachBack uses one shared policy for a problem-only Student and a PI-conditioned Teacher, generates paired trajectories, and updates the policy with preferences in both directions. Teacher construction, cross-scoring, and the balance between the two losses determine how supervision is generated and transferred.

3 METHOD

  • The shared policy acts as a problem-only Student or a PI-conditioned Teacher, and training pairs their trajectories before cross-condition preference updates.At test time, only the problem-only Student condition is used.Source: §3 · pp. 3–4, Figure 2 · p. 4

3.1 TEACHER CONSTRUCTION: REFERENCE-SUFFIX INJECTION AND PROMPT-STYLE CONTROL

  • The Teacher receives a variable-length reference suffix as optional, incomplete working notes, encouraging it to reconstruct the derivation, verify the notes, and produce a complete answer.Suffix length controls available privileged information, while prompt style changes how it is presented; the main runs use a suffix from the dataset reference.Source: §3.1 · p. 4

3.2 PAIRED ON-POLICY ROLLOUTS AND CROSS-SCORING

  • At each refresh, the fixed policy generates paired responses, which are replayed under both prompts for cross-scoring.Preference training retains the first K valid response tokens, focusing on solution direction, exploration, and the start of verification.Source: §3.2 · pp. 4–5

3.3 BIDIRECTIONAL PREFERENCE OPTIMIZATION

  • Under the problem-only condition, Teacher-origin responses are preferred; under the PI-conditioned condition, Student-origin responses feed compatible reasoning back into future Teacher targets.The reversed preference does not assume Student answers are more accurate and may reinforce errors or unproductive verbosity.Source: §3.3 · p. 5

3.4 JOINT CO-ADAPTATION AND TRAINING PROTOCOL

  • The weight λ balances Teacher-to-Student transfer and Teacher-side feedback; each refresh generates paired rollouts, computes both losses, and updates shared parameters.λ = 0 isolates Student transfer, λ = 1 isolates Teacher-side feedback, and intermediate values test co-adaptation; verifier outcomes do not filter pairs or rescale losses.Source: §3.4 · p. 6

4 EXPERIMENTAL PROTOCOL

The experiments evaluate TeachBack across mathematical and logical reasoning, then examine how Teacher construction and bidirectional training affect trajectory quality, preference optimization, and performance.

4.1 EXPERIMENTAL SETUP

  • The study trains three Qwen3 scales on DeepMath and evaluates math benchmarks and Knight-and-Knaves, comparing TeachBack with Base, GRPO, OPSD, SPIN, and PBSD.SPIN prefers dataset ground-truth reasoning traces; PBSD uses frozen privileged-Teacher trajectories; TeachBack generates both trajectories from the current shared policy and optimizes preferences under both prompt conditions.Source: §4.1 · p. 6, §4 · p. 6

4.2 MAIN RESULTS

  • TeachBack achieves the highest Math and KK averages at all three Qwen3 scales, improving over Base by 4.3/2.2/2.5 Math points and 6.1/3.9/7.5 KK points.The gains correspond to Qwen3-1.7B/4B/8B, respectively.Source: §4.2 · p. 6, Table 1 · p. 7

4.3 HOW PROMPT CONSTRUCTION CHANGES TEACHER TRAJECTORIES

  • Notes framing the reference suffix as optional evidence yields more exploration and less reference use than authoritative Reference framing; 25–50% suffix injection retains high accuracy versus longer suffixes.Notes has the highest average mean@4 at 74.21%, while Answer-only has the highest best mean@4 at 76.49%; Notes and Reference are close on both downstream metrics.Source: §4.3 · pp. 7–8, Figure 3 · p. 7

4.4 LOSS VALIDITY AND OPTIMIZATION STABILITY

  • At λ = 0.75, both preference margins remain positive, but the Teacher-side margin grows more, while diagnostic KL rises and the four prompt–response entropy levels differ.The Teacher-origin response’s score falls under the privileged prompt, so the stronger Teacher-side margin reflects stronger optimization of that preference.Source: §4.4 · p. 8, Figure 4 · p. 8

4.5 TEACHER-SIDE WEIGHT AND TRAINING DYNAMICS

  • With λ = 0.75 or 1, Student accuracy generally rises while Teacher accuracy stays near 94%–99%; final held-out Math mean@12 reaches 53.99% and 54.02%, respectively.At λ = 0 and 0.25, final held-out Math mean@12 is 45.44% and 47.45%; the 0.75 run has the highest validation peak and six-checkpoint average.Source: §4.5 · p. 9, Figure 5 · p. 9

5 CONCLUSION

TeachBack transfers privileged Teacher trajectories to a problem-only Student while using Student behavior to shape future Teacher rollouts. The conclusion reports improved aggregate mathematical and logical reasoning across tested model scales.

5 CONCLUSION

  • TeachBack combines Student learning from privileged Teacher trajectories with Student feedback that shapes future Teacher rollouts.The conclusion also reports that partial reasoning as optional working notes maintains high Teacher accuracy with less measured reference dependence than authoritative reference framing.Source: §5 · p. 9
  • The authors report using generative AI tools for theoretical proofs, baseline implementation, manuscript translation and editing, and literature identification.They state that they take responsibility for the final content.Source: AI USE STATEMENT · p. 10
  • The reproducibility statement points to appendices for algorithm details, prompt templates, experimental protocols, additional results, diagnostics, assumptions, and proofs.Section 3 describes the conditions, paired rollouts, and objectives; Section 4 documents experiments and evaluation.Source: REPRODUCIBILITY STATEMENT · p. 10

A PROMPT-CONSTRUCTION AUDIT AND IMPLEMENTATION DETAILS

The appendix audits how prompt construction affects Teacher behavior and validates prompt-style effects in fixed-Teacher OPSD. It also documents diagnostic definitions, prompt templates, and the bidirectional co-adaptation implementation.

A.1 PROMPT-AUDIT SETUP AND COMPLETE RESULTS

  • The prompt audit compares eight Teacher conditions and a problem-only Student baseline; Reference, Working Notes, and History use the same final 50% reasoning body.The audit uses Qwen3-4B on 256 DeepMath problems, with four rollouts per problem; Answer-only supplies no reasoning body, while Instruction supplies full reasoning.Source: §A.1 · p. 12, Table 2 · p. 12

A.2 DIAGNOSTIC DEFINITIONS

  • The scorecard measures verifier-based accuracy, response length, truncation, exploration density, and breadth across five reasoning-behavior categories.Exploration density is length-normalized per rollout, and breadth counts how many categories appear in a response.Source: §A.2 · p. 13
  • Reference dependence combines output-side copying of problem-exclusive privileged n-grams with explicit-mention scoring; strong attribution is reported separately.Answer-only content reference measures exact normalized answer agreement, so it is not directly comparable with reasoning-injection conditions.Source: §A.2 · pp. 13–14

A.3 DOWNSTREAM VALIDATION WITH FIXED-TEACHER OPSD

  • 74.21% Notes and 74.10% Answer-only average mean@4, versus 74.08% Reference and 73.32% History across steps 50–300.Notes has the highest average mean@4; Student entropy rises later under Reference and History but remains more controlled under Notes and Answer-only.Source: Table 4 · p. 14, §A.3 · pp. 14–15

A.2 DIAGNOSTIC DEFINITIONS

  • Exploration and reference metrics are surface-text proxies that can overcount intended behaviors or miss paraphrased use of privileged information.The lexical inventory groups exploration into five behavioral categories and distinguishes broad mention from strong attribution; these are observable wording, not latent states.Source: §A.2 · p. 14, Table 3 · p. 14

A.4 EXACT PROMPT TEMPLATES

  • The prompt audit removes the special </think> token from ground-truth reasoning because its inclusion was associated with premature termination in Qwen3 training runs.A boxed-answer notation had a smaller observed effect on reasoning continuation.Source: §A.4 · p. 15
  • Working Notes injects a selectable 25–100% ground-truth reasoning suffix as optional notes, while Reference and History frame or place the same 50% suffix differently.Answer-only provides the final answer without a reasoning body; Instruction provides the full solution and answer.Source: §A.4 · p. 15, §A.4.2 · p. 15, §A.4.3 · p. 16

B BIDIRECTIONAL CO-ADAPTATION ALGORITHM

  • Each co-adaptation step generates paired Teacher and Student responses from a shared policy snapshot, then updates the shared policy with Student-side and Teacher-side preference losses.Verifier outcomes are logged as diagnostics and do not filter preference pairs or gate either loss term.Source: B BIDIRECTIONAL CO-ADAPTATION ALGORITHM · p. 17

C.1 DATASETS AND EVALUATION SPLITS

The appendix specifies the mathematical and logical reasoning datasets, held-out splits, and shared evaluation protocol used for the main comparisons.

C.1 DATASETS AND EVALUATION SPLITS

  • Mathematical training uses DeepMath-103K with model-specific difficulty levels, and evaluation uses AMC 2023, AIME 2024, and AIME 2025.Qwen3-1.7B uses levels 5–7, Qwen3-4B levels 6–8, and Qwen3-8B levels 7–10.Source: §C.1 · p. 17
  • Knight-and-Knaves training and in-domain testing use 4–8-role problems, while out-of-domain testing uses 9-, 10-, and 11-role problems.Source: §C.1 · p. 17

C.2 PROTOCOL PARAMETERS

  • All methods share the held-out evaluation protocol; main Math results use mean@12, while KK uses greedy pass@1.The four-sample prompt audits and controlled diagnostics are separate experiments.Source: §C.2 · p. 17, Table 5 · p. 18

C.3 METHOD-SPECIFIC TRAINING SETTINGS

The appendix gives completion caps for the main runs and fixed-Teacher experiments, and explains why Math and KK use different Teacher-side loss weights.

C.3 METHOD-SPECIFIC TRAINING SETTINGS

  • Main TeachBack runs and fixed-Teacher OPSD prompt-transfer experiments use a 4,096-token completion cap; the Teacher-side weight sweep uses 8,192 tokens.Source: §C.3 · p. 18

C.4 TEACHER-SIDE WEIGHTS

  • Math uses λ = 0.75 for a larger privileged-role share, while KK uses λ = 0.25 to emphasize Student-side transfer and limit adaptation to repeated surface patterns.The paper motivates the distinction by Math's diverse solution paths and KK's more structured, repeatable reasoning traces.Source: §C.4 · p. 18

D CROSS-ARCHITECTURE MATH RESULTS

The paper evaluates TeachBack beyond the Qwen3 family on two additional model groups using the same mathematical reasoning evaluation convention.

D CROSS-ARCHITECTURE MATH RESULTS

  • DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B are evaluated on AMC 2023, AIME 2024, and AIME 2025.The passage states that the values are reported in Table 7 but does not provide them.Source: D CROSS-ARCHITECTURE MATH RESULTS · p. 18

E TEACHER ROLLOUT EVOLUTION ACROSS CHECKPOINTS

Across checkpoints, Teacher responses grow longer and more exploratory while measured reference overlap declines; lexical markers also show more uncertainty and backtracking.

E TEACHER ROLLOUT EVOLUTION ACROSS CHECKPOINTS

F THEORETICAL ANALYSIS OF BIDIRECTIONAL PREFERENCE OPTIMIZATION

The analysis gives an exact interpretation of the two reference-relative preference margins and characterizes the scope of that identity. It applies to fixed response pairs and does not imply positive margins for every pair.

F THEORETICAL ANALYSIS OF BIDIRECTIONAL PREFERENCE OPTIMIZATION

F.1 SETUP AND SCOPE

  • The margins compare Teacher- and Student-origin responses under their respective prompts, with positive values indicating preference relative to the reference policy rather than raw sequence-probability ordering.The analysis fixes sampled tokens, prompts, and response masks at a rollout refresh; derivatives are not through the sampling distribution across refreshes.Source: §F.1 · p. 20

F.2 AN EXACT INTERPRETATION OF THE TWO MARGINS

  • When both margins are positive, the privileged-context advantage of the Teacher-origin response over the Student-origin response is smaller than under the reference policy.The identity holds for fixed response pairs, including unequal-length pairs, with each response using its own length normalizer in both contexts.Source: §F.2 · p. 20
  • The sum of the two margins measures a reduction in the relative conditional advantage defined by the exact identity, but positive mean margins do not guarantee positive margins on every pair.The identity also holds in expectation under a fixed pair distribution for which the scores are integrable; it neither requires matching conditional distributions nor bounds their KL divergence.Source: §F.2 · p. 20

F.3 LOCAL COMPATIBILITY OF THE SHARED-PARAMETER UPDATES

A weighted combination of the Student and Teacher gradients can locally decrease both losses under stated conditions. The analysis does not guarantee improvement across rollout refreshes or for the optimizer used in practice.

F.3 LOCAL COMPATIBILITY OF THE SHARED-PARAMETER UPDATES

  • When the gradients are nonzero and have nonnegative inner product, every interior mixing weight gives strict descent in both losses for sufficiently small positive step sizes.For a negative inner product, both first-order changes are negative only within the specified weight interval, which is nonempty exactly when k^2 < ab.Source: §F.3 · p. 21
  • The local result analyzes ordinary gradient descent, not the full AdamW update, and guarantees neither a globally optimal λ nor improvement across rollout refreshes.With momentum, preconditioning, clipping, or weight decay, the relevant first-order diagnostic is the inner product of each branch gradient with the actual update.Source: §F.3 · p. 21
Loading 2610.29167v1…