# GKD: Generalized Knowledge Distillation for Auto-regressive Sequence Models

Rishabh Agarwal Nino Vieillard∗ Piotr Stanczyk† &Sabela Ramos† &Matthieu Geist
&Olivier Bachem & Google DeepMind denotes equal contribution. † denotes
infrastructure contribution. Correspondence to Nino<EMAIL_ADDRESS>and
Rishabh<EMAIL_ADDRESS>
###### Abstract

Knowledge distillation is commonly used for compressing neural networks to
reduce their inference cost and memory footprint. However, current
distillation methods for auto-regressive models, such as generative language
models (LMs), suffer from two key issues: (1) distribution mismatch between
output sequences during training and the sequences generated by the student
during its deployment, and (2) model under-specification, where the student
model may not be expressive enough to fit the teacher’s distribution. To
address these issues, we propose Generalized Knowledge Distillation (GKD). GKD
mitigates distribution mismatch by sampling output sequences from the student
during training. Furthermore, GKD handles model under-specification by
optimizing alternative divergences, such as reverse KL, that focus on
generating samples from the student that are likely under the teacher’s
distribution. We demonstrate that GKD outperforms commonly-used approaches for
distilling LLMs on summarization, machine translation, and arithmetic
reasoning tasks.

##  1 Introduction

Figure 1: Comparing GKD with common distillation methods on XSum summarization
(Narayan et al., 2018). We use the T5 models (Raffel et al., 2020) trained
with supervised fine-tuning (FT) as the student models for distillation.
Supervised KD and supervised FT use the XSum training dataset with ground-
truth summaries but KD can query the teacher to get probabilities while FT
does not. Furthermore, on-policy approaches sample summaries from the student
while ‘Mixed’ refers to uniform sampling from the ground-truth and student-
generated summaries. ImitKD (Lin et al., 2020) corresponds to ‘Mixed’ sampling
with Forward KL. GKD with reverse KL and generalized JS divergence (JSD), as
discussed in Section 2, outperforms other approaches. Figure 2: Number of
parameters in LLMs over the years.

Auto-regressive sequence models, such as language models (LMs), have shown
impressive capabilities in numerous tasks, where the key to this success is
often scaling the amount of training data as well as the number of model
parameters (Figure 2). However, scaling parameter count comes at a cost, and
the deployment of such state-of-the-art models in real-world is limited by
either their inference cost or memory footprint. Thus, a crucial goal for
practical use of large capable models, such as LLMs, is to compress them by
reducing their number of parameters, while retaining as much as possible of
their performance.

One of the prevalent techniques for compressing auto-regressive models is
knowledge distillation (Hinton et al., 2015). Distillation is the process of
training a model – the student – to replicate the knowledge of another model –
the teacher – on a specific set of tasks. Typically, the student has fewer
parameters than the teacher and as such, distillation can improve task-
specific performance while maintaining lower inference cost and memory
footprint than the teacher. Current distillation methods for auto-regressive
models either require generating a large amount of output sequences from the
teacher model, which can be quite costly depending on the teacher size, or a
fixed dataset of sequences that the teacher can label (_e.g._ , assign token-
level probabilities). However, using a fixed dataset can lead to distribution
mismatch between teacher-generated output sequences during training and the
sequences generated by the student auto-regressively during its deployment, a
well-known problem in imitation learning (Pomerleau, 1991; Ross & Bagnell,
2010). Furthermore, the common objective for distillation is to maximize the
likelihood of teacher-generated samples under the student’s distribution.
However, this objective may not be ideal due to model under-specification
(Huszár, 2015): the student model may not be expressive enough to fit the
teacher’s distribution and maximizing likelihood can lead to student-generated
samples that are unlikely under the teacher (_e.g._ , Figure 3).

Inspired by the recent success of reinforcement learning (RL) fine-tuning
approaches with auto-regressive language models, we mitigate the above issues
with supervised distillation by recognizing that we can view such models as
agents and as such, distillation as an imitation learning problem with a known
expert. In particular, instead of training the student using a fixed
distribution over outputs, we argue for using samples from the student’s
distribution itself during training, akin to RL and on-policy distillation
(Ross et al., 2011). Furthermore, to address model under-specification, we
argue that alternative objectives that focus on generating samples from the
student that are likely under the teacher’s distribution, such as reverse KL,
are more suitable for distilling auto-regressive models. Combining the above
ideas, we propose Generalized Knowledge Distillation (GKD), which generalizes
both on-policy and supervised distillation. We demonstrate that GKD
significantly outperforms commonly-used approaches for distilling large
language models on summarization (Figure 1), machine translation (WMT), and
arithmetic reasoning (GSM8K) tasks.

##  2 Background: KL-based Divergences

Figure 3: Mode _vs_ Mean-seeking KL under model under-specification. The plot
above shows the learned distribution when minimizing the forward and reverse
KL between a mixture distribution $P$ of 2 Gaussians with unit variance and a
unimodal Gaussian $Q_{\theta}$ with respect of $Q$. Reverse KL is mode-seeking
as minimizing it forces $Q_{\theta}$ to be zero where $P$ is zero and hence
makes it concentrate on one of the modes (last plot). However, forward KL is
mode-covering or mean-seeking minimizing it ensures that there is some mass
under $Q_{\theta}$ wherever there is some mass under $P$ (all plots). See Le
(2017) to replicate this plot.

