On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
Authors: Rishabh Agarwal
Nino Vieillard
Yongchao Zhou
Piotr Stanczyk
Sabela Ramos Garea
Matthieu Geist
Olivier Bachem
Venue: ICLR 2024 poster (2024)
Keywords: Language models
Distillation
RLHF
Primary area: generative models
TL;DR: Better distillation for autoregressive student models using on-policy student-generated data, which can be easily combined with RLHF.

Abstract: Knowledge distillation (KD) is widely used for compressing a teacher model to reduce its inference cost and memory footprint, by training a smaller student model. However, current KD methods for auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference. To address this issue, we introduce Generalized Knowledge Distillation (GKD). 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. Unlike supervised KD approaches, GKD also offers the flexibility to employ alternative loss functions between the student and teacher, which can be useful when the student lacks the expressivity to mimic the teacher's distribution. Furthermore, GKD facilitates the seamless integration of distillation with RL fine-tuning (RLHF). We demonstrate the efficacy of GKD for distilling auto-regressive T5 language models on summarization, translation, and arithmetic reasoning tasks.

Published as a conference paper at ICLR 2024
ON-POLICY DISTILLATION OF LANGUAGE MODELS: LEARNING FROM SELF-GENERATEI MISTAKES
Rishabh Agarwal
Nino Vieillard
Yongchao Zhou" 13
Piotr Stanczyk
Sabela Ramos
Matthieu Geist!
Olivier Bachem
'Google Deep.Mind
Mila
University of Toronto
ABSTRACT
Knowledge distillation (KD) is widely used for compressing teacher model reduce its inference cost and memory footprint, by training smaller student model. However, current KD methods for auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference. To address this issue, introduce Generalized Knowledge Distillation (GKD). Instead of solely relying on a fixed set of output sequences, GKI) trains the student its self-generated output sequences by leveraging feedback from the teacher on such sequences. Unlike supervised KD approaches, GKD also offers the flexibility to employ alternative loss functions between the student and teacher, which may be useful when the student lacks the expressivity mimic the teacher's distribution. Furthermore, GKD facilitates the seamless integration of distillation with RL fine-tuning of language models. We demonstrate the efficacy of GKD for distilling auto-regressive T5 language models for task-specific distillation on summarization, translation, and reasoning tasks, and agnostic distillation for instruction tuning.
INTRODUCTION
Supervised Fine-Tuning
Superv
SeqKD
GKD (On-policy)
XSum (Summarization)
WMT (Translation)
GSMBK (Reasoning w/ CoT)
Performance (ROUGE-2)
22
+111%
T5-XL (Teacher)
Performance (BEU)
+70%
Accuracy (w/ Calculator)
30
490%
25
20
27.5
18
27.0
16
10
26.0
24
25.5
T5-5mall (77M)
T5-Base (250M)
T5-Large (800M)
T5-Small
T5-Base
T5-Large (800M)
T5-Small
T5-8ase (250M)
TS-Large (800M)
(77M)
(250M)
(77M)
Student (params)
Student (params)
Student (params)
Comparing GKD with KD approaches across different student model sizes. We use the T5 models (Raffel et al., 2020) trained with supervised FT as students. We use supervised FT T5-XI (>3B params) as the teacher, whose performance is indicated by the horizontal line. Supervised KD and FT use ground-truth output sequences for training while SeqKD trains on output sequences generated by the teacher. On-policy GKD trains on output sequences sampled from the student. For GKD, we use JSD (0.1) WMT and forward KL on other tasks. For evaluation,
Figure 1:
greedy sampling for XSum and GSM8K and beam search for WMT.
Auto-regressive sequence models, such 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 (Kaplan et al., 2020). However, scaling parameter count comes at
"denotes equal contribution. denotes infrastructure contribution. Correspondence to Rishabh Agarwal <ri shabhagarwal @egggggaaaa@@@@gglllg @goog] e com> and Nino Vieillard <viei 1 ardd coms.

---

Published as a conference paper at ICLR 2024
cost, and the deployment of such models is limited by either their inference cost memory footprint. Thus, a crucial goal for practical use of large capable models is to compress them by reducing their parameter count, while retaining much as possible of their performance.
One of the prevalent techniques for compressing models is knowledge distillation (Hinton et al., 2015). Distillation is the process of training model the student 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 sequence models either require generating a fixed set of output sequences from the teacher model (Kim & Rush, 2016), which can be expensive, or a fixed dataset of sequences that the teacher can label by assigning token-level probabilities (Sanh al., 2019). However, using fixed dataset can lead to distribution mismatch between output sequences seen during training and the sequences generated by the student auto-regressively during inference, a well-known problem in imitation learning (Pomerleau, 1991; Ross & Bagnell, 2010). Furthermore, the common objective for distillation is to minimize the forward KL between the teacher and the student distributions. However, the student may not be expressive enough to fit the teacher's distribution, which can result in suudent-generatee samples that are unlikely to be generated by the teacher (e.g., Figure A.1
In this paper, we propose Generalized KD (GKD) to mitigate the above issues. First, we recognize that KD for auto-regressive sequence models can be viewed as an imitation learning problem with an interactive expert (Ross al., 2011). Using this insight, GKD trains the student its self- self-generated sequences that are on-policy, instead of a fixed set of outputs sequences, using teacher probabilities as expert labels on these sequences. Our idea is further supported by the recent success of fine-tuning large language models on their own output sequences (Ouyang et al., 2022; Singh et al., 2023). Furthermore, GKD provides the flexibility to optimize alternative divergence measures, such as reverse KL and generalized JSD (Section 2), that can use student's limited capacity to focus on generating samples that are likely under the teacher.
GKD unifies some existing KD methods for autoregressive LMs while instantiating new on-policy methods that substantially outperform prevalent approaches. In terms performance gains over the initial student from on-policy GKD, averaged across T5 student models of different sizes, we rela- tive gains of 2. on summarization, 1. on machine translation, and 1.9 x on arithmetic reasoning tasks, compared to the performance improvements achieved with baseline KD methods (Figure 1). Additionally, we exhibit GKD's efficacy in aask-agnottic distillation, resulting in 2% and 1% absolute accuracy improvement on the held-out BBH and MMLJU benchmark suites (Figure A.11).
Our key contributions are:
To tackle discrepancy during training and inference for auto-regressive LMs, we present GKD that leverages on-policy suudent-generate outputs for distillation, guided by the token-level teacher probabilities over these outputs. GKD substantially outperforms commonly-used methods in task-specific (Figure and task-agnostic KD (Figure 1I).).
We demonstrate that on-policy GKID can be seamlessly combined with RI. fine-tuning (e.g.. RLAIF) of language models, a combination that has not been previously explored (Figure 5).
Through a systematic evaluation of design choices in GKD, we offer practical insights about the importance of using muudent-generann on-policy output sequences during distillation and the task-dependent nature of optimal divergence between the student and the teacher.
2 PRELIMINARIES
Auto-regressive Generative Sequence Models. We denote the input and output sequence x, y respectively. Let V denote the vocabulary comprising of M tokens, Ven+1 = (y1, V2...., yn) denote the generated output sequence up the nh token, and Ly denote the length of sequence y. A token-level auto-regressive policy p(-lVen, 2) (0, 1 outputs next-token probability distribution p(-|2)
over all tokens in V. conditioned on the input x and output sequence yea- Furthermore, corresponds to sampled output sequence y given the input æ. For ease of notation, define pz), := p(yn yen, z). Auto-regressive generation involves predicting tokens one at time, based on the previously generated tokens. The probability of predicting th token yn. p(ya (z). determined exp(=./7)
where 2n is the logit score for the using a softmax with temperature 7: p(yn(x) EM exp(=1177)
2

---

Published as a conference paper at ICLR 2024
token Vn. Higher values of introduces randomness, while a lower value makes the output more deterministic by favoring the most probable words. During training, the student's 's temperature is kept at 1. For evaluation, use greedy sampling (7 0) or temperature sampling (y > 0).
KL-Based Divergences.
The divergence between two probability distributions is a measure of the similarity of the distributions, with KL. divergence a prevalent measure. The KL divergence between two discrete distributions P(C) and Q(C) given by: (xx(()))Q) E.ec P(c) log P(c) Q(e)'
The KL divergence is not symmetric: D + DKr, (QI As such, we refer to DxL. Pl)) as the forward KL while (Q]PP) as the reverse KL between and Q. Forward KL under an empirical data distribution corresponds to maximum likelihood, which we optimize in supervised learning. Given model capacity mismatch, when approximating P(C) using a distribution Qo(C). minimizing the reverse and forward KI. results in mean and mode-seeking behavior (Figure 1.
While KL. divergence can be unbounded, a well-known divergence that bounded even for probability distributions with disjoint supports is the generalized JSD Leesen--hannoo divergence). JSD(3)
interpolates between the forward and reverse KL using the bounded coefficient < 1:
(1- 3)Q) D (asp (PIl) = 3DRL (
+ (1- (
B)Q)
(1)
Huszár (2015) show that lim3-00 (JSD(D) (PIIQ)/3 Dx (Pll)). As such, gradients of JSD(3)
behave similarly to forward KL and reverse KL when /3 is close and respectively.
DISTILLATION FOR AUTO-REGRESSIVE SEQUENCE MODELS
Problem Setup. We are given two auto-regressive sequence models of different capacity, where Ps and Pr refers to the student and teacher respectively. We assume that the student has learnable parameters and pš is differentiable w.r.t t. We are also given a dataset of inputs X. Optionally, can also assume access dataset of input-output sequence pairs (X, Y). .. If not given, such dataset can be generated by sampling sequences from the teacher. For a divergence D, we define the discrepancy between token-level distributions of pr and ps as
(vlrl 1
D(pr(tlVen x) D( (l,n.))
Ly
for an input x and output sequence y. For example, using JSD(B) as D in equation 2 results in D DJSD(A) (Pr (I3) (w]=) ł, Σ, DUSD(3) (pr(-1g<n, x) (- #)).
Supervised FT. we are only given fixed dataset of ground-trutt truth output sequences but not query access the teacher policy, then a simple approach is to minimize the negative liglliiel of such sequences under the student policy: LSFT(0) (0) = E(x.)))(X,)))
Sequence-Leve KD (Kim & Rush, 2016). SeqKD) maximizes the likelihood of high probability sequences generated by the teacher, and can be viewed as supervised FT on teacher-gent seacher-generated outputs.
Supervised KD (Hinton et al., 2015; Sanh et al., 2019) is widely used technique where the student trained to imitate the token-level probability distributions of the teacher. The student ps is trained with the supervised objective Lsn over the target token-level probabilities of the teacher pr:
(f) = (xx.)) DxL (prlle3) (vl=)).
(3)
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.
3.1 GENERALIZED KNOWLEDGE DISTILLATION (GKD)
As discussed above, commonly-used KD approaches use fixed dataset of output sequences, either using ground-truth targets or seccher-geeertedd sequences. However, distilling auto-regressive student models using such approaches results in train-ine inin-infeeenee distribution mismatch. This is because the partial sequences encountered by the student during the regreseeres generation phase at inference can be quite different from the ones seen during the training phase. Since predictions at any step are contingent upon previous steps in auto-regressive models, this mismatch can have cascading effect
3

