PaperlayerView exact revision on arXiv

Source-linked AI summary

TinyBERT: Distilling BERT for Natural Language Understanding

Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, Qun Liu

arXiv:1909.10351v5cs.CLcs.AIcs.LG

TL;DR

Pretrained language models improve NLP performance but are costly to run on resource-restricted devices. TinyBERT introduces Transformer-specific distillation and applies it during both pretraining and task-specific learning, achieving strong GLUE performance with substantially lower model and inference costs.

  • Problem

    Pre-trained language models such as BERT are difficult to deploy on resource-restricted devices because they have many parameters and long inference times.

  • Method

    TinyBERT uses a Transformer-specific knowledge-distillation method and a two-stage framework covering pretraining and task-specific learning.

  • Results

    TinyBERT4 achieves more than 96.8% of teacher BERTBASE performance on GLUE while being 7.5x smaller and 9.4x faster on inference; TinyBERT6 performs on-par with BERTBASE.

  • Takeaways & Limitations

    TinyBERT provides an effective way to deploy BERT-based NLP models on edge devices while significantly reducing model size and inference time.

  • Takeaways & Limitations

    The general-distillation stage does not perform prediction-layer distillation.

Abstract

from arXiv · show

Language model pre-training, such as BERT, has significantly improved the performances of many natural language processing tasks. However, pre-trained language models are usually computationally expensive, so it is difficult to efficiently execute them on resource-restricted devices. To accelerate inference and reduce model size while maintaining accuracy, we first propose a novel Transformer distillation method that is specially designed for knowledge distillation (KD) of the Transformer-based models. By leveraging this new KD method, the plenty of knowledge encoded in a large teacher BERT can be effectively transferred to a small student Tiny-BERT. Then, we introduce a new two-stage learning framework for TinyBERT, which performs Transformer distillation at both the pretraining and task-specific learning stages. This framework ensures that TinyBERT can capture he general-domain as well as the task-specific knowledge in BERT. TinyBERT with 4 layers is empirically effective and achieves more than 96.8% the performance of its teacher BERTBASE on GLUE benchmark, while being 7.5x smaller and 9.4x faster on inference. TinyBERT with 4 layers is also significantly better than 4-layer state-of-the-art baselines on BERT distillation, with only about 28% parameters and about 31% inference time of them. Moreover, TinyBERT with 6 layers performs on-par with its teacher BERTBASE.

1 Introduction

TinyBERT addresses the deployment cost of pretrained language models by distilling BERT through a Transformer-specific method and a two-stage learning framework. The resulting compact models retain strong GLUE performance while reducing model size and inference time.

  • Motivation: Pre-trained language models have many parameters and long inference times, making deployment on resource-restricted edge devices difficult.Existing analyses indicate redundancy in these models, motivating computational and storage reduction while maintaining accuracy.
  • Approach: Knowledge distillation transfers knowledge from a large teacher network to a small student network by training the student to reproduce the teacher’s behaviors.The paper specializes this teacher-student approach for Transformer-based models, using BERT as the example large-scale pretrained model.
  • Approach: The proposed Transformer distillation fits embedding outputs, Transformer hidden states and attention matrices, and prediction-layer logits.These objectives target representations at multiple levels, including attention patterns associated with linguistic knowledge.
  • Approach: The two-stage framework performs Transformer distillation during both pretraining and task-specific learning to transfer general-domain and task-specific knowledge.This design addresses the need to distill BERT across its pretraining and downstream fine-tuning stages.
  • Results: More than 96.8% of teacher BERTBASE performance is achieved by TinyBERT4 on GLUE, with approximately 13.3% of the parameters and 10.6% of the inference time.TinyBERT6 performs on-par with teacher BERTBASE on GLUE.

2 Preliminaries

The preliminaries describe Transformer layers and knowledge distillation as the technical basis for TinyBERT. They emphasize self-attention representations as informative behaviors and identify effective behavior and loss functions across BERT’s training stages as the central distillation problem.

  • Transformer Layer: Transformer-based pretrained language models use layers with multi-head attention and feed-forward network sub-layers.Multi-head attention uses queries, keys, and values to compute attention-weighted representations, while the feed-forward network uses two linear transformations and a ReLU activation.
  • Transformer Layer: The attention matrix is computed from query-key compatibility and determines how values are weighted in the attention output.The key dimension acts as a scaling factor before the softmax weighting operation.
  • Knowledge Distillation: Knowledge distillation trains a small student network to mimic informative behaviors of a large teacher network.Behavior functions may be outputs of Transformer layers or intermediate representations such as the attention matrix.
  • Knowledge Distillation: The distillation objective minimizes a loss measuring the difference between teacher and student behavior functions over the training dataset.In this setting, the main design problem is selecting effective behavior and loss functions while supporting both pretraining and task-specific training.

