On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes

ArXiv: 2306.13649

🎯 Pitch

Training a student language model on its own self-generated outputs rather than on a fixed dataset can double distillation gains, slashing the performance gap to the teacher. This simple on-policy approach not only fixes the train–inference mismatch but also unlocks flexible divergences and seamless integration with RL fine-tuning—without needing a separate data generation phase.

---

1. Executive Summary

This paper proposes Generalized Knowledge Distillation (GKD), a framework for distilling auto-regressive language models that trains the student on its own self-generated output sequences using token-level feedback from the teacher, explicitly addressing the train-inference distribution mismatch that plagues supervised KD approaches. Evaluated on T5 models across summarization (XSum), machine translation (WMT), arithmetic reasoning (GSM8K), and instruction tuning (FLAN), GKD introduces two key degrees of freedom: the choice of output sequences—ranging from purely on-policy student-generated to purely fixed supervised data—and the choice of divergence between teacher and student distributions—spanning forward KL, reverse KL, and generalized JSD (interpolating between mean-covering and mode-seeking behavior). On-policy GKD yields average relative gains of 2.1× on summarization, 1.7× on translation, and 1.9× on reasoning over baseline KD methods, and further demonstrates that distillation can be seamlessly combined with RL fine-tuning (improving both factual consistency and summary quality on XSum), establishing that student-generated on-policy data consistently outperforms fixed datasets across tasks while the optimal divergence remains task-dependent (forward KL dominating on GSM8K with greedy decoding, reverse KL excelling on instruction tuning).

2. Context and Motivation

The Core Problem: Distribution Mismatch in Auto-Regressive Distillation

The fundamental problem this paper addresses is deceptively simple: when you train a smaller language model to imitate a larger one, the training conditions do not match the conditions the student will face during actual use. This mismatch arises from the auto-regressive nature of sequence generation. At inference time, a language model predicts tokens one by one, with each prediction conditioned on its own previously generated tokens. If any of those earlier tokens were improbable during training—or if the model makes an error early in the sequence—the subsequent tokens are generated from states the model never encountered during supervised learning. These errors compound, producing low-quality outputs that diverge progressively further from what the teacher would generate.

In the language of imitation learning, this is called exposure bias or distribution mismatch (Pomerleau, 1991; Ross & Bagnell, 2010). During standard supervised knowledge distillation, the student is trained exclusively on ground-truth sequences or teacher-generated sequences—states that represent expert trajectories. At inference time, however, the student navigates its own trajectory through token space, visiting intermediate states very different from those in the training corpus. The paper formalizes this: standard KD minimizes

$L_{SD}(\theta) = \mathbb{E}_{(x,y) \sim (X,Y)} \left[ D_{KL}\left(p_T \| p_S^\theta\right)(y|x) \right]$

where the expectation is over a fixed dataset $(X, Y)$. The student learns to mimic the teacher on the teacher's preferred paths, but never practices recovering from its own mistakes on paths it would actually take during deployment.

This is not a theoretical curiosity. The paper cites multiple empirical demonstrations that exposure bias leads to poor text generation (Zhang et al., 2019; Chiang & Chen, 2021; Arora et al., 2022). In the distillation setting specifically, a student trained via supervised KD has never seen what happens when it generates a slightly suboptimal token at step 3 and must continue generating tokens 4 through 50 from that compromised state. At inference time, this happens constantly. The gap between training and inference distributions is the central obstacle that GKD is designed to eliminate.

Why This Problem Matters Now

The problem's importance has grown dramatically with the scaling of language models. The paper situates itself in a landscape where scaling model parameters is the dominant path to improved capability (Kaplan et al., 2020), yet this scaling comes at a direct cost: deployment is limited by either inference cost or memory footprint. A model that requires a datacenter-scale GPU cluster to serve is economically and practically inaccessible for many applications. Knowledge distillation—compressing a large teacher into a smaller student—is therefore not an academic exercise but a practical necessity for putting capable language models into production.

The paper's motivation is reinforced by several trends that make robust distillation increasingly critical:

• On-device deployment: Running language models on phones, laptops, or embedded devices requires models orders of magnitude smaller than the largest available teachers. If distillation leaves performance on the table due to distribution mismatch, those deployments underperform relative to what is achievable with the same parameter budget.

• Self-improving systems: When smaller models are used to generate training data for themselves (a loop the paper explicitly connects to in discussing RL fine-tuning and self-distillation in Appendix A.1), the quality of their own outputs determines the ceiling of improvement. A model that generates poor-quality sequences during self-training cannot bootstrap itself effectively.

• The alignment era: With RLHF and RLAIF becoming standard post-training steps (Ouyang et al., 2022), distillation that can coexist with reinforcement learning—sharing the same on-policy sampling infrastructure—becomes architecturally valuable. The paper explicitly targets this synergy.

• Task-agnostic compression: As models like FLAN (Chung et al., 2022) demonstrate, a single model can handle hundreds of tasks via instruction following. Distilling such models requires robustness across diverse output distributions, where a one-size-fits-all supervised approach is particularly vulnerable to mismatch.

The practical stakes are clear: a $4\times$ smaller model that achieves the same performance as a larger one represents enormous cost savings at scale. But achieving that compression ratio requires solving the distribution mismatch problem, because the student's limited capacity means it cannot simply memorize the teacher's output distribution—it must learn to generalize, and generalization requires training on states it will actually encounter.

Where Prior Approaches Fall Short

The paper identifies specific limitations in existing distillation methods along three axes:

1. Supervised KD uses fixed data, guaranteeing mismatch.

Supervised KD (Hinton et al., 2015; Sanh et al., 2019)—by far the most widely used approach—trains the student to match the teacher's token-level probability distributions on a fixed dataset of input-output pairs. This produces a rich training signal (full distributions rather than just hard targets), but it fundamentally cannot address distribution mismatch because the student never trains on its own outputs. The sequences in the training set are "expert demonstrations" from either the teacher or human annotators. The student never experiences the consequences of its own errors during training, so it learns no recovery strategies. This is the train-inference discrepancy that the paper identifies as the root cause of supervised KD's limitations.

2. Sequence-level KD (SeqKD) swaps one fixed dataset for another.

SeqKD (Kim & Rush, 2016) addresses a different problem: it trains on high-probability sequences generated by the teacher, which can be more diverse than human-written references. However, this still uses a fixed set of output sequences. The student trains on states from the teacher's distribution, not its own. The distribution mismatch persists because the partial sequences encountered during the student's auto-regressive generation at inference time can be quite different from the teacher-generated sequences in the training set. SeqKD is expensive (requiring generation from the larger teacher model) without solving the fundamental problem.

3. ImitKD and f-distill partially recognize the problem but don't fully address it.

ImitKD (Lin et al., 2020) identifies the connection between distillation and imitation learning—an insight the paper builds on substantially. ImitKD samples sequences from both the student and a fixed dataset, mixing them during training. However, the paper argues ImitKD does not push the idea far enough in three ways: (a) it never explores purely on-policy data collection (its student data fraction never reaches 1.0, defaulting to a 50-50 mix), (b) it restricts itself to forward KL divergence, missing the opportunity to use mode-seeking divergences that better suit capacity-limited students, and (c) it doesn't integrate with RL fine-tuning. ImitKD can be viewed as a special case of GKD with forward KL and $\lambda = 0.5$.

f-distill (Wen et al., 2023) frames sequence-level KD as minimizing an f-divergence and proposes a tractable objective based on total variation distance. Like ImitKD, it uses mixed data but is restricted in divergence choice and doesn't explore full on-policy training. The paper shows empirically that both ImitKD and f-distill underperform on-policy GKD (Figures 2, 9).

4. The field lacks a unifying framework for distillation design choices.

Prior work treated the choice of training data (teacher-generated vs. ground-truth vs. mixed) and the choice of objective (forward KL vs. alternatives) as separate, independent decisions. There was no framework connecting these choices or enabling systematic exploration of their interactions. This fragmentation meant that when a particular approach failed, it was unclear whether the culprit was the data distribution, the divergence, or their interaction. The paper argues this obscures the underlying structure of the distillation problem.

5. Concurrent work on MiniLLM takes an RL approach but has limitations.

The concurrent work MiniLLM (Gu et al., 2023), published around the same time, also exploits the imitation learning connection and frames distillation as minimizing reverse KL at the sequence level using policy gradient. The paper acknowledges this connection but argues GKD is simpler and more stable because it does not backpropagate through the student's sampling process. MiniLLM requires multiple stabilizing tricks to handle high gradient variance, reward hacking, and generation length bias. GKD's approach—computing token-level divergences on sampled sequences without differentiating through the sampling step—is closer to supervised training and avoids these complications. Moreover, GKD is more general, supporting forward KL, reverse KL, and JSD, whereas MiniLLM is locked to reverse KL (which the paper shows is not always optimal—forward KL outperforms on GSM8K with greedy decoding).

6. No prior work combines distillation with RL fine-tuning.

The paper identifies a specific gap: while RL fine-tuning (RLHF, RLAIF) has become standard for aligning language models, and distillation has been studied extensively, no prior work performs both simultaneously. In standard RLHF, the policy is regularized to stay close to its initial (usually supervised fine-tuned) checkpoint. The paper proposes replacing this self-regularization with teacher-regularization—simultaneously maximizing a reward signal while staying close to a larger teacher model. This combination is natural from an optimization perspective but represents genuinely unexplored territory.

How This Paper Positions Itself

The paper frames distillation for auto-regressive models through a specific lens: it is an imitation learning problem with an interactive expert. This is not merely an analogy—it is the organizing principle that generates the entire GKD framework. In imitation learning, the fundamental challenge is that supervised learning on expert demonstrations produces policies that drift off the expert trajectory and then encounter states for which they have no training data. The solution in IL is on-policy data collection: let the learner act in the environment, observe where it goes, get expert feedback on those states, and retrain (Ross et al., 2011). The paper transplants this logic directly to distillation:

"Instead of solely relying on a fixed set of output sequences, GKD trains the student on its self-generated output sequences by leveraging feedback from the teacher on such sequences."

The "environment" is the auto-regressive generation process. The "expert" is the teacher model, which provides token-level probability distributions at every state the student visits. The "learner" is the student model, which generates sequences, receives feedback, and retrains.

This framing positions the paper not as proposing a single new method, but as unifying and generalizing existing approaches while instantiating new ones that substantially outperform them. The unification is explicit: supervised KD is GKD with $\lambda = 0$ and forward KL. SeqKD is supervised FT on teacher-generated outputs (a special case of the fixed-data path in GKD). ImitKD is GKD with forward KL and $\lambda = 0.5$. f-distill is a particular divergence choice within the GKD framework. The genuinely new territory is:

• Purely on-policy distillation ($\lambda = 1$), which has not been previously explored and consistently outperforms baselines.
• Alternative divergences (reverse KL, generalized JSD) in combination with on-policy data, which the paper shows is a powerful pairing—mode-seeking divergences complement on-policy training because the student's limited capacity makes it impossible to cover the full teacher distribution.
• Integration with RL fine-tuning, allowing simultaneous reward maximization and distillation.

The paper also positions itself within a broader trend: the increasing use of self-generated data for improving language models. It explicitly cites Ouyang et al. (2022) and Singh et al. (2023) as evidence that fine-tuning on model-generated outputs works, lending external credibility to the on-policy approach. The GKD framework provides a principled distillation-specific instantiation of this general idea.

A crucial design choice that distinguishes GKD from RL-based approaches like MiniLLM: gradients do not flow through the student's sampling process. The paper argues this makes training "stable and computationally efficient"—a practical consideration that matters when training models at scale. By treating the sampling step as a data collection mechanism rather than part of the computational graph, GKD avoids the high-variance gradient estimates and reward hacking that plague policy gradient methods, while still achieving the benefits of on-policy training.

Finally, the paper positions its empirical contribution around scalability of the on-policy idea. ImitKD experimented with students up to roughly 30M parameters (BERT-scale models). GKD demonstrates the approach with T5 models up to 800M parameters—roughly 26× larger—on tasks spanning summarization, translation, reasoning, and instruction tuning. This scaling dimension is important because distribution mismatch becomes more severe as the capacity gap between teacher and student widens; the fact that on-policy training helps more for the smallest students (38× smaller than the teacher) is consistent with the imitation learning theory that motivates GKD.

3. Technical Approach

3.1 Reader Orientation

The system is a training procedure —not a model architecture—that teaches a smaller language model (student) to replicate the behavior of a larger one (teacher) by having the student practice generating its own outputs and receiving token-by-token guidance from the teacher on those self-generated sequences. The problem it solves is that standard distillation trains the student only on "perfect" sequences from the teacher or dataset, so when the student actually generates text at inference time and makes small errors, it encounters unfamiliar situations and compounds its mistakes. The solution's shape is an on-policy loop: the student generates sequences → the teacher scores every token in those sequences → the student updates to better match the teacher on exactly the states it actually visits during generation → repeat, with the student's outputs improving over time.

3.2 Big-Picture Architecture (Diagram in Words)

The GKD system consists of four major components interacting in a training loop:

1. Student Model ($p_S^\theta$) —the smaller auto-regressive language model being trained. It has learnable parameters $\theta$ and generates output sequences token-by-token. At the start of training, it is already supervised fine-tuned on the target task (the paper uses a "warm-started" student, not randomly initialized).

2. Teacher Model ($p_T$) —the larger, frozen language model providing supervision. It takes an input $x$ and any partial output sequence $y_{<n}$ and outputs a full token-level probability distribution $p_T(\cdot|y_{<n}, x)$ over the vocabulary. This distribution is the "expert feedback" for every state the student visits.

3. Data Sampler —a mechanism that decides, for each training batch, whether to use (a) student-generated sequences (sampled from $p_S^\theta(\cdot|x)$ at temperature $\gamma = 1$) or (b) fixed dataset sequences drawn from $(X, Y)$. The mixing ratio is controlled by a hyperparameter $\lambda \in [0, 1]$, the fraction of on-policy student data. A random draw $u \sim \text{Uniform}(0, 1)$ at each step determines the source: if $u \leq \lambda$, the student generates fresh outputs; otherwise, sequences are sampled from the fixed dataset.

4. Divergence Computer —takes the teacher's token-level distribution and the student's token-level distribution at each position in a generated sequence, computes a chosen divergence $D(p_T \| p_S^\theta)$ between them, and produces a scalar loss. The divergence can be forward KL, reverse KL, or generalized JSD with parameter $\beta \in (0, 1)$. The loss is averaged across all tokens in the sequence and across the batch, then used for a gradient update.

Information flow per training step: Input prompts $x$ are drawn from $X$ → if $u \leq \lambda$, the student auto-regressively generates an output sequence $y \sim p_S^\theta(\cdot|x)$ (sampling, no gradient tracking through the sampling step); otherwise, a pair $(x, y)$ is drawn from the fixed dataset → the teacher computes log-probabilities at every token position in $y$ → the divergence computer computes the chosen $D(p_T \| p_S^\theta)(y|x)$ for each $(x, y)$ pair → the loss is averaged over the batch → an optimizer updates $\theta$. Critically, gradients do not backpropagate through the student's sampling process —the student's generated tokens are treated as data, not as part of the computational graph. This is the design choice that makes GKD stable (no high-variance policy gradients) while still providing the benefits of on-policy training.

3.3 Roadmap for the Deep Dive

To understand GKD thoroughly, I will explain the components in this order:

• First, the loss function formalism —what exactly is being minimized, how the token-level divergences are aggregated across a sequence, and what degrees of freedom the $D$ and $\lambda$ parameters introduce. This is the mathematical core of the framework.

• Second, the on-policy data generation mechanism —how the student generates training sequences, why gradients do not flow through sampling, what temperature is used, and how this connects to the imitation learning principle of DAgger (Ross et al., 2011).

• Third, the unified GKD objective —how the on-policy and supervised components combine into a single loss controlled by $\lambda$, and why this unification matters for practical deployment.

• Fourth, the choice of divergence $D$ —the forward KL (mode-covering), reverse KL (mode-seeking), and generalized JSD (interpolating between them), including the mathematical forms, intuitive behaviors, and task-dependent tradeoffs the paper observes.

• Fifth, the integration with RL fine-tuning —how on-policy GKD naturally combines with REINFORCE-style reward maximization to simultaneously improve task performance and optimize a non-differentiable reward, with the teacher replacing the standard self-regularization term used in RLHF.

• Sixth, the practical training loop and hyperparameters —the exact algorithm, including how $\lambda$, batch size, learning rates, and temperature interact, and what defaults the paper uses across tasks.