---

Published as a conference paper at ICLR 2024
Algorithm I Generalized Knowledge Distillation (GKD)
1: Given: Teacher model pr. Student Model ys: Dataset (X, Y) containing (input, output) pairs
2: Hyperparameters: Student data fraction À E [o. 11, Divergence 'D, Learning rate
3: for each step k = I,..., K do
4: Generate random value Uniform(0,1)
5: if u I À then
6: Sample imputs from X and generate outputs y (1()) obtain B ((rb. yo)
7: else
8: Sample batch of inputs and outputs from (X, Y)) to obtain B {(1b: wo))
9: end if
Update 9 to minimize LGKD: 0 + 1-1 1B E E(EWEB VoD(prllp)(vla)
end for
where error in prediction at early step can affect the future predictions, resulting in poor quality text generation. To address this mismatch, we draw heavily from imitation learning (IL). In particular, on-policy imitation approaches (e.g. Ross et al., 2011) iteratively collect sequences using the student policy, obtain expert labels for those sequences, and then retrain the student on this dataset. Despite their popularity in robotics and deep RL (Parisotto et al., 2015; Kelly et al., 2019; Agarwal et al., 2022), on-policy approaches are typically not used for distilling auto-regressive models.
Extending on-policy imitation to distillation, we present on-policy KD. When using on-policy data during distillation, the student receives token-specific feedback from the teacher's logits on the erroneous tokens in its self-generated output sequences. This enables a form of feedback loop akin to what we observe in RL, which helps minimize the train-inference distribution mismatch. Moreover, as the student evolves during training, the data it generates also improves in quality. Given an input I, the student generates the output sequence and imitates the teacher token-level distributions, pr (ynle (x), on intermediate states Ven- Specifically, the on-policy loss Lon is given by
(o) = Ex>x (gy)s(((>>) [PKz (prllP)) (x1=))]
(4)
where we do not backpropagate through the student's sampling distribution ps(-(a), similar to on-policy imitation. Not backpropagating through the sampling makes the training stable and computationally efficient. In on-policy KD, the training is done output sequences that the student likely to generate. During training, we use a temperature of y = to encourage diversity in student generated sequences. Moreover, given unlabeled input prompts, generating sequences using the student is computationally cheaper than the teacher, due to differences in their model sizes.
Building further upon on-policy KD, we unify supervised and on-policy approaches and propose more general approach, which we call Generalized KD (GKD). In GKD, we can choose both the divergence to optimize as well as the output sequences to train on. Specifically, we can optimize any divergence between the teacher and student token-level probability distributions. For output sequences, GKD uses a mixture of fixed dataset, either teacher-gent geneerr-eeeeraee or ground-truth, and on-policy student-generate sequences. Abstractly, GKD minimizes an objective of the form:
LGKD = (1 - AN + XE Ey--pt (m)) (P(prlleS
where D(pr, Ps) Ps) (ylx) is a divergence between teacher and student distributions (equation 2), and that controls the student data fraction, that is, the fraction of on-policy student-generated outputs. Akin to on-policy KD, we do not backpropagate gradients through the student's sampling process. On-policy and supervised KD are instantiations of GKD with divergence D set forward KI. and student data fractions À to 1 and 0 respectively. That said, GKD allows for other choices for the fraction À and the divergence, which explore in this work.
A E J0, 1] is a
As opposed to a randomly initialized student, we assume access to a student that can generate sequences of adequate quality, which the teacher can provide feedback upon. In our experiments, we start from student models that have undergone supervised FT. This is analogous to two-stage RLHF training, which is widely used for LMs, where we first run SFT followed by the online RL fine-tuning. As such, GKD can leverage hyperparameter tuming insights from RLHF and can be combined with RLHF with small compute overhead and additional hyperparameters.
Remark.
4

---

Published as a conference paper at ICLR 2024
Choice of Divergence in GKD. While forward KL is commonly-used for distillation, it requires the student to cover the entire support of the teacher token-level distribution Pr(-/een, Ø). In doing so, the student might end up assigning probability mass to tokens which have low probability under (e,,) I), which can result in hallucination and low-quality generations. When the student has much lower model capacity than the teacher, this issue is likely to happen with temperature sampling (e.g., Figure A.) Alternatively, 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 given input. Our experiments indicate that optimal divergence seems to be task-dependent. Overall, the diversity and performance trade-offs for particular task needs to be considered when choosing the GKD divergence (e.g., Figure 4, A.11).
3.2 RL FINE-TUNING + ON-POLICY GKD
In some tasks, it is plausible that distilling from a teacher model only provides a proxy to our main objective, which can also be non con-differentiabe. We can directly optimize this objective with reinforcement learning (RL). Conveniently, on on-policy GKD can be easily combined with RL fine- tuning from human (RLHF) Al feedback (RLAIF), it only requires output samples from the student. Indeed, consider that one wants to optimize the student policy for a scalar reward r, while staying close a teacher policy, then we get regularized RL fine-tuning objective of the form:
Exax (1 (1-a) a)) P1 [r(v)] -x Eyrrps (-1w) [pr (In) (v1=))
(5)
RL
Generalized On Policy Distillation
where a E [0, 1] controls the strength of the distillation loss compared to the RL objective. With 1, it will perform only distillation. The above objective allows us to maximize reward while improving other model capabilities via distillation, which can possibly reduce the "alignment tax" decrease in general model capabilities when aligning language models with human preferences (Ouyang al., 2022). We apply the above idea mitigate hallucination using RLAIF, while simultaneously improving downstream performance via distillation (Figure 5).
Remark. In RLHF or RLAIF, we typically use reverse KL constrain the learned policy to stay close to the initial policy. 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.
EXPERIMENTS
In this section, we evaluate GKD for distilling language models, a typical class of rego-segressive sequence models, abstractive summarization, machine translation, and arithmetic reasoning.
Student Teacher Models. Our experiments start from student and teacher models with different sizes, specifically open-sourced T5 models (Raffel et al., 2020), that are pretrained on the same datasets. We use supervised fine-tuned T5-XL (~> 3B params) as the teacher. For students, we use T5-small (77M params), T5-base (250M params), and T5-large (800M params), which are smaller than the teacher by factor of 38x, 12x and 3.8x respectively. See Appendix A.3 for more details.
GKD Variants. For choice of divergence D in GKD Algorithm 1, and three variants of (3): JSD (0.1), JSD (0.5) JSD (0.9). For student data fraction X, we try X = 1 (On-policy), À 0.5 (Mixed) and X 0 (Supervised). In particular, we are interested in the on-policy variants (A = 1), which have not been previously explored.
forward KL, reverse KL.
Baselines. We compare to the widely-used KD methods discussed in Section SeqKD and Supervised KD. We also evaluate ImitKD (Lin et al., 2020) and f-distill (Wen et al., 2023), which can be viewed as "mixed" data variants of GKD (A 0.5) with forward KL and total variation distance as divergence. All the baselines start from the same supervised fine-tuned student checkpoint as GKD.
4.1 CASE STUDY: ABSTRACTIVE SUMMARIZATION
We start by evaluating GKD on an abstractive summarization task of generating a summary that captures salient ideas of the input document. To do so, we the XSum dataset (Narayan et al., 2018), which consists of news articles paired with human- written summaries. Following PaLM (Chowdhery al., 2022), we evaluate performance using ROUGE-2 score (Lin, 2004) of predicted summaries
5

