[Paper Review] Distilling Knowledge Learned in BERT for Text Generation

Yen-Chun Chen, Zhe Gan+3 more|arXiv (Cornell University)|Nov 10, 2019

Topic Modeling58 references21 citations

TL;DR

This paper proposes a knowledge distillation framework that leverages BERT’s bidirectional context via a novel Conditional Masked Language Modeling (C-MLM) finetuning to improve autoregressive sequence-to-sequence models for text generation. By using the finetuned BERT as a teacher to provide soft probability targets, the student model learns global coherence, achieving state-of-the-art results on IWSLT German-English and English-Vietnamese translation benchmarks.

ABSTRACT

Large-scale pre-trained language model such as BERT has achieved great success in language understanding tasks. However, it remains an open question how to utilize BERT for language generation. In this paper, we present a novel approach, Conditional Masked Language Modeling (C-MLM), to enable the finetuning of BERT on target generation tasks. The finetuned BERT (teacher) is exploited as extra supervision to improve conventional Seq2Seq models (student) for better text generation performance. By leveraging BERT's idiosyncratic bidirectional nature, distilling knowledge learned in BERT can encourage auto-regressive Seq2Seq models to plan ahead, imposing global sequence-level supervision for coherent text generation. Experiments show that the proposed approach significantly outperforms strong Transformer baselines on multiple language generation tasks such as machine translation and text summarization. Our proposed model also achieves new state of the art on IWSLT German-English and English-Vietnamese MT datasets. Code is available at https://github.com/ChenRocks/Distill-BERT-Textgen.

Motivation & Objective

• To address the challenge of leveraging BERT’s bidirectional contextual knowledge for autoregressive text generation, which is inherently autoregressive and lacks future context.
• To improve global coherence in text generation by introducing sequence-level supervision through knowledge distillation from a finetuned BERT teacher model.
• To develop a modular, model-agnostic approach that enhances any sequence-to-sequence model without architectural constraints or parameter sharing.
• To achieve state-of-the-art performance on machine translation and abstractive summarization tasks using a lightweight, efficient distillation mechanism.

Proposed method

• Propose Conditional Masked Language Modeling (C-MLM), a finetuning objective that extends MLM by conditioning on input sequences, enabling BERT to adapt to sequence-to-sequence tasks.
• Fine-tune BERT on target generation datasets using C-MLM to produce soft probability distributions over next tokens, incorporating both left and right context.
• Use the finetuned BERT as a teacher model to generate soft-label logits for each token in training sequences, providing global sequence-level supervision.
• Train a student sequence-to-sequence model (e.g., Transformer) via a modified MLE loss that incorporates the BERT-generated soft targets as a regularization term.
• Introduce a distillation loss term, $\mathcal{L}_{\text{bidi}}$, that encourages the student model to match the BERT teacher’s probability distribution across the entire sequence.
• Decouple the teacher from inference, enabling fast decoding at test time while retaining the benefits of bidirectional context learned during training.

Experimental results

Research questions

• RQ1 Can BERT’s bidirectional context be effectively transferred to autoregressive text generation models through knowledge distillation?
• RQ2 Does incorporating global sequence-level supervision from BERT improve the coherence and fluency of generated text?
• RQ3 Can the proposed method achieve state-of-the-art performance on diverse text generation tasks without architectural constraints?
• RQ4 How does the distillation method perform across varying output lengths, especially in long-sequence generation?
• RQ5 Can the model generalize across different generation tasks such as machine translation and abstractive summarization?

Key findings

• The proposed method achieves a new state of the art on the IWSLT14 German-English translation dataset, outperforming strong Transformer baselines.
• On the IWSLT15 English-Vietnamese translation dataset, the model achieves a BLEU score of 27.85, surpassing the best baseline by 0.86 BLEU points.
• For long sequences (N > 24), the model shows consistent improvement over baselines, with gains of up to 10 BLEU points on longer translations in the IWSLT German-English set.
• Qualitative analysis shows that the model avoids errors in word choice (e.g., 'with' vs. 'at') by leveraging future context during training, leading to more coherent outputs.
• The distillation loss $\mathcal{L}_{\text{bidi}}$ significantly improves generalization, particularly in reducing overconfidence on one-hot labels and enhancing robustness to context shifts.
• The method is model-agnostic and compatible with various architectures, including Transformers and LSTMs, without requiring parameter sharing or architectural modifications.