In probability theory and statistics, the divergence between two probability
distributions is a measure of the similarity of the distributions. The most
common ones are the Kullback-Leibler (KL) divergence and the Jensen–Shannon
(JS) divergence. The KL divergence between two discrete probability
distributions $P(\mathcal{C})$ and $Q(\mathcal{C})$ is given by

$\mathcal{D}_{KL}(P\|Q)=\sum_{c\in\mathcal{C}}P(c)\log\frac{P(c)}{Q(c)}$ (1)

The KL divergence is not symmetric, that is,
$\mathcal{D}_{KL}(P\|Q)\neq\mathcal{D}_{KL}(Q\|P)$. As such, we refer to
$\mathcal{D}_{KL}(P\|Q)$ as the forward KL while
$\mathcal{D}_{RKL}(P||Q):=\mathcal{D}_{KL}(Q\|P)$ as the reverse KL between
$P$ and $Q$. Note that the forward KL under an empirical data distribution
corresponds to maximum likelihood, which we typically optimize in supervised
learning given a fixed dataset. Notably, when approximating $P(\mathcal{C})$
using a parameterized distribution $Q_{\theta}(\mathcal{C})$, minimizing the
reverse and forward KL under model under-specification results in mean and
mode-seeking behavior (Figure 3).

One issue with the KL divergence is that it is only finite when the support of
$P$ is contained in the support of $Q$. This means that if there is a point
$c$ where $Q(c)=0$ and $P(c)>0$, then
$\mathcal{D}_{KL}(P\|Q)\rightarrow\infty$. A well-known divergence that is
_bounded_ even for probability distributions with disjoint supports, albeit
not as common as the KL divergence, is the generalized Jensen-Shannon
divergence (JSD). JSD$(\beta)$ interpolates between the forward and reverse KL
using the bounded coefficient $0<\beta<1$:

$\mathcal{D}_{{JSD}[\beta]}(P\|Q)=\beta\mathcal{D}_{KL}\Big{(}P\Big{\|}\beta
P+(1-\beta)Q\Big{)}+(1-\beta)\mathcal{D}_{KL}\Big{(}Q\Big{\|}\beta
P+(1-\beta)Q\Big{)}$ (2)

Interestingly, it can be proved that $\lim_{\beta\to
0}\frac{\mathcal{D}_{{JSD}(\beta)}(P\|Q)}{\beta}=\mathcal{D}_{KL}(P\|Q)$
(Huszár, 2015). As such, JSD$(\beta)$ behaves similarly to forward KL for
small values of $\beta$. Similarly, JSD$(\beta)$ has similar behavior to
reverse KL for $\beta$ close to 1, since
$D_{{JSD}(\beta)}(P\|Q)=\mathcal{D}_{{JSD}(1-\beta)}(Q\|P)$.

##  3 Notation & Setup

We denote the input and output sequence as $x,y$ respectively. Let
$\mathbb{V}$ denote the input and output vocabulary comprising of $M$ tokens,
$y_{<n+1}=(y_{1},y_{2},\dots,y_{n})$ denote the generated output sequence up
to the $n^{th}$ token, and $L_{y}$ denote the length of sequence $y$. We
define the token-wise generative process as a deterministic Markov Decision
Process (MDP), where the initial state corresponds to the input context or
prompt $x$, the state at the $n^{th}$ token generation is the output sequence
generated thus far, $y_{<n}$, and the action space is defined over the
vocabulary $\mathbb{V}$.

A policy $p(.|y_{<n},x)\in(0,1)^{M}$ is a next-token probability distribution
over all tokens in $\mathbb{V}$, conditioned on the context $x$ and state
$y_{<n}$. Following this formulation, the policy is identical to a token-level
auto-regressive language model. Furthermore, $y\sim p(\cdot|x)$ corresponds to
a sampled output sequence $y$ given the input context $x$. For ease of
notation, we define $p(y_{n}|x):=p(y_{n}|y_{<n},x)$. Note that if we can not
compute logits but are given only samples from a policy (_e.g._ , accessing
GPT-3 (Brown et al., 2020) through its API), then $p(\cdot|y_{<n},x)$
corresponds to the empirical distribution on the sampled tokens $y_{n}$ given
$y_{<n}$ and input $x$.

Distillation Setup. We are given two auto-regressive models of different
capacity, where $p_{\text{S}}$ and $p_{\text{T}}$ refers to the student and
teacher policy respectively. We assume that the student is a neural network
with learnable parameters $\theta$ and $p_{\text{S}}^{\theta}$ is
differentiable w.r.t $\theta$. We are also given a dataset of input contexts
$X$. Optionally, we can also assume access to a dataset of input-output
sequence pairs $(X,Y)$. If not given, such a dataset can be easily generated
by sampling output sequences from the teacher, that is $\\{(x,y)\
\mathrm{s.t.}\,x\in X,y\sim p_{T}(\cdot|x)\\}$. For a divergence
$\mathcal{D}$, we define the discrepancy between token-level distributions of
$p_{T}$ and $p_{S}$ as

