PaperlayerView exact revision on arXiv

Source-linked AI summary

Sequence-Level Knowledge Distillation

Yoon Kim, Alexander M. Rush

arXiv:1606.07947v4cs.CLcs.LGcs.NE

TL;DR

NMT models often require very high capacity, making training and deployment on standard setups challenging. This paper applies word-level and sequence-level knowledge distillation, finding that sequence-level methods improve student performance and can enable faster greedy decoding.

  • Problem

    NMT systems typically require very high capacity, making training demanding and use on standard setups challenging.

  • Method

    The paper trains students on teacher-generated sequences from beam search and introduces sequence-level distillation variants, including interpolation between observed and teacher-generated data.

  • Results

    Sequence-level knowledge distillation outperforms word-level distillation on English →German, while sequence-level interpolation makes greedy decoding comparable to beam search on the original model: 19.6 versus 19.5.

  • Takeaways & Limitations

    Sequence-level distillation transfers knowledge at the sequence level and can improve translation quality while reducing reliance on beam search.

  • Takeaways & Limitations

    The interpolation training process doubles the training-data size and requires conditioning on observed and teacher-generated sequences that are often quite different.

Abstract

from arXiv · show

Neural machine translation (NMT) offers a novel alternative formulation of translation that is potentially simpler than statistical approaches. However to reach competitive performance, NMT models need to be exceedingly large. In this paper we consider applying knowledge distillation approaches (Bucila et al., 2006; Hinton et al., 2015) that have proven successful for reducing the size of neural models in other domains to the problem of NMT. We demonstrate that standard knowledge distillation applied to word-level prediction can be effective for NMT, and also introduce two novel sequence-level versions of knowledge distillation that further improve performance, and somewhat surprisingly, seem to eliminate the need for beam search (even when applied on the original teacher model). Our best student model runs 10 times faster than its state-of-the-art teacher with little loss in performance. It is also significantly better than a baseline model trained without knowledge distillation: by 4.2/1.7 BLEU with greedy decoding/beam search. Applying weight pruning on top of knowledge distillation results in a student model that has 13 times fewer parameters than the original teacher model, with a decrease of 0.4 BLEU.

1 Introduction

NMT can translate accurately but typically requires very large models, making training and deployment difficult. This paper applies knowledge distillation to compress NMT models and introduces sequence-level variants tailored to complete-sequence prediction.

  • NMT directly predicts each target word conditioned on the source sentence and previously generated target words.
  • 4×1000 and 16×512 LSTM configurations illustrate the very high capacity required for competitive NMT performance.These model sizes require cutting-edge training hardware and make standard deployment challenging.
  • Knowledge distillation trains a smaller student network to mimic a larger teacher by matching their outputs.
  • The paper evaluates standard word-level distillation and proposes two sequence-level variants that approximate the teacher’s sequence-level distribution.The sequence-level approach trains the student on a dataset generated by teacher beam search.
  • A 2×500 student roughly matches a 4×1000 teacher, while greedy decoding is 10 times faster with comparable performance.The experiments also compress a 2×500 model to 2×100 on a smaller dataset.

2 Background

The paper frames translation as sequence prediction with an attention-based encoder–decoder and reviews knowledge distillation as output-distribution matching. It then contrasts ordinary word-level training with teacher-guided sequence-level training.

  • Sequence-to-Sequence with Attention: NMT parameterizes p(t | s) with an encoder that reads the source and a decoder that generates target words sequentially.The paper uses an attentional architecture for English →German translation.
  • Knowledge Distillation: Knowledge distillation trains a smaller student to match the predictions of a previously trained, larger teacher network.For classification, the usual objective matches output probabilities rather than only observed labels.
  • Knowledge Distillation: Word-level knowledge distillation minimizes cross-entropy between student and teacher distributions for each target word.The comparison is illustrated as part of the word-level approach in Figure 1.

3 Knowledge Distillation for NMT

