Order-Only On-Policy Distillation

AI summary

Overview

Research area: Machine learning — specifically knowledge distillation for large language models, and within that, on-policy distillation (OPD).

Technical level: Intermediate. The paper assumes familiarity with how language models produce token probabilities, what a "temperature" setting does to those probabilities, and the basic teacher–student distillation setup.

Scope: The paper proposes Order-Only On-Policy Distillation (O³PD), a distillation objective that transfers only the ranking of a teacher model's preferred next tokens — not their probability values — and shows across six teacher–student pairs that this outperforms standard on-policy distillation on three competition-mathematics benchmarks.

What This Paper Is About

On-policy distillation trains a smaller "student" model on text it generates itself, rewarding it for matching the token probabilities of a larger "teacher" model at each step. The problem is that those probabilities mix two different things: which tokens the teacher likes, and how strongly it likes them. The second part reflects the teacher's own size and calibration, and a small student may simply be unable to reproduce it — and it also shifts whenever the teacher's temperature setting is changed. This paper asks whether the ordering alone is sufficient, and shows that training the student to reproduce only the teacher's ranking of its top-K candidates produces a better student than fitting the full probability distribution.

Key Contributions

1. A diagnosis of what on-policy distillation conflates. The authors argue that teacher token probabilities entangle token preference (the ordering) with teacher confidence (the values), and that the confidence component is tied to the teacher's capacity and calibration rather than to knowledge the student can or should absorb. They also note that teacher temperature rescales these values without changing the ordering, which makes standard OPD sensitive to an arbitrary hyperparameter choice.

2. The O³PD objective. At each prefix generated by the student itself, the student is trained to reproduce the order of the teacher's top-K tokens, with no fitting of the underlying probability values. Because temperature changes values but not order, the objective is invariant to teacher temperature and requires no temperature tuning.

3. A six-pair empirical comparison. Across six teacher–student pairs, O³PD beats standard OPD on the average of HMMT25, AIME24, and AIME25 in every single pair, by up to 2.9 points, with per-benchmark gains of up to 5.6 points on AIME25.

4. Evidence that ordinality, not the specific loss, is the source of the gain. Substituting other ranking-based losses produces similar improvements, isolating the benefit to ordinal supervision itself rather than to any one loss formulation.

Main Findings

• Ordering alone is enough, and is better. Training the student only on the teacher's top-K ordering outperforms training it on the teacher's full probability values across all six teacher–student pairs tested.

• Consistent, not cherry-picked, gains. O³PD wins in every one of the six pairs on the three-benchmark average (HMMT25, AIME24, AIME25), with improvements of up to 2.9 points.

• Largest single-benchmark gain on AIME25. Per-benchmark improvements reach up to 5.6 points on AIME25.

• The improvement is not an artifact of one particular loss function. Other ranking losses yield similar gains, which the authors use to argue that the benefit comes from ordinal supervision in general.

• Temperature tuning is eliminated. Because the objective depends only on the order of the teacher's top-K tokens, and temperature changes probability values without changing that order, O³PD needs no teacher-temperature selection — a hyperparameter that standard OPD is sensitive to.

• A counterintuitive analysis result. Measured against the teacher's probability distribution, O³PD actually departs further than OPD does — yet it reproduces the teacher's ranking more faithfully and produces a stronger student. The authors read this as evidence that the ranking, not the probability values, carries the knowledge on-policy distillation needs to transfer.

Methodology in Plain English

The setup is the standard on-policy distillation loop: the student model generates its own continuations, and at each prefix along those self-generated sequences the teacher model is queried for what it would have preferred next. Standard OPD then asks the student to match the teacher's probability for every candidate token as closely as possible.

O³PD changes only the target of that matching. At each student-generated prefix, the teacher's top-K candidate tokens are collected, and the student is trained to put those tokens in the same relative order the teacher puts them in — without being pushed toward the teacher's specific probability numbers. Since the ordering is unaffected by the teacher's temperature setting, the training signal stays the same no matter what temperature is used, so no temperature search is required.

The evaluation compares O³PD against standard OPD on three competition-mathematics benchmarks (HMMT25, AIME24, AIME25) using six different teacher–student pairings. To check that the result is about ordinality rather than about their specific loss design, the authors also swap in other ranking-based losses and confirm comparable gains. They then analyze the trained students by comparing their output distributions against the teacher's, both in terms of probability closeness and ranking fidelity.

