AIAny
Icon for item

Rethinking On-Policy Distillation of Large Language Models II: One Training Example

Studies on-policy distillation (OPD) at the data-minimal limit by training on a single query, measuring state coverage and alignment dynamics, and showing OPD is often data-overfed but algorithm-starved.

Introduction

Why this matters

Most work on OPD evaluates algorithmic variants or large datasets; this paper asks a counterintuitive question: how little data is enough? By training on a single prompt and analyzing the student’s rollouts, the authors expose that OPD’s apparent reliance on large datasets is often illusory — rollouts quickly cover most relevant states, but the student assimilates the teacher’s signals slowly, making OPD step-limited rather than data-limited.

Key Findings
  • A single semantically rich query reaches 71.5% of the states visited by full-data OPD, with most coverage appearing in the first ~100 steps. This shows rollouts can rapidly expose broad supervision from very little prompt data.
  • Adding diverse queries increases state coverage and validation accuracy monotonically; 16 semantically distinct queries reach 98.9% coverage and match full-data OPD performance.
  • Alignment (the rate at which student token distributions converge to the teacher) remains slow regardless of dataset size — even with fixed states it takes hundreds of optimization steps to absorb supervision. The paper summarizes this as “data-overfed but algorithm-starved.”
  • The state-coverage phenomenon generalizes across model families, multi-teacher OPD, content-light templates, and off-domain queries, suggesting task content and induced state coverage can diverge.
Who it's for and tradeoffs

Great fit if you want to understand why OPD yields large gains in practice and whether improving data collection or optimization would be more effective. The paper is valuable for researchers and engineers designing post-training pipelines, curriculum selection, or distillation schemes.

Look elsewhere if you need production deployment recipes or step-efficiency fixes — the paper diagnoses and measures the problem more than it provides a turnkey optimization that speeds up alignment. Code and experimental artifacts are provided as a reference implementation on GitHub, but practical step-efficiency solutions remain an open direction.

Where it fits

Positions OPD research toward algorithmic improvements (optimization, step efficiency, alignment strategies) rather than solely larger or more diverse prompt datasets. The work complements mechanistic analyses of token-level signals and offers concrete metrics (state coverage) to compare prompt sets and curriculum designs.

Information

  • Websitearxiv.org
  • OrganizationsAffiliation: University of Chinese Academy of Sciences, Affiliation: Tsinghua University, Affiliation: Northeastern University, Affiliation: University of Illinois Urbana-Champaign, Affiliation: Johns Hopkins University*Equal Contribution.†Project Lead.‡Corresponding Authors.
  • AuthorsZixuan Fu, Bingxiang He, Yuxin Zuo, Haohuan Huang, Jinqian Zhang, Ruhang Xiao, Cheng Qian, Qinyu Luo, Huan-ang Gao, Yudong Wang …
  • Published date2026/09/03

More Items

Internalizes context-induced performance gains into model parameters by generating prospective queries and distilling context-conditioned outputs into mergeable low-rank or structural feed-forward updates, balancing consolidation and retention across 4B–120B dense and MoE models.

Restores recoverability in on-policy distillation by having the teacher generate critiques of incorrect student rollouts, guide a critique-free student rewrite, and score the revised solution as privileged context; yields improved math-reasoning performance, largest gains on competition problems.

Explores asymmetric compute allocation in looped Mixture-of-Experts models by compressing the first loop visit's sequence and reallocating saved compute to enlarge hidden size, yielding consistent NLL reductions and predicted compute savings across 1.3B–7B configurations.