On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
1. Paper overview
In one sentence: GKD teaches a smaller language model on the prefixes it generates, improving distillation while adding sampling cost and a task-dependent choice of divergence. e02e04e05e08e09e11e13e15e19
| At a glance | What to know |
|---|---|
| Research problem | Source description A student trained only on reference or teacher-generated sequences may encounter different prefixes when generating its own text. Errors can then compound. A second problem is capacity mismatch: a smaller student may be unable to represent the teacher distribution well. GKD addresses where feedback is requested and how the resulting distributions are matched. e02e03e05 |
| Core mechanism | |
| A key reported result | GSM8K calculator-assisted arithmetic reasoning: +8.8 points with on-policy forward KL. Accuracy improvement in percentage points. Original test split; fine-tuned FLAN-T5-XL → Base; four-shot CoT and greedy decoding; final checkpoints averaged across three seeds. Mixed forward KL: +6.8; fixed-data forward KL: +4.7. Initial student: 10.16%; teacher: 27.9%. At fixed divergence, student-generated CoTs improve final-answer accuracy. The metric does not validate every reasoning step, and the teacher remains stronger in this Base comparison. e12e13e14 |
| Reading caution |
Core contributions
- Reader analysis
Distillation can regularize an RL objective toward a stronger teacher. The XSum experiment demonstrates an entailment–ROUGE tradeoff, rather than a universal improvement in factuality and reference overlap simultaneously. e06
Algorithm 1. The student chooses which histories receive teacher supervision. Original paper, p. 4 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Start at line 2: λ controls the fraction of training steps using student outputs, D selects the divergence, and η is the learning rate. Lines 4–8 draw u and choose the batch source. The inequality u ≤ λ leads to current student generation; the other branch retrieves fixed input–output pairs. At λ = 1, generation is always on-policy. Line 10 averages the divergence gradient across the batch and updates θ, the student parameters. Equation (2) supplies the operation hidden inside D: compare teacher and student next-token distributions at every shared output prefix, then average over the sequence length. e04e05
What it supports. The algorithm changes where the student practices. The teacher still supplies rich probability targets, but those targets are evaluated on histories the student is likely to encounter. Refreshing generations after parameter updates distinguishes this feedback loop from repeatedly fitting a fixed set of teacher completions.
Where the evidence stops. The sampling step is not differentiable training machinery here: page 4 explicitly stops gradients through the sampling distribution, although the pseudocode does not draw a stop-gradient symbol. Teacher feedback also presupposes an already useful initial student.
2. Motivation
2.1 The problem and the proposed response
A student trained only on reference or teacher-generated sequences may encounter different prefixes when generating its own text. Errors can then compound. A second problem is capacity mismatch: a smaller student may be unable to represent the teacher distribution well. GKD addresses where feedback is requested and how the resulting distributions are matched. e02e03e05
2.2 What this reading follows
A language model can learn to continue polished training examples yet struggle after its own imperfect token choices. GKD makes those self-generated histories the place where a teacher supplies feedback. The student chooses the trajectory; the teacher supplies a full next-token distribution at each prefix. This reading follows the sampling algorithm, then separates the effects of on-policy data, divergence and decoding across T5 summarization, translation and arithmetic reasoning. The results favor student-generated training sequences, but they do not identify one universally best loss. Appendix settings and reporting inconsistencies matter when interpreting the gains or designing a reproduction. e02e04e05e08e09e11e13e15e19
3. Research context
We place the paper in the collection through its world–action interface. The catalog labels and the reading’s assessment are shown separately.
| Catalog dimension | Recorded classification |
|---|---|
| Major category | Foundational work |
| Architecture | Not applicable |
| Prediction paradigm | Not applicable |
| Quadrant | Not applicable |
3.1 Evidence-based assessment
Supports the recorded classification
The recorded foundational training-optimization/distillation category is supported. World-action-model architecture, prediction paradigm and quadrant are not applicable: teacher and student are training participants, with token-distribution targets and no learned environment-dynamics/action interface. The imitation-learning connection does not establish a world model or robot-execution result. e04e05e22
This is the collection’s architectural analysis, not a new related-work survey. Benchmark comparisons and their protocols appear in Section 6.
4. Problem formulation
4.1 Inputs and outputs
| Inputs | Outputs |
|---|---|
|
|
4.2 Equations and their role
5. Method in detail
5.1 Why an imperfect continuation still provides a useful training example
Suppose a student has already generated a questionable prefix. Supervised KD on a fixed reference continuation may never expose the model to that prefix. GKD instead keeps the student's history and asks the teacher what next-token distribution it would assign there. The student then moves toward that distribution using the token-averaged divergence. This is different from treating the student's sampled words as correct hard labels: the sample determines the conditioning context, while the teacher determines the target probabilities. It is also different from differentiating a reward through the whole trajectory. The paper explicitly omits gradients through sampling. Repeating generation after each update makes the training distribution track the evolving student, but the formulation does not guarantee that a teacher can repair every badly corrupted history. e02e04e05
Table A.2. The reproducible unit includes the checkpoint, schedule and calculator-assisted protocol. Original paper, p. 16 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Read this as the distillation configuration, not the earlier supervised initialization. The main setup first fine-tunes FLAN-T5 on the supplied CoT dataset; this table then specifies 40,000 distillation steps. Its learning rate warms up over 2,000 steps and cools down linearly between steps 30,000 and 40,000. The lower rows distinguish maximum input/output lengths, the FLAN-T5 checkpoint family, teacher temperature 0.1 and greedy evaluation. The student still samples its training sequences at temperature one under Algorithm 1. The four-shot prompt and external calculator are specified elsewhere, so the table alone is not the entire evaluation recipe. e05e12e14e15
What it supports. A reproduction must preserve both the online learning rule and these task settings. Forty thousand updates, batch size 32, dropout 0.05 and the 512/320 sequence limits define this GSM8K run; replacing them with another task's defaults would change the experiment even if λ and the loss were identical.
Where the evidence stops. The table is a configuration record, not a measured performance table. It does not specify calculator parsing, software versions or a full hardware budget. Those gaps remain open implementation choices to document, not values to infer from the reported accuracy.
5.2 Separate the distribution of prefixes from the loss used on them
GKD exposes two independent design choices. Lambda determines which histories receive feedback; the divergence determines how strongly different teacher–student probability discrepancies are penalized. The heatmaps make this separation unusually easy to inspect: vertical comparisons hold the loss fixed, while horizontal comparisons hold the data mixture fixed. WMT favors JSD(0.1), GSM8K Base favors forward KL, and greedy XSum Large barely separates the five on-policy losses. Reader interpretation: the robust lesson is to test on-policy data, not to adopt a single divergence as a universal default. Capacity mismatch motivates mode-seeking alternatives, but task metrics and decoding affect their usefulness. The XSum diagnostic adds another condition: a gain in reference overlap can accompany lower diversity, which may or may not suit the intended application. e03e05e08e09e11e13
5.3 Distinguish a cheaper deployed student from cheaper training
The teacher participates in training by providing logits, whereas deployment evaluates the distilled student. This separation is the compression benefit GKD seeks; the paper does not provide a serving-latency benchmark. On-policy generation nevertheless adds training work, with the GSM8K appendix reporting overhead relative to fixed-output training. Reader interpretation: compare student quality at both equal updates and equal compute before concluding that a method is more efficient overall. The optional RL extension changes the objective again. It rewards outputs while penalizing divergence from the teacher, rather than merely keeping the model close to its initial student policy. In XSum, the entailment reward and distillation weight produce a tradeoff between factual-consistency proxy and ROUGE-2. That experiment concerns text alignment, not executed actions in an environment. e05e06e07e15e22
5.4 Training and inference
During training
Main task-specific experiments initialize from supervised fine-tuning, using a roughly 3B-parameter T5-XL teacher and 77M, 250M or 800M students. XSum/WMT use LM-adapted T5v1.1; GSM8K starts from FLAN-T5. Adafactor is used. The method assumes useful initial generations rather than a randomly initialized student. e05e07
Student sampling uses temperature 1. Teacher temperature is task-dependent; GSM8K uses 0.1. Forward KL is the teacher-to-student direction; reverse KL swaps that direction. JSD approaches forward-like behavior near β = 0 and reverse-like behavior near β = 1. e03e05e14
For RL integration, α weights the distillation penalty and 1 − α weights expected reward; α = 1 is pure distillation. XSum uses JSD(0.9) and reward from a T5-XXL textual-entailment classifier. Increasing α raises ROUGE-2 while reducing the entailment improvement in Figure 5. e06
During inference
The trained student generates autoregressively without teacher feedback. XSum is evaluated with greedy or temperature sampling, WMT with beam search, and GSM8K with greedy four-shot chain-of-thought prompting plus an external calculator. The calculator-assisted accuracy is not unaided reasoning accuracy. e08e09e11e12e14
5.5 Implementation flow
- Choose whose prefixes to visit
Algorithm 1 draws a uniform random value u for each training step. If u ≤ λ, the current student samples the batch outputs; otherwise the batch uses fixed ground-truth or teacher-generated outputs. Thus λ mixes batch sources, and λ = 1 uses only current student generations. e05
- Ask the teacher about the same history
For every generated prefix, teacher and student predict the next-token distribution over the vocabulary. Their divergence is averaged over output positions. Teacher feedback is a probability distribution, not simply a corrected answer or a reward assigned to the sampled token. e04e05
- Update the student and refresh its data
Only student parameters are updated. Gradients pass through the distribution-matching loss but not through the sampling distribution. The next iteration generates with the updated student. This is on-policy imitation; distillation alone does not use a REINFORCE estimator. e05
6. Experiments & results
Generalized Knowledge Distillation (GKD) improves a language-model student by querying its teacher on prefixes the student actually generates. It separates the choice of training sequences from the choice of token-distribution divergence. T5 experiments support using on-policy data, while showing task-dependent divergence preferences and extra sampling cost. This is a training method for text generation, with no demonstrated world-model control loop.
The source contains no architecture block diagram and no conventional benchmark-results table. Its method is presented as pseudocode and equations, its numerical results as plots and annotated heatmaps, and its four tables as hyperparameter configurations. This edition therefore uses original Algorithm 1 for the mechanism, original numerical heatmaps for results and ablations, and Table A.2 as a configuration table without presenting it as experimental performance. e23
6.1 Read the original evidence
Figure 7. On-policy chain-of-thought training improves calculator-assisted final answers. Original paper, p. 8 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. The rows again specify the source of training output sequences; 100% means all student-generated chains of thought during distillation. Columns choose how teacher and student distributions are compared along those chains. Unlike the XSum heatmap, each cell is an improvement over the starting student, measured in accuracy percentage points. Read the forward-KL column vertically: 4.7 for fixed data, 6.8 for mixed data and 8.8 for on-policy data. Section 4.3 specifies four-shot prompting, greedy generation and an external calculator on the original GSM8K test split; Appendix A.5 reports final checkpoints averaged over three seeds. e12e13e14e18
What it supports. The strongest displayed Base result uses on-policy forward KL, with an 8.8-point gain over the 10.16% starting student. Reverse KL also works well on-policy, at +8 points, but it is not the best column here. This contrasts with instruction tuning, where the source favors reverse KL.
Where the evidence stops. These scores concern final answers with calculator assistance, not the correctness of every intermediate reasoning step. The teacher scores 27.9%, so this Base result does not establish that the compressed student exceeds its teacher.
Figure A.13. With greedy decoding, the large student's divergence differences nearly disappear. Original paper, p. 16 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Start with the lower panel, which holds student size at T5-large. These cells are absolute ROUGE-2 scores, unlike the improvement cells in the WMT and GSM8K heatmaps. Moving up the forward-KL column changes the training data source from fixed to mixed to on-policy, giving 20.6, 20.9 and 21.2. Moving across the top row changes the divergence while keeping all training outputs student-generated. The upper panels repeat the experiment for smaller students, with separate color scales. The caption specifies greedy evaluation and teacher/student training temperatures of one; Section 4.1 identifies the validation split. e08e09e10e16
What it supports. The large student's on-policy row spans only 21.1–21.2 ROUGE-2 across all five divergences. Its initial score is 19.6 and the teacher scores 22. This result supports the data-source intervention more clearly than a strong preference for any one divergence under greedy generation.
Where the evidence stops. Figure A.12 uses temperature sampling and a different teacher temperature during training. Comparing the two figures therefore changes more than decoding alone. Their difference cannot isolate a pure evaluation-temperature effect; Figure 4 is the direct decoding tradeoff diagnostic.
6.2 Results and evaluation conditions
| Task & protocol | Reported result | Comparison & interpretation |
|---|---|---|
| XSum summarization with greedy decoding Validation split; supervised-fine-tuned T5-XL → T5-large; teacher/student training temperatures 1. | 21.2 with on-policy forward KL. ROUGE-2, higher is better | 20.6 with fixed-data forward KL; initial student 19.6; teacher 22. Changing the source of prefixes improves this comparison. At λ = 1, all five divergences give 21.1–21.2 for the large student; divergence differences are much smaller here than under temperature sampling. e09e10 |
| WMT14 English-to-German translation Validation split; fine-tuned T5-XL → T5-small/base; beam search; final checkpoints, three-seed averages. | +0.85 for small and +0.71 for base with on-policy JSD(0.1). BLEU-point improvement over the initial student | Fixed-data JSD(0.1): +0.28/+0.38; mixed JSD(0.1): +0.71/+0.48. Initial scores: 25.58/26.98; teacher: 28. The on-policy row wins each same-divergence comparison. These are absolute BLEU increments, not percentage gains; the heatmap supplies no uncertainty intervals. e11 |
| GSM8K calculator-assisted arithmetic reasoning Original test split; fine-tuned FLAN-T5-XL → Base; four-shot CoT and greedy decoding; final checkpoints averaged across three seeds. | +8.8 points with on-policy forward KL. Accuracy improvement in percentage points | Mixed forward KL: +6.8; fixed-data forward KL: +4.7. Initial student: 10.16%; teacher: 27.9%. At fixed divergence, student-generated CoTs improve final-answer accuracy. The metric does not validate every reasoning step, and the teacher remains stronger in this Base comparison. e12e13e14 |
| Held-out instruction-following benchmarks FLAN2021 distillation, intended FLAN-T5-XL → Base; 50K steps; few-shot exact-match accuracy averaged without task weighting. | Figure A.11 shows approximately +2 percentage points on MMLU and +1.1 on BBH for on-policy reverse KL; plot-read approximations. Change in average accuracy | Reverse KL exceeds the displayed supervised KD and ImitKD baselines; on-policy forward KL decreases MMLU accuracy. The introduction reverses the benchmark order of its rounded gains. The figure caption also says T5-large where the setup and panel titles say Base. Exact student identity is inconsistent within the source, so these gains remain qualified. e18e19 |
| Same-size GSM8K self-distillation FLAN-T5-large student and same-size teacher; teacher fine-tuned on GSM8K; test evaluation with calculator. | On-policy JSD(0.9) is the highest displayed bar in Figure A.10 and exceeds the teacher. Test accuracy | The caption reports teacher accuracy 20.5%; the displayed supervised KD baseline also exceeds the teacher. Surpassing the teacher is not exclusive to on-policy learning. No precise student score is transcribed from this plot. e21 |
6.3 Ablations and diagnostic examples
Read component removals and qualitative examples within their stated evaluation conditions.
Figure 4. Changing the loss and changing decoding temperature affect the same quality–diversity tradeoff. Original paper, p. 6 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Read the horizontal axis as generation similarity: moving right means higher Self-BLEU and less varied outputs. Moving up means better ROUGE-2 reference overlap. Each color and line style identifies an evaluation temperature γ, while marker shape identifies the distillation divergence. Compare shapes within one colored group to study the loss, and compare the same shape across colors to study decoding temperature. JSD's β, used in labels such as JSD(0.9), controls the loss mixture and is different from γ. At higher temperatures, divergence choice changes quality more visibly; the low-temperature groups cluster near high similarity. e03e08
What it supports. On-policy data does not eliminate the quality–diversity choice. Mode-seeking losses can improve ROUGE-2 under more stochastic decoding while making outputs more alike. Lowering the decoding temperature also sacrifices diversity and compresses the differences between losses. The plot therefore supports selecting a loss together with the intended generation setting.
Where the evidence stops. ROUGE-2 is not factuality. For the mode-seeking interpretation, follow page 5 and Figure A.16: page 3's mean/mode wording reverses their ordering. This illustration does not establish one divergence as optimal across tasks.
Figure 6. Translation benefits from on-policy prefixes across every displayed divergence. Original paper, p. 7 ↗
Excerpt from the authors’ paper; cropped without altering the figure or table.
How to read it. Each panel fixes the student size, with T5-XL supplying teacher probabilities. Columns vary the divergence; rows vary λ from all student-generated outputs at the top to fixed outputs at the bottom. Read vertically within a column to compare data sources while keeping the loss unchanged. Read horizontally to compare losses at a fixed data mixture. The printed cells are BLEU improvements over the initial student, not final scores. The left and right color scales differ, so compare numbers rather than shade across panels. Section 4.2 specifies the validation split; Figure 6’s caption specifies beam search and averaging over three seeds. e11e15
What it supports. JSD(0.1) gives the strongest on-policy gains: +0.85 BLEU for small and +0.71 for base. The corresponding fixed-data gains are +0.28 and +0.38. Every displayed column favors the all-student row over the fixed-data row, supporting the importance of visiting the student's own prefixes under this translation protocol.
Where the evidence stops. Three-seed averaging does not supply an uncertainty interval. The initial small/base BLEU scores are 25.58/26.98 and the teacher scores 28; neither gains nor color intensity should be interpreted as percentages or evidence of equal-compute superiority.
7. Analysis & limitations
7.1 What the evidence leaves open
Evidence is concentrated on T5-family text tasks. Audio, video and text-to-image extensions are proposed future work. ROUGE and BLEU measure reference overlap; the RL factuality signal is a model-based entailment proxy, not a human verification study. e06e08e11e22
Mode-seeking can reduce output diversity. The source also contains reporting inconsistencies: the KL prose on page 3 reverses the mean/mode ordering used on pages 5 and 18; Figure A.14 labels its right panel Large but its caption says Base. These ambiguities are not silently corrected. e03e08e20
Student sampling adds reported GSM8K overhead of approximately 1.8×, 2× and 2.2× for the three model-size ratios. Equal update counts therefore do not establish equal-compute superiority. e15
7.2 Questions for discussion
- Does refreshing student-generated prefixes still win when generation and teacher-query compute are matched?
- Does a divergence that improves calculator-assisted answer accuracy also improve the validity of intermediate reasoning?
8. Reproducibility audit
8.1 Requirements and known gaps
A GSM8K reproduction requires the roughly 5.3K PaLM-generated CoT pairs, FLAN-T5 checkpoints fine-tuned for 10K steps, and the supplied four-shot prompt. Distillation uses 40K steps, batch 32, learning rate 0.0003, dropout 0.05, input/output limits 512/320, 2K warmup and linear cooldown over 30K–40K steps. e12e14
Task settings must remain separate: XSum uses 40K distillation steps and lengths 1024/64, with learning-rate and teacher-temperature exceptions in Appendix A.4. WMT uses 100K steps and lengths 80/80. FLAN2021 uses 50K steps, batch 128 and learning rate 0.0001. e16e17e18
Appendix A.3 mentions eight TPUv3 cores for initial supervised fine-tuning, without a full per-experiment hardware budget. WMT beam settings are delegated to a cited paper; calculator parsing and software versions are not specified here. A faithful replication needs these choices documented before comparing scores. e12e15e17
8.2 Proposed reproduction checks
The following checks are proposals motivated by the paper. They have not been run as part of this reading.
Check 1: Does refreshing student prefixes explain the GSM8K advantage?
Reader-proposed check, not performed: start all runs from the same supervised-fine-tuned FLAN-T5-Base checkpoint and use the same teacher, prompts, forward KL and Table A.2 settings. Compare λ = 0, λ = 1 with fresh student generations, and a fixed cache generated once by the initial student. Evaluate calculator-assisted exact answer accuracy over three seeds, plus teacher–student divergence on a common held-out set of student prefixes. First match update counts; then match total generation, teacher-query and optimization compute. A fresh-generation advantage over both fixed conditions supports the adaptive-prefix explanation. If it disappears under matched compute or a fixed student cache, the proposed mechanism or efficiency benefit needs qualification. e04e05e12e13e14e15
Check 2: Separate teacher-temperature effects from decoding effects on XSum
Reader-proposed check, not performed: distill the same T5-large initialization from T5-XL with λ = 1, crossing forward KL versus JSD(0.9) with teacher temperatures 0.1 versus 1. Hold prompts, update count, learning rate and all remaining XSum settings fixed. Evaluate every trained checkpoint with greedy decoding and the four temperatures used in Figure 4; measure validation ROUGE-2 and Self-BLEU, repeating training seeds. This factorial comparison separates training-target sharpening from evaluation randomness, which differ together between Figures A.12 and A.13. If JSD's quality advantage remains at matched teacher temperature but shrinks under greedy decoding, the decoding-dependent interpretation gains support. A reversal driven mainly by teacher temperature would qualify it. e03e08e09e10e16
8.3 Reading coverage
Visual audit: Visually inspected the title/author page, method equations and Algorithm 1, all original figures (1–9 and A.10–A.16), and Tables A.1–A.4. Each of the six final crops was inspected separately with its axes, labels, legends or complete table rows retained. The algorithm's u ≤ lambda branch agrees with the mixture objective; the no-sampling-gradient rule is explicit in the surrounding formulation. Non-cropped pages include all cited checkpoint, compute, prompt, RL and instruction-tuning details. Source inconsistencies in KL prose and appendix labels are disclosed in the report. References on pages 10–12 were read as text; external material was not inspected. For WMT, page 7 contains both Section 4.2 (validation split) and Figure 6’s caption (beam search and three-seed averaging); page 17 supplies Appendix A.6 (final checkpoints).
PDF pages inspected for this edition: 1, 2, 3, 4, 5, 6, 7, 8, 9, 13, 14, 15, 16, 17, 18. Appendix coverage: reviewed.
Original figures and tables remain the work of the source’s authors. Extractions preserve their scientific content; any HTML wrapper layout is disclosed with each figure. The surrounding reading notes are our own.
Text reading scope & known omissions
- Abstract; Section 1 Introduction (pp. 1–2)
- Section 2 Preliminaries (pp. 2–3)
- Section 3 Distillation for Auto-Regressive Sequence Models; 3.1 GKD; 3.2 RL Fine-Tuning + On-Policy GKD (pp. 3–5)
- Section 4 Experiments; 4.1 XSum; 4.2 WMT; 4.3 GSM8K (pp. 5–8)
- Sections 5 Related Work and 6 Conclusion (p. 9)
- References (pp. 10–12)
- Appendix A.1 Self-Distillation; A.2 Task-Agnostic Distillation; A.3 T5 Models (pp. 13–14)
- Appendix A.4 XSum; A.5 GSM8K including the complete four-shot prompt (pp. 14–17)
- Appendix A.6 WMT; A.7 Instruction Tuning; A.8 Mode-Seeking vs Mode-Covering KL (pp. 17–18)
Outside the original text pass
- Text extraction does not reconstruct figure images; inspect the retained PDF for figures and equation/table layout.
- Separate supplemental material availability has not been fully verified.
- The extraction limitation above was addressed by inspecting original PDF pages 1–9 and 13–18, including all figures, all tables and Algorithm 1. References on pages 10–12 were read as text.
- Identity notes: the title page identifies this artifact as an ICLR 2024 conference paper. Its title and seven-author sequence match the catalog, except that it credits Sabela Ramos rather than the catalog form Sabela Ramos Garea. The supplied observedTitle preserves extraction spacing and capitalization. No revision number or revision date is stated; no other edition was supplied or compared.
- No code, externally linked checkpoints, cited papers or separate supplements were inspected. No training or evaluation was reproduced.
The visual audit above records the subsequent illustrated pass.
8.4 Traceable evidence
e01PDF p. 1, title, author block and conference header
The artifact is titled On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes and marked ICLR 2024. Credits list Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist and Olivier Bachem; affiliations are Google DeepMind, Mila and University of Toronto.
Go to primary source ↓e02PDF p. 2, Section 1, motivation paragraphs
Fixed output datasets create a mismatch with prefixes visited during student generation; limited student capacity motivates alternatives to standard forward-KL distillation.
Go to primary source ↓e03PDF p. 3, Section 2 and Eq. (1); p. 5, Choice of Divergence; p. 18, Appendix A.8/Figure A.16
Forward KL is D_KL(P||Q), reverse KL D_KL(Q||P). Generalized JSD uses beta P + (1-beta)Q and approaches forward/reverse-like gradients near beta zero/one. Page 3 pairs reverse/forward with mean/mode in a conflicting order; page 5 and Figure A.16 identify reverse KL as mode-seeking and forward KL as mode-covering.
Go to primary source ↓e04PDF p. 3, Section 3 Problem Setup, Eq. (2), Eq. (3), and baseline definitions
The models share input and prefix conditioning. Equation (2) averages token-distribution divergence over output length. SeqKD fits teacher-generated sequences; supervised KD uses teacher distributions on fixed sequences.
Go to primary source ↓e05PDF p. 4, Algorithm 1 lines 1–11; Eq. (4), boxed GKD objective and Remark
The u ≤ lambda branch samples student outputs, the other branch fixed pairs. Student parameters minimize the batch divergence. Sampling is not differentiated; student generation temperature is one. The main experiments assume a supervised-fine-tuned initial student.
Go to primary source ↓e06PDF p. 5, Section 3.2/Eq. (5); p. 7, Figure 5 and On-policy GKD with RL
Alpha weights distillation against expected reward. XSum combines JSD(0.9) with a T5-XXL NLI entailment reward; increasing alpha trades entailment improvement for ROUGE-2. RLEF* regularizes toward the initial student instead of the teacher.
Go to primary source ↓e07PDF p. 5, Section 4 Student/Teacher Models and Baselines; p. 14, Appendix A.3
Task-specific teacher is about 3B parameters; students are 77M, 250M and 800M. Baselines share initial fine-tuned checkpoints. XSum/WMT use LM-adapted T5v1.1; GSM8K uses FLAN-T5. Adafactor and task-specific fine-tuning stages are described.
Go to primary source ↓e08PDF p. 6, Figure 4, its caption and Section 4.1
XSum is evaluated on validation ROUGE-2. Figure 4 varies evaluation temperature 0.1, 0.3, 0.5 and 1.0 and five divergences for on-policy T5-small. Higher Self-BLEU means more similar generations; lower temperature reduces diversity and narrows quality differences.
Go to primary source ↓e09PDF p. 16, Figure A.13, T5-large heatmap and caption; p. 6, Section 4.1 evaluation split
Greedy XSum, training temperatures one: large-student forward KL scores are 21.2/20.9/20.6 for student-data fractions 100/50/0%. Its on-policy row is 21.2, 21.2, 21.1, 21.2, 21.1. Caption gives initial large student 19.6 and teacher 22; validation split is specified in Section 4.1.
Go to primary source ↓e10PDF p. 15, Figure A.12, heatmaps and caption
Temperature-sampled XSum evaluation uses gamma one and teacher training temperature 0.1. Large-student on-policy forward KL gives 19.2 and JSD(0.9) gives 20.3. These settings differ from the greedy protocol in Figure A.13.
Go to primary source ↓e11PDF p. 7, Figure 6, both heatmaps/caption and Section 4.2; p. 17, Appendix A.6
Section 4.2 specifies WMT14 en-de validation evaluation; Figure 6’s caption specifies beam search and three-seed averages. JSD(0.1) gains for small/base are 0.85/0.71 on-policy, 0.71/0.48 mixed, and 0.28/0.38 fixed-data. Initial BLEU is 25.58/26.98 and teacher BLEU 28. Appendix A.6 specifies final checkpoints and repeats the three-seed averaging protocol. The heatmaps do not show intervals.
Go to primary source ↓e12PDF p. 8, Section 4.3 Setup; p. 15, Appendix A.5
GSM8K uses four-shot CoT, original test answers checked with an external calculator, approximately 5.3K PaLM-generated training CoTs, and 10K-step fine-tuned FLAN-T5 initialization. Final distilled checkpoints are averaged across three seeds.
Go to primary source ↓e13PDF p. 8, Figure 7 heatmap and caption
Base-student accuracy gains for forward KL are 8.8, 6.8 and 4.7 for 100%, 50% and 0% student data. The on-policy reverse-KL gain is 8. Caption reports teacher/student starting accuracy 27.9/10.16.
Go to primary source ↓e14PDF p. 16, Table A.2; pp. 16–17, Few-shot CoT Prompt
GSM8K settings: 40,000 steps, batch 32, dropout 0.05, learning rate 0.0003, 2,000 warmup steps, linear cooldown 30,000–40,000, lengths 512/320, FLAN-T5, teacher temperature 0.1 and greedy evaluation. Four complete prompt examples are supplied.
Go to primary source ↓e15PDF p. 14, Appendix A.3 Computational cost of GKD
Initial supervised fine-tuning is described as a few hours on an eight-core TPUv3. GSM8K on-policy sampling overhead is approximately 1.8×/2×/2.2× versus fixed outputs for teacher-to-student size ratios 38×/12×/3.8×; a full per-experiment hardware or software inventory is absent.
Go to primary source ↓e16PDF p. 14, Appendix A.4 Learning rate sweep; p. 15, teacher-temperature paragraph and Table A.1
XSum distillation uses 40K steps, batch 32, no dropout and lengths 1024/64. Default learning rate is 0.0003; T5-small uses 0.001 except reverse KL uses 0.0003. Teacher temperature is one for greedy evaluation experiments and 0.1 for temperature-sampled results.
Go to primary source ↓e17PDF p. 17, Appendix A.6; p. 18, Table A.3
WMT uses 100K steps, batch 32, learning rate 0.0003, 5K linear warmup, zero dropout, lengths 80/80 and teacher temperature one. Beam-search hyperparameters are referred to Raffel et al. rather than enumerated.
Go to primary source ↓e18PDF p. 13, Appendix A.2; p. 14, Figure A.11; p. 18, Appendix A.7/Table A.4
FLAN2021 contains 5.36M examples from 62 tasks. Setup and panel titles specify XL-to-Base distillation, evaluated on held-out MMLU 57 tasks and BBH 23 tasks by unweighted few-shot exact match. Reverse-KL on-policy bars are about +2 and +1.1 points respectively. Training uses 50K steps, batch 128, LR 0.0001, no warmup/dropout and lengths 2048/256.
Go to primary source ↓e19PDF p. 2, introduction paragraph reporting task-agnostic gains; p. 14, Figure A.11 titles and caption; p. 13, Appendix A.2 Setup
The introduction gives rounded gains in BBH/MMLU order as 2%/1%, whereas the plot shows roughly 2 points MMLU and 1 point BBH. Figure A.11 caption calls the student T5-large, conflicting with Base in its panel titles and the setup.
Go to primary source ↓e20PDF p. 17, Figure A.14 right-panel title and caption
The right heatmap is labeled T5-XL → Large, while the caption describes that panel as T5-base. This mismatch prevents unqualified attribution from the caption alone.
Go to primary source ↓e21PDF p. 13, Appendix A.1 and Figure A.10
Same-architecture, same-size FLAN-T5-large self-distillation is tested on GSM8K. The JSD(0.9) on-policy bar is highest; all displayed distilled models exceed the teacher line. The caption reports teacher accuracy 20.5%.
Go to primary source ↓e22PDF p. 9, Sections 5–6; pp. 3–4, method formulation
The method concerns autoregressive language-model distillation and imitation-learning connections. Extensions to audio, video and text-to-image are future directions; no environment-dynamics model or robot-action execution is introduced.
Go to primary source ↓e23PDF p. 4, Algorithm 1; pp. 1, 6–8 and 13–18, Figures 1–A.16 and Tables A.1–A.4
The source presents its method as pseudocode and equations, empirical results as plots/heatmaps, and the four numbered tables as hyperparameter settings. It contains no architecture block diagram or conventional tabulated benchmark-results table.
Go to primary source ↓8.5 Primary sources
ON-POLICY DISTILLATION OF LANGUAGE MODELS : LEARNING FROM SELF -G ENERATED MISTAKES ↗
PDF · 9,941 extracted words
Source fingerprint
8d116b2d5dfa5151d8877b411e09de6421f90ff98e4ba000a972b5b0cd55c20f