---

Published as a conference paper at ICLR 2024
Supervised KD
Fdistill
GKD (SD(0099)
XSum TS-XL - T5-Small
XSum: TS-XL
Large
15
Supervised KD
22
Temperature Sampling
22
Greedy Sampling
ROUGE-2
14
GKD (Supervised)
ImitKD
21
13
GKD (On-policy)
ROUGE-2
20
20
12
11
18
18
0.5%
25% Training Dataset Fraction
100%
Figure Comparing GKD to baselines on distillation from TS-XL to T5-large on XSum. On-policy GKD
Figure 3: Scaling training data. We evaluate distilled TS-small using temperature sampling (7 = 1). GKD more data efficient than baselines.
variants generally outperform baselines.
Performance / Diversity Trade-off on XSum: Distilled T5-small
Performance (ROUGE-2)
16.5
14.5
GKD Divergence
12.5
Temperature (v)
Forward KL
JSD (0.9) Reverse KL
0.1
0.5
JSD (0.1)
10.5
0.3
1.0
JSD (0.5)
20
40
60
80
Generation Similarity (Self-BLEU)
Figure 4: Effect of Divergence on Performance and Diversity. Utilizing on-policy GKD with different divergences, evaluate the trade-off between the distilled student's i generation quality and diversity, by varying the sampling temperature. We quantify diversity using Self-BLEU (Zhu al., 2018), where indicates deterministic outputs and signifies maximum diversity. Transitioning from forward KL reverse through generalized JSD, leads to decreased diversity, attributed to the enhanced mode-seeking characteristic of the divergence. Mode-seeking divergences often yield superior quality, especially at high temperatures (? 1). Reducing the temperature curtails diversity while narrowing performance differences among divergences.
of 100
on the validation split XSum but observe similar trends in ROUGE-L and ROUGE-1. We use T5 models supervised fine-tuned on XSum students for distillation while the fine-tuned T5-XL the teacher. See Appendix A.4 for additional experimental details.
Comparison to baselines. First, explore how GKD compares widely-used KD approaches, namely SDqKD and Supervised KD, across different student model sizes. As shown in Figure 1, we observe consistent improvements with GKD, which demonstrates the scalability of GKD with respect to the student capacity. Notably, GKD allows us to surpass the few-shot performance of PaLM (540B) using 7000 smaller T5 model. We also compare GKD variants with ImitKD and f-distill, and evaluate performance with greedy sampling and temperature sampling (7 1) in Figure 2. On-policy GKD with JSD (0.9) outperforms these additional baselines both scenarios.
Data efficiency and scaling. To evaluate the efficiency and scalability of GKD, we distilled the T5-XL teacher using subsampled XSum training datasets: 1K (0.5%). 10K (5%), and 50K (25%) examples. We used T5-s T5-small as the student and report data scaling curves in Figure 3. Notably, on-policy GKD on the 5% subsampled dataset, without any ground-truth summaries, outperforms supervised KD and ImitKD with entire training dataset with ground-truth summaries.
GKD Ablations. We ablated different divergences and student data fractions in GKD for various student sizes in Figure A.12 and A.13. On-policy and mixed variants consistently outperform super- vised variants. Mode-secking divergences perform better when evaluation done using temperature sampling while the choice of divergence doesn't 't affect performance much with greedy sampling.
Choosing GKD Divergence. The divergence chosen for distillation is crucial in determining the trade-off between summarization quality and diversity. As the sampling temperature also be adjusted to balance summary quality and diversity, the optimal choice of divergence temperature- dependent. To understand this dependence, evaluate T5-small distilled using on-policy GKD)

---