The specific teacher and student models used in the six pairs, the value of K, and the training data composition are not reported in the text available here.

Why This Matters

Impact on research. The paper reframes a core assumption of on-policy distillation: that the goal is to match the teacher's distribution. If the transferable signal is the ordering rather than the values, then a large body of work built on forward KL-style probability matching is optimizing a target that includes information the student cannot use — and is doing so through a temperature hyperparameter that turns out to be avoidable. The finding that a student can be further from the teacher in probability space while being better is a useful corrective to the intuition that distributional closeness equals distillation quality.

Real-world applications:

• Compact reasoning models. Competition-mathematics benchmarks like AIME and HMMT are used as proxies for reasoning ability, so order-only distillation is directly relevant to building small models that retain strong multi-step reasoning.
• On-device and edge deployment. A smaller student that distills more effectively from a large teacher means capable models can run where the teacher cannot.
• Cost reduction at inference. Replacing a large teacher with a well-distilled small student lowers serving cost for any deployed assistant.
• Simplified training pipelines. Removing the teacher-temperature hyperparameter removes an entire tuning sweep from distillation workflows, which matters for teams repeatedly distilling across many teacher–student combinations.

Industry relevance. Distillation from a large proprietary or in-house teacher into a small deployable student is standard practice, and temperature selection is a routine source of engineering overhead. An objective that is invariant to that choice, and that empirically dominates the standard one across six pairings, is directly actionable for teams building production models under latency and cost constraints.

Future Directions

• Beyond competition mathematics. All three reported benchmarks — HMMT25, AIME24, and AIME25 — are mathematics benchmarks. Whether ordinal supervision wins as consistently on code generation, multilingual tasks, or open-ended instruction following is an open question.
• Choosing and varying K. The paper trains on the teacher's top-K tokens, but the sensitivity of the method to the choice of K, and whether an adaptive or curriculum-based K helps, is not addressed in the available text.
• Recovering useful magnitude information. The results suggest values are not needed, but they do not show that values are harmful in every regime. A hybrid objective that uses ordering as the primary signal and a heavily down-weighted magnitude term for near-ties is a natural next experiment.
• A theory of why ranking transfers. The paper demonstrates empirically that ranking fidelity correlates with student strength while probability closeness does not. A formal account of why the ordering of the teacher's preferences is the capacity-independent part of its distribution would sharpen both the method and its scope conditions.

Target Audience

Machine learning researchers and practitioners working on knowledge distillation and LLM compression, especially those running on-policy or sequence-level distillation pipelines and looking to remove temperature tuning. It is also relevant to anyone studying what actually transfers between models of different sizes, and to engineers who need to justify dropping probability-matching terms from a distillation loss. Readers without prior exposure to teacher–student distillation, token-level probability matching, and temperature will need background reading first.

Authors’ abstract

On-policy distillation (OPD) trains a student on its own generations by matching the teacher's token probabilities. These probabilities entangle two kinds of information: which tokens the teacher prefers, and how confident it is. The latter reflects the teacher's own capacity and calibration, which a smaller student may be unable to reproduce. Because the teacher temperature changes these values without changing the ordering, OPD is also sensitive to the temperature choice. We ask whether the teacher's ordering of candidate tokens is enough, and propose Order-Only On-Policy Distillation (O$^{3}$PD). At each student-generated prefix, O$^{3}$PD trains the student to reproduce the order of the teacher's top-$K$ tokens, without fitting their probability values. The objective is therefore invariant to the teacher temperature and needs no temperature tuning. Across six teacher--student pairs, O$^{3}$PD outperforms OPD on the average over HMMT25, AIME24, and AIME25 in every pair, by up to 2.9 points, with gains of up to 5.6 points on AIME25. Other ranking losses yield similar gains, so the benefit comes from ordinal supervision itself. Analysis shows that O$^{3}$PD departs further from the teacher's probabilities than OPD, yet reproduces its ranking more faithfully and yields a stronger student. The ranking, not the probability values, thus carries the knowledge that on-policy distillation needs to transfer.