3 Method

TinyBERT uses Transformer-specific distillation losses within a two-stage framework that transfers general-domain and task-specific knowledge from BERT. The method combines layer-level representation matching, task-specific distillation, and data augmentation to produce compact models with competitive performance.

  • 3.1 Transformer Distillation: Transformer distillation maps selected student layers to teacher layers and minimizes weighted losses over corresponding model-layer behaviors.The mapping includes embedding and prediction layers, while intermediate behavior functions can represent Transformer outputs or attention matrices.
  • 3.1 Transformer Distillation: Attention-based distillation matches multi-head attention matrices, while hidden-state distillation aligns Transformer-layer outputs through a learnable linear transformation.The attention target uses unnormalized matrices because experiments found faster convergence and better performance than fitting softmax outputs.
  • 3.1 Transformer Distillation: The method also distills embedding-layer representations and prediction-layer logits using mean-squared-error and soft cross-entropy objectives, respectively.Embedding and hidden-state dimensions may differ, so learnable transformations align student representations with teacher states.
  • 3.2 TinyBERT Learning: TinyBERT training has general and task-specific distillation stages that transfer pretrained BERT knowledge before focusing on augmented downstream-task data.General distillation uses original BERT and large-scale text, while task-specific distillation reuses Transformer distillation with a fine-tuned BERT teacher.
  • 3.2 TinyBERT Learning: Although the generally distilled model performs worse than BERT, the complementary two-stage framework with augmentation yields competitive downstream performance despite substantial model reduction.General distillation provides initialization, while task-specific distillation further improves the model by focusing on task-specific knowledge.
  • 3.2 TinyBERT Learning: General distillation omits prediction-layer distillation, whereas task-specific learning uses word-level augmentation to improve task-focused knowledge transfer.The augmentation procedure replaces single-piece words using BERT predictions and multi-piece words using GloVe similarity candidates.

4 Experiments

TinyBERT is evaluated across GLUE tasks, compared with compact BERT distillation baselines, and analyzed through ablations of its learning procedures and objectives.

  • Experimental setup: TinyBERT is evaluated on the GLUE benchmark, which includes single-sentence, sentence-similarity, and natural-language-inference tasks.The experiments use the official GLUE test set and single-task training.
  • Results: TinyBERT4 consistently outperforms BERTTINY across GLUE tasks, improving the average performance by 6.8%.The comparison indicates that the proposed knowledge-distillation framework improves small models across downstream tasks.
  • Results: TinyBERT4 exceeds 4-layer state-of-the-art distillation baselines by at least 4.4%, with approximately 28% of their parameters and 3.1x inference speedup.Against BERT4-PKD and DistilBERT4, the model combines higher performance with lower resource use.
  • Efficiency: TinyBERT4 is 7.5x smaller and 9.4x faster than BERTBASE while maintaining competitive GLUE performance.It also matches the average score of 24-layer MobileBERTTINY using only 38.7% FLOPs, although that comparison may not be fair.
  • Ablation studies: Removing any of general distillation, task-specific distillation, or data augmentation reduces performance, with task-specific procedures more helpful than general distillation.The ablation study finds comparable effects for task-specific distillation and data augmentation across four tasks.
  • Ablation studies: Removing Transformer-layer distillation drops performance from 75.6 to 56.3, while the proposed distillation objectives are complementary.The study attributes the drop to the importance of good student initialization during pretraining.
  • Ablation studies: Uniform layer mapping outperforms top- and bottom-layer strategies across tasks, while different tasks depend on knowledge from different BERT layers.Adaptive task-specific layer selection is left for future work.

5 Related Work

Related work places TinyBERT within broader efforts to compress pretrained language models, especially through knowledge distillation and layer-mapping strategies.

  • PLM compression: Pretrained language models can be compressed using low-rank approximation, weight sharing, knowledge distillation, pruning, or quantization.This paper focuses on knowledge distillation.
  • Layer mapping: Different TinyBERT layer-mapping strategies produce task-dependent outcomes, while uniform mapping performs better than top and bottom strategies across evaluated tasks.The paper leaves adaptive task-specific layer selection for future work.
  • Knowledge distillation for PLMs: Prior work distills task-specific BERT knowledge into BiLSTMs or smaller pretrained language models, including BERT-PKD and related methods.The cited approaches represent several existing directions for compressing pretrained language models.
  • Knowledge distillation for PLMs: TinyBERT4 is a 4-layer model that can achieve more speedup than some prior small pretrained language models.The comparison is stated specifically for the cited small PLM methods.