$\mathcal{D}\big{(}p_{\text{T}}\|p_{\text{S}}^{\theta}\big{)}(y|x):=\frac{1}{L_{y}}\sum_{n=1}^{L_{y}}\mathcal{D}\big{(}p_{\text{T}}(\cdot|y_{<n},x)\|p_{\text{S}}^{\theta}(\cdot|y_{<n},x)\big{)},$
(3)

for an input context $x$ and output sequence $y$. For example, setting
$\mathcal{D}$ in (3) to the reverse KL results in
$\mathcal{D}_{RKL}\big{(}p_{\text{T}},p_{\text{S}}^{\theta}\big{)}(y|x)=\frac{1}{L_{y}}\sum_{n}\mathcal{D}_{KL}\big{(}p_{\text{S}}^{\theta}(\cdot|y_{<n},x)\|p_{\text{T}}(\cdot|y_{<n},x)\big{)}$.

##  4 Distillation for Autoregressive Models

Supervised FT. If we are only given a fixed dataset of ground-truth output
sequences but not the teacher policy, then a simple approach is to minimize
the negative log-likelihood of such sequences under the student policy:
$L_{SFT}(\theta)=\operatorname{\mathbb{E}}_{(x,y)\sim(X,Y)}\big{[}-\log
p_{\text{S}}^{\theta}(y|x)\big{]}$.

Supervised KD (Hinton et al., 2015; Sanh et al., 2019) is a prevalent model
compression technique where a student policy is trained to imitate the token-
level probability distributions of a teacher policy. The student $p_{S}$ is
trained with the supervised objective $L_{SD}$ over the target token-level
probabilities of the teacher $p_{T}$:

$L_{SD}(\theta):=\operatorname{\mathbb{E}}_{(x,y)\sim(X,Y)}\Big{[}\mathcal{D}_{KL}\big{(}p_{\text{T}}\|p_{\text{S}}^{\theta}\big{)}(y|x)\Big{]},$
(4)

where the expectation is over the samples from the dataset. This supervised
objective results in a rich training signal by leveraging the full token-level
distribution of the teacher. Note that supervised KD reduces to supervised FT
if we only have access to the empirical distribution of teacher-generated
sequences and not the teacher policy.

Following Hinton et al. (2015), we use a softmax-temperature $\gamma$:
$p_{T}^{(v)}=\frac{\exp(z_{v}/\gamma)}{\sum_{i=1}^{N}\exp(z_{i}/\gamma)}$
where $\gamma$ controls the smoothness of the token-level teacher distribution
and $z_{v}$ is the logit score of this distribution for the vocabulary token
$v$. For simplicity, we set the temperature to 1 for the student both during
training and inference.

On-policy KD is a class of imitation learning approaches for knowledge
transfer from a teacher policy to another policy. In such approaches (_e.g._
Ross et al., 2011), we iteratively train a student policy using forward KL on
a non-stationary dataset with self-generated outputs. Given an input context
$x$, the student generates the output sequence $y$ and imitates the teacher
generated action distribution, $p_{T}(\cdot|y_{<n},x)$, on intermediate states
$y_{<n}$. Specifically, the on-policy loss $\mathcal{L}_{OD}$ on student-
generated outputs is given by

$L_{OnD}(\theta):=\operatorname{\mathbb{E}}_{x\sim
X}\Big{[}\operatorname{\mathbb{E}}_{y\sim
p_{\text{S}}(\cdot|x)}\big{[}\mathcal{D}_{KL}\big{(}p_{\text{T}}\|p_{\text{S}}^{\theta}\big{)}(y|x)\big{]}\Big{]},$
(5)

where we do _not_ backpropagate through the sampling distribution
$p_{\text{S}}(\cdot|x)$. On-policy distillation improves on supervised
distillation by training on a dataset that resembles the output distribution
the student policy is likely to generate. Despite the popularity of on-policy
distillation in imitation learning and RL (_e.g.,_ Parisotto et al., 2015;
Kelly et al., 2019; Agarwal et al., 2022), it is typically not used for
distilling auto-regressive models, with the exception of ImitKD (Lin et al.,
2020) that can be instantiated using GKD as we discuss below.

###  4.1 Generalized Knowledge Distillation (GKD)

Algorithm 1 Generalized Knowledge Distillation (GKD)

1: Given: Teacher model $p_{\text{T}}$, Student Model $p_{\text{S}}^{\theta}$,
Dataset $(X,Y)$ contain input contexts $x$ and possibly output sequences $y$.

2: Hyperparameters: Student data fraction $\lambda\in[0,1]$, Divergence
$\mathcal{D}$, Learning rate $\eta$

3: for each step $k=1,\ldots,K$ do

4: Generate a random value $u\sim Uniform(0,1)$

5: if $u\leq\lambda$ then

6: Sample batch of contexts from $X$ and for each $x$, sample one output
$y\sim p_{\text{S}}^{\theta}(\cdot|x)$ to obtain
$B=\\{(x_{b},y_{b})\\}_{b=1}^{B}$

7: else

8: Sample batch of contexts and outputs from $(X,Y)$ to obtain
$B=\\{(x_{b},y_{b})\\}_{b=1}^{B}$.

9: Update $\theta$ to minimize $L_{\mathrm{GKD}}$:
$\theta\leftarrow\theta-\eta\frac{1}{B}\sum_{(x,y)\in
B}\nabla_{\theta}\mathcal{D}(p_{\text{T}}\|p_{\text{S}}^{\theta})(y|x)$