This order builds from the abstract objective → the data generation mechanism → the full objective → the divergence choices → the RL extension → the concrete implementation, each step adding context that makes the next interpretable.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methodological contribution paper whose core idea is that knowledge distillation for auto-regressive models should be treated as an on-policy imitation learning problem, where the student trains on its own generated sequences using token-level teacher probabilities as expert feedback, and that the choice of divergence between student and teacher distributions is a critical free parameter that interacts with task characteristics.

---

The Token-Level Divergence Between Student and Teacher

The paper defines a generic divergence between two auto-regressive sequence models on a specific input-output pair $(x, y)$. This is the atomic unit of the loss—the thing that gets averaged over sequences and batches.

For any divergence $D$ between two probability distributions, the sequence-level discrepancy is defined as:

$D\left(p_T \| p_S^\theta\right)(y|x) := \frac{1}{L_y} \sum_{n=1}^{L_y} D\left(p_T(\cdot|y_{<n}, x) \| p_S^\theta(\cdot|y_{<n}, x)\right)$

where $L_y$ is the length of the output sequence $y$ in tokens, $y_{<n}$ denotes the sequence up to (but not including) the $n$-th token, $p_T(\cdot|y_{<n}, x)$ is the teacher's next-token probability distribution over the entire vocabulary given the input $x$ and the partial output $y_{<n}$, and $p_S^\theta(\cdot|y_{<n}, x)$ is the student's corresponding distribution.

What it computes: at each position $n$ in the output sequence, it computes the chosen divergence $D$ between two probability vectors of dimension $|V|$ (the vocabulary size)—the teacher's distribution and the student's distribution, both conditioned on the exact same prefix $y_{<n}$. It then averages these per-token divergences across all $L_y$ positions to produce a single scalar for the sequence. Operationally, this means that for a sequence of length 50, the loss is the mean of 50 individual divergence computations.

Why this form: the per-token decomposition is what makes the training signal "rich" compared to sequence-level approaches. The student receives feedback at every generation step, not just a single score at the end. This is critical because errors in auto-regressive generation are path-dependent—a mistake at token 5 affects all subsequent tokens. By providing supervision at every intermediate state $y_{<n}$ that the student actually visits (when $y$ is student-generated), the teacher can correct the specific tokens where the student deviates. The length normalization (dividing by $L_y$) ensures that longer sequences do not dominate the loss simply because they have more tokens. This is a standard practice in sequence modeling and prevents the optimization from being biased toward longer outputs.

When $D$ is instantiated as forward KL, this becomes:

$D_{KL}\left(p_T \| p_S^\theta\right)(y|x) = \frac{1}{L_y} \sum_{n=1}^{L_y} \sum_{v \in V} p_T(v|y_{<n}, x) \log \frac{p_T(v|y_{<n}, x)}{p_S^\theta(v|y_{<n}, x)}$

which is the standard supervised KD loss applied token-by-token. When $D$ is reverse KL, the arguments swap:

$D_{KL}\left(p_S^\theta \| p_T\right)(y|x) = \frac{1}{L_y} \sum_{n=1}^{L_y} \sum_{v \in V} p_S^\theta(v|y_{<n}, x) \log \frac{p_S^\theta(v|y_{<n}, x)}{p_T(v|y_{<n}, x)}$

This subtle swap has profound consequences for how the student allocates its limited capacity, which I will explain in the divergence section below.

---

On-Policy Data Generation: The Student Samples Its Own Training Sequences

The defining mechanism of on-policy GKD is that the student model generates the output sequences used for training. This is where the imitation learning connection becomes concrete.

The generation process. Given an input $x$ drawn from the dataset $X$, the student auto-regressively generates an output sequence $y$ by sampling from its own current policy:

$y \sim p_S^\theta(\cdot|x)$

At each step $n$, the student computes a next-token distribution $p_S^\theta(\cdot|y_{<n}, x)$ over the vocabulary, then samples a token from this distribution with temperature $\gamma = 1$:

$p_S^\theta(v|y_{<n}, x) = \frac{\exp(z_v / \gamma)}{\sum_{i=1}^{M} \exp(z_i / \gamma)}$

where $z_v$ is the logit score for token $v$, $M = |V|$ is the vocabulary size, and $\gamma = 1$ means the distribution maintains the full diversity of the model's predictions. A higher temperature produces more diverse outputs (flatter distribution); a lower temperature makes the output more deterministic (peaked around high-probability tokens). The paper uses $\gamma = 1$ during training "to encourage diversity in student generated sequences." This diversity is important because it exposes the student to a broader range of its own potential errors, providing richer training data.

The critical design choice: no backpropagation through sampling. The paper explicitly states: "we do not backpropagate through the student's sampling distribution $p_S(\cdot|x)$, similar to on-policy imitation." This means the student's generated tokens are treated as fixed data for the purpose of the loss computation. The teacher computes log-probabilities on whatever tokens the student happened to generate, the loss is computed, and gradients flow through the student's log-probability computation at those tokens—but there is no gradient path from the loss back through the discrete sampling operation that selected which token to generate.

Why this matters operationally. Discrete sampling from a categorical distribution is not differentiable (the operation is a hard argmax or a random draw, neither of which has a gradient). Policy gradient methods like REINFORCE work around this by using the score function estimator, which introduces high variance and requires variance reduction techniques. MiniLLM (Gu et al., 2023), the concurrent work, uses exactly this approach and "relies on a number of stabilizing tricks, to tackle high variance, reward hacking, and generation length bias." GKD avoids this entire complexity by treating the student's sampling step as uncoupled from the computational graph. The loss is still a function of $\theta$ through the student's log-probability computation at the generated tokens, so gradients flow—but only through the "evaluation" of those tokens, not through the "decision" of which tokens to generate.

Why this is stable and efficient. The computation is essentially supervised learning on dynamically generated data. At each training step: (1) sample sequences from the student (inference only, no gradients), (2) compute the divergence between student and teacher at each token position (gradients through the student's output layer), (3) update. There are no importance weights, no baseline corrections, no reward normalization—just a standard cross-entropy-like update on self-generated data. The paper argues this makes GKD "closer to supervised training" and therefore simple to implement and tune.

The imitation learning connection. This design directly mirrors DAgger (Dataset Aggregation; Ross et al., 2011), a foundational imitation learning algorithm. In DAgger, the learner executes its current policy in the environment to collect states, an expert labels those states with optimal actions, and the learner is retrained on the aggregated dataset of expert-labeled states. The key insight that DAgger introduced over simpler behavioral cloning is that the states the learner visits during its own execution are different from the states in the expert's demonstration trajectories, so the learner must be trained on its own state distribution to learn recovery behaviors. GKD applies this logic to auto-regressive distillation: the "states" are partial output sequences $y_{<n}$, the "actions" are next tokens, and the "expert" is the frozen teacher model. By continually generating fresh training data from the student's current policy throughout training, GKD ensures the student practices on the states it will actually encounter at inference time.

Computational cost considerations. Generating sequences from the student is cheaper than generating from the teacher because the student has fewer parameters. The paper reports that on GSM8K, "the computational overhead from student sampling is approximately 1.8×, 2× and 2.2× compared to sampling from a fixed dataset of outputs, for a student-teacher [size] ratio of 38×, 12× and 3.8×." This overhead is the cost of running student inference during training to generate on-policy data, versus simply loading pre-generated sequences from disk. The paper argues this overhead is acceptable because (a) it is small relative to the teacher's inference cost, (b) the performance gains outweigh the cost, and (c) "the majority of cost in real world use cases is due to serving cost at inference time and not due to fine tuning"—if deploying the model is already expensive, slightly more expensive training is a good trade.

---

The Generalized Knowledge Distillation (GKD) Objective

GKD unifies on-policy and supervised distillation into a single objective controlled by the student data fraction $\lambda$. The full GKD loss is:

$\mathcal{L}_{\text{GKD}}(\theta) := (1 - \lambda) \mathbb{E}_{(x,y) \sim (X,Y)} \left[ D(p_T \| p_S^\theta)(y|x) \right] + \lambda \mathbb{E}_{x \sim X} \left[ \mathbb{E}_{y \sim p_S(\cdot|x)} \left[ D(p_T \| p_S^\theta)(y|x) \right] \right]$

where $\lambda \in [0, 1]$ controls the fraction of training data that comes from on-policy student generation, $D$ is the chosen divergence (forward KL, reverse KL, or JSD($\beta$)), $(X, Y)$ is a fixed dataset of input-output pairs (either ground-truth or teacher-generated), and $X$ is the dataset of input prompts (used for on-policy generation, where outputs are not provided but generated by the student). In the second term, the inner expectation is over the student's sampling distribution $p_S(\cdot|x)$, and again, gradients do not flow through this sampling step.

What it computes. The loss is a convex combination of two terms. The first term, weighted by $(1 - \lambda)$, is the standard supervised distillation loss—it computes the divergence between teacher and student on a fixed dataset of known output sequences. The second term, weighted by $\lambda$, is the on-policy distillation loss—it has the student generate its own output sequences from the input prompts, then computes the divergence between teacher and student on those self-generated sequences. The total loss is the expectation of this combination over training batches.

Special cases that recover prior work:

• $\lambda = 0$, $D = D_{KL}$: pure supervised KD (Sanh et al., 2019), training only on fixed dataset sequences with forward KL divergence.
• $\lambda = 0$, with ground-truth outputs: supervised fine-tuning (if $D$ is forward KL, this is standard NLL minimization).
• $\lambda = 0$, with teacher-generated outputs: SeqKD (Kim & Rush, 2016).
• $\lambda = 0.5$, $D = D_{KL}$: ImitKD (Lin et al., 2020).
• $\lambda = 1$: purely on-policy GKD, the configuration the paper finds most consistently effective (denoted "on-policy GKD" in experiments).
• $\lambda = 0.5$: "mixed" GKD, using equal parts supervised and on-policy data.

Why this form. The mixing parameter $\lambda$ provides a practical continuum between supervised-only and on-policy-only training. This is valuable for two reasons. First, it allows practitioners to interpolate based on computational constraints—if on-policy generation is too expensive, reducing $\lambda$ reduces cost while still providing some on-policy benefit. Second, it allows the framework to naturally use available supervised data when it exists (the fixed dataset $(X, Y)$ might contain high-quality human annotations that are valuable to retain), while also leveraging the distribution-matching benefits of on-policy training. The paper's experiments systematically sweep $\lambda \in \{0, 0.5, 1.0\}$ across tasks to show that purely on-policy ($\lambda = 1$) typically performs best, but that even mixed training ($\lambda = 0.5$) substantially outperforms supervised-only ($\lambda = 0$).

The outer expectation over batches is implemented in practice as described in Algorithm 1: at each training step, a random draw $u \sim \text{Uniform}(0, 1)$ determines whether the batch comes from on-policy generation (if $u \leq \lambda$) or from the fixed dataset (otherwise). This stochastic mixing is equivalent to the expectation in the objective but is simpler to implement than producing exact fractions in each step.

---

The On-Policy Feedback Loop: Why Self-Generated Data Improves Training

The on-policy mechanism creates a virtuous cycle that the paper argues is central to GKD's performance:

Step 1 — Diverse error exposure. At the start of training, the warm-started student generates sequences of "adequate quality" but with errors. Because the student samples at temperature $\gamma = 1$, these errors are diverse—the student explores different failure modes rather than repeating the same deterministic mistakes. The teacher then provides token-level feedback on exactly those error states, showing the student what the correct distribution should have been at each position.

Step 2 — Targeted correction. The loss update moves the student's distribution closer to the teacher's at every state it visited, including the error states. This means the student learns not just "what the correct sequence looks like" (which it already learned during supervised FT) but "how to recover when I go off the correct path." This is the imitation learning insight—the student builds a policy that is robust to its own imperfections.

Step 3 — Improved generation. After the update, the student generates slightly better sequences. The teacher now provides feedback on a shifted distribution of states—states that are closer to the teacher's distribution because the student has improved.

Step 4 — Iterative refinement. Steps 2–3 repeat for thousands of steps. The student's output distribution progressively shifts toward the teacher's, and the on-policy data distribution tracks this improvement, meaning the student is always training on states that are challenging at its current capability level. This is in contrast to supervised KD, where the training data is static and cannot adapt to the student's evolving needs.

Evidence for this loop. The paper's results support this mechanism in two ways. First, the on-policy advantage is largest for the smallest students (38× smaller than teacher on XSum, Figure 1), where capacity mismatch is most severe and the student's own outputs are most likely to diverge from the teacher's distribution—consistent with the imitation learning prediction that on-policy training matters most when the learner's initial policy is far from the expert's. Second, the on-policy advantage persists even at scale, with instruction tuning on 5.36M examples showing measurable gains from on-policy training (Figure 10). Third, on GSM8K, increasing the on-policy data fraction beyond 25% consistently improves performance (Figure 8), showing a monotonic benefit rather than diminishing returns.

---

Choice of Divergence: Forward KL, Reverse KL, and Generalized JSD

GKD's second key degree of freedom is the divergence $D$ used to compare teacher and student distributions at each token position. The paper evaluates five choices: forward KL, reverse KL, JSD(0.1), JSD(0.5), and JSD(0.9). The choice matters because different divergences impose different penalties on the student when it cannot perfectly match the teacher.

Why the student cannot match the teacher perfectly. The student has fewer parameters than the teacher—in the paper's experiments, the student is 3.8× to 38× smaller. This capacity gap means the student's representable distribution family is strictly smaller than the teacher's. The student cannot simultaneously put high probability on every token that the teacher assigns high probability to, because its softmax bottleneck forces it to concentrate probability mass. The divergence choice determines which aspects of the teacher's distribution the student prioritizes fitting.

Forward KL: Mode-Covering

The forward KL divergence at a token position is:

$D_{KL}(p_T \| p_S^\theta) = \sum_{v \in V} p_T(v) \log \frac{p_T(v)}{p_S^\theta(v)}$

where I omit the conditioning on $(y_{<n}, x)$ for brevity. $p_T(v)$ is the teacher's probability for token $v$, and $p_S^\theta(v)$ is the student's probability for token $v$.

What it computes. For each token $v$ in the vocabulary, the term $p_T(v) \log(p_T(v) / p_S^\theta(v))$ measures how much the student's probability differs from the teacher's, weighted by the teacher's probability $p_T(v)$. If $p_T(v)$ is small, the term contributes little even if the student disagrees strongly. If $p_T(v)$ is large and $p_S^\theta(v)$ is small, the term contributes heavily (since $\log(1/\text{small}) \to \infty$ as $p_S^\theta(v) \to 0$). The sum over the vocabulary yields a non-negative scalar that is zero only when the two distributions are identical.

What behavior this encourages. Forward KL is "mode-covering" —it penalizes the student heavily for putting low probability on tokens the teacher assigns high probability to. To minimize the loss, the student must spread its probability mass across all tokens that have non-negligible probability under the teacher. This means the student learns a distribution that covers the full support of the teacher's distribution, even if doing so requires assigning some probability to tokens the teacher considers unlikely. In the extreme, if the teacher assigns probability 0.01 to each of 100 plausible tokens, the student must allocate at least some probability to all 100—it cannot concentrate entirely on the top 5 without incurring a large penalty from the 95 others.

Why this can be problematic. With limited capacity, the student cannot cover all teacher modes while also assigning high probability to the teacher's most likely tokens. The result is that the student's distribution becomes "spread thin" —it assigns non-trivial probability to many low-probability teacher tokens, which can lead to sampling those low-probability tokens during generation. The paper connects this to hallucination: "the student might end up assigning probability mass to tokens $v$ which have low probability under $p_T$, which can result in hallucination and low-quality generations."

When forward KL works well. The paper finds forward KL performs well on GSM8K with greedy decoding (Figure 7: forward KL achieves the highest accuracy improvement of 8.8 points for T5-Base). When evaluation uses greedy decoding ($\gamma \to 0$), the student always selects the single most probable token. In this regime, the exact distribution shape matters less than the relative ordering of high-probability tokens, and forward KL's broad coverage may help the student correctly identify which tokens should be in the top ranks.

Reverse KL: Mode-Seeking

The reverse KL divergence swaps the arguments:

$D_{KL}(p_S^\theta \| p_T) = \sum_{v \in V} p_S^\theta(v) \log \frac{p_S^\theta(v)}{p_T(v)}$

What it computes. Now the weighting is by $p_S^\theta(v)$, the student's probability, not the teacher's. The term $p_S^\theta(v) \log(p_S^\theta(v) / p_T(v))$ is large when the student assigns high probability to a token that the teacher assigns low probability to. If $p_S^\theta(v) > 0$ and $p_T(v) \approx 0$, the log-ratio becomes very large (since $\log(p_S^\theta(v) / \text{near-}0) \to \infty$), and the penalty is severe. Conversely, if the teacher assigns high probability to a token but the student assigns near-zero probability, the student's probability $p_S^\theta(v) \approx 0$ means the term contributes almost nothing—the student is not penalized for ignoring teacher modes as long as it doesn't claim they are important.

What behavior this encourages. Reverse KL is "mode-seeking" —it penalizes the student for putting probability mass where the teacher does not, but does not penalize the student for ignoring modes of the teacher that it cannot fit. The student is free to concentrate its limited probability mass on a subset of the teacher's high-probability tokens, as long as it avoids assigning mass to tokens the teacher considers unlikely. The student will "zero in" on one or a few modes of the teacher distribution and model those well, while effectively setting near-zero probability for other teacher modes—and this is not heavily penalized because those modes have $p_S^\theta(v) \approx 0$, multiplying the penalty to near zero.

Why this is useful for capacity-limited students. When the student is much smaller than the teacher, it cannot faithfully reproduce the teacher's full distribution. Reverse KL gives the student permission to specialize: it can focus on generating tokens that are highly likely under the teacher, avoiding the "spread thin" problem of forward KL. At generation time, this means the student is less likely to sample low-probability hallucinated tokens because it has learned to assign near-zero probability to them. The paper states: "mode-seeking divergences, such as reverse KL, prioritize the tokens where the teacher assigns high probability, which can avoid low-quality generations but at the expense of less diverse generations for a given input."

When reverse KL works well. The paper finds reverse KL substantially outperforms forward KL on instruction tuning evaluated with greedy decoding (Figure 10), where on-policy reverse KL achieves ~2% absolute accuracy improvement on MMLU vs. forward KL. The paper hypothesizes: "the efficacy of reverse KL in instruction tuning may stem from its mode-seeking nature as it ensures that the model zeroes in on the main intent or behavior specified by the instruction." In instruction following, there is often one core correct behavior, and modeling the full distribution of all plausible but suboptimal outputs is less useful than concentrating on the best ones.

The forward-vs-reverse trade-off visualized. The paper references Figure A.16, which shows a synthetic example of fitting a unimodal Gaussian to a mixture of two Gaussians using forward vs. reverse KL. Forward KL places the Gaussian to cover both modes (even though it fits neither well), while reverse KL selects one mode and fits it well, ignoring the other. This directly maps to the distillation setting: the teacher's distribution may have multiple plausible continuations at any prefix, and the student must choose between covering many poorly or one well.

Generalized JSD: Interpolating Between Forward and Reverse KL

The generalized Jensen-Shannon divergence with parameter $\beta \in (0, 1)$ is:

$D_{JSD(\beta)}(P \| Q) = \beta D_{KL}\left(P \| \beta P + (1 - \beta) Q\right) + (1 - \beta) D_{KL}\left(Q \| \beta P + (1 - \beta) Q\right)$

where $P$ is the teacher distribution and $Q$ is the student distribution.

What it computes. The JSD constructs an intermediate distribution $M = \beta P + (1 - \beta) Q$, which is a weighted average of the teacher and student distributions. It then computes the forward KL from the teacher to $M$ (weighted by $\beta$) plus the forward KL from the student to $M$ (weighted by $1 - \beta$). The result is symmetric in the sense that it doesn't "favor" either distribution—both are compared to the same reference distribution $M$.

What $\beta$ controls. The parameter $\beta$ interpolates between forward KL and reverse KL behavior:

• As $\beta \to 0$, the mixture $M$ approaches the student distribution $Q$, and the $D_{KL}(Q \| M)$ term vanishes. The remaining term $D_{KL}(P \| M) / \beta$ behaves like $D_{KL}(P \| Q)$ —forward KL. The paper cites Huszár (2015): $\lim_{\beta \to 0} D_{JSD(\beta)}(P \| Q) / \beta = D_{KL}(P \| Q)$.
• As $\beta \to 1$, the mixture $M$ approaches the teacher distribution $P$, and by symmetry the behavior approaches reverse KL.
• At $\beta = 0.5$, it is the standard symmetric JSD.

The three values tested in the paper are $\beta = 0.1$ (close to forward KL behavior), $\beta = 0.5$ (balanced), and $\beta = 0.9$ (close to reverse KL behavior). These span the continuum without testing every possible value.

Why JSD is useful. First, unlike forward or reverse KL, JSD is bounded even for distributions with disjoint support. This provides numerical stability and prevents the loss from exploding if the student assigns exactly zero probability to a token the teacher considers possible (which would make forward KL infinite). Second, the $\beta$ parameter provides a continuous knob for practitioners to tune between diversity (forward KL) and quality (reverse KL) without committing to either extreme. The paper's experiments show that the optimal $\beta$ is task-dependent: JSD(0.1) performs best on WMT translation with beam search, JSD(0.9) performs best on XSum summarization with temperature sampling, and forward KL ($\beta \to 0$ limit) performs best on GSM8K with greedy decoding.

Evidence for task-dependence (Figure 4). On XSum with temperature sampling ($\gamma = 1$), the mode-seeking divergences (JSD(0.9), reverse KL) produce higher ROUGE-2 scores than forward KL, but at the cost of lower diversity (higher Self-BLEU). As the sampling temperature is reduced, the diversity gap narrows and so does the performance gap between divergences. This makes sense: at low temperature, the student always selects the most probable token, so the shape of the distribution away from the mode matters less. At high temperature, the student samples from the full distribution, and forward KL's tendency to spread probability mass leads to sampling more low-quality tokens.

---

The Operational Training Loop: Algorithm 1 Deconstructed

The paper presents Algorithm 1 as the concrete implementation of GKD. Here is a detailed walkthrough:

Inputs:

• $p_T$: the frozen teacher model. It can be queried for token-level log-probabilities on any sequence but its parameters are never updated.
• $p_S^\theta$: the student model with learnable parameters $\theta$. It starts from a supervised fine-tuned checkpoint (not random initialization). The paper argues this two-stage approach—SFT followed by on-policy distillation—is "analogous to two-stage RLHF training, which is widely used for LMs."
• $(X, Y)$: a fixed dataset of (input, output) pairs. For summarization and translation, this is the standard training dataset. For GSM8K, it is teacher-generated CoT outputs. For instruction tuning, it is the FLAN2021 dataset (5.36M examples across 62 tasks).
• $\lambda \in [0, 1]$: student data fraction.
• $D$: the chosen divergence.
• $\eta$: learning rate.

Per-step procedure:

1. Data source selection (lines 3–9): Sample $u \sim \text{Uniform}(0, 1)$. If $u \leq \lambda$, the batch comes from on-policy generation; otherwise, from the fixed dataset.

2. On-policy path (line 6): Sample a batch of $B$ input prompts $x$ from $X$. For each $x$, generate an output sequence $y \sim p_S^\theta(\cdot|x)$ using the student's current policy at temperature $\gamma = 1$. The generation is auto-regressive: the student produces tokens one at a time, each conditioned on previously generated tokens, until an end-of-sequence token or maximum length. These $(x, y)$ pairs form the batch $\mathcal{B}$.

3. Supervised path (line 8): Sample a batch of $B$ input-output pairs $(x, y)$ directly from the fixed dataset $(X, Y)$. These sequences are "expert demonstrations"—either human-written references or teacher-generated outputs.

4. Loss computation (line 10): For each $(x, y)$ pair in the batch: (a) run the teacher model on $(x, y)$ to obtain token-level log-probabilities $p_T(\cdot|y_{<n}, x)$ at every position $n$, (b) run the student model on $(x, y)$ to obtain $p_S^\theta(\cdot|y_{<n}, x)$, (c) compute $D(p_T \| p_S^\theta)(y|x)$ using the chosen divergence (averaged across tokens as defined in Section 2), (d) average this loss over the $B$ pairs in the batch, (e) compute $\nabla_\theta$ of this average loss and update $\theta$ using the optimizer: $\theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L}_{\text{GKD}}$.

