Source-linked AI summary
SAPD: Step-Aligned Privileged Distillation
Tianle Wang, Jiayu Liu, Ruizhi Zhao, Ning Miao
TL;DR
On-policy post-training learns from student-generated trajectories but incurs costly rollout generation, raising whether fixed demonstrations can support competitive learning with better supervision. SAPD uses step-aligned privileged distributions on fixed trajectories; it outperforms supervised baselines on average, remains competitive with on-policy methods, and largely preserves coding performance while speeding training loops.
Problem
The paper asks whether fixed demonstrations can support competitive post-training without student rollouts, and whether supervision can provide informative preferences aligned with the reasoning decision being learned.
Source: §1 · p. 1Method
SAPD uses the reference solution's progression to align privileged, context-dependent distributional guidance with each reasoning transition on fixed trajectories.
Source: §1 · p. 2, §3 · pp. 3–4Results
SAPD outperforms SFT and label-smoothed SFT on average across mathematical reasoning benchmarks, remains competitive with GRPO and OPSD, largely preserves coding performance, and achieves approximately 2× training-loop speedups over on-policy baselines.
Source: §6 · p. 9Takeaways & Limitations
The findings suggest that carefully constructed supervision can make fixed demonstrations a competitive and efficient basis for reasoning post-training.
Source: §6 · p. 9
Abstract
from arXiv · showhide
On-policy post-training can improve large language models by learning from their own trajectories, but requires costly rollout generation. We ask whether fixed demonstrations can support competitive off-policy learning through better supervision. Our premise is that their usefulness depends not only on the training trajectories, but also on whether supervision provides informative preferences among continuations and connects this guidance to the reasoning decision being learned. We introduce Step-Aligned Privileged Distillation (SAPD), a rollout-free self-distillation method that turns demonstrations into step-aligned distributional supervision. Its key insight is to use the known progression of a reference solution to associate each reasoning transition with targeted privileged guidance, rather than treating the solution as undifferentiated context. On mathematical reasoning benchmarks, SAPD outperforms supervised fine-tuning and label smoothing on average while remaining competitive with on-policy reinforcement learning and self-distillation. Analyses support both the value of context-dependent distributional guidance and the benefit of aligning privileged information with the current step. SAPD also largely preserves out-of-domain coding performance and achieves approximately 2x training-loop speedups over the on-policy baselines. These findings suggest that carefully constructed supervision can make fully off-policy post-training a competitive and computationally efficient alternative. Our code is available at https://github.com/Miaow-Lab/SAPD.
1 INTRODUCTION
The introduction asks whether fixed demonstrations can support competitive post-training without costly student rollouts, and whether richer, step-aligned supervision can make them more useful. SAPD combines distributional guidance with step alignment and outperforms supervised baselines on average while remaining competitive with on-policy methods.
1 INTRODUCTION
- Fixed demonstrations could avoid costly student rollouts, but their supervision may need to be richer and more targeted to support competitive post-training.The introduction frames the question around distributional preferences among continuations and whether guidance should be localized to the reasoning step being learned.Source: §1 · p. 1
- One-hot imitation does not distinguish plausible alternative expressions from incompatible continuations, whereas predictive distributions can convey relative support beyond the demonstrated token.The paper hypothesizes that demonstration-informed distributions provide context-dependent preferences and uncertainty beyond hard targets or uniform smoothing.Source: §1 · pp. 1–2
- Using the full solution as privileged context leaves the teacher to determine which parts guide each continuation; the known step progression can instead align guidance to each local decision.The authors hypothesize this explicit correspondence provides more focused guidance than using the same global solution context throughout the trace.Source: §1 · p. 2
- SAPD associates each reference transition with corresponding privileged guidance, then distills context-dependent distributions on fixed reference trajectories without generating new training trajectories.Distributional supervision determines what preferences are transferred; step alignment determines which demonstrated information informs them.Source: §1 · p. 2
- SAPD outperforms SFT and label-smoothed SFT on average across three mathematical reasoning benchmarks while remaining competitive with GRPO and OPSD.It largely preserves out-of-domain coding performance and reports approximately 2× training-loop speedups over on-policy baselines.Source: §1 · p. 2
2 PRELIMINARIES
The preliminaries define the student and reference-solution setting, then contrast fixed-trajectory supervision with on-policy reward optimization and self-distillation. They also position SDFT as sharing the core paradigm of demonstration-conditioned on-policy distillation.
2 PRELIMINARIES
- The setup uses a dataset of problem-reference-solution pairs and an autoregressive student, with θ and ϕ denoting student and teacher parameters and t a token position.For a context and temperature T > 0, the notation distinguishes predictive distributions S and T.Source: §2 · p. 2
- SFT trains by teacher forcing on fixed reference trajectories, while label smoothing softens the target but assigns equal probability to every alternative token.Label smoothing replaces the one-hot target with a softened distribution and minimizes cross-entropy against it.Source: §2 · p. 2
- GRPO samples responses from the rollout policy, uses group-normalized verifiable rewards, and optimizes a clipped policy-gradient objective, optionally with a reference-policy KL penalty.Under outcome supervision, all tokens in a response share the same advantage; GRPO uses no learned value model.Source: §2 · p. 3
- OPSD distills a demonstration-conditioned teacher's distribution along student-generated trajectories; SDFT shares this core on-policy distillation paradigm.The cited OPSD objective is forward KL, omitting divergence clipping.Source: §2 · p. 3
3 STEP-ALIGNED PRIVILEGED DISTILLATION
SAPD keeps learning on fixed reference-prefix contexts while using the solution’s step structure to select privileged information for distributional targets. Teacher-forced training avoids autoregressive rollout generation.
3 STEP-ALIGNED PRIVILEGED DISTILLATION
- SAPD uses the reference solution’s step structure to pair each reference transition with privileged guidance, then distills context-dependent distributions on fixed trajectories.The reference trajectory supplies both training contexts and the structure that selects privileged information.Source: §3 · pp. 3–4
- Student prediction contexts use reference prefixes, with each solution partitioned into contiguous segments that preserve the original token sequence.The formulation allows different segmentation granularities, including an unsegmented solution.Source: §3 · p. 3
- For each token in a reference step, the teacher conditions on that corresponding step, while the student retains the complete reference prefix as its causal history.The step boundary limits teacher-side privileged information; it does not reset the student’s history.Source: §3 · p. 3
- A token-mean forward KL loss trains the student against aligned targets; both branches use teacher forcing, so training requires no autoregressive rollouts.The teacher is frozen at the initial checkpoint throughout training.Source: §3 · p. 4
4 EXPERIMENTS
The experiments evaluate SAPD’s mathematical reasoning, coding capability retention, supervision distributions, step alignment, and training cost. SAPD leads average mathematical accuracy while largely preserving coding performance; analyses support context-dependent guidance aligned with the current reasoning step.
4.1 EXPERIMENT SETUP
- Experiments use Qwen3-4B-Instruct-2507, train on OpenThoughts mathematics, and evaluate on three math benchmarks plus three coding benchmarks.The baselines are SFT, LS-SFT, GRPO, and OPSD; avg@12 is reported, with checkpoints selected by mean in-domain avg@k.Source: §4.1 · pp. 4–5
4.2 OVERALL PERFORMANCE AND CAPABILITY GENERALIZATION
- 47.87% average mathematical accuracy makes SAPD the strongest evaluated method; its coding OOD average is 56.68%, above SFT and LS-SFT and close to GRPO and OPSD.SAPD also produces the shortest average outputs among post-trained methods, at 9,927 tokens.Source: Table 1 · p. 5, §4.2 · pp. 5–6
4.3 BEYOND LABEL SMOOTHING: CONTEXT-DEPENDENT SUPERVISION
- In the equivalent-expression case comparing 44 = 256 with 64, SAPD assigns 0.395 to larger and 0.251 to bigger, but only 6.8 × 10−6 to contradictory less.When ten sparrows each peck x grains and their total exceeds 1100, the correct relation is 10x > 1100; SAPD assigns approximately unit probability to >. LS-SFT has reference probability 0.900 and entropy 1.518 in both examples.Source: §4.3 · p. 6
4.4 VALUE OF STEP-ALIGNED PRIVILEGED INFORMATION
- Full-solution guidance lowers average accuracy from 47.87% to 45.00%, while shuffled step guidance reaches 45.65%, 2.22 pp below SAPD.The ablations retain the off-policy reference contexts and objective while changing the teacher-side privileged information.Source: Table 2 · p. 7, §4.4 · p. 8
4.3 BEYOND LABEL SMOOTHING: CONTEXT-DEPENDENT SUPERVISION
- SAPD’s conditional non-reference KL decreases from approximately 1.07 to 1.02, while SFT rises to 1.47 and LS-SFT to 2.20.Full KL instead falls more under SFT than SAPD, illustrating that aggregate teacher matching does not establish better preservation of relative probabilities among alternatives.Source: §4.3 · pp. 6–7, Figure 2 · p. 7
4.5 EFFICIENCY OF ROLLOUT-FREE TRAINING
- 9.88 minutes of runtime (1.94× faster than OPSD; 2.23× than GRPO) accompanies SAPD’s approximately 65.3% AIME24 Avg@12 at 0.66 GPU-hours.Rollout generation dominates OPSD and GRPO runtime; SAPD spends most runtime on teacher computation and student updates, with the added teacher cost more than offset by avoiding rollouts.Source: Figure 3 · p. 8, §4.5 · pp. 8–9
5 RELATED WORK
Related work contrasts fixed-reference training with rollout-based approaches, distinguishes distributional guidance from token imitation, and connects SAPD’s supervision to the current reasoning step.
5.1 ON-POLICY AND OFF-POLICY LEARNING
- Unlike methods that supervise student-generated contexts, SAPD distills on fixed reference prefixes, avoiding autoregressive rollout generation.The related-work comparison includes GKD, MiniLLM, GRPO, OPSD (Zhao et al., 2026), and SDFT (Shenfeld et al., 2026).Source: §5.1 · p. 9
5.2 DISTRIBUTIONAL SUPERVISION
- SAPD transfers context-dependent next-token distributions at existing reference positions using privileged information and a frozen teacher conditioned on the current step.This differs from uniform label smoothing, which assigns equal mass to alternative tokens.Source: §5.2 · p. 9
5.3 LOCAL ALIGNMENT
- Reference-step boundaries select the privileged information used to construct each target while preserving the student’s complete reference history.SAPD does not require step-correctness annotations or a process reward model.Source: §5.3 · p. 9
6 CONCLUSION
SAPD applies step-aligned privileged information to fixed reference prefixes. On mathematical reasoning benchmarks it outperforms SFT and label-smoothed SFT on average, remains competitive with on-policy baselines, and largely preserves coding performance while speeding training loops.
6 CONCLUSION
- On mathematical reasoning benchmarks, SAPD outperforms SFT and label-smoothed SFT on average while remaining competitive with GRPO and OPSD.The conclusion also reports largely preserved out-of-domain coding performance and approximately 2× training-loop speedups over on-policy baselines.Source: §6 · p. 9
- SAPD largely preserves out-of-domain coding performance and achieves approximately 2× training-loop speedups over on-policy baselines.Source: §6 · p. 9
- Analyses support context-dependent preferences and current-step guidance as useful features of the method.The conclusion frames carefully constructed supervision as a competitive and efficient basis for reasoning post-training on fixed demonstrations.Source: §6 · p. 9
A TRAINING DETAILS
The training-details section specifies the label-smoothing baseline’s coefficient and the source of baseline training settings.
A TRAINING DETAILS
- LS-SFT uses the SFT configuration with label-smoothing coefficient ϵls = 0.1; baseline settings follow OPSD.Source: A TRAINING DETAILS · p. 13
B EVALUATION DETAILS
Evaluation follows the cited protocol, using specified sampling settings and answer extraction and grading procedures for mathematical reasoning benchmarks.
B EVALUATION DETAILS
- All models are sampled at temperature 1.0 and top-p 0.8, with top-k filtering disabled and generation capped at 38k tokens.The maximum total sequence length, including prompt and response, is 40k tokens.Source: B EVALUATION DETAILS · p. 13
- For mathematical reasoning benchmarks, evaluation follows the OPSD protocol for prompts, chat templates, answer extraction, and correctness grading.Source: B EVALUATION DETAILS · p. 13
- Models are instructed to reason step by step and place the final answer in \boxed{}; evaluation extracts the last boxed answer and checks it against ground truth.Correctness uses the implementation’s math verify-based grading function.Source: B EVALUATION DETAILS · p. 13
C TEACHER AND STUDENT PROMPT EXAMPLES
The example contrasts the privileged information available to the teacher with the student's more limited prompt. Both branches score the same reference continuation using teacher forcing.
C TEACHER AND STUDENT PROMPT EXAMPLES
- The teacher sees the complete current reference segment as privileged guidance, while the student does not see that second segment as privileged guidance.The student sees the question and first reference segment; the teacher's assistant prefix ends with a single newline.Source: C TEACHER AND STUDENT PROMPT EXAMPLES · p. 14
- Both branches score the same reference continuation, beginning with “First, the equation is x raised to the power of 2y”.At subsequent positions, both use the extended within-segment reference prefix.Source: C TEACHER AND STUDENT PROMPT EXAMPLES · p. 14
- Both branches use teacher forcing rather than sampling a continuation.Source: C TEACHER AND STUDENT PROMPT EXAMPLES · p. 14