In this work, we unify supervised and on-policy distillation approaches and
propose a more general approach, which we call Generalized Knowledge
Distillation (GKD), where we can choose the divergence to optimize as well as
the output distribution for the input training contexts. Specifically, we can
optimize any divergence on token-level teacher and student distributions on a
mixture of fixed dataset of output sequences, either teacher-generated or
ground-truth, and on-policy student-generated sequences. Abstractly, GKD
minimizes an objective of the form

$\boxed{L_{\mathrm{GKD}}(\theta):=(1-\lambda)\operatorname{\mathbb{E}}_{(x,y)\sim(X,Y)}\big{[}\mathcal{D}(p_{\text{T}}\|p_{\text{S}}^{\theta})(y|x)\big{]}+\lambda\operatorname{\mathbb{E}}_{x\sim
X}\Big{[}\operatorname{\mathbb{E}}_{y\sim
p_{\text{S}}(\cdot|x)}\big{[}\mathcal{D}(p_{\text{T}}\|p_{\text{S}}^{\theta})(y|x)\big{]}\Big{]}},$

where $D(p_{\text{T}},p_{\text{S}})(y|x)$ is a divergence between token-level
teacher and student distributions, as described in Eq (3), and
$\lambda\in[0,1]$ is a hyper-parameter that controls the student data
fraction, that is, the fraction of on-policy student-generated outputs. Akin
to on-policy distillation, we do not backpropagate gradients through the
student’s sampling process.

Remark. On-policy KD and supervised KD are special cases of GKD where we set
the divergence $\mathcal{D}$ to forward KL and student data fractions
$\lambda$ to $1$ and $0$ respectively. Furthermore, ImitKD (Lin et al., 2020)
can be viewed as GKD with forward KL as divergence and a non-increasing
schedule on $\lambda$, a simple choice being $\lambda=0.5$. That said, GKD
allows for other choices for the on-policy fraction $\lambda$ and the
divergence, which we explore in this work.

_Choice of Divergence_. When using forward KL for distillation, the student
model tries to cover the entire support of the token-level teacher
distribution $p_{\text{T}}(.|y_{<n},x)$, and in doing so might end up
assigning probability mass to tokens $y_{n}$ which have low probability under
$p_{\text{T}}(.|y_{<n},x)$. Since the student has typically much lower model
capacity than that of the student, this is likely to happen (_e.g._ , Figure
3). We posit that a reasonable choice for $\mathcal{D}$ in GKD might be the
reverse KL or JSD$(0.9)$, which has mode-seeking behavior.

_Choosing student data fraction_. It is well-known that using only a fixed
dataset of ground-truth outputs ($\lambda=0$) for training autoregressive
models, which is a common practice, can lead to train-test distribution
mismatch. This is because the distribution of partial sequences seen during
the auto-regressive generation phase can be very different from the ones
encountered during the model’s training. One simple way to mitigate this issue
is to sampling output sequences from the student during training itself, akin
to on-policy KD. Furthermore, if we only have access to an unlabeled dataset
$X$ of input contexts, generating sequences from the student rather than the
teacher is much less expensive due to the difference in their model sizes.
Particularly, when only using student-generated sequences ($\lambda=1$), GKD
generalizes on-policy KD to arbitrary divergences:

$L_{\mathrm{OnGKD}}(\theta):=\operatorname{\mathbb{E}}_{x\sim
X}\Big{[}\operatorname{\mathbb{E}}_{y\sim
p_{\text{S}}(\cdot|x)}\big{[}\mathcal{D}(p_{\text{T}}\|p_{\text{S}}^{\theta})(y|x)\big{]}\Big{]}.$

###  4.2 RL + On-policy GKD

In some tasks, it is plausible that the teacher only provides a proxy to our
main objective. With an RL approach, we can however directly optimize this
objective. Conveniently, a useful property of on-policy GKD is that it can
easily be combined with online RL fine-tuning as it only requires output
samples from the student. Indeed, consider that one wants to optimize the
student policy for a scalar reward signal $r$, while staying close to a
teacher policy, then we get a regularized RL fine-tuning objective of the
form:

$\operatorname{\mathbb{E}}_{x\sim X}\Big{[}(1-\alpha)\underbrace{E_{y\sim
p_{\text{S}}^{\theta}(\cdot|x)}\left[r(y)\right]}_{\text{RL
objective}}-\alpha\underbrace{\operatorname{\mathbb{E}}_{y\sim
p_{\text{S}}(\cdot|x)}\big{[}\mathcal{D}(p_{\text{T}}\|p_{\text{S}}^{\theta})(y|x)\big{]}}_{\text{Generalized
On-Policy Distillation}}\Big{]},$ (6)

where $\alpha\in[0,1]$ controls the strength of the distillation loss compared
to the RL objective. With $\alpha=1$, it will perform only distillation.
Following Roit et al. (2023), we use a REINFORCE-like policy gradient
algorithm for optimizing the RL loss.

Figure 4: Data efficiency and scaling on XSum. Mixed and On-policy GKD
variants with JSD(0.1) perform quite well with a small number of training
examples and have better scaling behavior than KD methods that use forward KL.
The T5-XL teacher obtains a ROUGE-2 score of 22. The results for 100% XSum
training dataset are shown in Figure 1. On-policy and Mixed GKD surpassses
supervised KD and ImitKD with 100% dataset using only 0.5% of the original
dataset.

