PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, Olivier Bachem

arXiv:2306.13649v3cs.LGcs.AIcs.CL

TL;DR

Autoregressive knowledge distillation can suffer when fixed training sequences differ from the student’s inference-time outputs. GKD trains on student-generated sequences with teacher feedback, supports alternative divergences, and achieves improvements across generation and task-agnostic distillation settings. The method also integrates with reinforcement-learning fine-tuning.

  • Problem

    Existing autoregressive distillation methods can create train–inference distribution mismatch because they rely on fixed ground-truth or teacher-generated sequences.

  • Method

    GKD mixes fixed and on-policy student-generated sequences while allowing arbitrary divergences between teacher and student token distributions.

  • Results

    GKD consistently outperformed commonly used distillation approaches across summarization, machine translation, and arithmetic reasoning, with 2% and 1% absolute accuracy gains on held-out BBH and MMLU.

  • Takeaways & Limitations

    GKD provides a distillation framework that uses student behavior, supports task-dependent divergence choices, and can be combined with reinforcement learning for sequence-level reward optimization.

  • Takeaways & Limitations

    Experiments assume a student already capable of generating sequences of adequate quality, typically after supervised fine-tuning.

Abstract

from arXiv · show

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 language models on summarization, translation, and arithmetic reasoning tasks, and task-agnostic distillation for instruction-tuning.

1 INTRODUCTION

Generalized Knowledge Distillation (GKD) addresses train–inference distribution mismatch by distilling from student-generated sequences and allows alternative divergences and reinforcement-learning integration. It outperforms commonly used distillation methods across task-specific and task-agnostic settings.

  • Motivation and approach: GKD trains students on self-generated, on-policy output sequences labeled with teacher token probabilities rather than relying only on fixed sequences.This targets the mismatch between sequences used during training and those generated during autoregressive inference.
  • Motivation and approach: GKD permits alternative divergences, including reverse KL and generalized JSD, to accommodate limited student capacity.The choice of divergence can let students focus on teacher-likely samples instead of exactly matching the teacher distribution.
  • Results: GKD unifies existing distillation methods while introducing on-policy variants that substantially outperform prevalent approaches.The framework covers both fixed-sequence and student-generated training regimes.
  • Results: Relative gains over baseline KD improvements averaged 2.1× on summarization, 1.7× on machine translation, and 1.9× on arithmetic reasoning.These gains are reported across T5 students of different sizes.
  • Results: GKD improves task-agnostic distillation by 2% absolute accuracy on BBH and 1% on MMLU.The improvements are reported on held-out benchmark suites.
  • Extensions: On-policy GKD can be combined with RL fine-tuning to optimize sequence-level rewards alongside teacher distillation.The paper identifies this combination as relevant to RLHF training.

2 PRELIMINARIES

The paper introduces autoregressive sequence-model notation and contrasts divergence objectives used to compare teacher and student token distributions. Forward and reverse KL differ in direction and optimization behavior, while generalized JSD provides a bounded interpolation.

  • Autoregressive sequence models: Autoregressive models predict tokens sequentially, conditioning each next-token distribution on the input and previously generated prefix.A sampled output sequence is drawn from the model conditioned on the input.
  • Autoregressive sequence models: Temperature γ controls sampling randomness: higher values increase randomness, while lower values favor deterministic high-probability tokens.Training keeps student temperature at 1; evaluation uses greedy or temperature sampling.
  • KL-based divergences: KL divergence measures differences between discrete distributions and is asymmetric, so DKL(P∥Q) is forward KL while DKL(Q∥P) is reverse KL.The direction determines which distribution appears first in the divergence.
  • KL-based divergences: Under capacity mismatch, minimizing forward KL produces mean-seeking behavior, whereas reverse KL produces mode-seeking behavior.Forward KL under an empirical data distribution corresponds to maximum likelihood.
  • KL-based divergences: Generalized JSD is bounded even for distributions with disjoint supports and interpolates between forward and reverse KL through 0 < β < 1.Its gradients resemble forward KL near β = 0 and reverse KL near β = 1.

3 DISTILLATION FOR AUTO-REGRESSIVE SEQUENCE MODELS

