Variant 03 · complete reading and print edition
The Prefix Observatory
A causal field guide for keeping trajectory source, divergence direction, feedback aperture, and gradient estimator in their own lanes.
“On-policy distillation” is easy to compress into a false shortcut: student samples, therefore reverse KL. The canonical GKD formulation makes the two choices independent. One choice determines where the visited completion states came from; another determines how the teacher and student distributions are compared at those states.
Sourced In GKD, the student-data fraction and the divergence are separate inputs to the algorithm, and the experiments cross source fractions with forward KL, reverse KL, and generalized JSD. “On-policy” therefore names trajectory provenance, not the order of the KL arguments. See Agarwal et al., §3.1 and §4. OPD-CL-003
Sourced A fixed prompt set is compatible with on-policy completions. GKD draws a prompt x from a dataset, then draws the continuation y from the current student conditioned on that prompt. The prompt is fixed data; the generated completion can still be on-policy. OPD-CL-006
01 · Separate the axes
Ask two questions, in order
Who produced the visited path?
A fixed demonstration supplies a predetermined continuation. A student-generated path comes from the active student policy. The source belongs to the causal history that produced the string.
Rows: fixed trajectory or student-generated trajectory.
Which distribution weights the comparison?
At the selected prefix, forward KL uses the teacher as numerator and weight. Reverse KL uses the student. This changes the loss without changing the path’s history.
Columns: forward DKL(T∥S) or reverse DKL(S∥T).
A prefix state is the pair (x, y<t). Sampling a token moves the autoregressive generator to a new prefix. Once a path is selected, teacher and student are both evaluated at the same displayed prefix. The teacher query supplies supervision; it does not retroactively become the source of the prefix.
Sourced GKD Eq. 2 conditions the teacher and student next-token distributions on the same prefix, while Eq. 4 identifies student-sampled continuations as on-policy. GKD §3, Eqs. 2–4. OPD-CL-001 OPD-CL-002
02 · Read the calculation
Just enough KL
For positive normalized categorical distributions P and Q, KL divergence is:
OBS-CL-001One token contribution may be negative. The complete supported KL sum is non-negative. Therefore the signed terms below are printed alongside their local sum; no bar length or isolated term should be mistaken for the divergence.
The static temperature snapshots use this probability-space form of logit temperature scaling:
OBS-CL-002The declared path reduction is an unweighted mean across the two visited positions:
03 · Static observatory
The filled canonical state
Predict first: a student-generated completion with DKL(teacher∥student) belongs in which quadrant?
Answer: student-generated / forward KL. “Student-generated” answers the source question. “Teacher ∥ student” answers the divergence question. Neither label determines the other.
x → A → xA
The fixed demonstration selects A. Its second state is xA. No categorical draw determines this route.
Toy provenance Fixed teacher-generated demonstration completion.
x → B → xB
At τ = 1, the root student distribution is (0.75, 0.25). The fixed draw u = 0.8 falls after the cumulative mass for A, so it selects B.
Toy provenance Student-generated completion from the active student policy.
Complete normalized probability tables
The teacher and student use the same vocabulary {A, B}. All canonical distributions are shown, including the state not visited by the currently selected route.
| Prefix | Policy | P(A) | P(B) | Sum |
|---|---|---|---|---|
x | Student | 0.750000 | 0.250000 | 1.000000 |
x | Teacher | 0.500000 | 0.500000 | 1.000000 |
xA | Student | 0.300000 | 0.700000 | 1.000000 |
xA | Teacher | 0.800000 | 0.200000 | 1.000000 |
xB | Student | 0.600000 | 0.400000 | 1.000000 |
xB | Teacher | 0.900000 | 0.100000 | 1.000000 |
Filled 2×2 canonical matrix
| Trajectory source | Forward KL · T∥S |
Reverse KL · S∥T |
|---|---|---|
| Fixed | x → xA0.338976 nats |
x → xA0.356749 nats |
| Student-generated | x → xB0.185065 nats |
x → xB0.221025 nats |
Signed terms, local sums, and reductions
Each signed token term is the weighted contribution required by the selected KL order. The local sum adds the two token terms; the path mean averages the two unrounded local sums.
| Source · direction | Visited state | Numerator | A term | B term | Local sum | Two-state mean |
|---|---|---|---|---|---|---|
| Fixed · forward | x | Teacher | −0.202733 | +0.346574 | 0.143841 | 0.338976 |
| Fixed · forward | xA | Teacher | +0.784663 | −0.250553 | 0.534111 | |
| Fixed · reverse | x | Student | +0.304099 | −0.173287 | 0.130812 | 0.356749 |
| Fixed · reverse | xA | Student | −0.294249 | +0.876934 | 0.582685 | |
| Student · forward | x | Teacher | −0.202733 | +0.346574 | 0.143841 | 0.185065 |
| Student · forward | xB | Teacher | +0.364919 | −0.138629 | 0.226289 | |
| Student · reverse | x | Student | +0.304099 | −0.173287 | 0.130812 | 0.221025 |
| Student · reverse | xB | Student | −0.243279 | +0.554518 | 0.311239 |
04 · Representative counterfactuals
Temperature changes probabilities, not definitions
The active tempered student is both the behavior policy that supplies the fixed draw and the student distribution inside the printed loss. These are static snapshots of the same deterministic rule—not separate measured runs. Teacher distributions remain those in the canonical table.
τ = 0.5
Student draw: u = 0.8 selects A because root P(A)=0.900000.
| State | A | B | Sum |
|---|---|---|---|
x | 0.900000 | 0.100000 | 1.000000 |
xA | 0.155172 | 0.844828 | 1.000000 |
Both paths spell x → xA, but one is fixed and one was student-generated.
Means: fixed F 0.767361 · fixed R 0.665405 · student F 0.767361 · student R 0.665405.
τ = 1
Student draw: u = 0.8 selects B because root P(A)=0.750000.
| State | A | B | Sum |
|---|---|---|---|
x | 0.750000 | 0.250000 | 1.000000 |
xA · fixed | 0.300000 | 0.700000 | 1.000000 |
xB · student | 0.600000 | 0.400000 | 1.000000 |
Paths: fixed xA; student-generated xB.
Means: fixed F 0.338976 · fixed R 0.356749 · student F 0.185065 · student R 0.221025.
τ = 2
Student draw: u = 0.8 selects B because root P(A)=0.633975.
| State | A | B | Sum |
|---|---|---|---|
x | 0.633975 | 0.366025 | 1.000000 |
xA · fixed | 0.395644 | 0.604356 | 1.000000 |
xB · student | 0.550510 | 0.449490 | 1.000000 |
Paths: fixed xA; student-generated xB.
Means: fixed F 0.189680 · fixed R 0.213047 · student F 0.164676 · student R 0.220648.
05 · Follow the computation
The update begins after the draw
Derived teaching decomposition The six stops below separate operations that are often compressed into one phrase. They are a faithful reading aid for canonical GKD, not terminology quoted from the paper.
-
Sample or select the completion
The current student supplies a categorical draw, or a fixed demonstration supplies its predetermined path.
-
Record the visited prefix
The selected token determines the second state:
xAorxB. -
Evaluate both policies there
Teacher and student return next-token distributions at exactly the same recorded prefix.
-
Compute the local divergence
Direction chooses which distribution weights each log ratio.
-
Reduce across positions
This toy takes the unweighted mean of the two local state divergences.
-
Update student parameters θ
Gradients flow through the local student distributions under the selected divergence.
Sourced method distinction MiniLLM is not merely “the reverse column” of this tokenwise toy. It defines a sequence-level reverse KL, derives a likelihood-ratio policy gradient with future log-ratio returns, then decomposes the estimator into an exact full-vocabulary single-step term and a sampled long-horizon term. Its practical estimator also uses teacher–student mixture sampling. See Gu et al., §2.1, Eq. 1 and §2.2, Eqs. 2–3. OPD-CL-007
06 · A third, separate dimension
How wide is the feedback aperture?
The 2×2 answers where the state came from and which divergence is intended. It does not say how much distribution information crosses the supervision boundary. Full-vocabulary, sampled-token, and top-k feedback are three different payloads; switching among them does not move the source/divergence cell.
OBS-FEEDBACK-001
An invented top-k non-identifiability witness
Toy Two teacher distributions have the same retained top-2 payload and the same omitted mass, but different hidden tails.
| Policy | A | B | C | D | Sum |
|---|---|---|---|---|---|
| Student S | 0.550000 | 0.250000 | 0.150000 | 0.050000 | 1.000000 |
| Teacher T1 | 0.400000 | 0.200000 | 0.300000 | 0.100000 | 1.000000 |
| Teacher T2 | 0.400000 | 0.100000 | 0.300000 | 0.200000 | 1.000000 |
Mode 1 · Full vocabulary
The exact calculation sees every probability. The table exposes all four signed contributions and their local sum.
| Direction · teacher | A term | B term | C term | D term | Full KL |
|---|---|---|---|---|---|
Forward T1∥S | −0.127381 | −0.044629 | +0.207944 | +0.069315 | 0.105249 |
Forward T2∥S | −0.127381 | −0.091629 | +0.207944 | +0.277259 | 0.266192 |
Reverse S∥T1 | +0.175150 | +0.055786 | −0.103972 | −0.034657 | 0.092306 |
Reverse S∥T2 | +0.175150 | +0.229073 | −0.103972 | −0.069315 | 0.230935 |
Mode 2 · One sampled token
For the selected order P ∥ Q, draw one token from the numerator distribution and inspect the unweighted log ratio. Repeated numerator-distributed draws have KL as their expectation:
OBS-CL-004| Direction | Sampling distribution | Selected token | One log-ratio observation | Corresponding weighted exact term |
|---|---|---|---|---|
Forward T1∥S | Teacher T1 | C | +0.693147 | +0.207944 |
Reverse S∥T1 | Student S | B | +0.223144 | +0.055786 |
Not equivalent: the unweighted one-draw log ratio is not the weighted token contribution and is not the full KL. Another draw can be negative. The sampling distribution must be named; a reverse-KL student sample cannot silently stand in for an unweighted forward-KL sample.
Mode 3 · Teacher top-k
The retained teacher candidates are {A, C}. Both teachers expose the exact same payload:
| Teacher | A | C | Omitted mass over B,D |
|---|---|---|---|
| T1 | 0.400000 | 0.300000 | 0.300000 |
| T2 | 0.400000 | 0.300000 | 0.300000 |
If—and only if—we explicitly redefine the objective on retained support, the renormalized distributions are:
| Policy | Conditional A | Conditional C | Sum |
|---|---|---|---|
| Student S | 0.785714 | 0.214286 | 1.000000 |
| Teacher T1 | 0.571429 | 0.428571 | 1.000000 |
| Teacher T2 | 0.571429 | 0.428571 | 1.000000 |
| Declared conditional objective | A term | C term | Conditional KL |
|---|---|---|---|
Forward conditional T1∥S | −0.181974 | +0.297063 | 0.115090 |
Reverse conditional S∥T1 | +0.250214 | −0.148532 | 0.101682 |
Sourced implementation example Current Tinker documentation defines one top-k soft-target rule by renormalizing teacher probabilities over retained candidates and using them as cross-entropy weights. That is one explicit implementation contract, not a universal definition of top-k distillation. Tinker, “Top-K distillation”. OBS-CL-006
Sourced engineering example The 2025 Thinking Machines report instead samples student trajectories, queries teacher log probabilities on those sampled tokens, and passes per-token reverse-KL-style advantages to an RL importance-sampling update. It states that its reported experiments did not use top-k logit distillation. Implementation and pseudocode. OBS-CL-005
07 · Keep neighboring methods distinct
Same neighborhood, different addresses
This is a boundary map, not a leaderboard. Similar path sources do not imply identical supervision or gradient estimators.
| Method | Training trajectory / prefix source | Supervision | Objective or update distinction | Source status |
|---|---|---|---|---|
| Canonical GKD on-policy endpoint | Student completion from fixed prompts | Teacher and student full next-token distributions at the same prefixes | Chosen tokenwise divergence; discrete draw detached in the paper’s estimator | Peer-reviewed ICLR 2024 |
| SeqKD mode approximation | Fixed teacher beam-search sequences | Teacher-generated target tokens | Student likelihood on a fixed teacher-generated dataset | Peer-reviewed EMNLP 2016 |
| ImitKD | Mixed initial/static and student-explored states | Teacher correction; token and distribution variants | Imitation-learning algorithm; its pool may include older student generations | Peer-reviewed EMNLP 2020; “mixed” is GKD’s cross-paper interpretation |
| MiniLLM | Sequence samples under a teacher–student mixture in the practical estimator | Teacher/student log-density ratios plus an exact local term | Sequence-level reverse KL with a policy-gradient decomposition | Peer-reviewed ICLR 2024; current author manuscript v6 |
| Scalar-reward RL | Often student trajectories | Scalar task or environment reward | Policy-gradient objective; a teacher token distribution is not required | General category; GKD Eq. 5 keeps its reward and distillation terms separate |
| Thinking Machines 2025 implementation | Student trajectories | Teacher and student log probabilities on sampled tokens | Reverse-KL-style per-token advantages in an RL importance-sampling update | First-party engineering report, not peer-reviewed |
Sourced SeqKD’s implemented mode approximation uses a teacher beam-search sequence as fixed student training data. See Kim and Rush, §3.2 and Fig. 1. OBS-CL-007
Sourced · qualified ImitKD lets the student explore and queries teacher correction on visited states, but its training-state mixture is not simply the pure, always-current λ=1 GKD endpoint. See Lin et al., abstract and §§2.2–3, plus GKD §4’s unifying classification. OBS-CL-008
Sourced GKD can be combined with a separate scalar-reward RL term; Eq. 5 keeps reward and distillation objectives distinct. On-policy state collection does not make every method ordinary scalar-reward RL. OBS-CL-009
08 · Transfer
Classify both axes
Both answers are fully expanded in this edition. Treat them as worked transfer prompts, not as an assessment instrument.
Reverse KL on fixed xA
Prompt: A learner uses reverse KL on a fixed teacher-generated xA prefix. Is the divergence reverse KL, and is the state sampling on-policy? Explain both axes separately.
Expanded answer
Divergence axis: yes, it is reverse KL because the student is the numerator and weight in DKL(S∥T).
Source axis: no, the training state is not student on-policy. The teacher-generated route to xA was fixed before the local comparison.
Why the labels coexist: evaluating the student and teacher at xA does not alter how xA entered the training batch. A reverse objective cannot rewrite trajectory provenance.
Forward KL on student xB
Prompt: A learner samples B from the student, queries the teacher at xB, and computes forward KL. Is this on-policy distillation? Explain why divergence direction does not change the state source.
Expanded answer
Source axis: yes, under canonical GKD’s state-provenance meaning. The current student produced B, so xB is a student-visited prefix.
Divergence axis: it is forward KL because the teacher weights ln(T/S) in DKL(T∥S).
Why the labels coexist: the teacher labels the already visited xB state; it did not generate the path. Forward KL changes the local comparison, not the trajectory history.
09 · Production boundary
What the toy does not settle
Before accepting an “OPD” label for a real implementation, ask:
- Which policy, dataset, or mixture produced the completion trajectories?
- Are teacher and student evaluated at the same exact prefixes?
- Which divergence or sequence-level objective is optimized, with what reduction?
- Is feedback full-vocabulary, sampled-token, or top-k—and under which explicit estimator?
- Does gradient flow only through local distributions, or also through a likelihood-ratio trajectory estimator?
- How are EOS, padding, variable sequence length, temperature, and tokenizer mismatch handled?
- How fresh is the behavior policy relative to the student being updated?
Limitations
- Unresolved The fixture assumes a shared two-token vocabulary. It does not model tokenizer mismatch or vocabulary projection.
- Unresolved It has exactly two visited positions and no EOS, padding, masking, variable-length normalization, batching, or zero-probability edge case.
- Unresolved It does not model stale behavior policies, asynchronous sampling, importance correction, optimizer dynamics, or convergence.
- Unresolved Teacher probabilities are supervision signals, not guarantees of factual correctness or calibration.
- Unresolved Dense tokenwise feedback does not by itself identify which earlier token caused a later outcome; it is not causal process supervision.
- Toy The numerical values validate only the declared arithmetic. They do not measure quality, compute savings, factuality, stronger-teacher benefit, or empirical superiority.
- Unresolved No learner study has tested whether this page improves understanding, retention, or transfer.
Primary sources and inspected versions
- Rishabh Agarwal et al., “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.” Peer-reviewed ICLR 2024 proceedings record and conference paper; author manuscript v3, 17 January 2024. Key locators: §2 “KL-Based Divergences”; §3 Eq. 2; §3.1 Eq. 4, sentence immediately after Eq. 4, Algorithm 1, and boxed objective; §4 “GKD Variants.”
- Yuxian Gu et al., “MiniLLM: Knowledge Distillation of Large Language Models.” Peer-reviewed ICLR 2024 proceedings record and conference paper; current author manuscript v6, 31 January 2026. Key locators: §2.1 Eq. 1; §2.2 Eqs. 2–3; practical teacher–student mixture and importance-weight estimator in Eqs. 4–5; length normalization in Eq. 6.
- Yoon Kim and Alexander Rush, “Sequence-Level Knowledge Distillation.” Peer-reviewed EMNLP 2016 ACL Anthology record and paper. Key locators: §3.2 and Fig. 1 for the beam-search mode approximation and fixed teacher-generated targets.
- Alexander Lin et al., “Autoregressive Knowledge Distillation through Imitation Learning.” Peer-reviewed EMNLP 2020 ACL Anthology record and paper. Key locators: abstract and §§2.2–3 for prefix states, learner exploration, oracle correction, and the mixed training dataset.
- Kevin Lu and Thinking Machines Lab, “On-Policy Distillation.” First-party engineering report, 27 October 2025; the page also notes later implementation updates. Key locators: “Loss function: reverse KL,” “Implementation,” and pseudocode steps 2–4. This is not a peer-reviewed paper.
- Tinker documentation, “Top-K distillation.” First-party software documentation. Key locator: the retained-candidate renormalization formula and example. This defines one implementation, not a universal top-k objective.