##  5 Experiments

In this section, we evaluate GKD for distilling LLMs, a widely-studied class
of auto-regressive models, on a variety of language tasks, namely abstractive
summarization, machine translation and arithmetic reasoning.

Student / Teacher Models. In all experiments, we are given two models: a
student and a teacher, pretrained on the same datasets, but with different
model sizes. We focus on the pretrained T5 v1.1 models (Raffel et al., 2020)
of different sizes. For all experiments, we use supervised fine-tuned T5-XL
($\sim 3$B params) as the teacher. For student models, we use T5-small (77M
params), T5-base (250M params), and T5-large (800M params), which are smaller
than the teacher by a factor of $38\times$, $12\times$ and $3.8\times$
respectively.

GKD Variants. For choice of divergence $\mathcal{D}$ in GKD, we use Forward
KL, Reverse KL and two variants of JSD$(\beta)$, JSD$(0.1)$ and JSD$(0.9)$.
For student data fraction $\lambda$, we try $\lambda=0$ (Supervised),
$\lambda=1$ (On-policy) and $\lambda=0.5$ (Mixed). In particular, we are
interested in the mixed and on-policy variants with reverse KL and JSD, which
have not been previously explored. We compare to the following baselines:

  * •

Supervised FT: See Section 4 It directly fine-tunes the student on either
ground-truth or teacher-generated outputs.

  * •

Supervised KD: See Eq (4), GKD with $\lambda=0$ and forward KL as divergence
$\mathcal{D}$.

  * •

On-policy KD: See Eq (5), GKD with $\lambda=1$ and forward KL as divergence.
Note that purely on-policy distillation hasn’t been employed for auto-
regressive models.

  * •

ImitKD (Lin et al., 2020): GKD with $\lambda=0.5$ and forward KL as
divergence.

###  5.1 Case Study: Abstractive Summarization

Training Setup. We start by evaluating GKD on an abstractive summarization
task of generating a concise summary that captures salient ideas of the input
document. To do so, we use the XSum dataset (Narayan et al., 2018), which
consists of news articles paired with human-written summaries. For evaluating
summarization performance, we report ROUGE scores (Lin, 2004) of predicted
summaries on the validation dataset split of XSum. Following PaLM (Chowdhery
et al., 2022), we mainly use ROUGE-2 but observe similar trends in ROUGE-LSum
and ROUGE-1. We use T5 models fine-tuned on XSum as the base models for
distillation while the T5-XL model as the teacher. We report ROUGE-2
performance after a fixed number of training steps with batch size of 32.

Figure 5: Trade-off between reward maximization and distillation performance
(ROUGE-2) on XSum. We report performance improvements relative to the original
student model (T5-base). Following Roit et al. (2023), we use the binary
entailment score from a NLI classifier as the reward while $\alpha$ controls
the strength of the on-policy GKD loss with JSD (0.1). As $\alpha$ increases,
ROUGE-2 increases while entailment score improvement decreases. For
comparison, we show the relative performance of the $12\times$ larger teacher
model. RLEF* corresponds to RLHF-like method from Roit et al. (2023), where
the student is regularized towards the original student model itself instead
of the teacher. As expected, on-policy GKD + RL achieves higher ROUGE-2
compared to the RLEF* and obtain much better entailment reward compared to the
teacher.

Data efficiency and scaling. To evaluate the efficiency and scaling
characteristics of various methods, we perform knowledge distillation by using
the teacher only with subsampled XSum training datasets containing the first
1K $(0.5\%)$, 10K $(5\%)$ and 50K$(\sim 25\%)$ examples. For this experiment,
we use T5-small as the student and report the results in Figure 4. Notably,
on-policy GKD with JSD$(0.9)$ on the 0.5% dataset, without any ground-truth
summaries, outperforms supervised KD and ImitKD with 100% of the training
dataset with ground-truth summaries (Figure 1). Furthermore, as we scale the
dataset size from 0.5% to 25%, on-policy KD results in a $6\times$ larger
improvement than supervised KD and ImitKD.

Varying student model capacity. Next, we explore how GKD compares to prior
distillation approaches across different student model sizes. As shown in
Figure 1, we observe consistent improvements with GKD, which demonstrates the
generalizability of GKD. Furthermore, we find that GKD variants with reverse
KL or JSD$(0.1)$ substantially outperform approaches that use the forward KL.
Interestingly, GKD allows us to nearly match the few-shot performance of PaLM
(540B) using a $7000\times$ smaller model (T5-small).

Figure 6: Ablating GKD variants on XSum. We distill T5-XL to T5-Small (77M
params), which obtain a ROUGE-2 score of 7.96 and 22 respectively. For the
distilled student, we report improvement in ROUGE-2. On-policy GKD with JSD
(0.9) performs the best.

Figure 7: Effect of teacher temperature. We plot the evaluation performance on
XSum for the distilled student, which uses a temperature of 1, for several
values of the teacher temperature. Interestingly, contrary to supervised KD,
we find that temperature values that are less than 1 perform quite well.