Published as a conference paper at ICLR 2024
RL Fine-Tuning + Distillation
Improvement in Factual Consistency
RLEF*
45.0%
(A Entailment
a
0.05
35.0%
0.1
25.0%
a
0.25 0.5
Teacher (T5-XL)
15.0%
3
4
7
Improvement in Summary Quality (AROUGE-2)
Figure 5: RLAIF On-policy GKD. We show the trade-off between reward maximization and summarization performance on XSum. Wc report improvements relative to the original T5-base student. Following Roit et al. (2023), we use the textual entailment score from a T5-XXL NLI classifier as the reward. controls the strength of the on-policy GKD loss with JSD (0.9). As increases, ROUGE-2 increases while improvement in factual consistency decreases.
comparison. we show relative performance RLEF* corresponds to RLAIF method from Roit et al. (2023), where the student is regularized towards the original student model itself instead of the teacher. On-policy GKD RL achieves higher ROUGE-2 compared RLEF" while generating more factually consistent summaries compared to the teacher.
12x larger T5-XL teacher.
Student Data Fraction (A)
BLEU Improvement: TS-XL
Small
BLEU Improvement: T5-XL
Base
6.7
0.8
100%
0.55
0.85
0.77
0.79
0.46
0.7
100%
0.55
0.71
0.58
0.58
0.47
0.6
50%
0.5
0.71
0.57
0.43
0.38
0.5
50%
0.52
0.48
0.54
0.4
0.3
02
0%
0.08
0.28
0.26
0.29
0.27
0.22
0.38
0.17
0.11
0.02
..1
Forward J&b (0. JSD (0.5)
JSD (0.9)
Reverse
Forward jáb (0 (0.1)
JSD (0 10D3DD (0.9)
Reverse
Divergence (D)
Divergence (D)
Figure 6: Varying student data fraction and divergence in GKD WMT + de. For evaluation, heam search and report the improvement in BLEU score of distilled student relative the original student. Results are averaged across three seeds. We observe that using only student-generater output samples outperform other GKD variants. We use the TS-XL (~3BB params) supervised fine-tuned on WMT as the teacher, which obtains a BLEU score of 28. (Left)) We use T5-small (77M params) the student, which obtain BLEU score of 25.58. (Right) Student corresponds to T5-basc (250M params) with BLEU score of 26.98.
with different divergences. As shown in Figure 4, certain divergences, like JSD (0.5) and JSD (0.9), offer better quality but less diversity at high temperatures. However, as temperature decreases, the difference in quality among divergences narrows, while diversity also drops.
On-policy GKD with RL. In summarization, we want moder-generated summaries to be factually consistent with their input documents. However, distillation alone might not improve factual consis- tency as even large models halluncinate and generate inconsistent summaries. Recently, Roit et al. (2023) mitigate hallucination on summarization tasks by using RL with textual entailment feedback as the reward (RLEF), as faithful summaries must be textually entailed from their input documents. Inspired by their success, explore combining RL fine-tuning using a REINFORCE-HE with on- on-policy GKD, as described in Section 3.2. As shown in Figure 5, GKD with RL fine-tuning substantially improves factual consistency compared to the teacher model while obtaining large improvements in summarization quality for the distilled student model.
objective
4.2 MACHINE TRANSLATION
To evaluate GKD beyond summarization. we consider the task on translating English to German using WMT14 en-de (Bojar et al., 2014). We report performance on the validation split using the BLEU score, which measures the similarity of machine-translale We use supervised fine-tume mee-tueed T5-XL as the teacher with a softmax- cettert-ee 1.0 (BLEU score of 28). See Appendix A.6 for additional experimental details.
text to high quality reference translations.
7

---

Published as a conference paper at ICLR 2024
Student Data Fraction (A)
Accuracy Improvement T5 XL - Base
GSM8K: Flan TS-XL Small
100%
7.6
7
6.7
Test Accuracy
11
Forward KL
Reverse KL
50%
6.2
6.4
6.9
7.0
9
665
-60
8
0%
4.7
5.1
5.1
5.7
55
Forward Kbb (0.
ISD (0.5)
JSD (0.9)
Reverse KL
0%
25%
50%
75%
100%
Student Data Proportion (A)
Divergence (D)
Figure 7: Ablating GKD on GSM8K. We distill fine- tuned T5-XL T5-Base, which obtain accuracy of 27.9 and 10.16 with greedy sampling.
Figure 8: Varying on-policy data GSM8K. As increase fraction of suuent-generated data beyond 25%, performance typically improves.
TS-XL Small
Flan
30
Flan T5-XL
GSMBK Test Accuracy
0-shot dovincl 002* (3758)
Few PaLM (5408)
28
On-policy
ImitKD
16
(Supervised)
(4-shot- D-shot-CoTTio
f-distill
8
Supervised KD
14
24
SegKD
12
Figure 9: Distillation on GSM8K with few-shot CoT prompting. On-policy GKD substantially outperform other approaches. As reference, provide GPT-3 davinci-002 results well PaLM (540B) results (without a calculator). We use forward KL and reverse KL respectively for on-policy and supervised GKD.
Results. Figure 1, A.15 show that on-policy GKD outperforms commonly commonly-used KD approaches. Furthermore, we ablate GKD variants using T5-small and T5-base students in Figure 6. We observe that generalized JSD divergences perform better than forward or reverse KI. but their performance gap reduces when using larger student. Moreover, purely on on-policy and mixed data distributions consistently outperform GKD variants only using fixed supervised dataset, showing the importance of generating sequences from the student. This finding aligns with our results on XSum.
4.3 ARITHMETIC REASONING
Wei et al. (2022) show that reasoning abilities only appear to emerge in LLMs with at least several billions parameters, making KD important for improving reasoning abilities of smaller models. To this end, we evaluate GKD GSM8K (Cobbe et al., 2021), high-quality dataset of grade school math word problems requiring multi-step logical inference. Here, we explore GKD in conjunction with chain-of-the tf-toutt (CoT) (Wei al., 2022), common approach improve reasoning abilities of LLMs by prompting them to produce intermediate reasoning steps before giving the final answer.
Setup. We perform few-shot prompting by prepending the math problems in GSM8K with the first CoT input-output exemplars from Wei et al. (2022). For evaluation, we report accuracy on the 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 al. (2022), resulting around 5.3K (problem, CoTs) pairs the original training split of GSM8K. We use Flan-T5 models (Chung et al., 2022) supervised fine-tuned for 10K steps on the above CoT dataset starting point for distillation. We use the fine-tuned FLAN T5-XL as the teacher, which obtains test accuracy of 27.9. See additional experimental in Appendix A.5.
Results. We first ablate GKD variants and report results Figure and A.14. We observe that when using only the fixed CoT dataset or mixing it with student- -eneertteeeerrted Co'Ts, performance consistently falls short of using solely the student-genent similar to findings XSum with greedy sampling. Notably, reverse KL also performs well, especially when training using only a fixed dataset. Additionally, Figure 8 shows that performance consistently improves the proportion of on on-policy data increases, provided that at least 25% of the data is on-policy. Moreover, we demonstrate that on on-policy GKD have superior performance compared to baseline KD methods, across all student sizes, shown in Figure 9.
CoTs Furthermore, forward KL performs quite well,
8

---

Published as a conference paper at ICLR 2024
5 RELATED WORK
Knowledge distillation. Supervised KD (Buciluă et al., 2006; Hinton et al., 2015) is a classic approach and has been successfully used for distilling outo-regressive models (Sanh et al., 2019). Another approach for distilling such models is sequence-level KD (Kim & Rush, 2016). On-policy GKD substantially outperforms supervised KD and SeqKD (Figure 1). Other KD approaches train the student 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 distillation and imitation learning, and purely supervised approach can suffer from train- inference mismatch, also known as exposure bias (Ranzato et al., 2015; Bengio et al., 2015). While He et al. (2019) argue that this mismatch may not be critical, several papers demonstrate that exposure bias leads to poor text generation (Zhang et al., 2019; Chiang & Chen, 2021; Arora et al., 2022).
ImitKD (Lin et al., 2020) identifies this connection by sampling sequences from both the student and a fixed dataset but does not push the idea further. Unlike GKD, ImitKD does not explore purely on-policy data collection, nor does it integrate RL fine-tuning. Moreover, ImitKD keeps the forward KL at the token level, which is not necessary when one has access to the teacher's s log- peg-robbbbilities, rather than just samples. Furthermore, GKD demonstrates the scalability of the idea, handling student models roughly 26x larger than those explored by ImitKD. ImitKD) can be viewed as GKD with forward KL and a non-increasing schedule on X, a simple choice being = 0.5. More recently, f-distill (Wen et al., 2023) formulates sequence-level KD as minimizing an f-divergence and propose an tractable objective based on total variation distance between the token-level student and teacher distributions. In essence, both ImitKD and f-distill are specific instances of GKD, which demonstrate lead to worse empirical results than on-policy GKD (Figure 2, 9).
The concurrent work MiniLLM (Gu et al., 2023) also exploits the link to imitation and frame distillation as an RL problem. In particular, Minil.M optimizes reverse KI. between the teacher and the student at the sequence level (while likelihood maximization is the forward one) using policy gradient approach. However, 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 also more general as it can also be used with other divergences such as forward KL. or JSD, which can perform better than reverse KL (Figure 6, 7).
RL fine-tuning. There are now numerous examples of language models being fine-tuned with RI., be the reward optimizing for some metric (Wu et al., 2018) or learned using human feedback (Ouyang et al., 2022). In these approaches, it is typical to regularize the RL fine-tuned model towards the initial (usually supervised fine-tuned) model. However, as far we know, we are the first to perform distillation and RL fine-tuning at the same time (Figure 5). If it may seem natural, it is quite different from an optimization perspective, as it changes the regularization towards the initial policy towards the teacher policy, and we show empirically that it is a viable approach.
Distillation with reasoning traces rationales. Chain-of -houught prompting (Nye et al., 2021; 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 with CoT prompts for fine-tuning the student (Magister et al., 2022; Ho et al., 2022; Hsich 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- on-policy distillation with various divergences. It shows the versatility of GKD, and improves upon the purely supervised approaches, as seen in our results on GSM8K (Figure 9).
KD, by extending the teacher dataset
CONCLUSION
In this work, we proposed GKD to address the train-inference distribution mismatch and model underspecification when distilling language models. GKD consistently outperformed commonly distillation approaches on several generative tasks. We further showed that GKD can be combined with reinforcement learning to optimize a sequence-level reward in addition to distilling the know wedge of large teacher model, which we believe can improve the widely-used RLHF training phase for language models. One interesting direction would be extending GKD to auto regressseee sequence models for audio (Radford et al., 2023), video (Villegas et al., 2022) and text-to-imagg to-image generation (Yu et al., 2022). We hope that our work will be valuable for researchers and practitioners who are working on improving performance and efficiency of generative auto-regressive sequence models.

---

Published as a conference paper at ICLR 2024
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Reincarnating reinforcement learning: Reusing prior computation accelerate progress. Advances in Neural Information Processing Systems, 35:28955-28971, 2022.
Kushal Arora, Layla Asri, Hareesh Bahuleyan, and Jackie Chi Kit Cheung. Why exposure bias matters: An imitation learning perspective of error accumulation in language generation. arXiv preprint ar arXiv:2204.01 171, 2022.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
Ondiej 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.
Cristian Buciluă, Rich Caruana, and Alexandru Niculescu-Mizzi. Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 535-541, 2006.
Ting-Rut -ui Chiang and Yun-Nung Chen. Relating neural text degeneration to exposure bias. arXiv preprint arXiv:2 087058 2021.
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.023111 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:22 1416, 2022.
pntructioo--inetuneed
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Hcewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, al. Training verifiers to solve math word problems. arXiv preprint 1rXiv:2110.114 14168, 2021.
Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. Knowledge distillation of large language models. arXiv preprint arXiv:2306.0855 08543, 2023.
Tianxing He, Jingzhao Zhang, Zhiming Zhou, and James Glass. Exposure bias versus self- recovery: Are distortions really incremental for autoregressive text generation? arXiv preprint arXiv:1905.1000 10617, 2019.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.025: 2015.
Namgyu Ho, Laura Schmid, and Se- Young Yun. Large language models are reasoning teachers. arXiv preprint arXiv:2212.1007 2022.
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.
Ferene Huszár. How (not) to train your generative model: Scheduled sampling, likelihood, adversary? arXiv preprint arXiv: 111.v:151..0510,, 2015.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tenybert: Distilling bert for natural language understanding. Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4163-4174, 2020.
10

---

Published as a conference paper at ICLR 2024
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint wrXhv:2011.0831,, 2020.
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.
Yoon Kim and Alexander M Rush. Sequence-level knowledge distillation. arXiv preprint arXiv:166..0780 079477 2016.
Tuan Anh Le. Reverse vs forward kl, December 2017. URL html.
uk/notes https//veeeeeeereeeeeeeeeeeeeeee
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 2rXiv:2210.06777 2022.
Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. Autoregressive knowledge distillation through imitation learning. arXiv preprint arXiv:2009. 07253, 2020.
Chin- Yew Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 14, 2004.
Lucie Charlotte Magister, Jonathan Mallinson, Jakub Adamek, Eric Malmi, and Aliaksei Severyn. Teaching small language models to reason. arXiv preprint arXiv:2212.08410, 2022.
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.02777, 2023.
Shashi Narayan, Shay Cohen, and Mirella Lapata. Don't give me the details, just the sum- mary! topic-aware convolutional neural networks for extreme summarization. arXiv preprint 1rXiv:1808.08745 2018.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski. Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114 2021.
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744,202
Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. arXiv preprint arXiv: 15100632,, 2015.
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1):88-9) 97, 1991.
Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever. Robust speech recognition via large-scale weak supervision. In International Conference on Machine Learning, pp. 28492-28518 PMLR, 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-551) 2020.
Mare' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. arXiv preprint 151iv: 2015.
11

---

Published as a conference paper at ICLR 2024
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 27X1v:2306.00116. 2023.
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.
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 artificial intelligence and statistics, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01 2019.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp.
PMLR, 2018.
Avi Singh, John D Co- Re-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil. Peter Liu, James Harrison, Jaehoon Lee, Kelvin Xu, Aaron Parisi, et al. Beyond human data: Scaling self-training for problem-solving with language models. arXiv preprint wrxiv:23/2.06588 2023.
Ruben Villegas, Mohammad Babaeizadeh, Pieter-Jan Kindermans, Hemnnn Moraldo, Han Zhang, Mohammad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan. Phenaki: Variable length video generation from open domain textual description. arXiv preprint drXiv:2210.02399, 2022.
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self- attention distillation for task-agmostic compression of pre-trained transformers. Advances in Neural Information Processing Systems, 33:5776-5788, 2020.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022.
Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. f-divergence minimization for sequence-level knowledge distillation. arXiv preprint arXiv:2307. 15190, 2023.
Lijun Wu, Fei Tian, Tao Qin, Jianhuang Lai, and Yin Liu. A study of reinforcement learning for neural machine translation. arXiv preprint arXiv: 10w.::8688.0..6,, 2018.
Junho Yim, Donggyu Joo, Jihoon Bae, and Junmo Kim. A gift from knowledge distillation: Fast optimization, network minimization and transfer learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4133-4141, 2017
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Bureu Karagol Ayan, al. Scaling autoregressive models for content- rich text-to-image generation. Transactions on Machine Learning Research, 2022.
Wen Zhang, Yang Feng, Fandong Meng, Di You, and Qun Liu. Bridging the gap between training and inference for neural machine translation. In Anna Korhonen, David Traum, and Lluís Märquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019. URL https: : ac haa/aaaaannnoooooooooooooooooooa org/ 101//199-1222
Yaoming Zhu, Sidi Lu, Lei Zheng, Jiaxian Guo, Weinan Zhang, Jun Wang, and Yong Yu. Texygen: A benchmarking platform for text generation models. In The 41st international ACM SIGIR conference on research & development in information retrieval, pp. 1097-1100, 2018.
12

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 SHLF-DISTILLATION
Self-Distil (GSM8K): FLAN TS-Large
Accuracy (w/ Calculator)
24
Teacher: FLAN La-Large (STT)
20
16
Student: FLAN La-Large
Supervised KD
GKD (FKL)
ImitkD
f-distill
GKD (SD(0.9))
Figure A010: Sclf-Distilation on GSM8K. Here, GKD corresponds to on-policy GKD (A = 1). On-policy GKD variants outperform other approaches including supervised KD. The teacher FLAN T5-Large supervised fine-tuned GSM8K and achieves accuracy of 20.5 %% while the student FLAN T5-large (not trained on GSM8K) obtains an accuracy 14.4% on the test set.
Here, we investigate whether GKD works for self-distillation (Yim et al., 2017), where we want to transfer knowledge from a teacher model student model with the same architecture and size. To investigate this, we consider self-distillation on GSM8K with FLAN-TS large as the student and teacher, where the teacher supervised fine-tuned GSM8K. As shown in Figure A.10, self-distilled students surpass surpasses the teacher's performance on the test set. Moreover, distillation using student-generated data outperforms supervised KD, with on-policy GKD performing the best.
A.2 TASK-AGNOSTI DISTILLATION: INSTRUCTION TUNING
While task task-specific distillation provides optimized performance for predefined tasks, which is often crucial for deployment purposes, task-agnostic distillation offers a compelling alternative in scenarios where the exact nature of the task is not known beforehand and can vary during deployment. As highlighted by Sanh et al. (2019), the allure of task-agnostic distillation lies in its efficiency: once distilled, a model can be re-purposed for multiple downstream tasks via prompting or fine-tuning.
Setup. To study task-agnostic KD, we focus on instruction tuning (Chung et al., 2022). Our aim is to enhance the distilled model's proficiency to handle diverse tasks presented in the form of instructions. To achieve this, employ the FLAN T5-XL model
teacher and distill its knowledge into the FLAN T5-Base, as introduced by Chung et al. (2022). Our distillation process utilizes the comprehensive FLAN2021 instruction tuning dataset, which boasts 5.36 million examples spanning 62 distinct language understanding and generation tasks. For hyperparameter details, see Table A.4.
Evaluation. To gauge the versatility of task-agnostic model, it is essential test it across a diverse set of tasks. In line with Chung et al. (2022), we evaluate our distilled TS-base student on two held-out benchmark suites: (1) MMLU (Massive Multitask Language Understanding) includes exam questions from 57 tasks such as mathematics, history, law, and medicine, and (2) BBH (BIG-Bench Hard) includes 23 tasks from BIG-Bench for which PaLM 540B (Chowdhery et al., 2022) performs below average human raters. For performance, we report the distilled model's 's ability to directly predict the answer via standard few-shot prompting, averaged across tasks in MMLU and BBH.
We report the performance of distilled checkpoints obtained after 50K training steps for various methods in Figure 1. We find that on-policy GKD with reverse KL substantially outperforms supervised KD and ImitKD. Notably, in the context of instruction tuning, we find that using reverse KL performs much better than forward KL. We hypothesize that the efficacy of reverse KL in instruction tuning may stem from its mode-seeking nature it ensures that the model zeroes in on the main intent or behavior specified by the instruction. As a result, the model might prioritize core behaviors over less relevant details, leading better performance held-out tasks.
Results.
13

---

Published as a conference paper at ICLR 2024
GKD (On Forward KL)
Supervised
GKD (On- palnppllyyy Reverse
ImitKD
Supervised GKD (Reverse KL)
MMLU (57 tasks): FLAN TS-XL Base
BBH (23 tasks): FLAN T5 -XL
Base
Improvement In Avg. MMLU Accuracy (%)
2.0
Avg. BBH Accuracy uurayy (%)
1.0
1.5
Improvement in
0.8
1.0
0.6
0.5
0.0
0.2
0.5
0.0
Distillation Method
Distillation Method
Figure A.11: aask-agnostic Distillation on FLAN (Chung et al., 2022). On-policy GKD with reverse KL outperforms other approaches. The evaluation metric both the MMLU and BBH benchmark suites is all tasks. These few-shot prompted accuracy (exact match), where take unweighted average run SeqKD due to evaluation benchmarks are held-out (not included the distillation data). Here, its computational inefficiency for generating data from the teacher during training. The teacher FLAN TS-XL achieves an accuracy of 52.4% on MMLU and 41% on BBH, while the student T5-large model obtains an accuracy of 35.6% MMLU and 31.25% on BBH.
A.3 T5 MODELS
As base checkpoints, start from LM-adapted T5v1.1 models. These LM-adapted models initialized from T5v1. 1 and trained for an additional 100K steps on the LM objective discussed in the T5 paper (Raffel al., 2020). These checkpoints open-soureed https: //t////onnn console. cloud. google..ooooooooooooooooooo... com/ woom/stoore weooaagebbowwsee 5 data/ppetrained_.eo...
In our experiments, initialize both the student and teacher models for distillation by running further supervised FT on the original training dataset, as described below:
XSum. For small, base, large and XL models, use LM-Adapted T5v1.1 models super- vised fine-tuned for 100K, 50K, 38k and 8K steps respectively.
WMT. For small, base, large and XL models, we use LM-Adapted T5v1.1 models super- vised fine-tuned for 250K, 250K, 110k and 50K steps respectively.
GSM8K. All models were supervised fine-tuned starting from FLAN-T5 models on the Palm-540 generated CoT dataset for 10K steps.
Similar to T5 and FLAN-T5, experiments the Adafactor optimizer (Shazeer & Stern, 2018).
Computational cost GKD. All methods including basclines start from the supervised fine-tuned student checkpoint, which requires training for few hours on the smallest TPUv3 (8 cores). On GSMSK, the computational overhead from student sampling is approximately 1.8 x, 2x and 2.2x compared sampling from fixed dataset of outputs, for student-teacher ratio of 38x, 12x and 3.8x. For RLHF + GKD, the computational overhead is somewhat small as we are only running inference get teacher logits instead of student logits.
Moreover, the majority of cost in real world use cases is due to serving cost at inference time and due to fine tuning. Concretely, if might also be too expensive to serve this model to users (which may range from tens of thousands billionn). Overall, performance benefits from on-policy GKD might be worth the compute cost, especially when combined with RLHF.
costly to sample from the student during fine-tuning,
A.4 XSUM
Learning rate sweep. We performed sweep over the learning rate in {0.0001, 0.0003. 0.001 and found 0.0003 to be the best for T5-base, T5-large and 0.001 leads to best performance for T5-small. As such, by default
an LR of 0.0003 except when reporting results for T5-small, which uses 0.001. We found that reverse KL was more sensitive to higher LRs and we default to using 0.0003 for all models when using reverse KI.
14

---

Published as a conference paper at ICLR 2024
Teacher Softmax-temperature When using greedy sampling for evaluation, set the teacher temperature to 1. However, when reporting student performance with temperature sampling (? = 1), as done in Figures and 3, set teacher temperature to 0.1 for the student.
Table A.I: Hyperparameter Details for experiments XSum.
Hyperparameter
Value
Training Steps
40,000
Batch size
Eval Split
Validation
Dropout
0.0
Learning Rate (LR)
0.0003
LR Warmup Steps
2,000
LR Cooldown (Begin, End)
(30,000, 40,000)
Warmup Schedule
Linear (from 0 to LR)
Cooldown Schedule
Linear (from LR 0)
Max Input Length
1024
Max Output Length
Evaluation
Greedy & Temp. Sampling
GKD Ablations. We ablated different divergences and student data fractions in GKD for various student sizes in Figure A.12 and A.13. On-policy and mixed variants consistently outperform super- vised variants. Mode-seeking divergences perform better when evaluation done using temperature sampling while the choice of divergence doesn't 't affect performance much with greedy sampling.
Student Data Fraction (A)
ROUGE-2 2: TS-XL
TS sm-small
Student Fraction (A)
ROUGE-22 2: TS-XL TS-base
33.3
100%
13.2
14.5
100%
17.5
10.0
17.8
50%
13.1
15
15.5
15.2
14.5
50%
16.4
17.5
17.8
17.9
17.9
17.4
172
0%
13.1
14.3
14.9
15
14.4
0%
16.7
17.1
17.3
17.4
17.6
16.8
Forward
bb
1)
JSD (0.5)
JSD (0.9)
Reverse KL
Forward
J&b
(0. 15D
JSD (0.9)
Reverse KL
Divergence (D)
Divergence (D)
ROUGE-2: T5-XL T5-large
Student Data Fraction
100%
19.2
19.9
20.1
20.3
19.8
20.2
50%
19.3
19.9
20
19.8
19.3
15.1
0%
18.6
19.1
19.3
19.3
19.1
19.2
18.
Forward
(0.
DD (0.5)
1SD
Reverse KL
Divergence (D)
Figure A.12: Ablating GKD XSum with evaluation via temperature sampling (7 1). We distill supervised fine-tuned T5-XL model to different sized student models. Here, we evaluate using temperature sampling and teacher temperature to 0.1 during training. In the plots above, we report the ROUGE-2 score of the student post distillation. On-policy GKD approaches with reverse KL and JSD (0.9) performs the best, while forward KI. performs poorly.
A.5 GSM8K
For training, the CoT outputs generated from Palm-540B by Magister et al. (2022). We report accuracy on the original test split of the GSMSK dataset (Cobbe et al., 2021). We use checkpoints at the end of training after distillation for reporting results, which are averaged across 3 seeds.
15

---

Published as a conference paper at ICLR 2024
Student Data Fraction (A)
ROUGE-2: T5-XL TS-small
ROUGE-2:
TS-XL T5-base
100%
16.4
15.6
Student Data Fraction
100%
18.7
18.7
16.6
244
18.6
:4.
:3.8
50%
18.3
18.4
23.4
18.
0%
15.2
15.1
0%
Forward Jáb (0. JSD (0. JSD (0.9)
Reverse
Forward JSD
(0.
Divergence Reverse KL
[SD (0.
JSD (0.9)
Divergence (D)
Student Data Fraction (A)
ROUGE-2: T5-XL + T5-large
21.2
100%.
21.2
21.2
21.1
21.2
21.1
203
50%
20.9
21
21
20.0
20.6
0%
20.6
20.5
20.3
20.5
20.2
20.4
Forward Jbb
(0. 10.11
JSD
(0.5)
JSD (0.9)
Reverse KL
Divergence (D)
Figure A.13: Ablating GKD on XSum with evaluation via greedy sampling. We distill supervised fine-tuned T5-XI. model to different sized student T5 models. Here, evaluate using greedy sampling and set student and teacher temperature
during training. When evaluated using greedy sampling, the teacher model obtains ROUGE-2 score of 22 while respectively. In the plots above, we report the ROUGE-2 score of the student post distillation. Ou-policy GKD approaches performs the best, with small differences among different divergences. Furthermore, on-policy and mixed variants strongly outperform supervised variants.
student TS-small, base and large models obtain score 13.4, 17.9. 19.6
Table A.2: Hyperparameter Details for experiments on GSM8K.
Hyperparameter
Value
Training Steps
40,000
Batch size
32
Evaluation Split
Test
Dropout
0.05
Learning Rate (LR)
0.0003
LR Warmup Steps
2,000
LR Cooldown (Begin, End)
(30,000, 40,000)
Warmup Schedule
Linear (from 0 to LR)
Cooldown Schedule
Linear (from LR to 0)
Max Input Length
512
Max Output Length
320
Checkpoint
Flan-T5
Teacher softmax temperature
0.1
Evaluation
Greedy Sampling
Few-shot CoT Prompt. Here, specify the 4-shot CoT prompt used for the experiments:
Q: There are 15 trees in the grove. Grove workers will plant trees in the grove today. After they are done, there will be 21 trees. How many trees did the grove workers plant today?
A: There are 15 trees originally. Then there were 21 trees after some more were planted. So there must have been 21 15 6. The answer is 6.
Q: If there are 3 cars in the parking lot and more cars arrive, how many cars are in the parking lot?
A: There originally cars. 2 more arrive. 3 =. The answer is 5.
16

---

Published as conference paper at ICLR 2024
Student Data Fraction (A)
Accuracy Improvement: TS-XL Small
Student Data Fraction (
Accuracy Improvement: TS-XL- Large
100%
6.6
4.7
3.8
4.9
5.5
100%
3.5
50%
4.2
3.8
3.6
3.2
4.5
50%
4.0
0%
2.5
3.1
2.9
2.9
3.8
3.5
0%
3.9
3.9
3.7
4.5
Forward J&b (0.1)
JSD (0.5)
JSD (0.9)
Reverse KL
Forward 5b (0.1)
JSD (0.5)
JSD (0.9)
Reverse KL
Divergence (D)
Divergence (D)
Figure A. 14: Ablating GKD variants on GSM8K with 4-shot CoT. For evaluation, we use greedy sampling and report improvement in test accuracy of the student after distillation. Results are averaged across three seeds. Using only student-generated output samples typically outperform other GKD variants. We the supervised fine-tuned T5-XI which obtains an accuracy of 4.585. (Right) Student corresponds to T5-base base with an accuracy of 20.5.
teacher, which obtains an accuracy 27.9. (Left) We use T5-small as the student,
Q: Leah had chocolates and her sister had 42, If they ate 35, how many pieces do they have left in total?
A: Originally, Leah had 32 chocolates. Her sister had 42. So in total they had 32 42 74, After eating 35, they had 74 35 = 39. The answer is 39.
Q: Jason had 20 lollipops. He gave Denny some lollipops. Now Jason has 12 lollipops. How many lollipops did Jason give to Denny?
A: Jason started with 20 lollipops. Then had 12 after giving some Denny. So he gave Denny 20 12 8. The answer is 8.
A.6 WMT
For evaluation, we use beam search with same hyperparameters as Raffel et al. (2020). We report performance of the final checkpoints obtained after training. To reduce the variance in results, report the results averaged across 3 seeds.
WMT on-de: T5-XL + Small
WMT enide: T5-XL Base
0.8
BLEU imprevement
0.6
0.4
0.4
BLEU
0.2
ImikKD
GKD
ImikD
GKD
Figure A.15: Comparing GKD to ImitKD and f-distill WMT. Here, GKD corresponds to best performing variant WMT, with À (on-policy) and JSD (0.1). On-policy GKD leads to 53% higher BLEU improvement over ImitKD and 162% over f-distill, averaged across small and base models.
17

---

Published as a conference paper at ICLR 2024
Table A.3: Hyperparameter Details for WMT en-de experiments.
Hyperparameter
Value
Training Steps
100,000
Batch size
32
Seqio Task Name
"wmt 12t ende v003
Eval Split
Validation
Dropout
0.0
Warmup Steps
5,000
Warmup Schedule
Linear (from 0 LR)
Learning Rate (LR)
0.0003
Input Length (Tokenized)
80
Output Length (Tokenized)
80
Teacher softmax temperature
1.0
Evaluation
Beam Search
A.7 INSTRUCTION TUNING
Table A.4: Hyperparameter Details for FLAN Instruction Tuning.
Hyperparameter
Value
Training Steps
50,000
Batch size
128
Task
FLAN2021
Dropout
0.0
Warmup Schedule
No Warmup
Learning Rate (LR)
0.0001
Input Length (Tokenized)
2048
Output Length (Tokenized)
256
Teacher softmax temperature
1.0
Evaluation
Greedy Sampling
A.8 MODE-SERKING VS MODE-COVERING KI.
argming KLIPI
argming KLIQ (IP)
Figure A.16: Mode-sccking vs Model-covering KL with capacity mismatch. We show the learned distribution Qo when minimizing the forward and reverse KL. w.r.t between mixture distribution P and unimodal Gaussian Qs. Reverse KI. is mode-seeking it forces Qo be where zero and hence makes concentrate on one the modes (last plot). However, forward KL mode-covering as ensures that there some mass under Qa wherever there some mass under P. See Le (2017) to replicate this plot.
18


Reviews:
[{"review_id": "dSXampbRhz", "reviewer": "Reviewer_Qwi5", "summary": "This paper proposes GKD for auto-regressive LM distillation, aiming to address the distribution mismatch between the output sequence during training and inference. GKD trains the student on its self-generated output sequences. GKD significantly improves student’s performance over commonly used KD methods for autoregressive LMs.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper proposes a unified framework for existing KD methods within the context of on-policy imitation learning.\n2. GKD can be easily combined with RLHF. This allows the student to learn capabilities from the teacher while aligning with humans.\n3. The authors investigate different divergence objectives and student data fractions on various tasks, providing practical insights.", "weaknesses": "1. GKD has limited technical novelty upon ImitKD/f-distill. ImitKD/f-distill adopts similar imitation learning ideas to GKD, except that GKD makes $\\lambda$ tunable and explores different divergences. The latter is more of a hyper-parameter choice as it highly depends on target tasks and decoding strategies. \n2. To verify the initial claim that GKD alleviates the train-test distribution mismatch issue, the authors may consider include an experiment with a mismatched test domain.\n3. Missing comparison between GDK and ImitKD/f-distill in NMT.\n4. GKD seems to be computationally expensive due to the need of fine-tuning a good student initialization and sampling from student distribution. What is the computational overhead?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hrt4bfDFaA", "reviewer": "Reviewer_YEFj", "summary": "This paper presents Generalized Knowledge Distillation (GKD), a method that executes knowledge distillation on output sequences generated by both the student and teacher models. The training of the student model utilizes the teacher model's probabilities on the sampled sequences. The proposed method enhances the performance of the student model over standard KD methods across various tasks, including summarization and machine translation.", "questions": "It seems that the GKD method uses the finetuned student as the initialization during task-specific KD. For the other baselines, do you use the same settings?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed GKD method is a straightforward and efficient approach for the knowledge distillation of autoregressive models. It demonstrates substantial enhancements across a range of tasks. The paper also introduces a series of ablation studies.", "weaknesses": "The paper mainly focuses on task-specific distillation, which requires finetuning the student on the specific tasks first. The two-stage training also brings in additional hyperparameters that need to be adjusted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BQKfAdOPdB", "reviewer": "Reviewer_Jqut", "summary": "The paper introduces a novel approach called Generalized Knowledge Distillation (GKD) to address distribution mismatch issues in knowledge distillation for auto-regressive sequence models. GKD trains the student model using self-generated output sequences with guidance from the teacher model. It offers the flexibility to employ alternative loss functions between the student and teacher and enables the seamless integration of distillation with reinforcement learning fine-tuning. GKD significantly improves the performance of student models compared to traditional knowledge distillation methods. The paper's contributions include the introduction of GKD, the use of alternative divergence measures, and the combination of distillation with reinforcement learning fine-tuning for the popoular large language models (LLMs).", "questions": "1. In line 10 of Algorithm 1, the gradient calculation of $L_{GKD}$ is incorrect, right? Because that just corresponds to the gradient of a special case of supervised KD when $\\lambda$ is 0;\n\n2. I don't quite understand that since the student model itself is smaller and has poorer performance, the quality of the sequences generated by the student model is also not good. So why, judging from the experimental results, is it helpful to introduce output sequences from the student in the training set? Can the authors give a reasonable and detailed explanation?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "For auto-regressive sequence models, the proposed GKD aims to tackle the distribution mismatch problem between sequences seen during teacher-student training and those generated by the student during inference. GKD, in contrast to traditional methods, trains the student on self-generated sequences with guidance from the teacher.\n\nThe paper's strengths are as follows:\n\n1. The paper provides extensive experimental results, demonstrating the effectiveness of GKD in significantly improving student model performance, particularly in tasks like summarization, machine translation, and arithmetic reasoning. The approach is supported by solid empirical evidence;\n\n2. The paper is well-structured, providing clear explanations and insights into the proposed GKD method. The concepts are effectively communicated, making them accessible to the reader;\n\n3. GKD has the potential to enhance knowledge distillation for auto-regressive sequence models, contributing to the development of more efficient and cost-effective models. The paper's findings hold significance in the fields of natural language processing and model compression.", "weaknesses": "1. The GKD method has been popular in some related works (Parisotto et al., 2015; Kelly et al., 2019; Agarwal et al., 2022), but it has not been applied to distilled autoregressive models, so the idea of ​​this work is not very novel and original;\n\n2. Although the paper reports significant improvements in performance over baseline KD methods, a more extensive evaluation could strengthen the claims. It would be helpful to see comparisons on a wider range of datasets and tasks to demonstrate the generalizability of GKD;\n\n3. The contributions are not very clear, the authours should explicitly state the key contributions in a dedicated section. This would help readers quickly grasp the paper's core contributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QS07sFNBh7", "reviewer": "Reviewer_HokL", "summary": "The authors propose generalized knowledge distillation for auto-regressive language models. The authors argue that current form of knowledge distillation does not cater for auto-regressive language models; a student model faces discrepancy between the samples seen in the training and those generated by itself in inference. Inspired by imitation learning, the authors propose on-policy KD, where the student generates a sequence and \"imitates\" the teacher. In addition, the authors illustrate the proposed on-policy KD can be easily coupled with RL finetuning.", "questions": "- It would be interesting to see if the proposed KD works on self-knowledge distillation, where a teacher and a student are of same size. \n- In Appendix Figure A.13, why does the temperature of the teacher bring significant difference in the performance? Temperatures are hardly set to 0.2, but the student works best when temp set to 0.2.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow. \n- The proposed distillation method is generalized knowlede distillation that can optimze any divergence between the teacher and student token-level probability distributions. \n- The authros illustrate effectiveness of the propose KD on diverse tasks, such as summarization, machine translation, and arithmetric reasoning.\n- The proposed method can be easily combined with RLAIF.", "weaknesses": "- The authors argue that \"auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference\", but this argument requires a thorough citations. There is a contradictory opinion as in [1], where the authors argue that the discrepancy may not be critical.\n- The organization can be improved. The authors allocate much portion of this paper on different divergences, yet the use of different divergence seems trivial.\n- (minor) in the section 4, \"In particular, we are interested in the on-policy variants (λ = 0), which have not been previously explored.\" should be modified to \\lambda=1.\n\n[1] Exposure Bias versus Self-Recovery: Are Distortions Really Incremental for Autoregressive Text Generation?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)