Source-linked AI summary
Autoregressive Knowledge Distillation through Imitation Learning
Alexander Lin, Jeremy Wohlwend, Howard Chen, Tao Lei
TL;DR
Deep autoregressive architectures improve language generation but make inference costly for real-world deployment. The paper introduces imitation-based knowledge distillation, which uses imitation learning to address exposure bias. On translation and summarization, it outperforms other distillation methods and produces faster students.
Problem
Growing model size and Transformer decoding complexity make state-of-the-art autoregressive models costly to deploy, while static-data training creates exposure bias.
Source: §1 · pp. 1–2Method
ImitKD trains a student in an imitation-learning framework, treating the teacher as an oracle that corrects the student at every time step.
Source: §1 · p. 2, §3.1 · p. 3Results
ImitKD consistently outperforms other distillation algorithms on translation and summarization, with students 1.4 to 4.8 BLEU/ROUGE points above scratch-trained models and up to 14 times faster than the teacher.
Source: §1 · p. 2Takeaways & Limitations
Imitation-based distillation compresses large autoregressive models into smaller, faster counterparts while retaining strong generation performance.
Source: §1 · p. 2, §7 · pp. 8–9The approach assumes the teacher can provide better actions than the student on the student’s own states.
Source: §3.1 · p. 3
Abstract
from arXiv · showhide
The performance of autoregressive models on natural language generation tasks has dramatically improved due to the adoption of deep, self-attentive architectures. However, these gains have come at the cost of hindering inference speed, making state-of-the-art models cumbersome to deploy in real-world, time-sensitive settings. We develop a compression technique for autoregressive models that is driven by an imitation learning perspective on knowledge distillation. The algorithm is designed to address the exposure bias problem. On prototypical language generation tasks such as translation and summarization, our method consistently outperforms other distillation algorithms, such as sequence-level knowledge distillation. Student models trained with our method attain 1.4 to 4.8 BLEU/ROUGE points higher than those trained from scratch, while increasing inference speed by up to 14 times in comparison to the teacher model.
1 Introduction
Autoregressive models are powerful for language generation but increasingly costly to deploy because model scale and Transformer decoding slow inference. ImitKD addresses this through imitation-learning-based distillation, outperforming competing methods while producing substantially faster students.
1 Introduction
- Model growth and Transformer self-attention have increased autoregressive inference costs, especially on CPUs and low-resource devices.The paper identifies both larger architectures and more expensive step-by-step decoding as deployment barriers.Source: §1 · p. 1
- Knowledge distillation transfers information from a pretrained teacher to a smaller student whose architecture can differ substantially from the teacher.This flexibility supports compression for memory and speed requirements.Source: §1 · p. 1
- Static teacher-generated training data creates exposure bias because students train on teacher states but generate from their own states at inference.The resulting mismatch can reduce generation quality.Source: §1 · pp. 1–2
- Up to 14× faster inference results from compressing deep, high-performing Transformer teachers into shallow RNN students.The method is especially suited to reducing inference cost while preserving generation quality.Source: §1 · p. 2
- 1.4 to 4.8 BLEU/ROUGE points separate ImitKD-trained students from models trained without a teacher.This improvement is reported across translation and summarization generation metrics.Source: §1 · p. 2
- ImitKD consistently outperforms word-level and sequence-level knowledge distillation algorithms on the reported generation tasks.The algorithm is motivated by an imitation-learning perspective on autoregressive distillation.Source: §1 · p. 2
2 Background
The paper formalizes autoregressive distillation as learning a student policy to match a pretrained teacher over generated sequences. It interprets SeqKD as behavioral cloning and motivates dynamic teacher interaction to address exposure bias.
2.1 Autoregressive Distillation
- Autoregressive distillation trains a student model to generate sequences well by minimizing loss relative to a pretrained teacher.The objective uses next-token discrepancy under a chosen distribution of source-target contexts.Source: §2.1 · p. 2
2.2 Distillation as Imitation Learning
- Autoregressive generation forms a T-step Markov decision process where partial sequences are states, next tokens are actions, and transitions append tokens.The vocabulary is the action space, and the transition is deterministic.Source: §2.2 · p. 2
- Imitation learning treats the student as a learner and the high-performing teacher as an oracle whose policy the student learns to match.This avoids requiring an objective token-level reward for generation quality.Source: §2.2 · p. 2
2.3 SeqKD as Behavioral Cloning
- SeqKD is behavioral cloning: the teacher generates approximate mode sequences that become a static training dataset for the student.The student is trained on teacher tokens with a negative log-likelihood objective.Source: §2.3 · p. 3
- SeqKD is simple and adds no training overhead beyond creating the teacher-generated dataset.This is its principal practical advantage over vanilla supervised learning.Source: §2.3 · p. 3
- Behavioral cloning suffers exposure bias because the student never trains on its own states, allowing step-by-step errors to compound during testing.The cited analysis describes regret relative to the oracle as quadratic in the time horizon T.Source: §2.3 · p. 3
- Dynamic teacher interaction lets the teacher query and correct states generated by the student, offering a way to ameliorate exposure bias.This motivates mixing oracle and learner state distributions rather than relying only on a static dataset.Source: §2.1 · p. 2, §2.3 · p. 3
3 Imitation-Based Distillation Algorithm
ImitKD frames autoregressive distillation as imitation learning: the student explores its own states while the teacher acts as an oracle that corrects each next-token decision. The algorithm mixes initial data with student-generated sequences, supports practical dataset replacement and annealed sampling, and accelerates training through batched generation.
3.1 Design Principles and Rationale
- ImitKD trains the student on its own state distribution while mixing student-generated examples with an initial dataset to mitigate exposure bias and cold start.This strategy builds on DAgger’s oracle-learner mixture, which is associated with improved regret under stated convexity assumptions.Source: §3.1 · p. 3
- A strong Transformer teacher can act as an oracle by correcting student generations at each time step and producing better completions for the student’s partial states.The preliminary IWSLT study compares teacher completions with shallow-RNN generations using greedy-decoding BLEU scores.Source: §3.1 · p. 3, Table 1 · p. 4
3.2 The ImitKD Algorithm
- The ImitKD algorithm samples from a data mixture, generates student sequences, replaces training examples, queries teacher targets, and returns the best validation policy.The loss can use the teacher’s optimal next token or the full teacher-student cross entropy.Source: §3.2 · p. 4
3.3 Data Mixture Selection and Annealing
- Sequence-level data replacement uses either the original supervised dataset or SeqKD data, mixes in student-generated samples, and can follow an exponential annealing schedule.This avoids teacher generation at every iteration because the teacher only computes the batched loss; β_i = r^i/I controls the schedule.Source: §3.3 · p. 4
3.4 Speeding Up Training
- Training efficiency improves by generating B · M examples in parallel once every M iterations, with small M values such as 2–8 reported to preserve final performance.The trade-off is that losses may use examples generated by parameters from up to M iterations earlier.Source: §3.4 · p. 5
4 Related Work
Related work addresses autoregressive generation through sequence-level distillation, behavioral cloning, and methods that generate during training. The paper distinguishes ImitKD by using a dynamic teacher oracle to define supervision for student-generated states.
Sequence-level distillation
- SeqKD trains students on teacher-generated datasets with standard negative log-likelihood and has been widely used in autoregressive and non-autoregressive translation.The cited applications include machine translation, low-resource settings, and ensemble distillation.Source: §4 · p. 5
Generation during training
- Generation-based methods such as MIXER, Beam Search Optimization, SEARNN, and scheduled sampling address training-time generation or exposure bias through sequence metrics, local losses, or altered contexts.These methods differ in how they construct supervision once generated sequences depart from static training data.Source: §4 · p. 5
Dynamic oracle supervision
- Without a dynamic queryable oracle, generated states can lack a consistent target, whereas distillation can use the teacher oracle to support imitation learning.The paper positions ImitKD as addressing exposure bias in autoregressive distillation rather than targeting non-autoregressive translation.Source: §4 · p. 5
5 Experimental Setup
The experiments evaluate autoregressive distillation across translation and summarization settings using Transformer teachers and recurrent students. They compare supervised, SeqKD, and ImitKD variants with BLEU, ROUGE, and perplexity metrics.
Datasets
- The evaluation covers IWSLT 2014 German→English translation, WMT 2016 English→German translation, and CNN/DailyMail abstractive summarization.The datasets use established train, validation, and test splits with task-specific BPE or BERT tokenization.Source: §5 · p. 5
Models
- State-of-the-art Transformers serve as teachers, while SRU recurrent neural networks serve as students across tasks.Additional Transformer, GRU, and LSTM students are trained on IWSLT to examine the method across architectures.Source: §5 · pp. 5–6
Variants
- The variants compare vanilla supervised learning, SeqKD, and ImitKD, including original-data or SeqKD initialization and optimal-token versus full cross-entropy losses.Table 2 distinguishes where each variant obtains its next-token target and whether it uses the full teacher-student distribution.Source: Table 2 · p. 6, §5 · p. 6
Evaluation
- Translation is evaluated with BLEU, summarization with ROUGE-1, ROUGE-2, and ROUGE-L, and all tasks also report perplexity.The checkpoint with the highest validation BLEU or ROUGE-1 is selected for test evaluation.Source: §5 · p. 6
6 Results
Across translation and summarization, ImitKD consistently improves student performance over distillation and vanilla baselines while substantially reducing inference cost. Its gains persist across architectures, data regimes, and sequence lengths, with particular benefits for recurrent students and longer sequences.
IWSLT De→En Translation
- 1.4 to 4.8 BLEU points: ImitKD outperforms all baselines on IWSLT across three student architectures and three evaluation metrics.The 3-layer SRU with ImitKD + Full slightly exceeds the teacher, while ImitKD beats SeqKD by up to 1.4 BLEU.Source: §6 · p. 6, Table 3 · p. 7
Different Student Architectures
- ImitKD outperforms baselines across all tested 2-layer student architectures, with especially large gains for recurrent models.These results use validation BLEU1 with beam size 1.Source: §6 · p. 6, Table 4 · p. 7
WMT En→De Translation
- ImitKD remains competitive with SeqKD on WMT, performs best with the Full variant, and is especially effective when training data is reduced.The method avoids creating a teacher-modified dataset, while low-data results show much stronger gains over baselines.Source: §6 · pp. 6–7, Table 5 · p. 7
CNN/DailyMail Summarization
- On CNN/DailyMail, ImitKD + Full brings the best student within 1 ROUGE point of the teacher and beats baselines on every ROUGE metric.The comparison uses a 6-layer Transformer teacher and a 2-layer SRU student with beam size 5.Source: §6 · p. 7, Table 6 · p. 7
Size and Speed Analysis
- 4–14 times faster: shallow recurrent students substantially reduce model size and CPU inference time while minimizing performance loss relative to deep Transformer teachers.Recurrent students are also 2–3 times faster than similarly sized Transformer students.Source: §6 · p. 7, Table 7 · p. 8
Performance Analysis at Different Lengths
- Across decoding-length bins, ImitKD improves BLEU at every length and shows its largest gains on longer sequences, unlike vanilla and SeqKD models whose quality degrades more noticeably.The analysis uses two-layer IWSLT models and bins sequences in widths of 20; the pattern supports addressing exposure bias.Source: §6 · pp. 7–8, Figure 1 · p. 8
Additive Effect of Fine-Tuning
- ImitKD models retain superior performance after SeqInter fine-tuning, although SeqInter enables vanilla models to close the gap with SeqKD.The comparison concerns IWSLT student architectures evaluated with beam size 1.Source: §6 · p. 8, Table 8 · p. 8
7 Conclusion
The work introduces an imitation-learning-inspired distillation technique for compressing large autoregressive models into smaller, faster counterparts, with empirical success over popular baselines. Future directions include stronger imitation-learning algorithms, imitation-based fine-tuning, and applying ImitKD to large language models for transfer learning.
7 Conclusion
- Imitation-based knowledge distillation compresses large, cumbersome autoregressive models into smaller and faster counterparts.The technique is inspired by imitation learning.Source: §7 · p. 8
- The method achieves empirical success over popular baselines on several natural language generation tasks.Source: §7 · p. 9
- Future work may extend ImitKD with advanced imitation-learning algorithms, imitation-based fine-tuning, and compression of large language models for transfer learning.Source: §7 · p. 9
A.1 DAgger Algorithm
DAgger minimizes action-discrepancy loss over state trajectories and iteratively aggregates states encountered by a mixed policy. The procedure returns the best policy selected on a validation set.
A.1 DAgger Algorithm
- DAgger minimizes action-discrepancy loss between an oracle and policy learner over a distribution of T-step state trajectories.The objective uses D as the trajectory distribution and ℓπ∗(s, π) as the action-discrepancy loss.Source: §A.1 · p. 12
- DAgger initializes a policy randomly, runs the environment using a mixed policy, and samples state trajectories.Source: §A.1 · p. 12
- The algorithm appends sampled states to the training dataset before continuing aggregation.Source: §A.1 · p. 12
- DAgger returns the policy that performs best on a validation set.Source: §A.1 · p. 12
A.2 Implementation Details
The implementation details specify model architectures, datasets, optimization settings, evaluation procedures, and decoding choices for the experiments. They cover recurrent and Transformer students across translation, summarization, and length-dependent analyses.
Model architecture
- RNN-based models use bidirectional encoders with hidden dimension N/2 and left-to-right decoders with hidden dimension N.Source: §A.2 · p. 12
Evaluation
- BLEU and ROUGE are evaluated with the NLTK and py-rouge libraries, respectively.Source: §A.2 · p. 12
IWSLT
- The IWSLT experiments use an 8-layer Transformer teacher and 2- or 3-layer SRU or 2-layer Transformer students, with top-K sampling during ImitKD training.The final mixing rate is r = 0.005 and K = 5.Source: §A.2 · p. 12
WMT
- The WMT experiments use Fairseq data, a pretrained Transformer-large teacher, and a 4-layer SRU student trained for 15 epochs.The final mixing rate is tuned over r ∈ {0.5, 0.1, 0.005}.Source: §A.2 · pp. 12–13
CNN/DailyMail
- CNN/DailyMail uses a 6-layer Transformer-base teacher and a 2-layer SRU student with greedy decoding during ImitKD training.All ImitKD variants use final mixing rate r = 0.1 and M = 4.Source: §A.2 · p. 13
Size and Speed Analysis
- CPU generation times are averaged over test-set examples and measured on a 2019 MacBook Pro with a 2.6GHz 6-core Intel Core i7 processor.Source: §A.2 · p. 13
Performance Analysis at Different Lengths
- Length-dependent IWSLT performance is measured with greedy decoding by grouping decoded sequences into six length bins and computing each bin’s BLEU score.Source: §A.2 · p. 13
Additive Effect of Fine-Tuning
- Fine-tuning experiments use SeqInter data generated with beam size K = 5, Adam optimization, 20K iterations, and validation every 1K iterations.Source: §A.2 · p. 13