Source-linked AI summary
Rethinking On-Policy Distillation of Large Language Models II: One Training Example
Zixuan Fu, Bingxiang He, Yuxin Zuo, Haohuan Huang, Jinqian Zhang, Ruhang Xiao, Cheng Qian, Qinyu Luo, Huan-ang Gao, Yudong Wang, Zhiyuan Liu, Ning Ding, Chaojun Xiao
TL;DR
Existing OPD research has emphasized algorithmic dynamics, leaving how training data shapes OPD unclear. This paper studies the data-minimal case through one-shot OPD and explains its behavior using state coverage and the student’s absorption rate. One query recovers most of full-data OPD’s gain, while 16 semantically diverse queries match full-data training across single- and multi-teacher settings.
Problem
Prior OPD work largely studies algorithmic behavior, leaving how training data and the algorithm interact insufficiently examined.
Method
The paper trains OPD on single and semantically diverse query sets, measuring rollout state coverage, validation accuracy, and the rate of teacher–student alignment.
Results
One query recovers most of full-data OPD’s gain across domains and model families, while 16 semantically diverse queries reach 98.9% state coverage and match full-data training, including in MOPD.
Takeaways & Limitations
OPD is data-overfed but algorithm-starved: its rollouts expose broad supervision quickly, while the student’s absorption rate increasingly limits training progress.
Takeaways & Limitations
State coverage is a semantic proxy based on full-data rollouts with equally weighted clusters, and the absorption-rate mechanism remains open; the MOPD study uses three domains and one teacher per domain.
Abstract
from arXiv · showhide
On-policy distillation (OPD) combines student-generated rollouts with dense token-level supervision from a teacher. Existing work has mainly studied its algorithmic behavior, leaving the role of training data unclear. We examine this role at the data-minimal limit by training on a single query. One-shot OPD keeps improving for hundreds of steps and recovers most of full-data OPD's gain across task domains and model families. We explain this result through the states visited during training and the rate at which the student aligns with the teacher. We measure \emph{state coverage}, the fraction of the states full-data OPD visits that a query set's rollouts reach. A single query already reaches \(71.5\%\), most of it within the first 100 steps. Adding semantically distinct queries raises coverage and validation accuracy together, until 16 queries reach \(98.9\%\) and match full-data training. Yet alignment slows at a similar pace whether OPD trains on one query or the whole dataset, and even a fixed set of states takes hundreds of steps to absorb. OPD is therefore data-overfed but algorithm-starved. Its rollouts quickly expose broad supervision, while the student absorbs that supervision increasingly slowly. The state-coverage result extends to multi-teacher OPD, where 16 semantically diverse queries per domain match full-data MOPD. As a further stress test, content-light templates and off-domain WildChat queries also approach the real-query baseline. Task content and induced state coverage can therefore come apart. We hope these findings direct future work toward the step efficiency of OPD, and prompt a re-examination of the data and the mechanisms behind its recent successes in frontier post-training.
1. Introduction
OPD’s data requirements and optimization dynamics can be separated by training on a single query. One query reaches much of the state space visited by full-data OPD, while absorption of the resulting supervision slows regardless of query count.
- Motivation: OPD combines student-generated rollouts with dense teacher supervision at every visited prefix, unlike outcome-level reward in RLVR.The student samples its own trajectories, and the teacher supplies a full next-token distribution at each visited state.
- One-shot OPD: One-shot OPD keeps improving for hundreds of steps and achieves substantial gains across four task domains and three model families, including on queries the student never solves.The single-query setting disentangles data supply from the rate of supervision absorption.
- State coverage: 71.5% of the states visited by full-data OPD are covered by one query, while 16 semantically diverse queries cover 98.9% and match full-data training.State coverage measures the fraction of semantic clusters reached by a query set’s rollouts among states visited by full-data OPD.
- Optimization dynamics: The student’s absorption rate declines similarly with one query and all 17k queries, leaving much of the available supervision unexploited.The absorption rate is the proportion of the remaining teacher–student gap closed by one update.
- Beyond one-shot OPD: Sixteen semantically diverse queries per domain suffice to match full-data multi-teacher OPD, while off-domain WildChat prompts and content-free templates approach real-query performance.These results suggest that induced reasoning states, rather than task content alone, determine much of an input’s usefulness.
- Implications: OPD is data-overfed but algorithm-starved: data curation should prioritize the states inputs induce, alongside methods that increase absorption speed.The paper frames future progress around state-level query selection and more step-efficient training.
2. Preliminaries
The preliminaries define OPD as on-policy alignment between student and teacher distributions on student-visited states. They introduce sampled-token and top-k estimators, together with metrics for gap recovery, token overlap, and residual disagreement.
- Notation: A state is the autoregressive context s_i = (x, y_<i), where x is the input and y_<i is the response prefix before token position i.The student and teacher define next-token distributions over a shared vocabulary, and trajectories pair inputs with autoregressively sampled responses.
- On-Policy Distillation: OPD samples trajectories from the student and aligns it with the teacher on the prefixes the student actually visits.Its objective can also be viewed as dense KL-constrained reinforcement learning with a teacher-induced token-level reward.
- On-Policy Distillation: The per-token advantage is the teacher log-probability minus the student log-probability for the emitted token at a visited state.This estimator evaluates the correction only at the token sampled by the student.
- On-Policy Distillation: The top-k estimator averages the teacher–student log-probability difference over the student’s k most likely tokens, renormalized within that set.It trades the distribution tail for lower variance while estimating the same divergence.
- Evaluation metrics: Gap recovery normalizes student improvement by the initial teacher–student score gap, while full-data recovery normalizes it by the full-data OPD gain.A full-data recovery ratio near 100% indicates that a reduced query set nearly matches full-data OPD’s gain.
- Alignment diagnostics: Top-k token overlap measures agreement between the student’s and teacher’s high-probability token sets, while overlap-token advantage measures their weighted log-probability difference on shared tokens.The overlap-token advantage is on a smaller, non-comparable scale because it averages over top-k entries and weights them by student probability.
3. The One-Shot Phenomenon
One-shot OPD recovers most of full-data OPD’s improvement across model families and task domains, and remains effective across query and rollout conditions. In mathematics, it reaches 68.5 versus 69.8 for full-data OPD at step 300, while comparable gains extend to code, instruction following, and agentic tool use.
- Cross-domain robustness: One query recovers most of the teacher–student gap across task domains and model families.The effect was evaluated across mathematics, code generation, instruction following, and agentic tool use, using three student–teacher families.
- Mathematical reasoning: 68.5 versus 69.8: one-shot OPD approaches full-data OPD on mathematics at step 300.This corresponds to recovering 69% of the teacher–student gap and 87% of full-data OPD’s gain at that step.
- Training dynamics: At step 1000, one-shot OPD reaches 68.4 versus 72.1 for full-data OPD and recovers 72% of its gain.Beyond step 300, the two curves remain within a band of about 3 points, while alignment metrics follow the same process under both settings.
- Cross-domain robustness: 73%, 66%, and 64%: one-shot OPD recovers these respective teacher–student gaps in code generation, instruction following, and agentic tool use.These results extend the one-shot effect beyond mathematical reasoning.
- Model-family robustness: 77.1 to 85.5, 28.2 to 40.2, and 70.8 to 82.4: final averaged scores increase across the three student–teacher families.The respective families are R1-Distill-1.5B, Llama-3B-It, and OLMo-7B-It-DPO.
- Query and rollout robustness: One-shot OPD remains effective across easy, medium, and hard queries, including a hard query the student never solves.Tightening the response-length cap and lowering rollout temperature also preserve the gain.
4. Data Perspective: Abundant States
OPD consumes visited states rather than queries, so one query can expose much of the supervision reached by full-data training. State coverage and validation accuracy both rise with semantically diverse queries, reaching near-full-data levels with a small query set.
- State coverage: OPD trains on states produced by query–response prefixes, so a single query can yield tens of thousands of supervised states.Each sampled response creates a state at every token position, and 64 rollouts per update provide substantial supervision even for one query.
- Measuring coverage: State coverage measures the fraction of shared representation-space clusters reached by a setting’s rollouts, rather than the raw number of visited prefixes.Clusters are formed from teacher signatures of full-data OPD states and assigned using PCA followed by K-means with K=200.
- State coverage: 71.5% state coverage is reached by one-shot OPD at step 300, with 65.9% already reached by step 100.The run adds only 5.6 percentage points over the following 200 steps, indicating sharply diminishing discovery of new clusters.
- State coverage: One-shot OPD raises validation accuracy from 59.1 to 66.9, compared with 70.8 for full-data OPD at step 300.The comparison uses full-data OPD as the benchmark for both state coverage and validation accuracy.
- Diversity expands coverage: 16 semantically distinct queries reach 98.9% state coverage and match full-data OPD validation accuracy.Controlled comparisons show that semantic diversity, rather than query count alone, raises coverage and validation accuracy; off-policy accuracy also rises with distinct trajectories.
5. Algorithm Perspective: Slow Alignment
OPD’s alignment continues throughout training, but each update absorbs a shrinking fraction of the remaining teacher–student gap. This slowdown persists across query counts and remains even when training states are fixed.
- Alignment dynamics: The student keeps aligning with the teacher, but absorption slows steadily, requiring hundreds rather than tens of steps.Distance falls throughout the run, while the absorption-rate curves flatten on a logarithmic scale.
- Alignment dynamics: 78% to 84% of each run’s step-30 distance is removed by step 300 across training-set sizes.The absorption rate falls in much the same way for one, four, sixteen, and all DAPO-Math-17k queries.
- Alignment dynamics: Alignment pace is largely a property of OPD rather than query count, although learning rate changes how many steps the process takes.The reported similarity concerns the way distance falls, not the total number of steps under every learning rate.
- Fixed-state ablation: Fixed training states still support steadily improving OPD for about 200 steps before accuracy stops increasing.The off-policy run reuses 64 trajectories sampled from the initial student, while alignment metrics also change over hundreds of steps.
- Fixed-state ablation: The prolonged run therefore cannot be explained by the continued arrival of fresh on-policy states alone.Holding states fixed leaves the run comparably long, suggesting that the fixed states themselves contain supervision the student absorbs gradually.
6. One-Shot Extends to Multi-Teacher OPD
The query-diversity result extends to multi-teacher OPD: selecting semantically diverse queries within each domain can match full-data MOPD. Full-data MOPD itself nearly matches separate full-data OPD.
- Multi-teacher extension: 16 semantically diverse queries per domain are sufficient to match full-data MOPD.The setting covers mathematical reasoning, code generation, and instruction following with domain-specific teachers.
- Multi-teacher extension: Full-data MOPD raises average validation accuracy from 43.5 to 52.8, recovering 79% of the teacher–student gap.Separate full-data OPD reaches 53.8, establishing full-data MOPD as a close reference.
- Query diversity: Increasing from 1 to 16 queries per domain raises average validation accuracy from 50.1 to 52.9 at step 300.Full-data MOPD reaches 52.8, so 16-shot MOPD recovers 101% of its gain.
- Query diversity: The 16-shot result is consistent across domains, with gain-recovery proportions of 93% for mathematics, 136% for code generation, and 109% for instruction following.The corresponding proportions are calculated separately for each domain.
- State-space interpretation: The findings support a state-space account in which diverse queries induce complementary states within each domain.This account explains why a small per-domain query set can match full-data MOPD.
7. Discussion
OPD can benefit from inputs that contain little or no task content when they still induce states exposing useful teacher–student gaps. Compared with one-shot RLVR, OPD extracts substantially more learning signal from the same query because its supervision is dense and token-level.
- Content-light and off-domain inputs: Content-light templates and off-domain WildChat inputs track the real-query baseline, raising the three-benchmark average from 59.1 to 69.8.All three conditions reach that level on one-third to one-half of their rollout tokens, but a template that immediately closes the block collapses into unusable meta-level replies.
- Content-light and off-domain inputs: Task content is not the only source of useful OPD signal because inputs also act as state generators.Explicit domain content changes final performance by about one point in these experiments.
- OPD versus RLVR: The two methods therefore value inputs differently: RLVR requires verifiable, varied outcomes, while OPD requires trajectories that expose teacher–student gaps.OPD input value is not tied directly to verifiability or difficulty.
- OPD versus RLVR: Over 1000 steps, OPD closes 72% of its gap to the teacher and gains more than twice as much validation accuracy as RLVR on the same query.The advantage remains when methods are matched by rollout tokens rather than training steps.
- OPD versus RLVR: RLVR’s signal weakens when outcome groups become nearly unanimous, whereas OPD continues learning from local teacher–student gaps after the query is solved.OPD’s remaining limitation is gradual shrinkage of those gaps rather than loss of outcome variation.
8. Related Work
Prior work explains OPD through its objectives, update geometry, and conditions for success, while parallel research studies data-efficient reasoning and synthetic or unsupervised post-training data. This paper differs by examining how inputs and induced states shape dense OPD supervision.
- On-policy distillation: OPD research has studied objective and update geometry, including dense per-token targets, token-level policy-gradient views, and localized versus dense update structure.These analyses contrast OPD’s supervision with sequence-level outcome rewards and SFT-style updates.
- On-policy distillation: A second OPD research thread examines when the method succeeds or fails.The supplied passage introduces this thread as complementary to work on objective and update geometry.
- Data-efficient reasoning post-training: Data-efficient reasoning work reduces training data through one-shot RLVR, example-impact ranking, prefix scheduling, and high-value subset selection.These approaches span sparse outcome-reward RL and broader post-training data selection.
- Data-efficient reasoning post-training: Unlike sparse outcome rewards or reference solutions, OPD supplies dense per-token targets at each visited prefix and can remain effective on queries the student never solves.This distinguishes the paper’s data-efficiency setting from related RL and SFT approaches.
- Synthetic and unsupervised data: Synthetic-data work reduces dependence on curated human-written inputs by synthesizing instructions or removing ground-truth labels from reinforcement learning.The related directions include Self-Instruct, Evol-Instruct, Magpie, and label-free RL approaches.
9. Conclusion, Limitations, and Future Work
The paper concludes that OPD is data-overfed but algorithm-starved: few queries induce most relevant states, while declining absorption rates limit training speed. It identifies state-aware data selection and more step-efficient training as priorities, while noting limits of the coverage proxy and current MOPD scale.
- Conclusion: One query already covers most of the state space visited by full-data OPD, while 16 semantically diverse queries match full-data training in single-domain OPD and MOPD.
- Conclusion: The absorption rate keeps falling at the same pace on one query as on all 17k, so the optimizer rather than training-set size sets run duration.
- Limitations: Coverage is a semantic proxy anchored to full-data rollouts, equally weighting clusters regardless of visitation frequency or remaining teacher signal.
- Limitations: The MOPD evidence uses three domains with one teacher each, leaving its extension to more teachers unresolved.
- Future Work: State coverage should guide query selection, while future efficiency work could reuse batches under trust regions or weight tokens by remaining teacher signal.
A.1. Experimental Setup and Training Queries
The setup uses domain-specific teachers and standardized OPD defaults, with mathematical queries selected by initial student difficulty. Instruction-following uses a separately trained teacher, and its full-data baseline is a sampled dataset subset.
- Instruction-following setup: The instruction-following teacher starts from DeepSeek-R1-Distill-Qwen-1.5B and is post-trained with GRPO on synthetic tasks with programmatically checkable constraints.
- Data convention: For instruction following, “full-data” means a sampled UltraData-SFT-2605 subset, and the one-shot query is randomly drawn from that subset.
- Training queries: Mathematical queries are selected as easy, medium, and hard using pretraining rollout pass rates of 8/8, 4/8, and 0/8.
- Hyperparameters: Mathematical one-shot runs use K=16, while all other experiments use the sampled-token objective under the default OPD settings.
B.1. Training-State Coverage
Training-state coverage measures how much of a full-data reference state space query rollouts reach. Coverage rises with semantically diverse queries, remains robust across constructions, and complements distance diagnostics that expose one-shot’s heavier distant-state tail.
- Coverage construction: State coverage is estimated from rollouts collected across optimization steps and compared against clustered full-data reference and held-out states.
- Robustness: Coverage ordering is stable across cluster counts, reference constructions, and state positions, with 16-shot coverage remaining close to full data.
- State extraction: Final-token extraction lowers coverage for few-query settings, whereas eight relative positions sample states throughout responses of varying lengths.
- Distance diagnostics: One-shot states have a heavier upper tail of distances from full-data states, but AUROC 0.552 indicates they are not separable from full-data states overall.
- Diversity: 16 semantically diverse queries reach 70.9 validation accuracy at step 300, versus 69.9 for 16 queries from one cluster.
- Scheduling: At step 600, sequential and mixed two-query schedules reach 68.8 and 69.1, both above the shared one-shot baseline of 67.4.
- Alignment measurements: Table 7 summarizes the remaining step-30 distance at step 300 and the absorption-rate decline between steps 50 and 200.
C.2. The Learning Rate Rescales the Step Axis
Learning rate changes how quickly OPD moves through its alignment trajectory but does not change the trajectory’s slowdown shape. Related experiments compare input conditions, token accounting, entropy, and multi-teacher settings across domains.
- Learning-rate effect: A larger learning rate reaches a given teacher–student distance sooner, with the three runs separated by roughly their learning-rate ratios.
- Learning-rate effect: Rescaling time as u=(lr/10^-6)t narrows the run spread from about 2.2× to about 1.1×, showing that learning rate mainly rescales the step axis.
- Learning-rate effect: The rescaled curves still bend flatter, so the learning rate changes run duration without removing the observed slowdown.
- Code generation: The code-generation setup evaluates R1-Distill-1.5B with DeepCoder-1.5B-Preview on LiveCodeBench v6 under avg@3.
- Code generation: Across code-generation input conditions, final accuracies differ by at most 1.7 points, while rollout-token use differs substantially: 277M for full-data baseline versus 18M for the base template.