Why gradient flow works this way. On the on-policy path, when the student generates $y \sim p_S^\theta(\cdot|x)$, the sampling operation is not differentiated. The generated tokens $y_1, y_2, \ldots, y_{L_y}$ are treated as fixed when computing the loss. However, the loss is still a function of $\theta$ because it involves $\log p_S^\theta(y_n|y_{<n}, x)$—the student's log-probability of the generated token at each position. The gradient $\nabla_\theta$ flows through this log-probability computation, teaching the student to increase the probability of tokens the teacher considers good (for forward KL) or decrease the probability of tokens the teacher considers bad (for reverse KL). But there is no gradient through the auto-regressive sampling mechanism itself—no credit assignment for "the student chose to generate token A which led to state B which led to loss C." This is the crucial simplification that makes GKD stable.

Temperature settings. During training, the student samples at $\gamma = 1$ (full diversity) to generate on-policy data. The teacher's temperature during training is task-dependent: for XSum with temperature sampling evaluation, the teacher temperature is set to 0.1 (making the teacher more deterministic, which provides sharper supervision); for GSM8K, the teacher temperature is 0.1; for WMT, it is 1.0. The student's training temperature is always 1.0, and the optimizer sees the student's logits at temperature 1.0 (standard practice for cross-entropy-based training).

Why this algorithm is simple. There are no importance sampling corrections, no trust region constraints, no advantage estimators, no value function baselines, no KL penalties between consecutive policies. The entire mechanism is: generate data from the current policy, compute a token-level divergence against the teacher, update. This simplicity is explicitly contrasted with MiniLLM, which requires policy gradient estimators, reward normalization, and length bias correction.

---

Integration with RL Fine-Tuning: The Regularized Objective

On-policy GKD's sampling infrastructure naturally supports reinforcement learning, because both require generating sequences from the student and scoring them. The paper proposes a combined objective that simultaneously maximizes a scalar reward $r(y)$ and distills from the teacher:

$\mathbb{E}_{x \sim X} \left[ (1 - \alpha) \mathbb{E}_{y \sim p_S^\theta(\cdot|x)} [r(y)] - \alpha \mathbb{E}_{y \sim p_S(\cdot|x)} \left[ D(p_T \| p_S^\theta)(y|x) \right] \right]$

where $\alpha \in [0, 1]$ controls the strength of distillation relative to reward maximization. When $\alpha = 1$, this is pure on-policy GKD. When $\alpha = 0$, it is pure RL fine-tuning with no distillation. Intermediate values blend both objectives.

What it computes. For each input $x$, the student generates an output sequence $y$. This single generation is used for both the reward term and the distillation term—a practical efficiency. The reward $r(y)$ is a scalar (e.g., a factual consistency score from an NLI classifier), and its gradient is estimated using REINFORCE (or a similar policy gradient method). The distillation term $D(p_T \| p_S^\theta)(y|x)$ is the same token-level divergence used in standard GKD. The total loss is the weighted combination, with the distillation term subtracted because it is minimized (while reward is maximized).

Why this combination matters. The paper identifies a specific problem: standard RL fine-tuning for alignment (RLHF/RLAIF) often suffers from an "alignment tax"—the reward-optimized model performs worse on general capabilities because it overfits to the reward signal. In standard RLHF, the model is regularized to stay close to its initial (SFT) checkpoint using a reverse KL penalty. The paper's innovation is replacing this self-regularization with teacher-regularization: instead of penalizing deviation from the student's own past self, penalize deviation from a better teacher model. This means the student can improve its general capabilities (via distillation) even as it optimizes for the specific reward.

Concrete example on XSum (Figure 5). The paper combines on-policy GKD (with JSD(0.9)) with RLAIF using textual entailment feedback as the reward (following Roit et al., 2023). The reward is the entailment score from a T5-XXL NLI classifier—summaries that are textually entailed by the input document receive higher reward, encouraging factual consistency. The distillation teacher is T5-XL (12× larger than the T5-base student). Results show a clear trade-off controlled by $\alpha$:

• At $\alpha = 0.05$ (heavy RL, light distillation): +45% improvement in factual consistency but only +3 ROUGE-2.
• At $\alpha = 0.5$ (balanced): +25% factual consistency improvement and +7 ROUGE-2.
• The student with $\alpha = 0.5$ achieves higher ROUGE-2 than the RLEF baseline (which regularizes to the student's own initial checkpoint) while simultaneously generating more factually consistent summaries than the 12× larger teacher—a dual improvement that neither RL nor distillation alone achieves.

Implementation note for RL practitioners. The paper suggests: "If one wants to only make slight modifications to existing RL fine-tuning workflows, we recommend using reverse KL or JSD(0.9) when integrating GKD with RL." This is because reverse KL is already the standard divergence used in RLHF for the KL penalty toward the initial policy. Using reverse KL for the teacher-regularization term means the only change to existing RLHF code is swapping the reference model from the initial student checkpoint to the teacher model.

---

Hyperparameters, Configurations, and Practical Defaults

The paper provides detailed hyperparameter tables for each task in the Appendix. Here are the key configurations, organized by task.

XSum (Summarization) — Table A.1:

• Training steps: 40,000
• Batch size: 32
• Learning rate: 0.0003 default (0.001 for T5-small; reverse KL is "more sensitive to higher LRs" and uses 0.0003 for all models)
• Warmup: 2,000 steps, linear from 0 to LR
• Cooldown: from step 30,000 to 40,000, linear from LR to 0
• Optimizer: Adafactor (standard for T5 models; Shazeer & Stern, 2018)
• Dropout: 0.0
• Max input length: 1024 tokens
• Max output length: 64 tokens
• Teacher temperature: 1.0 when evaluating with greedy sampling; 0.1 when evaluating with temperature sampling
• Student training temperature: 1.0
• GKD divergences tested: forward KL, reverse KL, JSD(0.1), JSD(0.5), JSD(0.9)
• Student data fractions tested: $\lambda \in \{0, 0.5, 1.0\}$

WMT (Translation) — Table A.3:

• Training steps: 100,000
• Batch size: 32
• Learning rate: 0.0003
• Warmup: 5,000 steps, linear
• Dropout: 0.0
• Max input/output length: 80 tokens each
• Teacher temperature: 1.0
• Evaluation: beam search (same parameters as Raffel et al., 2020)
• Results averaged over 3 seeds

GSM8K (Arithmetic Reasoning) — Table A.2:

• Training steps: 40,000
• Batch size: 32
• Learning rate: 0.0003
• Warmup: 2,000 steps
• Cooldown: step 30,000 to 40,000
• Dropout: 0.05
• Max input length: 512 tokens
• Max output length: 320 tokens (to accommodate chain-of-thought reasoning)
• Teacher temperature: 0.1
• Base checkpoints: Flan-T5 (not plain T5)
• Training data: CoT outputs generated by PaLM-540B (from Magister et al., 2022), ~5.3K (problem, CoT) pairs
• Evaluation: greedy sampling with external calculator
• Few-shot prompt: 4 CoT exemplars from Wei et al. (2022)

FLAN Instruction Tuning — Table A.4:

• Training steps: 50,000
• Batch size: 128
• Learning rate: 0.0001
• Warmup: none
• Dropout: 0.0
• Max input length: 2048 tokens
• Max output length: 256 tokens
• Teacher temperature: 1.0
• Dataset: FLAN2021 (5.36M examples, 62 tasks)

A critical design pattern across all tasks: the student is always initialized from a supervised fine-tuned checkpoint, never from scratch. The paper argues this is necessary because the student must "generate sequences of adequate quality, which the teacher can provide feedback upon." A randomly initialized student would produce near-random sequences, and the teacher's token-level distributions on those sequences would provide a very weak training signal—the states visited would be too far from any reasonable distribution for the teacher's feedback to be meaningful. The supervised FT warm-start ensures the student starts from a policy that is already task-competent, and GKD refines it to better match the teacher while reducing distribution mismatch.

Data efficiency design. On XSum, the paper evaluates GKD with training dataset fractions of 0.5% (1K examples), 5% (10K), 25% (50K), and 100% (full dataset). On-policy GKD on the 5% subset "without any ground-truth summaries, outperforms supervised KD and ImitKD with entire training dataset with ground-truth summaries" (Figure 3). This is a strong result: the on-policy mechanism allows the student to effectively generate its own training signal from the teacher, reducing dependence on large supervised datasets.

---

Design Choices and Their Justifications: A Summary

• No backpropagation through sampling: avoids high-variance policy gradients, reward hacking, and length bias that require stabilizing tricks in approaches like MiniLLM. Makes GKD behave like supervised learning.

• Student data fraction $\lambda$ as a mixing parameter: unifies supervised and on-policy distillation into a single framework, recovers prior work as special cases, and allows practitioners to trade off between computational cost and distribution-matching benefits.

• Divergence choice as a free parameter: recognizes that optimal divergence is task-dependent and temperature-dependent, providing flexibility without committing to a single approach. Forward KL for mode-covering (diversity, useful with greedy decoding), reverse KL for mode-seeking (quality, useful with temperature sampling or instruction following), JSD for smooth interpolation.

• Warm-started student (supervised FT before GKD): ensures initial student-generated sequences are of sufficient quality for the teacher to provide meaningful feedback. Analogous to the SFT → RLHF pipeline, making integration natural.

• Single generation serving both RL reward and distillation loss: computational efficiency when combining RL and distillation, using one forward pass to support both objectives.

• Teacher temperature tuned per task: recognizes that the "sharpness" of teacher supervision should match the evaluation regime—sharper (lower temperature) when evaluating with temperature sampling, more diffuse (temperature 1.0) when evaluating with beam search or greedy decoding.

• Adafactor optimizer throughout: consistency with the T5 training recipe, sublinear memory cost suitable for large models.

4. Key Insights and Innovations

Innovation 1: Framing Distillation as On-Policy Imitation Learning Changes the Data Generation Strategy, Not Just the Objective

The paper's most fundamental conceptual move is recognizing that knowledge distillation for auto-regressive models is structurally identical to imitation learning with an interactive expert, and that the core pathology—distribution mismatch between training and inference—has a known solution from the IL literature that had been almost entirely overlooked in distillation research. This is not merely an analogy but a diagnostic reframing that reveals why standard supervised KD inherently underperforms and what must change to fix it.

Before GKD, the dominant assumption in distillation was that the training data—whether human-written references or teacher-generated sequences—was fundamentally the right thing to train on, and the challenge was in how to extract richer supervision from that data (e.g., token-level probabilities instead of hard labels in Supervised KD; Sanh et al., 2019). ImitKD (Lin et al., 2020) recognized the IL connection but treated it as a data-augmentation idea—mix in some student-generated sequences with the fixed dataset, using forward KL throughout. This partially addressed distribution mismatch but did not question the assumption that the fixed dataset was necessary or that forward KL was the right divergence.

GKD's reframing is qualitatively different: it says the student should train primarily or entirely on its own generated sequences because those are the states it will encounter at inference time. The fixed dataset is optional—useful when available but not the core training signal. This inverts the standard KD workflow from "train on expert data, occasionally mix in student data" to "train on student data, optionally supplement with expert data." The paper shows this inversion works empirically: on-policy GKD ($\lambda = 1$) consistently outperforms mixed ($\lambda = 0.5$) and supervised ($\lambda = 0$) variants across summarization, translation, reasoning, and instruction tuning (Figures 1, A.12, A.13, A.14, 6, 7). On GSM8K, Figure 8 shows a monotonic trend: "performance consistently improves as the proportion of on-policy data increases, provided that at least 25% of the data is on-policy."

Beyond the practical inversion, the IL framing provides an intellectual genealogy that explains why GKD works and when it should be expected to work best. In imitation learning, the benefit of on-policy training is proportional to how far the learner's initial policy is from the expert's—a near-perfect initial policy visits states very similar to expert demonstrations, so supervised training suffices; a weak initial policy diverges quickly and needs on-policy correction. This predicts that on-policy GKD should help more for smaller students (larger capacity gap from teacher) and for harder tasks. The paper's results support this: on XSum (Figure 1), moving from supervised KD to on-policy GKD provides larger relative gains for T5-small (77M, 38× smaller than teacher) than for T5-large (800M, 3.8× smaller)—the capacity gap determines the need for on-policy correction. This is a diagnostic principle rather than just an empirical observation: it tells practitioners when investing in on-policy training is worth the computational overhead.

This reframing is a fundamental shift, not incremental. It changes the default answer to "what data should I train my student on?" from "the best available supervised data" to "the student's own outputs," and it provides a theoretical framework (DAgger-style on-policy aggregation) for understanding why.

---

Innovation 2: Divergence Choice Is a First-Class Design Dimension, Not a Fixed Convention, and Its Optimal Value Is Task-Dependent

The field of knowledge distillation has overwhelmingly used forward KL divergence as the default objective—it is the natural choice when training on fixed data because it corresponds to maximum likelihood under the teacher's distribution. GKD's second conceptual contribution is demonstrating that this default is neither necessary nor always optimal, and that the choice of divergence should be treated as a tuneable design parameter that controls a fundamental trade-off between generation quality and diversity.

Prior work treated the divergence as fixed: Supervised KD uses forward KL (Sanh et al., 2019); ImitKD uses forward KL; SeqKD uses hard likelihood maximization; MiniLLM (Gu et al., 2023) uses reverse KL but locks the entire method to it without exploring alternatives. f-distill (Wen et al., 2023) recognized that different f-divergences are possible but proposed total variation distance as a specific tractable choice rather than systematically exploring the space. No prior work treated divergence selection as an empirical question whose answer varies by task, evaluation regime, and student capacity—or provided a conceptual framework for understanding why different divergences behave differently.

GKD surfaces this as a first-class design dimension by evaluating forward KL, reverse KL, and three JSD variants ($\beta \in \{0.1, 0.5, 0.9\}$) across four tasks and two evaluation regimes. The results reveal clear task-dependent optimality:

• GSM8K with greedy decoding: forward KL performs best (Figure 7, 8.8 accuracy improvement vs. 6.9 for reverse KL), because greedy decoding only cares about the relative ordering of top tokens, and forward KL's broad coverage helps the student correctly rank all high-probability teacher tokens.
• WMT with beam search: JSD(0.1) performs best (Figure 6, 0.71 BLEU improvement vs. 0.55 for forward KL), suggesting a slight mode-seeking preference helps translation quality.
• XSum with temperature sampling ($\gamma = 1$): mode-seeking divergences (JSD(0.9), reverse KL) outperform forward KL (Figure 2, 4), because at high temperature the student samples broadly from its distribution, and forward KL's tendency to spread probability mass leads to sampling low-quality tokens.
• FLAN instruction tuning: reverse KL substantially outperforms forward KL (Figure 10, ~2% absolute improvement on MMLU), because instruction following requires the model to zero in on the core intended behavior rather than modeling all plausible continuations.

The quality-diversity trade-off revealed in Figure 4 is the paper's most illuminating diagnostic of this dimension. As the divergence moves from forward KL → JSD(0.5) → JSD(0.9) → reverse KL, the student's generated summaries become less diverse (higher Self-BLEU) but higher quality (higher ROUGE-2) at high sampling temperatures. At low temperatures, the gap narrows because the sampling mechanism itself enforces mode-seeking behavior regardless of the training divergence. This trade-off is not captured by any single metric and requires practitioners to make an explicit choice based on their deployment needs.

This contribution is a conceptual refinement of the distillation problem formulation rather than a new algorithm. It establishes that "what objective" is not a given but a design choice that interacts with task characteristics and evaluation protocols. The practical implication—that practitioners should sweep divergences as part of their hyperparameter tuning, just as they sweep learning rates—is straightforward but had not been articulated before.

---

Innovation 3: Distillation and RL Fine-Tuning Are Not Competing Paradigms but Naturally Complementary Objectives That Can Be Jointly Optimized

Prior to GKD, distillation and RL fine-tuning (RLHF/RLAIF) existed in separate workflows with different goals: distillation compressed a teacher into a student, while RL fine-tuning aligned a policy with a reward signal. In standard RLHF, the learned policy is regularized to stay close to its own initial (SFT) checkpoint using a reverse KL penalty—a self-referential constraint that prevents reward over-optimization but provides no mechanism for the policy to improve its general capabilities beyond what the SFT model already knew. The paper identifies a gap that had not been previously recognized: this self-regularization term can be replaced with teacher-regularization, turning the RL fine-tuning step into a simultaneous distillation and reward optimization process.

The conceptual move is recognizing that on-policy GKD and policy-gradient RL share identical infrastructure—both require sampling sequences from the student and computing some signal on those sequences. The on-policy GKD loss provides a token-level signal from the teacher, while the RL loss provides a sequence-level signal from the reward function. Combining them requires no additional sampling, no architectural changes, and no new hyperparameters beyond the mixing coefficient $\alpha$. The combined objective (Equation 5) is a convex combination that smoothly interpolates between pure distillation ($\alpha = 1$) and pure RL ($\alpha = 0$).

The significance of this combination goes beyond convenience. The paper demonstrates on XSum summarization (Figure 5) that joint optimization achieves a Pareto improvement—a point that improves on both objectives simultaneously relative to baselines—which is rare in multi-objective optimization. Specifically:

• RLEF (Roit et al., 2023), which regularizes to the student's own initial checkpoint, improves factual consistency by ~35% but yields minimal ROUGE-2 gains.
• Pure on-policy GKD improves ROUGE-2 but does not directly target factual consistency.
• Joint optimization with $\alpha = 0.5$ achieves both: +25% factual consistency improvement and +7 ROUGE-2, producing summaries that are more factually consistent than the 12× larger teacher while also being higher quality than the RLEF baseline.

This is a conceptual advance because it reframes the standard RLHF pipeline's regularization term not as a necessary evil to prevent reward hacking, but as an opportunity to inject additional capability improvement. The standard RLHF KL penalty toward the initial policy serves only to constrain—it says "don't change too much." Replacing it with a KL penalty toward a teacher says "change, but change toward something better." This transforms RL fine-tuning from a pure alignment step into a capability-plus-alignment step, potentially reducing the "alignment tax" that Ouyang et al. (2022) documented.

Beyond the specific XSum result, this innovation opens a design space for practitioners: any RL fine-tuning workflow that currently includes a KL penalty toward the initial policy can be augmented by substituting a teacher model, with the mixing coefficient $\alpha$ controlling the trade-off between reward optimization and capability transfer. The paper's recommendation to use reverse KL or JSD(0.9) for the distillation term when integrating with RL (since reverse KL is already the standard choice for the regularization penalty) makes adoption straightforward—it is a one-line change from kl_divergence(current_policy, initial_policy) to kl_divergence(current_policy, teacher_policy).

---

Innovation 4: The On-Policy Advantage for Distillation Is Robust Across Tasks, Scales, and Evaluation Protocols—Establishing It as a Reliable Principle Rather Than a Fragile Trick

A recurring pattern in machine learning is that methods which work well in controlled experiments prove brittle when applied across diverse settings—sensitive to hyperparameters, model scale, or task characteristics. The paper's fourth contribution is an empirical demonstration of robustness that elevates on-policy distillation from an interesting idea to a reliable principle. The evidence spans four dimensions:

Across tasks (summarization, translation, reasoning, instruction tuning). On-policy GKD outperforms supervised KD baselines on every task tested: 2.1× relative gain on XSum, 1.7× on WMT, 1.9× on GSM8K (Figure 1), and ~2% absolute gain on MMLU for instruction tuning (Figure 10). The consistency is notable because these tasks have fundamentally different output structures—free-form summaries, constrained translations, chain-of-thought reasoning traces, and diverse instruction-following responses—yet the on-policy mechanism benefits all of them.

Across student scales (77M to 800M parameters). Figure 1 shows monotonic improvements across T5-small, T5-base, and T5-large for XSum and WMT, and Figure 9 shows the same for GSM8K. The on-policy advantage does not disappear at larger student sizes, even though the capacity gap to the teacher shrinks—suggesting the mechanism addresses something fundamental about auto-regressive generation rather than a pathology specific to extreme compression.

Across evaluation protocols (greedy, temperature sampling, beam search). On XSum, on-policy GKD outperforms baselines under both greedy decoding and temperature sampling (Figures 2, A.12, A.13). On WMT, it outperforms under beam search (Figures 6, A.15). On GSM8K, it outperforms under greedy decoding with external calculator (Figures 7, 9). The method does not require tuning for a specific decoding strategy to work.

Across data regimes. On XSum (Figure 3), on-policy GKD trained on just 5% of the training data—without any ground-truth summaries—matches or exceeds supervised KD trained on the full dataset with ground-truth summaries. This is a striking result: the on-policy mechanism effectively generates its own training signal by having the teacher score student-generated outputs, reducing dependence on large supervised datasets. This has practical implications for domains where annotated data is scarce but an unlabeled input distribution and a teacher model are available.

Beyond performance numbers, the paper includes negative and boundary results that strengthen the robustness claim by showing the method's limits are well-behaved:

• Self-distillation (Appendix A.1, Figure A.11): distilling a model into itself using on-policy GKD improves performance over the original, and on-policy outperforms supervised KD, showing the mechanism is not dependent on a capacity gap.
• ReST$^{EM}$ -style RL optimization of the revision model (Appendix K, Figure 16, from the reference example): though this is from the companion paper's revision experiments, it shows that not all on-policy methods work—naïvely applying RL to self-training can hurt performance, making GKD's stability (no backprop through sampling) a meaningful design choice.
• The GSM8K data fraction sweep (Figure 8): on-policy data fractions below 25% perform worse, suggesting there is a threshold below which the on-policy signal is too sparse, but above which it reliably helps.

This robustness is a practical contribution of substantial value. It tells practitioners that on-policy GKD is not a method requiring careful per-task tuning to work—it is a default that consistently improves over supervised alternatives across a wide range of realistic distillation scenarios. The fact that the paper demonstrates this with open-source T5 models and publicly available datasets makes the claim verifiable and the method immediately adoptable.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper uses four primary datasets: (1) XSum (Narayan et al., 2018) for abstractive summarization, consisting of news articles paired with human-written summaries, with evaluation on the validation split; (2) WMT14 en-de (Bojar et al., 2014) for machine translation from English to German, evaluated on the validation split; (3) GSM8K (Cobbe et al., 2021) for arithmetic reasoning, a dataset of grade-school math word problems requiring multi-step logical inference, evaluated on the test split; and (4) the FLAN2021 instruction tuning dataset (Chung et al., 2022) for task-agnostic distillation, containing 5.36 million examples spanning 62 language understanding and generation tasks, with evaluation on the held-out MMLU (57 tasks) and BBH (23 tasks) benchmark suites. For XSum data efficiency experiments, the paper additionally uses subsampled training sets at 0.5% (1K), 5% (10K), and 25% (50K) of the full dataset.

• Base model(s). All experiments use the T5 model family (Raffel et al., 2020), specifically T5v1.1 checkpoints that have undergone LM-adaptation (an additional 100K steps of language model training). The teacher is always a supervised fine-tuned T5-XL with approximately 3B parameters. The students are T5-small (77M, 38× smaller), T5-base (250M, 12× smaller), and T5-large (800M, 3.8× smaller). For GSM8K, the base checkpoints are Flan-T5 (Chung et al., 2022) rather than plain T5, with the teacher being a supervised fine-tuned Flan T5-XL. The choice of the T5 family is motivated by the availability of multiple scales pretrained on the same data, enabling clean capacity-gap comparisons, and by the models being open-sourced and representative of contemporary encoder-decoder architectures.

• Metrics. For summarization on XSum, the primary metric is ROUGE-2 (Lin, 2004), which measures bigram overlap between predicted and reference summaries; the paper notes similar trends in ROUGE-L and ROUGE-1. For translation on WMT, the metric is BLEU score (case-sensitive), which measures n-gram precision of machine-translated text against reference translations; evaluation uses beam search with the same hyperparameters as Raffel et al. (2020), and results are averaged across three seeds to reduce variance. For arithmetic reasoning on GSM8K, the metric is test accuracy with an external calculator: the model's final answer (extracted from its chain-of-thought output) is compared to the ground-truth answer using the grading function from Cobbe et al. (2021). For instruction tuning, the metric is few-shot prompted accuracy (exact match), reported as an unweighted average across all 57 MMLU tasks and all 23 BBH tasks. For XSum diversity analysis in Figure 4, the paper additionally reports Self-BLEU (Zhu et al., 2018) as a diversity measure, where a score of 100 indicates fully deterministic outputs and 0 indicates maximum diversity. For the RL + distillation experiments in Figure 5, factual consistency improvement is measured as the change in textual entailment score from a T5-XXL NLI classifier relative to the original student baseline.

• Baselines. The paper compares against five established methods:

  1. Supervised FT: standard negative log-likelihood training on ground-truth output sequences (no teacher).
  2. SeqKD (Kim & Rush, 2016): sequence-level knowledge distillation, which trains the student via supervised FT on output sequences generated by the frozen teacher model.
  3. Supervised KD (Hinton et al., 2015; Sanh et al., 2019): token-level knowledge distillation using forward KL divergence on a fixed dataset of input-output pairs (either ground-truth or teacher-generated).
  4. ImitKD (Lin et al., 2020): an approach that samples sequences from both the student and a fixed dataset with a non-increasing schedule on the student data fraction, using forward KL divergence. The paper notes ImitKD can be viewed as GKD with forward KL and $\lambda = 0.5$.
  5. f-distill (Wen et al., 2023): sequence-level KD formulated as minimizing an f-divergence, specifically using total variation distance between token-level student and teacher distributions. Like ImitKD, this uses mixed data but with a different divergence.

  All baselines start from the same supervised fine-tuned student checkpoint as GKD, ensuring fair comparison. For the RL experiments, the baseline is RLEF (Roit et al., 2023), which uses REINFORCE with a textual entailment reward while regularizing the policy toward its own initial checkpoint (rather than toward the teacher).

• Generation budget / compute accounting. The paper does not use a universal "generation budget" metric in the style of the reference example. Instead, all methods including GKD train for the same number of steps and use the same batch size per task (40K steps with batch size 32 for XSum and GSM8K; 100K steps with batch size 32 for WMT; 50K steps with batch size 128 for instruction tuning). The computational overhead of on-policy GKD comes from the student needing to auto-regressively generate output sequences during training (rather than simply loading pre-computed sequences from disk). The paper quantifies this overhead on GSM8K: "the computational overhead from student sampling is approximately 1.8×, 2× and 2.2× compared to sampling from a fixed dataset of outputs, for a student-teacher [size] ratio of 38×, 12× and 3.8× respectively." The paper argues this overhead is acceptable because (a) student inference is cheaper than teacher inference, (b) the majority of real-world cost is at serving time, not fine-tuning time, and (c) for RLHF + GKD, the overhead is reduced since teacher logits are already being computed. Fair comparison between GKD variants is maintained by controlling total training steps and batch size; the on-policy overhead affects wall-clock time but not the theoretical optimization budget.

• Cross-validation / statistical protocol. There is no cross-validation or statistical significance testing reported in the paper. For GSM8K, results are averaged across 3 seeds (as noted in Figure 6 and Appendix A.4). For WMT, results are also averaged across 3 seeds (as noted in Figure 6 and Appendix A.5). For XSum and instruction tuning, the paper does not specify the number of seeds or report any variance estimates. The paper selects the best-performing checkpoint at the end of training for reporting results on GSM8K; for instruction tuning, results are reported at 50K steps for all methods to ensure comparable training budget. The lack of confidence intervals or standard deviations across most results is a limitation.

Main Quantitative Results

Summarization on XSum

Headline result. On-policy GKD substantially outperforms all baseline KD methods across student sizes, with the advantage largest for the smallest student. As shown in Figure 1 (left panel), T5-large distilled with on-policy GKD achieves a ROUGE-2 score of approximately 21.5, compared to approximately 19.5 for supervised KD—a relative improvement of roughly 10%. For T5-base, on-policy GKD achieves approximately 17.8 vs. approximately 16.8 for supervised KD. For T5-small, on-policy GKD achieves approximately 15.0 vs. approximately 13.8 for supervised KD. The paper characterizes these gains as "2.1× on summarization... compared to the performance improvements achieved with baseline KD methods," though the exact computation of this multiplier is not spelled out—it appears to represent the ratio of improvement over the supervised FT student baseline (GKD improvement divided by baseline KD improvement).

Comparison to additional baselines. Figure 2 compares on-policy GKD variants with ImitKD and f-distill for T5-XL → T5-large distillation on XSum. Under temperature sampling ($\gamma = 1$), on-policy GKD with JSD(0.9) achieves approximately 21.0 ROUGE-2, outperforming ImitKD (approximately 19.5) and f-distill (approximately 19.0). Under greedy sampling, GKD (forward KL) achieves approximately 21.2, also outperforming ImitKD (~20.5) and f-distill (~20.3). Supervised KD achieves approximately 20.8 under greedy and approximately 19.8 under temperature sampling. These results establish that purely on-policy GKD variants outperform methods that use mixed (student + fixed) data.

Data scaling curves. Figure 3 evaluates data efficiency by training on subsets of the XSum training data (0.5%, 5%, 25%, 100%) with T5-small as student. On-policy GKD achieves approximately 11.5 ROUGE-2 with only 0.5% of the data (1K examples) and approximately 14.0 with 5% (10K examples). Critically, on-policy GKD on the 5% subset without any ground-truth summaries outperforms both supervised KD and ImitKD trained on the full dataset with ground-truth summaries (both at approximately 13.8). This is a striking result: the on-policy mechanism effectively substitutes for a 20× larger supervised dataset while not requiring human reference outputs. Supervised GKD ($\lambda = 0$, forward KL) tracks closely with supervised KD, confirming that the on-policy data is the source of the data efficiency gain, not the GKD framework itself when operated in supervised mode.

GKD ablations across student sizes and evaluation regimes. Appendix Figures A.12 (temperature sampling evaluation) and A.13 (greedy sampling evaluation) provide comprehensive heatmaps for all combinations of divergence × student data fraction × student size.

Under temperature sampling (Figure A.12, teacher temperature 0.1):

• For T5-small: purely on-policy ($\lambda = 1$, bottom row of each heatmap) with JSD(0.9) achieves 15.5 ROUGE-2, vs. 13.2 for supervised ($\lambda = 0$) with forward KL. Reverse KL with $\lambda = 1$ achieves 14.5. The gap between best on-policy and best supervised configurations is approximately +2.3 ROUGE-2 points.
• For T5-base: on-policy JSD(0.9) achieves 18.2, vs. 16.7 for supervised forward KL. JSD(0.5) and reverse KL with $\lambda = 1$ also outperform all supervised variants.
• For T5-large: on-policy JSD(0.9) achieves 20.3, vs. 19.2 for supervised forward KL. The gap narrows as student capacity increases—the 38× smaller T5-small benefits more (+2.3) than the 3.8× smaller T5-large (+1.1)—consistent with the imitation learning prediction that larger capacity gaps create larger distribution mismatch.
• Mode-seeking divergences (reverse KL, JSD(0.9)) consistently outperform forward KL when evaluation uses temperature sampling, because temperature sampling draws from the full distribution and forward KL's tendency to spread probability mass leads to sampling low-quality tokens.

Under greedy sampling (Figure A.13, teacher temperature 1.0):

• The overall pattern is similar (on-policy > mixed > supervised) but the gaps between divergences largely disappear.
• For T5-small: on-policy forward KL achieves 16.6 ROUGE-2, vs. 15.2 for supervised forward KL. Reverse KL with $\lambda = 1$ achieves 15.6—actually worse than forward KL, showing that reverse KL's mode-seeking behavior is less beneficial when evaluation is greedy.
• For T5-base: on-policy variants achieve 18.5-18.8, supervised achieve 17.9-18.2.
• For T5-large: all on-policy variants hover around 21.1-21.2, vs. supervised at 20.2-20.6.
• With greedy decoding, the student always selects the single most probable token, so the exact shape of the distribution away from the mode matters little—explaining why forward KL, reverse KL, and JSD converge to similar performance.

Divergence-quality-diversity trade-off. Figure 4 provides a detailed analysis of how divergence choice affects the quality-diversity Pareto frontier for T5-small distilled with on-policy GKD. At high sampling temperature ($\gamma = 1.0$), forward KL produces the most diverse outputs (Self-BLEU ≈ 20) but the lowest ROUGE-2 (approximately 20.0). Moving through JSD(0.1) → JSD(0.5) → JSD(0.9) → reverse KL progressively reduces diversity while improving quality: reverse KL achieves ROUGE-2 ~21.5 but Self-BLEU ~65. As sampling temperature is reduced from 1.0 to 0.1, the diversity of all methods decreases and performance differences among divergences narrow substantially. At $\gamma = 0.1$, forward KL achieves approximately 21.5 ROUGE-2 with Self-BLEU ~85, while reverse KL achieves approximately 21.8 with Self-BLEU ~95. The practical implication is clear: if deploying with high-temperature sampling (for creative generation), choose reverse KL or JSD(0.9); if deploying with greedy or low-temperature sampling, the divergence choice matters little and forward KL is a safe default.

GKD with RL fine-tuning. Figure 5 reports the trade-off between summarization quality and factual consistency when combining on-policy GKD (JSD(0.9)) with RLAIF. The rewards are textual entailment scores from a T5-XXL NLI classifier (Roit et al., 2023). Using T5-base as the student (12× smaller than the T5-XL teacher), the results show:

• RLEF baseline (regularizing to the student's own initial checkpoint): approximately +35% factual consistency improvement, approximately +3 ΔROUGE-2.
• GKD + RL with $\alpha = 0.05$ (heavy RL, light distillation): approximately +45% factual consistency improvement, approximately +3 ΔROUGE-2—much better factual consistency than the teacher but minimal summarization quality gain.
• GKD + RL with $\alpha = 0.1$: approximately +38% factual consistency improvement, approximately +4.5 ΔROUGE-2.
• GKD + RL with $\alpha = 0.25$: approximately +32% factual consistency improvement, approximately +6 ΔROUGE-2.
• GKD + RL with $\alpha = 0.5$ (balanced): approximately +25% factual consistency improvement, approximately +7 ΔROUGE-2—summaries are more factually consistent than the 12× larger teacher and higher quality than the RLEF baseline. This is the Pareto-dominating point.
• Teacher (T5-XL): 0 factual consistency improvement (reference point), 0 ΔROUGE-2 (reference point).

The dotted diagonal line in Figure 5 shows the trade-off frontier—as $\alpha$ decreases, the student optimizes more aggressively for the entailment reward at the expense of ROUGE-2 improvement. The key finding is that this frontier lies substantially above the RLEF* baseline, meaning GKD's teacher-regularization enables simultaneous improvement on both metrics rather than forcing a zero-sum trade-off.

---

Machine Translation on WMT

Headline result. On-policy GKD with JSD(0.1) achieves the best BLEU improvement over baselines for WMT14 en-de translation. As shown in Figure 1 (middle panel), T5-large distilled with on-policy GKD reaches approximately 27.2 BLEU vs. approximately 26.7 for supervised KD (teacher T5-XL achieves 28.0). For T5-base, on-policy GKD yields approximately 27.2 vs. approximately 26.9 for supervised KD. For T5-small, on-policy GKD yields approximately 26.4 vs. approximately 25.9 for supervised KD. The paper reports "1.7× on machine translation" relative improvement over baseline KD methods.

Full ablation heatmaps. Figure 6 provides comprehensive results for all GKD variants on WMT with beam search evaluation:

• T5-small (left heatmap, baseline BLEU 25.58): on-policy ($\lambda = 1$, bottom row) JSD(0.1) achieves +0.71 BLEU improvement, the best configuration overall. On-policy forward KL achieves +0.55. Supervised ($\lambda = 0$, top row) reverse KL achieves +0.08—near-zero improvement. Across all divergences, on-policy or mixed data consistently outperforms purely supervised data. The best performing divergence is JSD(0.1), which is closest to forward KL behavior, followed by forward KL itself, with reverse KL performing worst (as low as +0.02 with $\lambda = 0$).
• T5-base (right heatmap, baseline BLEU 26.98): on-policy JSD(0.1) achieves +0.71 BLEU improvement, matching the T5-small result. On-policy forward KL achieves +0.55. The gap between best and worst configurations is smaller for T5-base than T5-small—JSD(0.1) with $\lambda = 0.5$ achieves +0.52 and even supervised forward KL manages +0.52, suggesting the larger student is less sensitive to both data distribution and divergence choice. Reverse KL still underperforms (as low as +0.02 with $\lambda = 0$) but the penalty is less severe than for T5-small.

A key observation from the heatmaps: the optimal divergence for WMT is different from XSum. On WMT, JSD(0.1) (close to forward KL) performs best, while on XSum with temperature sampling, reverse KL and JSD(0.9) perform best. This reinforces the paper's claim that optimal divergence is task-dependent—translation with beam search rewards mode-covering behavior (diversity of candidate translations), while summarization with temperature sampling rewards mode-seeking behavior (avoiding low-quality tokens).

Comparison to ImitKD and f-distill. Appendix Figure A.15 reports a direct comparison on WMT. For T5-small, on-policy GKD achieves approximately +0.64 BLEU improvement vs. approximately +0.42 for ImitKD and approximately +0.24 for f-distill—GKD provides a 53% and 162% relative improvement over ImitKD and f-distill respectively. For T5-base, on-policy GKD achieves approximately +0.55 BLEU improvement vs. approximately +0.35 for ImitKD and approximately +0.20 for f-distill. The consistent ranking (GKD > ImitKD > f-distill) across both student sizes confirms that purely on-policy data with an appropriately chosen divergence is superior to mixed-data approaches.

---

Arithmetic Reasoning on GSM8K

Headline result. On-policy GKD with forward KL achieves the largest accuracy improvement for arithmetic reasoning with chain-of-thought prompting, evaluated with greedy decoding and an external calculator. Figure 1 (right panel) shows T5-large (Flan-T5 base) distilled with on-policy GKD achieving approximately 27.5% test accuracy vs. approximately 25% for supervised KD—the teacher Flan-T5-XL achieves 27.9%. For T5-base, on-policy GKD achieves approximately 17% vs. approximately 13% for supervised KD. For T5-small, on-policy GKD achieves approximately 11.5% vs. approximately 9% for supervised KD. The paper reports "1.9× on arithmetic reasoning tasks" relative improvement over baseline KD methods.

Full ablation heatmap. Figure 7 reports GKD variants for Flan-T5-XL → Flan-T5-Base (student baseline accuracy 10.16%, teacher 27.9%):

• On-policy ($\lambda = 1$, bottom row) forward KL achieves +8.8 accuracy improvement, the best configuration overall. On-policy reverse KL achieves +8.0. On-policy JSD(0.9) achieves +6.9.
• Supervised ($\lambda = 0$, top row) reverse KL achieves +6.9—this is notably strong, tied with on-policy JSD(0.9). Supervised forward KL achieves +4.7.
• A striking pattern: reverse KL performs well across all data fractions (+5.0 to +8.0), while forward KL degrades sharply with less on-policy data (+4.7 at $\lambda = 0$, +6.8 at $\lambda = 0.5$, +8.8 at $\lambda = 1$). This suggests forward KL benefits more from on-policy data than reverse KL on GSM8K—perhaps because forward KL's mode-covering behavior is more vulnerable to distribution mismatch, making on-policy correction more critical.
• JSD(0.1) and JSD(0.5) perform in between, with JSD(0.1) closer to forward KL and JSD(0.9) closer to reverse KL.

Effect of on-policy data proportion. Figure 8 varies the on-policy student data fraction from 0% to 100% while keeping divergence fixed (forward KL and reverse KL separately). With forward KL, accuracy improves monotonically as the on-policy fraction increases beyond 25%: at 0%, the student achieves approximately 7.5% test accuracy; at 25%, approximately 8.3%; at 50%, approximately 9.0%; at 75%, approximately 9.8%; at 100%, approximately 10.5%. The biggest jump occurs between 0-25% (the student needs a minimum amount of on-policy data to see benefit) and between 50-100% (continued monotonic improvement). With reverse KL, the trend is less monotonic and the overall performance is lower at most fractions. This provides direct evidence that the on-policy data fraction is not just a binary choice—more on-policy data continues to help up to $\lambda = 1$.

Comparison to baselines across student sizes. Figure 9 reports test accuracy for all KD methods across Flan-T5-small, base, and large, with external calculator:

• Flan-T5-small (baseline ~4.6%, teacher 27.9%): on-policy GKD achieves ~11.5%, ImitKD achieves ~9.5%, supervised KD achieves ~8.5%, f-distill achieves ~7.5%, SeqKD achieves ~7%.
• Flan-T5-base (baseline ~10.2%): on-policy GKD achieves ~19%, ImitKD achieves ~15%, f-distill achieves ~13%, supervised KD achieves ~12.5%, SeqKD achieves ~11.5%.
• Flan-T5-large (baseline ~20.5%): on-policy GKD achieves ~27.5% (nearly matching the teacher's 27.9%), ImitKD achieves ~26%, f-distill achieves ~24.5%, supervised KD achieves ~25%, SeqKD achieves ~23.5%.
• Reference lines are provided: 0-shot davinci-002 (175B) at ~12%, few-shot PaLM (540B, without calculator) at ~18%. On-policy GKD with Flan-T5-base surpasses the 540B PaLM model's few-shot performance.

The consistent ranking (on-policy GKD > ImitKD > supervised KD > f-distill > SeqKD) across all three student sizes reinforces the robustness of the on-policy advantage for reasoning tasks. The gap is largest for the smallest student (+7 percentage points over SeqKD for T5-small) and narrows for T5-large (+4 percentage points), consistent with the imitation learning prediction.

Self-distillation. Appendix Figure A.11 reports self-distillation results where the student and teacher have the same architecture and size (Flan-T5-Large, ~800M parameters), with the teacher being supervised fine-tuned on GSM8K (20.5% accuracy) and the student starting from a non-fine-tuned Flan-T5-Large (14.4% accuracy). On-policy GKD with forward KL achieves ~25%, on-policy GKD with JSD(0.9) achieves ~24.5%, f-distill achieves ~23%, ImitKD achieves ~22%, supervised KD achieves ~21%. All self-distilled students surpass the teacher's performance (20.5%), demonstrating that the on-policy mechanism provides benefits beyond simple compression—it improves performance even without a capacity gap.

---

Task-Agnostic Distillation: Instruction Tuning

Headline result. On-policy GKD with reverse KL substantially outperforms supervised KD and ImitKD for task-agnostic instruction tuning, evaluated on held-out benchmark suites. Figure 10 reports results for Flan-T5-XL → Flan-T5-Base after 50K training steps on the FLAN2021 dataset (5.36M examples):

• MMLU (57 tasks): On-policy GKD with reverse KL achieves approximately +2.0 percentage points improvement in average accuracy over the initial student. On-policy GKD with forward KL achieves approximately -0.25 (a slight degradation). Supervised GKD with reverse KL achieves approximately +1.0. Supervised KD achieves approximately +0.25. ImitKD achieves approximately -0.5 (worse than the initial student). The initial Flan-T5-Base achieves 35.6% on MMLU; the teacher Flan-T5-XL achieves 52.4%.
• BBH (23 tasks): On-policy GKD with reverse KL achieves approximately +1.0 percentage point improvement. On-policy GKD with forward KL achieves approximately +0.1. Supervised GKD with reverse KL achieves approximately +0.5. Supervised KD achieves approximately +0.25. ImitKD achieves approximately -0.5 (again worse than the initial student). The initial Flan-T5-Base achieves 31.25% on BBH; the teacher achieves 41%.

The divergence effect is dramatically different from previous tasks: reverse KL is clearly superior to forward KL for instruction tuning. The paper hypothesizes this is because instruction following requires the model to "zero in on the main intent or behavior specified by the instruction" rather than modeling all plausible completions—mode-seeking behavior aligns with the task structure. Forward KL's tendency to spread probability mass across many plausible outputs appears actively harmful in this setting (on-policy forward KL performs worse than the initial student on MMLU).

Note: SeqKD is not run for instruction tuning "due to its computational inefficiency for generating data from the teacher during training"—the teacher would need to generate 5.36M sequences, which is prohibitively expensive for the 3B-parameter T5-XL.

---

Ablation Studies and Robustness Checks

Student data fraction sweep ($\lambda \in \{0, 0.5, 1.0\}$): Systematically evaluated across XSum (Figures A.12, A.13), WMT (Figure 6), and GSM8K (Figures 7, A.14). The consistent finding is that $\lambda = 1$ (purely on-policy) matches or outperforms $\lambda = 0.5$ (mixed) and $\lambda = 0$ (supervised) across all tasks, with the gap largest when (a) the student is smallest relative to the teacher and (b) the evaluation uses stochastic decoding rather than greedy search. On GSM8K with forward KL, the monotonic benefit of increasing on-policy data is especially clear in Figure 8, where accuracy increases from ~7.5% at 0% on-policy to ~10.5% at 100% on-policy. An important threshold appears around 25% on-policy data—below this, the benefit is small or nonexistent, suggesting a minimum amount of on-policy training is needed for the student to experience its own error distribution.

Divergence choice sweep (forward KL, reverse KL, JSD(0.1), JSD(0.5), JSD(0.9)): Evaluated across all tasks. The optimal divergence is confirmed to be task-dependent: forward KL dominates on GSM8K with greedy decoding (Figure 7, +8.8 accuracy improvement vs. +8.0 for reverse KL); JSD(0.1) (near forward KL) dominates on WMT with beam search (Figure 6, +0.71 BLEU improvement vs. +0.55 for forward KL); mode-seeking divergences (JSD(0.9), reverse KL) dominate on XSum with temperature sampling (Figure A.12) and on instruction tuning (Figure 10, reverse KL +2.0 vs. forward KL -0.25 on MMLU). On XSum with greedy sampling (Figure A.13), divergence choice matters little—all on-policy divergences within ~0.5 ROUGE-2 of each other for each student size. This confirms the paper's framework that divergence should be treated as a tuneable hyperparameter.

Teacher softmax temperature sensitivity: Not systematically ablated but varies across tasks in the reported configurations. On XSum, teacher temperature is set to 1.0 when evaluating with greedy sampling but 0.1 when evaluating with temperature sampling (Figure 2, A.12 vs. A.13). On GSM8K, teacher temperature is 0.1 (Table A.2). On WMT, teacher temperature is 1.0 (Table A.3). On instruction tuning, teacher temperature is 1.0 (Table A.4). The paper does not provide an ablation over teacher temperature values or explain the rationale for these choices beyond the statement in Appendix A.3 that for XSum with temperature sampling evaluation, "we set teacher temperature to 0.1 for the student." The dependence of results on this hyperparameter is unknown—if teacher temperature significantly affects the quality of on-policy supervision, this would be an important practical consideration.

Student size scaling: Evaluated on XSum (T5-small, base, large; Figure 1 left), WMT (T5-small, base, large; Figure 1 middle), and GSM8K (Flan-T5-small, base, large; Figures 1 right, 9). On-policy GKD outperforms baselines at every scale, but the relative advantage decreases as student capacity increases. On XSum, the gap between on-policy GKD and supervised KD narrows from ~1.5 ROUGE-2 for T5-small (38× smaller) to ~0.8 ROUGE-2 for T5-large (3.8× smaller). On GSM8K, the gap narrows from ~4 percentage points for T5-small to ~2.5 percentage points for T5-large. This scaling behavior is consistent with the imitation learning motivation: distribution mismatch is most severe when the capacity gap is large, so on-policy training provides the greatest benefit for the smallest students.

Training dataset size: Evaluated on XSum with T5-small at {0.5%, 5%, 25%, 100%} of the full training set (Figure 3). On-policy GKD on 5% of the data (~10K examples) matches supervised KD on 100% of the data. On-policy GKD on 25% (~50K examples) achieves approximately 14.8 ROUGE-2, outperforming supervised KD on 100% (approximately 13.8). This suggests that on-policy GKD is not just more efficient but also can compensate for limited supervised data—a finding with practical implications for domains where high-quality reference outputs are scarce but input prompts are abundant.

Self-distillation: Evaluated on GSM8K with Flan-T5-Large where student and teacher share the same architecture and size, but the teacher is fine-tuned on GSM8K while the student is not (Appendix Figure A.11). On-policy GKD with forward KL achieves approximately 25% accuracy, surpassing the teacher's 20.5% by a substantial margin. This demonstrates that on-policy distillation provides benefits even when there is no capacity gap, contradicting the intuition that distillation only helps because the teacher is larger. The mechanism may involve the teacher's supervised fine-tuning providing a better target distribution, with on-policy training allowing the student to effectively bootstrap from this target.

RL + distillation mixing coefficient ($\alpha$): Evaluated on XSum with T5-base (Figure 5) at $\alpha \in \{0.05, 0.1, 0.25, 0.5\}$. The trade-off between ROUGE-2 improvement and factual consistency improvement is smooth and controllable. At $\alpha = 0.5$, the student achieves both better factual consistency than the teacher and better ROUGE-2 than the RLEF baseline. At $\alpha = 0.05$, factual consistency improvement is maximized (+45%) but ROUGE-2 improvement is minimal (+3). The smoothness of the trade-off curve suggests $\alpha$ is a reliable control parameter that practitioners can tune based on their relative preference for task quality vs. reward optimization.

Evaluation with different decoding strategies: Greedy vs. temperature sampling results are reported for XSum (Figures 2, A.12 vs. A.13). The divergence effect is pronounced with temperature sampling ($\gamma = 1$) but largely disappears with greedy decoding. This is an important practical finding: if deploying with greedy decoding, the divergence choice is not critical and forward KL (the simplest and most common choice) is as good as any alternative. If deploying with stochastic decoding, divergence selection matters substantially and should be tuned.

RL fine-tuning with self-regularization vs. teacher-regularization: The comparison between RLEF* (regularize to initial student) and GKD + RL (regularize to teacher) in Figure 5 shows that teacher-regularization enables substantially higher ROUGE-2 at any given level of factual consistency improvement. This is a direct ablation of the regularization target, confirming that the teacher provides a better reference distribution than the student's own initial checkpoint.

---

Critical Assessment

Claim: On-policy GKD substantially outperforms commonly-used KD approaches. This claim is well-supported across four tasks, three student scales per task, and multiple evaluation protocols. The evidence is consistent and the margins are substantial—the paper does not present borderline improvements that could disappear with different random seeds. However, there are important caveats:

• Single model family. All experiments use T5 (or Flan-T5) encoder-decoder models. It is unknown whether the on-policy advantage generalizes to decoder-only architectures (GPT-style), which dominate current LLM development. The auto-regressive generation mechanism is the same in principle, but decoder-only models have different attention patterns and may exhibit different sensitivity to distribution mismatch. Experiments on LLaMA, Mistral, or Pythia would substantially strengthen the generality claim.

• Single teacher per task. Each task uses exactly one teacher (T5-XL or Flan-T5-XL). The paper does not explore how teacher quality affects on-policy distillation—does a better teacher provide more useful on-policy feedback, or does a weak teacher actually benefit more from the on-policy mechanism because its supervised outputs are lower quality? The self-distillation result (Figure A.11) provides a partial answer by showing that even a same-capacity teacher helps, but a systematic sweep over teacher quality is missing.

• No statistical uncertainty for most results. Only WMT and GSM8K report averaging over multiple seeds (3 seeds each). XSum results (Figures 1, 2, 3, 4, A.12, A.13) and instruction tuning results (Figure 10) report point estimates without any indication of variance. Given that XSum evaluation is on the validation split and GSM8K on the test split, and that batch sizes are moderate (32-128), there is likely non-trivial run-to-run variation. The claim of "substantial" improvement would be stronger with confidence intervals.

• Checkpoint selection unclear for XSum. For GSM8K, the paper states "We use checkpoints at the end of training after distillation for reporting results," and for instruction tuning, "We report the performance of distilled checkpoints obtained after 50K training steps for various methods." For XSum, no explicit checkpoint selection protocol is described—it is not clear whether the reported numbers are from the final checkpoint, the best validation checkpoint, or an average over the last N checkpoints. If the best validation checkpoint was selected, the reported gains may include an optimistic selection bias relative to methods that are less sensitive to checkpoint choice.

• The 2.1×, 1.7×, 1.9× multipliers are computed relative to baseline KD improvements, not absolute performance. The paper states "In terms of performance gains over the initial student from on-policy GKD, averaged across T5 student models of different sizes, we see relative gains of 2.1× on summarization." This means if baseline KD improves ROUGE-2 by +2.0 and on-policy GKD improves by +4.2, the relative gain is 2.1×. This metric inflates the apparent advantage when baseline improvements are small—for example, if initial performance is already close to the teacher, even a tiny absolute improvement can produce a large multiplier. The absolute gains are more informative: on XSum with T5-small, on-policy GKD achieves ~15.0 vs. supervised KD ~13.8 (an absolute +1.2 ROUGE-2); on GSM8K with Flan-T5-base, ~19% vs. ~12.5% (+6.5 percentage points). These are practically meaningful but less dramatic than the multipliers suggest.

Claim: On-policy GKD can be seamlessly combined with RL fine-tuning. This claim is supported by the XSum experiment (Figure 5), which is genuinely novel—no prior work combined distillation with RL fine-tuning. However:

• Single task, single reward model. The combination is demonstrated only on XSum with a textual entailment reward. It is unknown whether the approach works with other reward types (human preference models, automated metrics like BLEU or ROUGE, safety classifiers) or on other tasks. The reward model's quality and calibration likely affect the interaction with the distillation term—a poorly calibrated reward might dominate or be dominated by the distillation signal in unpredictable ways.

• The RL objective uses REINFORCE, but the paper provides no implementation details for this component. There is no description of how the policy gradient is estimated, whether advantage normalization or baseline subtraction is used, or what learning rate is applied to the RL term vs. the distillation term. The $\alpha$ parameter controls the relative weight, but if the two loss terms have different scales or variance, $\alpha$ alone may be insufficient for stable training. This makes the result harder to reproduce than the pure distillation experiments.

• Only one student size tested for RL + distillation (T5-base). The interaction between student capacity, distillation benefit, and RL benefit is unexplored. The claim of "seamless" combination would be stronger with results across multiple student scales.

Claim: Optimal divergence is task-dependent and interacts with evaluation protocol. This is one of the paper's most insightful claims and is well-supported by the data—forward KL is best on GSM8K (greedy), JSD(0.1) is best on WMT (beam search), reverse KL/JSD(0.9) is best on XSum (temperature sampling), reverse KL is best on instruction tuning (greedy). However:

• The paper provides no predictive theory for which divergence to use on a new task. The post-hoc rationalizations (mode-seeking for instruction following, mode-covering for translation diversity) are plausible but unfalsifiable within the paper's scope. A practitioner approaching a new task has no principled way to choose a divergence without running the same sweep that the paper did.

• Only three specific JSD $\beta$ values are tested (0.1, 0.5, 0.9). The space between these values is unexplored—it is possible that JSD(0.3) outperforms both JSD(0.1) and JSD(0.5) on WMT. The paper establishes that $\beta$ matters but does not characterize the shape of the performance-vs-$\beta$ curve.

• The instruction tuning result (Figure 10) is from a single training run at 50K steps. There is no learning curve showing whether the reverse KL advantage over forward KL emerges early, persists throughout training, or widens/narrows over time. If forward KL catches up with longer training, the task-dependence claim would need qualification.

Claim: On-policy data is more important than divergence choice. This is implicit in the results—within any given divergence, $\lambda = 1$ typically outperforms $\lambda = 0$, and the best supervised configuration is usually worse than the worst on-policy configuration on XSum (Figures A.12, A.13) and WMT (Figure 6). However:

• On GSM8K, supervised reverse KL (+6.9) outperforms on-policy JSD(0.9) (+6.9, tied) and on-policy JSD(0.1) (+4.2) in Figure 7 left. This means there exist divergence-data combinations where the divergence matters more than the data source. The "on-policy is more important" heuristic has exceptions.

• On instruction tuning, supervised GKD with reverse KL (+1.0) outperforms on-policy GKD with forward KL (-0.25) by a large margin (Figure 10). For this task, the divergence choice dominates the data source choice. The claim should be qualified: on-policy data is more important when using an appropriate divergence for the task, but a bad divergence can negate the on-policy benefit entirely.

Missing experiments that would strengthen the paper:

• Scaling to larger models: The largest teacher is 3B parameters. It is unknown whether GKD's on-policy advantage persists or diminishes with much larger teachers (e.g., 70B, 540B), where the teacher's distribution is significantly richer and the student's capacity to approximate it may be a harder bottleneck than distribution mismatch.

• Decoding strategy ablation within tasks: The XSum experiments show results with both greedy and temperature sampling, but WMT uses only beam search and GSM8K uses only greedy decoding. It would be valuable to know whether the optimal divergence for WMT changes if evaluating with temperature sampling, or whether GSM8K results are robust to decoding strategy.

• Interaction between $\lambda$ and training duration: The paper fixes total training steps and varies $\lambda$. But on-policy training generates new data every step while supervised training reuses fixed data. For a fixed step budget, on-policy training sees more unique training examples (because they are generated on-the-fly). A fairer comparison might match on total unique tokens seen rather than total steps—this would isolate the benefit of on-policy data from the benefit of data diversity.

• Ablation of warm-start quality: All experiments start from a supervised fine-tuned student. How does the quality of this warm-start affect GKD's benefit? A poorly warm-started student might benefit more from on-policy training (consistent with IL theory) or might be too unstable to generate useful self-training data.

• Computational cost in wall-clock time: The paper reports the sampling overhead as a multiplier on per-step cost (1.8× to 2.2× on GSM8K), but does not report total wall-clock training time for any experiment. A practitioner choosing between GKD and a simpler baseline needs to know whether the performance gain justifies the time investment in absolute terms, not just relative terms.

Summary assessment: The experiments provide strong evidence that on-policy GKD improves over supervised KD across a diverse set of tasks and student scales, with particularly compelling results on data efficiency (matching full-dataset supervised KD with 5% of the data) and on the task-dependence of optimal divergence. The combination with RL fine-tuning, while demonstrated only on a single task, represents a genuine proof-of-concept for a novel capability. The main weaknesses are the restriction to a single model family (T5), the lack of statistical uncertainty quantification for most results, and the absence of a predictive framework for divergence selection on new tasks. The paper establishes on-policy distillation as a reliable principle but leaves open the question of how practitioners should choose divergences without running expensive sweeps, and whether the findings transfer to the much larger decoder-only models that dominate current practice.

6. Limitations and Trade-offs

The Cost of Difficulty Estimation Is Unaccounted for in the Headline Gains

The assumption or constraint. The entire compute-optimal framework from the companion paper rests on estimating prompt difficulty before allocating the inference budget. The method for doing so — generating 2048 samples per question and averaging PRM final-answer scores — is extraordinarily expensive. As the authors acknowledge in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The paper reports efficiency gains (up to 4× over best-of-N) by comparing the compute spent after difficulty is known, without including the cost of learning the difficulty in the first place.

The consequence. In a real deployment, total cost equals difficulty estimation plus strategy execution. The difficulty estimation step alone — 2048 sample generations and PRM scoring per prompt — consumes more compute than the largest test-time budgets studied (256–512 generations). If difficulty estimation costs, say, 2048 generations and the compute-optimal strategy uses 64 generations, the total cost is 2112 generations, whereas a naive best-of-256 baseline costs 256 generations. The 4× efficiency claim becomes misleading: the compute-optimal strategy may actually be more expensive in total when the estimation overhead is amortized. The paper's results should be interpreted as an upper bound on achievable efficiency — valid only in the limit where difficulty can be predicted nearly for free.

What evidence exists in the paper. The paper explicitly acknowledges this gap in Section 3.2 and Figure 4 caption context, but provides no experiments measuring the estimation-accounted efficiency. There is no ablation on the number of samples needed for difficulty estimation (is 2048 necessary, or would 128 suffice?). There is no comparison of total compute (estimation + execution) between compute-optimal and baseline approaches. The difficulty estimation cost is simply excluded from all reported metrics.

Mitigation status. The paper identifies this as "a key avenue for future work" (Section 3.2), suggesting training models to predict difficulty directly from question text, but develops no such model and provides no evidence that difficulty can be predicted cheaply enough to preserve the efficiency gains. Until this is addressed, the 4× figure should be understood as a theoretical potential rather than a realized deployment gain.

---

All Results Are on a Single Benchmark with a Single Model Family

The assumption or constraint. Every experiment in the paper uses the MATH benchmark (500 test questions of competition-level mathematics) and the PaLM 2-S* model family. The paper states (Section 4) that it "believe[s] this model is representative of the capabilities of many contemporary LLMs," but provides no cross-model or cross-benchmark validation.

The consequence. Several aspects of the findings could be specific to the MATH-PaLM combination rather than general properties of test-time compute scaling:

• PRM over-optimization behavior depends on the verifier's calibration on PaLM 2-S* outputs. A model with different output distributions (e.g., more diverse token-level entropy) may show different vulnerability to search exploitation.
• Revision model training depends on the base model's ability to learn from in-context incorrect examples. The 38% correct-to-incorrect reversion rate (Section 6.1) might be higher or lower for other model families.
• Difficulty-dependence patterns (beam search hurting easy problems, revisions helping easy problems) might not generalize to non-symbolic reasoning domains. MATH consists exclusively of problems with well-defined symbolic answers — it is unclear whether code generation, logical reasoning, or open-ended QA would exhibit the same difficulty-strategy interactions.
• The 5-difficulty-bin structure is derived from the base model's pass@1 distribution on MATH. On a benchmark with different difficulty characteristics (e.g., mostly very easy or mostly very hard problems), the bin assignments and resulting compute-optimal policy would differ.

The 500-question test set, split into 5 quintiles of ~100 questions each, further split by two-fold cross-validation, means strategy selection is based on ~50 questions per fold per bin. This is a small sample for optimizing over a discrete set of strategy hyperparameters (search algorithm, beam width, revision chain length, etc.), and the computed-optimal policy may have high variance.

What evidence exists in the paper. None. There are no experiments on any benchmark other than MATH, and no experiments with any model other than PaLM 2-S*. The paper does not report confidence intervals on the compute-optimal scaling curves, so the variance from the small test set is unquantified.

Mitigation status. Not addressed. The authors acknowledge the single-benchmark limitation implicitly by focusing claims on MATH specifically, but the broader framing (e.g., "compute-optimal test-time scaling") implies generality that is not supported by the experimental design. Replication on additional benchmarks and model families is left entirely to future work.

---

The Larger-Model Baseline in FLOPs-Matched Comparisons Is Artificially Weak

The assumption or constraint. Section 7 compares PaLM 2-S* with compute-optimal test-time scaling to a model with approximately 14× more parameters. This larger model uses: (1) parameter-only scaling with fixed training data, departing from compute-optimal pretraining where both data and parameters scale (Hoffmann et al., 2022), and (2) greedy decoding with no test-time compute augmentation of its own.

The consequence. Both design choices weaken the pretraining baseline, inflating the apparent advantage of test-time compute:

• Parameter-only vs. Chinchilla-optimal scaling. A model trained with 14× more total FLOPs allocated optimally between parameters and data would likely outperform a model where only parameters are scaled. The paper explicitly acknowledges this choice (Section 7): "We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." This means the reported advantages — e.g., +27.8% relative improvement on easy-medium questions at low inference-to-pretraining ratios — may shrink or reverse against a properly compute-optimal larger model.
• Greedy decoding for the larger model. The 14× larger model gets no majority voting, no best-of-N, no search, no revision chaining — zero test-time compute augmentation. Yet the paper's central claim is about the substitutability of test-time compute for pretraining compute. A fairer comparison would give the larger model some test-time compute budget (the same budget that the smaller model receives?), or at minimum a modest best-of-N. Without this, the comparison is between "small model + optimized inference" and "large model + minimal inference," which conflates two effects: the value of test-time compute and the value of optimizing its allocation.

What evidence exists in the paper. The parameter-only scaling and greedy decoding choices are explicitly stated in Section 7. No ablation is provided with a Chinchilla-optimal pretraining baseline or with the larger model receiving test-time compute. The FLOPs accounting formulas (Section 7) include only pretraining and inference FLOPs, not the cost of training a PRM or revision model for the smaller model, which further advantages the test-time compute condition.

Mitigation status. The paper transparently acknowledges the parameter-only scaling choice and frames it as leaving the compute-optimal pretraining comparison to future work. The greedy decoding choice for the larger model is not discussed as a limitation. Both choices make the FLOPs-matched results a lower bound on pretraining's advantage — the true substitutability of test-time compute for pretraining is likely weaker than reported, potentially substantially so on hard problems.

---

Verifier Over-Optimization Is a Hard Ceiling, and the Compute-Optimal Policy Only Mitigates, Not Solves, It

The assumption or constraint. All search-based test-time compute methods rely on a learned verifier (PRM or ORM) to score candidate solutions. The verifier is imperfect, and aggressive search optimization amplifies its imperfections, eventually degrading rather than improving performance. This phenomenon is the central bottleneck preventing unbounded improvements from additional test-time compute.

The consequence. The compute-optimal policy mitigates over-optimization by routing easy problems away from aggressive search (using best-of-N instead of beam search for difficulty bins 1–2), but it cannot overcome the underlying limitation. On medium-difficulty problems where beam search is deployed (bins 3–4), over-optimization still limits the scaling ceiling — the beam search curves in Figure 3 (right) flatten and eventually decline as budget increases, meaning performance stops improving well before the budget is exhausted. On the hardest problems (bin 5), no method helps, but the failure mode is different (lack of capability, not verifier over-optimization).

What evidence exists in the paper. The evidence is extensive and detailed:

• Figure 3 (left): Lookahead search — the most powerful optimizer — paradoxically performs worst at matched generation budgets because its extra per-step cost reduces effective beam count while also over-optimizing the verifier signal.
• Figure 3 (right): Beam search on easy problems (bin 1) decreases in accuracy as budget increases from 4 to 256 generations, the clearest signature of verifier exploitation.
• Appendix M: Qualitative examples show search producing degenerate outputs that score highly under the PRM but are incorrect — repetitive low-information steps, overly short 1–2 step solutions, and solutions that exploit the PRM's blind spots.

The compute-optimal policy improves the efficiency of budget allocation but does not prevent over-optimization from eventually dominating. The scaling curves in Figure 4 for compute-optimal search flatten around 39–40% accuracy at 256 generations, suggesting a hard ceiling determined by verifier quality. Improving the verifier (through better training data, adversarial robustness, ensembles) would likely shift this ceiling upward, but the paper does not explore verifier improvement as an intervention.

Mitigation status. The paper identifies over-optimization as the primary bottleneck (Sections 5.3, 8) and the compute-optimal policy can be viewed partly as a strategy for staying below the over-optimization threshold per difficulty level. But the fundamental problem — how to build verifiers robust to aggressive search — is not addressed. The paper frames improving verifier robustness as a key direction for future work rather than as something solved within the current framework.

---

Hard Problems Remain Essentially Unsolved, and Test-Time Compute Cannot Compensate for Fundamental Capability Gaps

The assumption or constraint. The paper's approach assumes the base model already has some non-trivial probability of generating a correct solution for a given prompt. Test-time compute can amplify this probability — guiding search toward correct solutions, refining nearly-correct answers — but it cannot create capability where none exists.

The consequence. Across all methods — search, revisions, compute-optimal combinations — the hardest questions (difficulty bin 5) show near-zero improvement regardless of compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budget levels. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5%, well below the 14× larger model's performance at any inference-to-pretraining ratio. The paper is candid about this in the Section 7 takeaway:

"On the hardest questions (bins 4–5), pretraining is almost always more effective."

This means the approach offers no path forward for problems genuinely outside the base model's capability range. For a production system where a meaningful fraction of queries fall into this regime, investing in larger pretraining is the only viable option — no amount of clever inference-time optimization will help. The practical implication is that organizations must understand their query difficulty distribution to decide whether test-time compute scaling or pretraining scaling is the better investment.

What evidence exists in the paper. The evidence is consistent and unambiguous across every figure that breaks down results by difficulty bin. Bin 5 accuracy is near-zero under all conditions. The paper correctly identifies this as a fundamental limitation and does not overclaim.

Mitigation status. The paper explicitly acknowledges this boundary (Section 7 and the difficulty-bin analyses throughout). There is no attempt to solve the hard-problem regime within the current framework, and the authors do not suggest any approach for doing so — the limitation is presented as inherent to the concept of test-time compute scaling. Future work would need to address capability acquisition (via pretraining, retrieval, or tool use) rather than capability amplification, which is outside the scope of test-time compute optimization.

---

Sequential Revisions Introduce a Latency Tax That the Paper Does Not Account For

The assumption or constraint. The paper measures test-time compute in "generations" — the number of complete solution samples produced. This is a reasonable proxy for total FLOPs but ignores wall-clock latency. Sequential operations cannot be parallelized, while parallel operations can be executed simultaneously given sufficient hardware.

The consequence. The compute-optimal policy, particularly for easy-to-medium problems, favors sequential strategies — long revision chains (up to 64 sequential steps) or beam search with serial expansion. A strategy that allocates, say, 64 generations as 64 sequential revisions takes approximately 64× the wall-clock time of a strategy that runs 64 parallel samples simultaneously. For latency-sensitive applications — interactive assistants, real-time decision-making, any user-facing system — this serial dependency may be prohibitive regardless of accuracy advantages. The compute-optimal policy might select a strategy that is optimal in FLOPs but unacceptable in latency, and the paper provides no framework for incorporating a latency constraint into the allocation optimization.

What evidence exists in the paper. The paper does not measure or discuss wall-clock time, latency, or throughput for any experiment. The "computation cost" discussed in Section 3.2 and Appendix A.2 refers only to total floating-point operations (or number of generations), not to the serial vs. parallel nature of those operations. The revision model's sequential chain structure (Section 6) is described as generating revisions one at a time, with each revision conditioned on the previous one — an inherently serial process. The paper notes that the revision model generates up to 64 sequential revisions (Figure 6), but does not quantify the latency cost of this serial dependency relative to parallel sampling.

Mitigation status. Not addressed. The paper treats "generation budget" as the sole resource constraint, implicitly assuming either that latency is irrelevant or that sufficient parallelism is available to mask serial dependencies (e.g., processing many queries simultaneously in a batch). For practitioners deploying in latency-constrained settings, this is a significant gap: the compute-optimal strategy must be re-derived with both FLOP and latency constraints, which could dramatically change the optimal sequential-to-parallel ratio, especially for easy problems where the paper's policy heavily favors sequential revisions.

7. Implications and Future Directions

How This Work Changes the Landscape

GKD reframes knowledge distillation for auto-regressive language models from a static data problem to a dynamic interaction problem. This is not a new architecture, a new loss function in the narrow sense, or a new training trick—it is a methodological reframing that changes what practitioners consider the default answer to "what should I train my student on." Before GKD, the standard workflow was: collect a high-quality dataset (human references or teacher-generated outputs), train the student to mimic the teacher's probabilities on that fixed dataset, deploy. GKD's central demonstration is that this workflow systematically underperforms because it ignores the distribution mismatch between the states the student sees during training (expert trajectories) and the states it encounters during its own auto-regressive generation at inference time (self-generated trajectories with accumulated errors). The fix—having the student generate its own training sequences and receive token-level teacher feedback on those self-generated states—is conceptually simple but represents a genuine shift in perspective: the student's own outputs are not just an evaluation target but the primary training medium.

The magnitude of this shift is substantial but foundational rather than revolutionary. The paper does not claim that supervised KD is obsolete—in fact, GKD's unified framework explicitly retains the ability to use fixed datasets through the $\lambda$ mixing parameter, and the paper shows that combining on-policy and supervised data ($\lambda = 0.5$) often performs nearly as well as purely on-policy ($\lambda = 1$). Rather, GKD establishes that on-policy data should be the default starting point, with supervised data as a useful supplement, inverting the prior practice. This is analogous to how DAgger (Ross et al., 2011) changed imitation learning from "train on expert demonstrations" to "collect learner rollouts, get expert labels on those rollouts, retrain, repeat"—not by showing expert demonstrations are useless, but by demonstrating that the learner's own state distribution is the more critical training signal.

Reconciling prior contradictions. The paper resolves a tension that had been brewing in the distillation literature. On one side, supervised KD (Sanh et al., 2019) and SeqKD (Kim & Rush, 2016) demonstrated that distilling from a larger teacher could improve student performance, establishing distillation as a reliable technique. On the other side, research on exposure bias (Zhang et al., 2019; Chiang & Chen, 2021; Arora et al., 2022) showed that auto-regressive models trained on fixed datasets suffer from cascading errors at inference time because they never practice recovering from their own mistakes. These findings were not contradictory—they were describing different aspects of the same phenomenon—but the field lacked a framework for understanding that supervised KD's success despite exposure bias means it is leaving substantial performance on the table. GKD quantifies this gap: across summarization, translation, reasoning, and instruction tuning, on-policy training consistently outperforms supervised training on the same teacher with the same student architecture, with the largest gains for the smallest students (38× smaller than teacher on XSum, Figure 1) where distribution mismatch is most severe. The prior apparent success of supervised KD was real but incomplete—it represented a lower bound on what distillation could achieve, not the ceiling.

Research directions that become more attractive. Several lines of inquiry gain momentum from GKD's findings:

• On-policy data generation for other compression methods. Pruning, quantization, and architecture search typically use fixed calibration datasets to guide compression decisions. GKD's demonstration that training on self-generated data improves student quality suggests a parallel question: should the student's own outputs be used as calibration data for determining which weights to prune or how to quantize? If the student's deployment-time input distribution differs from the training distribution, compression decisions made on static data may be suboptimal. This direction is newly tractable because GKD provides a template for generating on-policy calibration data from the student itself.

• RLHF pipeline redesign. The standard RLHF workflow is SFT → reward model training → PPO with KL penalty toward the initial SFT policy. GKD's demonstration that the KL penalty target can be profitably replaced with a teacher model (Section 3.2, Figure 5) opens a direct path to modifying this pipeline: SFT → teacher distillation (on-policy GKD) → reward model training → PPO with KL penalty toward the teacher. This would mean the RL fine-tuning step simultaneously aligns the model with human preferences (via the reward) and improves its general capabilities (via teacher-regularized distillation), potentially reducing or eliminating the alignment tax documented by Ouyang et al. (2022). The combination is computationally natural because on-policy GKD and PPO both require sampling from the current policy—they can share the same forward passes, with GKD using the teacher's log-probabilities and PPO using the reward model's scores.

• Self-improving systems with teacher guidance. The paper's self-distillation result (Appendix A.11, Figure A.11)—where a model distilled from a fine-tuned copy of itself using on-policy GKD surpasses the teacher's performance—is particularly provocative. It suggests that on-policy distillation can serve as an optimization mechanism, not just a compression mechanism. The student generates outputs, the teacher (a better-trained version of the same architecture) scores them, and the student improves beyond the teacher by learning from its own exploration. This is conceptually similar to the self-play dynamics that have driven progress in game-playing AI, but applied to language model training. The result is preliminary (single task, single model size) but suggests a research program around iterative on-policy distillation cycles.

Research directions that become less attractive. GKD's results also suggest diminished returns for certain approaches:

• Developing ever-more-complex fixed-data augmentation strategies for distillation. If the primary bottleneck is distribution mismatch rather than data quantity or diversity, then sophisticated data curation, filtering, or synthesis pipelines that still produce static datasets will hit the same fundamental ceiling as supervised KD. The paper's data efficiency result (Figure 3)—on-policy GKD on 5% of the XSum training data outperforms supervised KD on 100%—directly demonstrates that the data source (on-policy vs. fixed) matters more than the data volume. Researchers working on distillation data pipelines should consider whether their efforts would be better spent implementing on-policy generation.

• Pure sequence-level distillation methods that don't use token-level teacher feedback. SeqKD trains on teacher-generated sequences but uses hard likelihood maximization rather than token-level probability matching. GKD shows that SeqKD is substantially worse than on-policy GKD across student sizes on GSM8K (Figure 9, SeqKD is the lowest-performing method for every student size). The token-level signal from the teacher—showing the student not just "this is the correct token" but "here is the full distribution over what could have been correct, and here is how your distribution differs"—appears essential for the on-policy mechanism to work effectively. Methods that discard this rich signal in favor of sequence-level objectives are swimming upstream.

• Approaches that rely on forward KL as the only divergence. The paper's most consistent empirical finding across tasks is that the optimal divergence is task-dependent and can differ substantially from forward KL. On XSum with temperature sampling, mode-seeking divergences (reverse KL, JSD(0.9)) substantially outperform forward KL (Figure 4). On instruction tuning, reverse KL provides a ~2% absolute accuracy gain on MMLU while forward KL actually degrades performance relative to the initial student (Figure 10). On WMT, JSD(0.1) outperforms forward KL (Figure 6). On GSM8K with greedy decoding, forward KL is best (Figure 7). This pattern means that any method that hard-codes forward KL—including most existing distillation implementations—is leaving performance on the table for some tasks and evaluation regimes. The divergence should be treated as a tuneable hyperparameter, and frameworks that don't support this (including ImitKD and standard supervised KD) are architecturally limited.

---

Follow-Up Research This Work Enables

1. Characterizing the divergence-task interaction with a predictive theory. The paper demonstrates that optimal divergence is task-dependent—forward KL for GSM8K greedy decoding, reverse KL for instruction tuning, JSD(0.1) for WMT beam search, JSD(0.9) for XSum temperature sampling—but provides only post-hoc explanations. A direct follow-up would systematically vary task properties and measure which divergence performs best, aiming to build a predictive model. For example: does the optimal divergence correlate with the entropy of the teacher's output distribution? With the number of valid outputs per input? With whether evaluation uses stochastic or deterministic decoding? A concrete experiment would take 10–20 diverse tasks (spanning classification, generation, reasoning, translation, summarization, code generation), run the full GKD divergence sweep on each, and regress optimal divergence against measurable task properties. If a reliable predictor emerges—e.g., "mode-seeking divergences are preferred when the teacher's average token entropy exceeds threshold H"—practitioners could select divergences without expensive per-task sweeps. The paper's existing data already provides five data points (XSum greedy, XSum temperature, WMT beam, GSM8K greedy, FLAN instruction tuning); expanding to 15–20 would enable meaningful analysis.

2. Scaling GKD to GPT-style decoder-only models at current LLM scales. All experiments use T5 encoder-decoder models with the teacher capped at 3B parameters. The most pressing external validity question is whether on-policy GKD's benefits persist with much larger decoder-only architectures (e.g., distilling LLaMA-70B into LLaMA-7B, or Mixtral 8×7B into a dense 7B model). Decoder-only models have different attention patterns (causal self-attention throughout vs. encoder-decoder cross-attention) and may exhibit different sensitivity to distribution mismatch. A strong follow-up would replicate the core result—on-policy GKD vs. supervised KD across student sizes—on a decoder-only family with a teacher in the 30B–70B range, on at least two tasks (summarization and reasoning, to cover both the forward-KL-preferred and reverse-KL-preferred regimes). The experiment would answer: (a) does the on-policy advantage scale to much larger absolute model sizes, (b) does the optimal divergence pattern replicate, and (c) is the computational overhead of student sampling still acceptable when student inference already costs tens of billions of FLOPs per sequence? A negative result—on-policy GKD providing minimal benefit at scale—would force a refinement of the theory: perhaps distribution mismatch is primarily a small-model phenomenon, or perhaps decoder-only architectures handle it differently.

3. Cheap difficulty estimation via learned difficulty predictors for compute-optimal allocation. The companion paper's compute-optimal test-time scaling framework (the reference example analyzed in prior sections) is bottlenecked by the cost of difficulty estimation: 2048 sample generations per question. GKD's on-policy mechanism suggests a potential solution: train a lightweight "difficulty predictor" model using distillation from the PRM's difficulty assessments, then use that predictor to allocate test-time compute. Concretely: generate difficulty labels for a large set of prompts using the expensive 2048-sample PRM method, then fine-tune a small classifier (e.g., T5-small or a lightweight BERT variant) to predict the difficulty quintile directly from the prompt text, using on-policy GKD with the PRM as teacher to ensure the predictor is robust to its own inference-time distribution. A strong evaluation would compare the performance of compute-optimal test-time scaling using (a) oracle difficulty, (b) the expensive PRM-based difficulty, (c) a cheap learned predictor, and (d) no difficulty adaptation (uniform best-of-N), measuring both total compute (estimation + execution) and final accuracy. If the learned predictor achieves difficulty estimates within 10% accuracy of the PRM-based method at <1% of the compute cost, the compute-optimal framework becomes immediately practical for deployment. This research direction is directly enabled by GKD's demonstration that on-policy distillation from a frozen verifier works—the difficulty predictor is the student, the PRM is the teacher, and the training data is generated on-policy from the predictor as it learns.

4. Iterative on-policy self-distillation cycles for capability amplification. The self-distillation result (Appendix A.11)—where on-policy GKD from a fine-tuned teacher of the same architecture yields a student that surpasses the teacher—appears once in the paper but is not explored systematically. A natural follow-up would test whether this process can be iterated: train a model on a task via supervised FT (generation 0), use it as teacher to distill a fresh copy of the same architecture via on-policy GKD (generation 1), then use the generation-1 model as teacher to distill generation 2, and so on. The question is whether performance continues to improve, plateaus, or eventually degrades (due to amplification of model-specific biases or diversity collapse). This is analogous to the self-play dynamics in Silver et al. (2017) but for language tasks rather than games. A concrete experiment would run 5–10 iterations on GSM8K (where the self-distillation gain is cleanly measurable) and on XSum (to test whether the effect transfers to generative tasks), measuring accuracy/ROUGE at each iteration. If performance monotonically improves for 3+ iterations, this would establish on-policy self-distillation as a general capability amplification method independent of model compression. A negative result—performance peaking at iteration 1 and then declining—would suggest that the benefit comes from the teacher's supervised FT providing a one-time "correction signal" that cannot be bootstrapped further.

5. Understanding and mitigating the 38% correct-to-incorrect reversion rate in revision models. The companion paper's revision model suffers from a specific pathology: approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect answers in subsequent steps (reference example, Section 6.1). This occurs because the revision model is trained only on incorrect-to-correct trajectories, so it has no training signal for "the current answer is already correct, do nothing." GKD's framework provides a natural diagnostic and potential fix. First, diagnose: train a revision model with on-policy GKD where the teacher is the base model (not a fine-tuned revision teacher) and the divergence is reverse KL. If the model is mode-seeking toward the base model's high-probability continuations, it may learn to preserve correct answers because the base model already assigns high probability to them—the reverse KL penalty for changing a correct answer would be large. Second, fix: construct training data that includes correct-to-correct trajectories (where the model is shown a correct answer in context and trained to reproduce it unchanged), and use GKD's $\lambda$ parameter to mix these with the standard incorrect-to-correct trajectories. A strong evaluation would report the reversion rate (fraction of correct answers changed to incorrect) before and after the intervention, along with overall task accuracy. This work is directly enabled by GKD's unification of on-policy training with flexible divergence choice—the revision model training is a special case of on-policy distillation where the "teacher" is either the ground-truth correct answer or a verifier, and GKD shows how to optimize the divergence and data mixture for this setting.

6. Stress-testing the on-policy advantage: when might on-policy data hurt? The paper consistently finds that $\lambda = 1$ (purely on-policy) matches or outperforms $\lambda < 1$, but all experiments start from a supervised fine-tuned student of at least "adequate quality." This leaves open a critical boundary condition: if the student's initial policy is very poor—generating mostly nonsensical or irrelevant outputs—does on-policy training still help, or does it create a destructive feedback loop where the student trains on low-quality states, learns to produce more low-quality outputs, and diverges further? This question matters for distillation scenarios where no warm-start is available (e.g., distilling into a randomly initialized small model, or distilling a model trained on English into a model for a low-resource language). A concrete stress test would start GKD from students of deliberately varying quality: (a) fully supervised fine-tuned (standard warm start), (b) partially fine-tuned (1/10th the steps), (c) few-shot prompted with no fine-tuning, (d) randomly initialized. For each, measure whether on-policy GKD improves over the initial student and over supervised KD. The imitation learning theory predicts a U-shaped or threshold effect: on-policy training should help when the student is good enough that its errors are informative (states where the teacher can provide useful corrections) but may hurt when the student is so poor that its states are effectively noise. Finding this threshold would provide crucial practical guidance for when GKD can be applied without a warm-start phase, and would connect GKD more tightly to the IL literature (where DAgger is known to require an initial policy that is at least somewhat competent).

---

Practical Applications and Downstream Use Cases

1. Cost-efficient deployment of large language models via on-policy distillation with small students. Organizations serving language models at scale face a direct economic tension: larger models provide better quality but cost proportionally more per query. GKD offers a concrete recipe for narrowing this gap. The paper shows that a T5-small student (77M parameters, 38× smaller than the teacher) distilled with on-policy GKD achieves ROUGE-2 of ~15.0 on XSum, compared to ~13.8 for supervised KD—a ~9% relative improvement at the same deployment cost. On GSM8K, a Flan-T5-base student (250M parameters) with on-policy GKD achieves ~19% accuracy (matching the few-shot performance of PaLM 540B without a calculator, Figure 9), while the 3B teacher achieves 27.9%. The implication is that an organization could serve summarization or reasoning queries using a model that costs 12–38× less per token than the largest available teacher, while recovering 70–80% of the teacher's performance, by investing in on-policy distillation during training. The training overhead (1.8–2.2× on GSM8K per the paper's estimates) is a one-time cost; the inference savings compound with every query served. For a deployment serving millions of queries per day, this translates to substantial infrastructure savings.

2. Self-improving data annotation pipelines that reduce dependence on human references. The paper's data efficiency result on XSum (Figure 3) has direct implications for domains where human-annotated output sequences are scarce or expensive. On-policy GKD trained on just 5% of the XSum training data (~10K examples) without any ground-truth summaries matches supervised KD trained on the full dataset (~200K examples) with ground-truth summaries. This means an organization with a strong teacher model, a large corpus of unlabeled input prompts, and a small set of human-annotated examples can bootstrap a high-quality student model: (a) use the small annotated set to supervised fine-tune both the teacher and a warm-start student, (b) run on-policy GKD using the unlabeled prompts (student generates outputs, teacher scores them), (c) deploy the distilled student. The unlabeled prompts serve as the input distribution for on-policy generation, and the teacher provides all the supervision—no additional human annotation is needed. This is immediately applicable to domains with abundant unlabeled text but limited annotation budgets: legal document summarization, medical report generation, customer support response drafting, and low-resource language tasks where parallel corpora are limited but monolingual text is plentiful.

3. Mitigating the alignment tax in RLHF pipelines. The standard RLHF workflow (SFT → reward model → PPO with KL penalty to SFT policy) is known to reduce general model capabilities even as it improves alignment with human preferences—the alignment tax (Ouyang et al., 2022). GKD's demonstration (Figure 5) that replacing the self-regularization KL penalty with a teacher-regularization penalty enables simultaneous improvement in task quality and reward optimization provides a practical path to reducing this tax. In a production RLHF pipeline, the modification is minimal: instead of computing KL(current_policy || initial_sft_policy), compute KL(current_policy || teacher_policy) where the teacher is a larger, more capable model (which may already exist in the organization's model registry). On XSum with a T5-base student and T5-XL teacher, this substitution yields summaries that are both more factually consistent than the 12× larger teacher and higher quality than the self-regularized RLEF baseline (Figure 5). For organizations already running RLHF, this is a low-implementation-cost change that could improve both alignment and capability metrics simultaneously, reducing or eliminating the need to trade one off against the other.

4. Task-agnostic model compression for multi-task serving. The instruction tuning results (Figure 10) demonstrate that on-policy GKD with reverse KL improves a distilled model's performance on held-out benchmarks—tasks the model was not explicitly distilled on. The Flan-T5-Base student distilled with on-policy reverse KL gained ~2% absolute accuracy on MMLU (57 diverse tasks) and ~1% on BBH (23 challenging tasks) over the initial student, while supervised KD and ImitKD provided smaller or even negative gains. This is significant for organizations deploying a single model to handle a wide variety of user requests (the standard "model-as-a-service" paradigm). A single distilled model that is 12× smaller than the teacher but retains more of its broad capabilities reduces serving costs without requiring task-specific distillation for every possible use case. The practical pipeline: take a large instruction-tuned model (the teacher), its instruction tuning dataset, and a smaller pre-trained model; run on-policy GKD with reverse KL; deploy the resulting student as a general-purpose endpoint. The reverse KL choice is critical here—the paper shows forward KL actually degrades MMLU performance relative to the initial student, so using the wrong divergence would produce a model worse than not distilling at all.

---

When to Prefer This Method

The paper does not frame GKD as a strict alternative to named competing methods with explicit decision boundaries. Rather, it presents GKD as a unifying framework that recovers existing methods as special cases ($\lambda = 0$ gives supervised KD; forward KL with $\lambda = 0.5$ gives ImitKD) while instantiating new, better-performing configurations (purely on-policy, alternative divergences). The empirical results consistently show that the new configurations—particularly on-policy with task-appropriate divergence—outperform the special cases that correspond to prior methods. The paper's positioning implies:

• If you are currently using supervised KD or SeqKD: switching to on-policy GKD with the same divergence (forward KL) provides a consistent improvement with the same teacher and student, at the cost of 1.8–2.2× training time overhead (per the GSM8K estimate). The improvement is largest for small students relative to the teacher.

• If you are currently using ImitKD or f-distill: these are already GKD variants with suboptimal choices ($\lambda = 0.5$, specific divergences). Moving to $\lambda = 1$ and tuning the divergence for your task and evaluation protocol provides additional gains, as shown in Figures 2, 9, and A.15.

• If you are doing RLHF: replacing the self-regularization KL penalty with a teacher-regularization KL penalty (using GKD's combined objective, Equation 5) can improve both reward optimization and general capability, as shown on XSum (Figure 5). Start with $\alpha = 0.5$ and reverse KL or JSD(0.9) for the distillation term, then tune $\alpha$ based on the reward-quality trade-off.

The paper does not identify conditions under which supervised KD or SeqKD would be preferred over on-policy GKD. The closest it comes to a boundary is the acknowledgement that GKD requires a warm-started student capable of generating "sequences of adequate quality" (Section 3.1)—if no such warm-start is available, supervised training is necessary to reach the threshold where on-policy training becomes beneficial. This threshold is not quantified in the paper.