Paper Feed

Issue 36 · Project 06 GitHub AI / ML ✓ read

Thinking-Space/One-Shot-OPD

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

TL;DR: A Tsinghua group trained students via on-policy distillation (OPD) on a single query and recovered most of the gain of training on the full 17k-query dataset; 16 semantically diverse queries matched full-data training outright. Their diagnosis — OPD is "data-overfed but algorithm-starved" — is a genuinely useful reframing of where the bottleneck in distillation-based post-training actually sits, and the repo is a runnable veRL-based implementation of OPD, multi-teacher OPD, and a GRPO baseline.

The claim, and why it matters

OPD is the recipe behind a lot of recent frontier post-training: the student generates its own rollouts, and a frozen teacher provides dense token-level supervision (here, reverse KL as the reward). Everyone has studied its algorithmic behavior; almost nobody has asked how much data it actually needs. This is Part II of a series ("Rethinking On-Policy Distillation", Part I from April 2026), and it goes to the data-minimal limit: one training query.

The headline findings from the abstract:

  • One-shot OPD keeps improving for hundreds of steps and recovers most of full-data OPD's gain, across task domains and model families.
  • A single query induces 71.5% state coverage relative to full-data OPD, most of it within the first 100 steps. Sixteen semantically distinct queries reach 98.9% coverage and match full-data validation accuracy.
  • Alignment with the teacher slows at the same rate whether you train on 1 query or 17k — even a fixed set of states takes hundreds of steps to absorb.
  • Stress tests: content-light templates and off-domain WildChat queries approach the real-query baseline. Task content and induced state coverage come apart.

The explanatory frame: because rollouts are student-generated, even one prompt fans out into a broad distribution of states, and the teacher supervises every token of every rollout. So data diversity is quickly saturated — the constraint is the rate at which SGD absorbs that dense supervision, i.e., step efficiency of the algorithm, not query count. This flips the usual instinct ("post-training quality = curated data at scale") for the OPD regime specifically, and echoes the one-example-RLVR results from 2025 but with a mechanistic account (state coverage vs. absorption rate).

1 query (or 16, or 17k) student rollouts → visited states 71.5% coverage from 1 query (most in first 100 steps) 98.9% from 16 queries coverage saturates fast: data is "overfed" teacher reverse-KL absorbed slowly, 100s of steps algorithm is "starved": same slow rate at 1 or 17k queries
The paper's mechanism: student rollouts from even one prompt quickly expose broad supervision (state coverage), while alignment to the teacher proceeds slowly regardless of dataset size — so the bottleneck is step efficiency, not data.

What's in the repo

Code, not checkpoints or data. It's a vendored fork of veRL 0.9 with launcher scripts:

  • grpo.sh (rule-verifier RL baseline), opd.sh (single teacher, reverse-KL token reward), mopd.sh (multiple teachers routed per-row by an ability column).
  • DOMAIN picks teacher + parquets (math/code/fc/if); MODE picks data variant (full/oneshot/template — the last has the student write its own inputs). Only math ships a one-shot corpus.
  • Interesting knobs: LOG_PROB_TOP_K switches from sampled-token OPD to top-K token rewards with several support-set strategies (only_stu, union, intersection, …) and token weighting modes — a small design space for the reward itself.
  • Eval protocols: MATH-500/AMC/AIME (avg@16), LiveCodeBench v6, Multi-IF (external multi-turn harness in eval/), BFCL v3.

What's missing: no datasets (rebuilt via data/prep/ from public sources), no released weights, and the README quotes state-coverage numbers but no benchmark accuracies — those are in the paper (arXiv 2609.04172). Which model families / teacher-student pairs were tested isn't stated in the README either, so verify generality in the paper before extrapolating.

What to expect

Reference config is a single 8×H100/A100-80GB node, with the teacher worker sharing the actor's GPU pool. Stack is bleeding-edge (Python 3.12, vLLM 0.28, torch 2.13, transformers 5.10), Apache-2.0. This is a research artifact for reproducing the paper, not a distillation library — but the OPD/MOPD recipes on top of veRL are clean enough to adapt.

One honest caveat: "one query recovers most of the gain" is a claim about this reverse-KL OPD setup at these scales; the stress tests (templates, off-domain WildChat prompts nearly matching real queries) are striking but also suggest the gains here may be partly about generic distribution-sharpening toward the teacher rather than task learning. The authors themselves frame the takeaway as "go work on step efficiency," which seems like the right read.

Try it

git clone https://github.com/Thinking-Space/One-Shot-OPD.git
cd One-Shot-OPD
export MODEL_BASE=/your/models       # student + teacher checkpoints
export CONDA_ENV_BIN=/your/env/bin   # vllm >= 0.18, transformers 5.x
# rebuild data per data/README.md, then:
MODE=oneshot bash recipe/opd.sh      # one-shot OPD (math)

Paper: arXiv 2609.04172. If you only skim one thing, look for the state-coverage-vs-accuracy plot and the alignment-rate curves — that's where the "data-overfed, algorithm-starved" argument lives.