On-policy GKD with RL. For summarization, we’d like model-generated summaries
to be factually consistent with their input documents. However, distillation
alone might not be for improving factual consistency of the student model as
even large teachers generate factually inconsistent summaries. Recently, Roit
et al. (2023) mitigate this issue by using RL with textual entailment feedback
(RLEF) as faithful summaries must be textually entailed from their
corresponding input documents. Inspired by their success, we explore combining
RL optimization with on-policy GKD, as shown in (6). As shown in Figure 5,
this approach can substantially improve factual consistency over the teacher
model while obtaining large improvements in ROUGE-2 score for the distilled
student model.

Ablating GKD variants. We now evaluate the impact of different KL-based
divergences and student data fractions in GKD for distilling T5-XL to a
T5-small student. As shown in Figure 7, forward KL performs poorly compared to
other divergences while on-policy variants typically perform at least as well
or better than the mixed and supervised variants.

Varying Teacher Temperature. Furthermore, contrary to conventional wisdom in
supervised learning, Figure 7 demonstrates that setting teacher’s softmax-
temperature lower than 1 can sometimes lead to large performance improvements.

Figure 8: Varying on-policy student data fraction and divergence in GKD on WMT
English $\rightarrow$ German. We report the BLEU score improvement of
distilled student relative to the original student model. We observe that
using student-generated samples as well as mode-seeking divergences, that is
reverse KL and JSD (0.9), outperforms using only teacher-generated data or
using forward KL. We use the T5-XL ($\sim$3B params) supervised fine-tuned on
WMT as the teacher, which obtains a BLEU score of 28. (Left) Student
corresponds to T5-base (250M params) with a BLEU score of 27.30. GKD with
reverse KL and 50-50 split between student-generated and ground-truth
translations performs the best. (Right) We use T5-small (77M params) as the
student, which obtain a BLEU score of 25.52. On-policy GKD with reverse KL or
JSD (0.9) perform the best.

###  5.2 Machine Translation

To evaluate GKD beyond summarization, we consider WMT14 en-de (Bojar et al.,
2014), a machine translation task that requires rewriting text in English into
German, preserving the content, semantics and style of the input. We report
performance on the original test split using the BLEU score, which measures
the similarity of machine-translated text to a set of high quality reference
translations. We use fine-tuned T5-Base and T5-small models as students while
T5-XL as the teacher with a softmax-temperature of $0.1$.

Results. As shown in Figure 8, we observe that mode-seeking KL-based
divergences typically outperform their mean-seeking counterparts. Furthermore,
using on-policy and mixed data distributions consistently outperforms GKD
variants only using a fixed supervised dataset, showing the importance of
generating on-policy output sequences from the student. These findings are
consistent with what we observed on XSum. However, reverse KL seem to perform
better than JSD $(0.9)$, while for XSum their ranking seems to be reversed.

Figure 9: Distillation results on GSM8K with 4-shot chain-of-thought (CoT)
prompting. On-policy GKD variants outperform other distillation approaches. As
a reference, we provide UL2 results with CoT + calculator as well as LaMDA
results with CoT (w/o calc). The teacher FLAN T5-XL achieves a test accuracy
of 27.1. SFT corresponds to the approach of Magister et al. (2022).

###  5.3 Arithmetic Reasoning

Wei et al. (2022) show that reasoning abilities only appear to emerge in LLMs
with at least several billions parameters, making knowledge distillation an
important area of research for improving reasoning abilities of smaller
models. To this end, we evaluate GKD on natural language math problems
requiring multi-step logical inference. In particular, we use GSM8K (Cobbe et
al., 2021), a high quality linguistically diverse dataset of 8.5K grade school
math word problems. Furthermore, we explore GKD in conjunction with chain-of-
thought (CoT) prompting (Wei et al., 2022), a common approach to improve
reasoning abilities of LLMs, by simply augmenting the input contexts with the
given prompt.

Training Setup. We perform few-shot prompting by augmenting the math problems
in GSM8K with the first 4 CoT exemplars from Wei et al. (2022), which we use
as input contexts. For evaluation, we report accuracy on the GSM8K test split
by checking whether the target answer matches the final answer given an
external calculator, akin to Cobbe et al. (2021). For supervised training, we
use the CoT outputs generated by Magister et al. (2022) with Palm 540B and
filtered using the target answer, resulting in around 5.3K (problem, CoTs)
pairs in the original training split of GSM8K. Magister et al. (2022) show
that supervised fine-tuning (SFT) on this dataset substantially improves test
accuracy on GSM8K. As such, we use Flan-T5 models (Chung et al., 2022) fine-
tuned for 10K steps on the above CoT dataset as a starting point for
distillation. We use the supervised fine-tuned FLAN T5-XL model as the
teacher,

Figure 10: Ablating GKD variants on GSM8K.

which obtains a test accuracy of 27.1.

Results. As shown in Figure 9, we observe that on-policy GKD methods perform
quite well compared to other approaches. Notably, in contrast to our XSum
results, we notice that GKD with forward KL (on-policy KD) performs quite
well. However, GKD with reverse KL outperform the on-policy KD for the FLAN
T5-base and T5-large student models. Also, we observe that mixing student-
generated outputs and supervised CoT dataset typically underperforms using
only the on-policy student-generated outputs (Figure 10).

##  6 Related work