This section frames autoregressive distillation as vulnerable to train–inference distribution mismatch and introduces GKD, which combines fixed and student-generated sequences with flexible divergences. On-policy training supplies teacher feedback on the student's own outputs and can also integrate with RL fine-tuning.

  • Existing Distillation Approaches: GKD includes supervised fine-tuning, sequence-level KD, and supervised KD as related fixed-sequence approaches or special cases.Sequence-level KD maximizes the likelihood of high-probability teacher-generated sequences, while supervised KD imitates the teacher's token-level distributions.
  • Motivation: Fixed ground-truth or teacher-generated sequences can create train-inference distribution mismatch during autoregressive student generation.Errors early in generation can cascade into later predictions and reduce text quality.
  • On-Policy Knowledge Distillation: On-policy KD trains students on self-generated sequences and uses teacher logits to provide token-specific feedback on erroneous tokens.The student is trained on outputs it is likely to generate, while sampling remains outside backpropagation for stability and efficiency.
  • Generalized Knowledge Distillation: GKD unifies supervised and on-policy distillation by mixing fixed ground-truth or teacher-generated sequences with student-generated sequences.It allows the choice of both the training sequences and the divergence between teacher and student token-level distributions.
  • Generalized Knowledge Distillation: The student data fraction λ controls the share of on-policy outputs, with λ=1 yielding on-policy KD and λ=0 yielding supervised KD under forward KL.GKD also permits intermediate λ values and alternative divergences.
  • Distillation with RL Fine-Tuning: GKD can combine with RLHF or RLAIF by adding a distillation loss to reinforcement learning, with α controlling the distillation-loss strength.For integration with existing RL fine-tuning workflows, the authors recommend reverse KL or JSD (0.9).

4 EXPERIMENTS

Experiments evaluate GKD across summarization, translation, and arithmetic reasoning, varying student data sources, divergences, sampling temperatures, and integration with RL. Across tasks, on-policy or mixed student-generated data generally outperforms fixed supervised data, while divergence effects depend on task and decoding conditions.

  • Experimental setup: GKD is evaluated on abstractive summarization, machine translation, and arithmetic reasoning using T5 students of varying sizes and a supervised fine-tuned T5-XL teacher.Students include T5-small, T5-base, and T5-large, which are 38×, 12×, and 3.8× smaller than the teacher.
  • Abstractive summarization: On-policy GKD with JSD (0.9) outperforms SeqKD, Supervised KD, ImitKD, and f-distill under both greedy and temperature sampling on XSum.The comparison spans multiple student model sizes and demonstrates scalability with respect to student capacity.
  • Abstractive summarization: 5% of the XSum training data without ground-truth summaries lets on-policy GKD outperform supervised KD and ImitKD trained on the full dataset.The 5% setting uses 10K examples and T5-small as the student.
  • Ablations and divergence choice: On-policy and mixed variants consistently outperform supervised variants, while mode-seeking divergences help more with temperature sampling than with greedy sampling.For temperature-dependent divergence choice, JSD (0.5) and JSD (0.9) offer better quality but less diversity at high temperatures; lowering temperature narrows quality differences and reduces diversity.
  • RL fine-tuning: On-policy GKD combined with RL fine-tuning improves factual consistency relative to the teacher while producing large gains in summarization quality for the student.In Figure 5, increasing the GKD-loss strength raises ROUGE-2 while reducing the improvement in factual consistency.
  • Translation and arithmetic reasoning: On WMT en →de and GSM8K, purely student-generated data generally outperforms fixed or mixed data, and increasing on-policy data beyond 25% typically improves performance.For instruction tuning, on-policy GKD with reverse KL substantially outperforms supervised KD and ImitKD, while generalized JSD performs better than KL variants in some translation settings.

5 RELATED WORK

Prior work includes supervised and sequence-level distillation, imitation-learning connections, and RL fine-tuning, while GKD differs through on-policy data, divergence flexibility, and joint RL optimization.

  • Knowledge distillation: Supervised and sequence-level KD are established approaches for distilling autoregressive models.Supervised methods can suffer from train-inference mismatch, also called exposure bias.
  • Imitation learning: ImitKD connects distillation with imitation learning but does not use purely on-policy collection or integrate RL fine-tuning.Unlike GKD, it retains token-level forward KL.
  • Imitation learning: MiniLLM frames distillation as an RL problem using sequence-level reverse KL and policy gradients, whereas GKD avoids backpropagation through sampling.The passage also characterizes GKD as simpler and more stable, while supporting forward KL, reverse KL, and JSD.
  • RL fine-tuning: RL fine-tuning commonly regularizes toward the initial model, while the paper jointly distills from a teacher and optimizes a sequence-level reward.The authors describe this as changing regularization from the initial policy to the teacher policy.
  • Reasoning traces or rationales: Reasoning-trace distillation extends teacher data with chain-of-thought prompts but remains supervised.The passage notes that other enhanced prompts could also be considered.
  • Application to speculative decoding: Later work applies GKD to speculative decoding to improve draft-target alignment and inference speedup.

6 CONCLUSION