6 Conclusion and Future Work

The paper presents Transformer-specific distillation and a two-stage TinyBERT framework that reduces BERTBASE size and inference time while retaining competitive performance.

  • Conclusion: The proposed method combines Transformer-based distillation with general and task-specific distillation stages for TinyBERT.The framework performs distillation during both pretraining and fine-tuning.
  • Conclusion: TinyBERT achieves competitive performance while significantly reducing BERTBASE model size and inference time.The authors identify deployment on edge devices as an effective application of the approach.
  • Future work: Future work includes transferring knowledge from wider and deeper teachers and combining distillation with quantization or pruning.These directions target further compression and broader teacher architectures.

A More Comparisons on GLUE

Additional GLUE comparisons evaluate TinyBERT6 against prior methods under matched architecture and evaluation settings, with results reported on GLUE development tasks.

  • Comparison setup: Table 5 compares TinyBERT with other baselines on the GLUE development set using task-specific metrics including Matthews correlation, Pearson, and Spearman correlation.The table is intended to provide direct comparisons with prior reported results.
  • Results: TinyBERT6 outperforms all compared baselines under the same architecture and evaluation-method settings.The authors describe this as further confirmation of TinyBERT’s effectiveness.

B Results on SQuAD v1.1 and v2.0

TinyBERT is evaluated on SQuAD v1.1 and v2.0 as token-level sequence-labeling tasks. It consistently outperforms 4-layer and 6-layer baselines, although better QA-TinyBERT remains future work.

  • Task setup: SQuAD v1.1 and v2.0 are treated as sequence-labeling tasks that predict each token’s probability of being an answer-span start or end.Prediction-layer distillation uses the original training dataset rather than the augmented dataset used for GLUE tasks.
  • Results: TinyBERT consistently outperforms both 4-layer and 6-layer baselines on the question-answering tasks.The result indicates that the framework also works for token-level labeling.
  • Scope: The authors leave construction of a better QA-TinyBERT for future work because question answering depends on subtler knowledge than sequence-level GLUE tasks.This dependence increases the difficulty of knowledge distillation.

C Initializing TinyBERT with BERTTINY

The framework first performs General Distillation during pretraining, then uses the resulting TinyBERT to initialize Task-specific Distillation. Experiments compare this approach with direct BERTTINY initialization.

  • Two-stage framework: General Distillation transfers teacher knowledge from intermediate layers during pretraining before downstream Task-specific Distillation begins.The resulting general TinyBERT initializes the student for downstream task-specific learning.
  • Evaluation: Table 7 reports dev-set results for General Distillation and Task-specific Distillation, with TD denoting task-specific distillation without data augmentation.The table compares methods at the pretraining stage.
  • Initialization comparison: BERTTINY(+TD) performs worse than BERTTINY on MRPC and CoLA when task-specific distillation follows direct BERTTINY pretraining.The authors conjecture that mismatched intermediate representations are further disturbed by supervision from fine-tuned BERTBASE.
  • Ablation results: General Distillation achieves comparable results with BERTTINY, 61.1 versus 63.9, even without MLM and NSP pretraining tasks.Task-specific distillation then boosts TinyBERT by continuing to learn task-specific knowledge from fine-tuned BERTBASE.

D GLUE Details

The GLUE details section defines the benchmark tasks used to evaluate language understanding across entailment, question equivalence, similarity, and paraphrase identification.

  • Entailment and equivalence: MNLI is a crowd-sourced entailment classification task that labels a premise-hypothesis pair as entailment, contradiction, or neutral.The task uses pairs of premise and hypothesis sentences.
  • Entailment and equivalence: QQP consists of question pairs from Quora, with the goal of determining whether the questions are semantically equivalent.
  • Similarity and paraphrase: STS-B evaluates the similarity of sentence pairs from news headlines and other domains using scores from 1 to 5.
  • Similarity and paraphrase: MRPC is a paraphrase-identification dataset in which systems determine whether two sentences are paraphrases.
Loading 1909.10351v5…