Knowledge distillation. Supervised KD (Buciluǎ et al., 2006; Hinton et al.,
2015) is now a classic approach for model compression and has been
successfully used for distilling auto-regressive models (Sanh et al., 2019).
Another natural approach for distilling auto-regressive models is sequence-
level KD (Kim & Rush, 2016), which maximizes the log-likelihood of the student
on sequences generated by the teacher. Other approaches rather train the
student to match different quantities obtained from the teacher, such as
hidden states (Jiao et al., 2020) or attention scores (Wang et al., 2020).
However, none of these approaches make the connection between sequence-level
distillation and imitation learning, and a purely supervised learning approach
can lead to distribution drifts, that is exposure bias (Arora et al., 2022).

The first approach to make this connection is ImitKD from Lin et al. (2020),
inspired by Dagger (Ross et al., 2011). It samples trajectories from both the
student and the teacher to reduce the exposure bias, but does not push the
idea further and keep the forward KL at the token level. This is not necessary
when one has access to the teacher’s log-probabilities, rather than just
samples. In essence, ImitKD is a special case of the proposed GKD, which we
show to generally lead to worse empirical results in our experiments.

More recently, the concurrent work on MiniLLM (Gu et al., 2023) also exploits
the link to imitation while also recognizing the limitation of forward KL.
They frame distillation as an RL problem with the reward being the reverse KL
between the teacher and the student, this being equivalent to the reverse KL
at the sequence level (while likelihood maximiziation is the forward one), and
solve it using a policy gradient approach. This is close to what we propose,
but we argue that GKD is simpler and more stable, being closer to supervised
training, since it does not backpropagate through the student’s sampling
process. Indeed, MiniLLM relies on a number of stabilizing tricks, to tackle
high variance, reward hacking, and generation length bias. GKD is also more
general, in the sense that we also consider the generalized Jensen-Shannon
divergence, which sometimes provides better results.

RL finetuning. There are now numerous examples of large language models being
fine-tuned with RL, be the reward optimizing for some metric (Wu et al.,
2018), or learned using human feedback (Ziegler et al., 2019). In these
approaches, it is typical to regularize the RL finetuned model towards the
initial (usually supervised fine-tuned) model. However, as far as we know, we
are the first to perform distillation and RL finetuning at the same time. If
it may seem natural, it is quite different from an optimization perspective,
as it changes the regularization towards the initial policy to towards the
teacher policy, and we show empirically that it is a viable approach.

Knowledge Distillation with Reasoning Traces. CoT prompting (Wei et al., 2022)
has recently demonstrated that LLMs can solve complex reasoning tasks, step by
step, just by prompting. This idea was quickly adapted to knowledge
distillation, by extending the teacher dataset with CoT prompts for fine-
tuning the student (Magister et al., 2022; Ho et al., 2022; Hsieh et al.,
2023). The distillation is still done in a supervised way, and other kind of
enhanced prompts could be considered (Li et al., 2022; Mukherjee et al.,
2023). We adopt the same approach, but combine it with on-policy distillation
with various divergences, for the first time to the best of our knowledge. It
shows the versatility of GKD, and improves upon the purely supervised
approach, as seen in our results on GSM8K (Figure 9).

##  7 Conclusion

In this work, we proposed generalized knowledge distillation (GKD) for auto-
regressive models that addresses the challenges of distribution mismatch and
model under-specification. We investigated performance of GKD on three natural
language generation tasks: abstractive summarization, machine translation, and
arithmetic reasoning. Our method consistently outperformed more commonly-used
knowledge distillation baselines. We further showed that the approach can be
combined with reinforcement learning to optimize a sequence-level reward in
addition to distilling the knowledge of a large teacher model. We believe that
our method will be a valuable resource for researchers and practitioners who
are working on improving performance of small auto-regressive models.

## Acknowledgments

We are thankful to Johan Ferret, Léonard Hussenot and Ramki Gummadi for
providing feedback on an early draft of this work. We are especially thankful
to Nikola Momchev, and Sertan Girgin, for key contributions to our research
infrastructure and general support for infrastructure-related queries. We’d
also like to acknowledge Aleksandra Faust, Olivier Pietquin, Dale Schuurmans,
Aaron Courville, Léonard Hussenot, Robert Dadashi, Adrien Ali Taiga, and Max
Schwarzer for helpful discussions.