The paper presents GKD for reducing train-inference mismatch in autoregressive language-model distillation and reports consistent gains across generation tasks, with simultaneous RL fine-tuning as an additional capability.

  • GKD addresses train-inference distribution mismatch when distilling autoregressive language models.
  • GKD consistently outperformed commonly used distillation approaches on abstractive summarization, machine translation, and arithmetic reasoning.
  • GKD can be combined with reinforcement learning to optimize a sequence-level reward while distilling knowledge from a large teacher model.
  • The paper identifies extending GKD to autoregressive audio, video, and text-to-image generation as future work.

A.1 SELF-DISTILLATION

Self-distillation transfers knowledge between same-size, same-architecture models; on GSM8K, on-policy GKD variants outperform supervised KD and can surpass the teacher.

  • On-policy GKD variants outperform other approaches, including supervised KD, and self-distilled students surpass the teacher on the test set.
  • Self-distillation transfers knowledge between teacher and student models with the same architecture and size.
  • On GSM8K, the teacher FLAN-T5 large was supervised fine-tuned on the task, while the student was not trained on GSM8K.

A.2 T5 MODELS

The experiments use LM-adapted T5 checkpoints with task-specific supervised initialization and a shared Adafactor optimizer; GKD adds measurable sampling cost during fine-tuning.

  • Base checkpoints: Base checkpoints are LM-adapted T5v1.1 models trained for an additional 100K language-modeling steps.
  • Initialization: Student and teacher models are initialized with further supervised fine-tuning on each task’s original training data.
  • Task setups: XSum, WMT, and GSM8K use task-specific supervised fine-tuning schedules and starting checkpoints.XSum and WMT use LM-Adapted T5v1.1 models, while GSM8K starts from FLAN-T5 models trained on a generated chain-of-thought dataset.
  • Optimization: The experiments use the Adafactor optimizer, consistent with T5 and FLAN-T5.
  • Computational cost: 1.8×, 2×, and 2.2× overhead is reported for student sampling on GSM8K at student-teacher ratios of 38×, 12×, and 3.8×.For RLHF + GKD, overhead is described as somewhat small because only teacher logits require inference.
  • Computational cost: Real-world cost is dominated by inference serving rather than fine-tuning, and the authors judge on-policy GKD’s performance benefits potentially worth its compute cost.

A.3 XSUM

The XSum ablations compare divergences, data mixtures, and training settings across student sizes. On-policy and mixed GKD variants consistently outperform supervised variants, with mode-seeking divergences especially helpful under temperature sampling.

  • Training settings: A learning-rate sweep selects 0.0003 for T5-base and T5-large, while T5-small uses 0.001; reverse KL defaults to 0.0003 for all models.
  • Training settings: Teacher temperature is 1 for greedy evaluation and 0.1 when reporting student performance with temperature sampling.
  • Training settings: Table A.1 provides the hyperparameter details for the XSum experiments.
  • GKD Ablations: On-policy and mixed GKD variants consistently outperform supervised variants across the XSum ablations.
  • GKD Ablations: Mode-seeking divergences perform better with temperature-sampling evaluation, whereas divergence choice has little effect under greedy sampling.

A.4 GSM8K

The GSM8K experiments evaluate distilled students on four-shot chain-of-thought arithmetic reasoning, using generated training outputs and averaged results across three seeds. The ablations compare GKD variants under greedy sampling.

  • Experimental setup: GSM8K training uses chain-of-thought outputs generated from PaLM-540B, and reported accuracy is averaged across three seeds.
  • Experimental setup: Table A.2 provides the hyperparameter details for the GSM8K experiments.
  • Experimental setup: The experiments use a specified four-shot chain-of-thought prompt for evaluation.
  • GKD Ablations: Ablations on GSM8K use greedy sampling and report improvement in student test accuracy after distillation, averaged across three seeds.
  • GKD Ablations: Using only student-generated output samples typically outperforms other GKD variants.
  • GKD Ablations: The supervised fine-tuned T5-XL teacher obtains 27.9 accuracy, while the T5-small student obtains 4.585 before the reported improvement.

A.5 WMT

The WMT experiments compare GKD with ImitKD and f-distill, alongside specified evaluation and hyperparameter settings. On-policy GKD achieves substantially larger BLEU improvements than both comparison methods.

  • Evaluation: WMT evaluation uses beam search with the same hyperparameters as Raffel et al. (2020), and results are averaged across three seeds.
  • Results: On-policy GKD produces 53% higher BLEU improvement than ImitKD and 162% higher improvement than f-distill, averaged across small and base models.
  • Method: For WMT, GKD uses the best-performing variant with λ = 1 and JSD (0.1).
  • Evaluation: Table A.3 lists hyperparameter details for the WMT English-German experiments.
  • Related settings: Table A.4 lists hyperparameter details for FLAN instruction tuning.
  • Divergence interpretation: Under capacity mismatch, reverse KL is mode-seeking while forward KL is mode-covering for a mixture distribution approximated by a unimodal Gaussian.
Loading 2306.13649v3…