The paper applies knowledge distillation to NMT at both word and sequence levels, using teacher-generated sequences to approximate otherwise intractable sequence-level objectives. Sequence-level interpolation selects a beam candidate that balances teacher probability with similarity to the observed target.

  • Overview: NMT knowledge distillation trains a smaller student to learn from a larger teacher, and the paper explores three applications of this technique.The approaches include word-level distillation, sequence-level distillation, and sequence-level interpolation.
  • Word-Level and Sequence-Level Distillation: Word-level distillation transfers local teacher word distributions, but sequence-level distillation instead matches the teacher’s distribution over complete target sequences.Sequence-level matching is motivated by error propagation from earlier decisions during test-time generation.
  • Sequence-Level Knowledge Distillation: Sequence-level distillation approximates the teacher’s sequence distribution with its mode, using beam search to obtain the teacher-generated target sequence.The mode is used because finding it exactly is intractable; a beam of size 5 captures more teacher probability mass than a beam of size 1 in the reported experiments.
  • Sequence-Level Interpolation: Sequence-level interpolation selects a beam sequence that is close to the observed target and has high probability under the teacher, using a similarity function such as BLEU.The method is a single-sequence approximation inspired by local updating and uses smoothed sentence-level BLEU in the described implementation.
  • Sequence-Level Interpolation: The interpolation rationale models the observed target as a noisy version of an unobserved true sequence, making a mixture-distribution mode potentially different from both the gold and teacher-decoded sequences.This provides a motivation for selecting a beam candidate near the observed target rather than training on both observed and teacher-generated sequences.

4 Experimental Setup

The experiments evaluate distillation methods on high- and low-resource translation tasks with teacher–student LSTM pairs. They compare word-level and sequence-level training variants, including beam-search-generated data, interpolation data, fine-tuning, and combinations.

  • Datasets: The study evaluates English→German as a high-resource task and Thai→English as a low-resource task.The English→German data comes from WMT 2014, while the Thai→English data comes from IWSLT 2015.
  • Datasets and Models: The English→German teacher is a 4×1000 LSTM, with 2×300 and 2×500 LSTM students.The training set contains 4m sentences, and the vocabulary retains the 50k most frequent words.
  • Methods and Evaluation: The experiments evaluate tokenized BLEU and compare baseline, word-level distillation, teacher-generated sequence data, sequence-level interpolation, fine-tuning, and combined variants.Word-level distillation uses α ∈{0.5, 0.9}, with α = 0.5 reported as better; Seq-KD uses beam size K = 5, while Seq-Inter uses K = 35.
  • Implementation: Sequence-level interpolation data for English→German is generated on approximately 50% of the training set for efficiency.The interpolation procedure fine-tunes a pretrained model with learning rate 0.1.

5 Results and Discussion

Sequence-level distillation improves translation quality and concentrates probability near the teacher’s mode, making greedy decoding effective. Weight pruning further reduces model size with limited BLEU loss, while word embeddings remain a major parameter bottleneck.

  • Translation results: Seq-KD outperforms Word-KD on English →German and matches it on Thai →English, while combining both yields further gains for 2 × 300 and 2 × 100 models.The gains from combining them do not extend to the 2 × 500 model, suggesting local and sequence-level transfer provide complementary information mainly in smaller students.
  • Distribution concentration: 16.9% of the probability mass belongs to the approximate argmax for the 2 × 500 Seq-KD model, versus 0.9% for the baseline on English →German.The authors use this concentration around the teacher’s mode to explain why Seq-KD models have more peaked distributions.
  • Distribution concentration: 7.6% of the distribution is accounted for by the greedily decoded sequence under Seq-Inter, providing a compromise between Seq-KD and the broader baseline distribution.The paper links this concentration to the success of greedy decoding.
  • Translation results: Seq-KD improves English →German despite perplexity rising from 8.2 for the baseline 2 × 500 model to 22.7 for Seq-KD.The Seq-KD model gains +4.2 BLEU with greedy decoding and +1.4 BLEU with beam search.
  • Decoding speed: 10 times faster decoding is achieved by the student with greedy decoding than by the teacher with beam search, with comparisons spanning GPU, CPU, and smartphone settings.Table 2 reports source words translated per second across these hardware platforms; the 4 × 1000 model could not be opened on the smartphone.
  • Weight pruning: 80% weight pruning produces a model with 13× fewer parameters than the original teacher and only a 0.4 BLEU decrease.The pruning procedure removes weights with the lowest absolute values, then retrains on Seq-KD data and fine-tunes toward Seq-Inter data.

6 Related Work

Prior compression work includes pruning, knowledge transfer across architectures, and alternative parameter-reduction techniques such as factorization, regularization, binarization, and weight sharing.

  • Prior pruning methods remove weights or neurons using criteria such as Hessian approximations or parameter magnitude.
  • Knowledge transfer has also been studied across network architectures, including transferring from RNNs to non-recurrent networks and from LSTMs to CNNs.
  • Other compression approaches use low-rank factorization, sparsity-inducing regularization, weight binarization, or weight sharing.

7 Conclusion

The paper introduces sequence-level knowledge distillation variants that improve over word-level distillation and considers translation as a test case for broader sequence-to-sequence model compression.

  • The two sequence-level knowledge distillation variants improve over standard word-level knowledge distillation for NMT.
  • The methods may help train smaller models for other sequence-to-sequence tasks beyond translation.
Loading 1606.07947v4…