## References

  * Agarwal et al. (2022) Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare.  Reincarnating reinforcement learning: Reusing prior computation to accelerate progress.  _Advances in Neural Information Processing Systems_ , 35:28955–28971, 2022.
  * Arora et al. (2022) Kushal Arora, Layla El Asri, Hareesh Bahuleyan, and Jackie Chi Kit Cheung.  Why exposure bias matters: An imitation learning perspective of error accumulation in language generation.  _arXiv preprint arXiv:2204.01171_ , 2022.
  * Bojar et al. (2014) Ondřej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, et al.  Findings of the 2014 workshop on statistical machine translation.  In _Proceedings of the ninth workshop on statistical machine translation_ , pp. 12–58, 2014.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Buciluǎ et al. (2006) Cristian Buciluǎ, Rich Caruana, and Alexandru Niculescu-Mizil.  Model compression.  In _Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining_ , pp. 535–541, 2006.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_ , 2022.
  * Chung et al. (2022) Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al.  Scaling instruction-finetuned language models.  _arXiv preprint arXiv:2210.11416_ , 2022.
  * Cobbe et al. (2021) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al.  Training verifiers to solve math word problems.  _arXiv preprint arXiv:2110.14168_ , 2021.
  * Gu et al. (2023) Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang.  Knowledge distillation of large language models.  _arXiv preprint arXiv:2306.08543_ , 2023.
  * Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.02531_ , 2015.
  * Ho et al. (2022) Namgyu Ho, Laura Schmid, and Se-Young Yun.  Large language models are reasoning teachers.  _arXiv preprint arXiv:2212.10071_ , 2022.
  * Hsieh et al. (2023) Cheng-Yu Hsieh, Chun-Liang Li, Chih-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alexander Ratner, Ranjay Krishna, Chen-Yu Lee, and Tomas Pfister.  Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes.  _arXiv preprint arXiv:2305.02301_ , 2023.
  * Huszár (2015) Ferenc Huszár.  How (not) to train your generative model: Scheduled sampling, likelihood, adversary?  _arXiv preprint arXiv:1511.05101_ , 2015.
  * Jiao et al. (2020) Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu.  Tinybert: Distilling bert for natural language understanding.  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , pp. 4163–4174, 2020.
  * Kelly et al. (2019) Michael Kelly, Chelsea Sidrane, Katherine Driggs-Campbell, and Mykel J Kochenderfer.  Hg-dagger: Interactive imitation learning with human experts.  In _2019 International Conference on Robotics and Automation (ICRA)_ , pp. 8077–8083. IEEE, 2019.
  * Kim & Rush (2016) Yoon Kim and Alexander M Rush.  Sequence-level knowledge distillation.  _arXiv preprint arXiv:1606.07947_ , 2016.
  * Le (2017) Tuan Anh Le.  Reverse vs forward kl, December 2017.  URL https://www.tuananhle.co.uk/notes/reverse-forward-kl.html.
  * Li et al. (2022) Shiyang Li, Jianshu Chen, Yelong Shen, Zhiyu Chen, Xinlu Zhang, Zekun Li, Hong Wang, Jing Qian, Baolin Peng, Yi Mao, et al.  Explanations from large language models make small reasoners better.  _arXiv preprint arXiv:2210.06726_ , 2022.
  * Lin et al. (2020) Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei.  Autoregressive knowledge distillation through imitation learning.  _arXiv preprint arXiv:2009.07253_ , 2020.
  * Lin (2004) Chin-Yew Lin.  Rouge: A package for automatic evaluation of summaries.  In _Text summarization branches out_ , pp. 74–81, 2004.
  * Magister et al. (2022) Lucie Charlotte Magister, Jonathan Mallinson, Jakub Adamek, Eric Malmi, and Aliaksei Severyn.  Teaching small language models to reason.  _arXiv preprint arXiv:2212.08410_ , 2022.
  * Mukherjee et al. (2023) Subhabrata Mukherjee, Arindam Mitra, Ganesh Jawahar, Sahaj Agarwal, Hamid Palangi, and Ahmed Awadallah.  Orca: Progressive learning from complex explanation traces of gpt-4.  _arXiv preprint arXiv:2306.02707_ , 2023.
  * Narayan et al. (2018) Shashi Narayan, Shay B Cohen, and Mirella Lapata.  Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization.  _arXiv preprint arXiv:1808.08745_ , 2018.
  * Parisotto et al. (2015) Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov.  Actor-mimic: Deep multitask and transfer reinforcement learning.  _arXiv preprint arXiv:1511.06342_ , 2015.
  * Pomerleau (1991) Dean A Pomerleau.  Efficient training of artificial neural networks for autonomous navigation.  _Neural computation_ , 3(1):88–97, 1991.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551, 2020.
  * Roit et al. (2023) Paul Roit, Johan Ferret, Lior Shani, Roee Aharoni, Geoffrey Cideron, Robert Dadashi, Matthieu Geist, Sertan Girgin, Léonard Hussenot, Orgad Keller, et al.  Factually consistent summarization via reinforcement learning with textual entailment feedback.  _arXiv preprint arXiv:2306.00186_ , 2023.
  * Ross & Bagnell (2010) Stéphane Ross and Drew Bagnell.  Efficient reductions for imitation learning.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pp. 661–668. JMLR Workshop and Conference Proceedings, 2010.
  * Ross et al. (2011) Stéphane Ross, Geoffrey Gordon, and Drew Bagnell.  A reduction of imitation learning and structured prediction to no-regret online learning.  In _Proceedings of the fourteenth international conference on artificial intelligence and statistics_ , pp. 627–635. JMLR Workshop and Conference Proceedings, 2011.
  * Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  _arXiv preprint arXiv:1910.01108_ , 2019.
  * Wang et al. (2020) Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou.  Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers.  _Advances in Neural Information Processing Systems_ , 33:5776–5788, 2020.
  * Wei et al. (2022) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou.  Chain of thought prompting elicits reasoning in large language models.  _arXiv preprint arXiv:2201.11903_ , 2022.
  * Wu et al. (2018) Lijun Wu, Fei Tian, Tao Qin, Jianhuang Lai, and Tie-Yan Liu.  A study of reinforcement learning for neural machine translation.  _arXiv preprint arXiv:1808.08866_ , 2018.
  * Ziegler et al. (2019) Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving.  Fine-tuning language models from human preferences.  _arXiv preprint arXiv:1909.08593_ , 2019.