Source-linked AI summary
Distilling Task-Specific Knowledge from BERT into Simple Neural Networks
Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, Jimmy Lin
TL;DR
Large pretrained language models are difficult to deploy because of their size and inference cost, motivating smaller alternatives. The paper distills BERT into single-layer BiLSTM models, using rule-based augmentation for NLP transfer data, and finds ELMo-comparable results with substantially lower resource use. The evidence also shows that the distilled models remain below deep transformer performance.
Problem
Large models such as BERT and GPT-2 pose deployment challenges in resource-restricted and real-time systems because of parameter count and inference inefficiency.
Method
The paper distills task-specific knowledge from fine-tuned BERT into a single-layer BiLSTM and a siamese BiLSTM for sentence-pair tasks, using rule-based textual augmentation.
Results
Across sentence classification and matching tasks, the distilled model achieves results comparable to ELMo while using around 100 times fewer parameters and 15 times faster inference.
Takeaways & Limitations
The results suggest that shallow BiLSTMs have greater representation power for natural-language tasks than previously thought.
Takeaways & Limitations
The distilled model does not outperform deep transformer models, performing 4–7 points worse on average.
Abstract
from arXiv · showhide
In the natural language processing literature, neural networks are becoming increasingly deeper and complex. The recent poster child of this trend is the deep language representation model, which includes BERT, ELMo, and GPT. These developments have led to the conviction that previous-generation, shallower neural networks for language understanding are obsolete. In this paper, however, we demonstrate that rudimentary, lightweight neural networks can still be made competitive without architecture changes, external training data, or additional input features. We propose to distill knowledge from BERT, a state-of-the-art language representation model, into a single-layer BiLSTM, as well as its siamese counterpart for sentence-pair tasks. Across multiple datasets in paraphrasing, natural language inference, and sentiment classification, we achieve comparable results with ELMo, while using roughly 100 times fewer parameters and 15 times less inference time.
1 Introduction
Large neural language models improve task performance but create deployment challenges because of their size and inference cost. This paper transfers BERT’s task-specific knowledge to shallow BiLSTMs, using rule-based augmentation to support distillation in NLP.
- Motivation: Large models such as BERT and GPT-2 can be undeployable on resource-restricted devices because of their parameter counts and inference inefficiency.The paper also motivates compression through slowing hardware scaling.
- Approach: The proposed approach distills task-specific knowledge from BERT into a shallow bidirectional LSTM.Knowledge distillation uses a larger teacher model and a smaller student model, and can transfer knowledge across architectures.
- Approach: Rule-based textual augmentation constructs an unlabeled transfer set because obtaining additional task-specific NLP samples is difficult.The augmented samples are not fluent natural-language sentences, but experiments show they work well for distillation.
- Results: Across sentence classification and matching tasks, distillation significantly outperforms training the original simpler network alone.The evaluation covers three tasks in sentence classification and sentence matching.
- Results: The shallow BiLSTM achieves results comparable to ELMo while using around 100 times fewer parameters and 15 times faster inference.The paper presents it as a state-of-the-art small model for neural NLP.
2 Related Work
Earlier NLP architectures were trained on task-specific data, while contextual language models introduced broad pretrained representations and strong task performance. Knowledge distillation offers a way to transfer knowledge from large models to smaller students with different architectures.
- Earlier NLP architectures: Generic CNN, recurrent, and recursive architectures support sentence classification and matching but are trained only on data from a particular task.These architectures represent earlier task-specific approaches in NLP.
- Pretrained representations: ELMo and BERT learn deep contextualized or language representations from large-scale language-modeling data and improve performance across multiple NLP tasks.BERT is reported to achieve state-of-the-art results on eleven NLP tasks.
- Model compression: Model-compression methods reduce model size, but irregular weight sparsity can prevent highly optimized computation routines.The cited compression pipeline achieved a 40-times reduction in model size without hurting accuracy.
- Knowledge distillation: Knowledge distillation transfers knowledge from a large teacher to a smaller student at the output level, allowing the student to use a different architecture.This flexibility supports studying whether shallow neural networks retain representation power for language understanding.
3 Our Approach
The approach transfers BERT’s task-specific output knowledge to minimally engineered BiLSTM students through logits-based distillation and rule-based augmentation. It covers single-sentence and sentence-pair classification while using synthetic unlabeled examples to improve transfer.
- Distillation procedure: The distillation procedure combines a logits-regression objective with an augmented transfer dataset.The student is trained with the distillation objective alongside cross-entropy against labeled or teacher-predicted targets.
- Model architecture: BERT is the teacher, while the student is a single-layer BiLSTM with a ReLU fully connected classifier and softmax output.For sentence pairs, the student uses shared BiLSTM encoder weights and a concatenate–compare operation before classification.
- Distillation objective: MSE penalizes differences between student and teacher logits, while cross-entropy trains against one-hot ground-truth or teacher-predicted labels.The authors report that MSE performed slightly better than cross-entropy with soft targets in preliminary experiments.
- Data augmentation: The augmentation method masks words, replaces words with same-part-of-speech alternatives, and samples n-grams from training examples.Masking reduces teacher confidence around altered words, POS replacement perturbs semantics, and n-gram sampling drops other words more aggressively.
- Data augmentation: Synthetic examples are generated by independently sampling word-level operations, optionally applying n-gram sampling, and appending results to the unlabeled dataset.The procedure is repeated per example with duplicate samples discarded; sentence pairs are augmented by cycling through each sentence and both sentences.
4 Experimental Setup
The experiments use fine-tuned BERT logits and synthesized examples to train distilled BiLSTMs, evaluated on representative GLUE tasks with specified model settings.
- Teacher and distillation setup: BERTLARGE is fine-tuned without data augmentation, then used to provide logits for the original and synthesized training examples.Four learning rates, {2, 3, 4, 5} × 10^-5, are evaluated with Adam, selecting the best validation model.
- Datasets: The study evaluates one dataset from each of GLUE’s three categories: SST-2, MNLI, and QQP.These represent single-sentence classification, inference, and similarity/paraphrase tasks, respectively.
- Student model and optimization: The BiLSTM uses 150 or 300 hidden units and a ReLU layer with 200 or 400 units, selected by validation performance.The models use 300-dimensional word2vec and multichannel embeddings.
- Student model and optimization: AdaDelta with learning rate 1.0 and ρ = 0.95 is used for optimization, with batch sizes of 50 for SST-2 and 256 for MNLI and QQP.Augmentation probabilities are fixed across datasets, with 20 iterations for SST-2 and 10 for MNLI and QQP.
- Baselines: Baseline comparisons include BERT, GPT, ELMo-based BiLSTM models, and previously reported BiLSTM systems.BERT is described in base and large variants, while GPT is unidirectional and uses previous context at each time step.
5 Results and Discussion
Distillation improves the shallow BiLSTM across the evaluated tasks and yields competitive quality with substantially better efficiency, though deep transformers remain more accurate.
- Model Quality: 1.9–4.5 points of absolute improvement are achieved against the base BiLSTM through logit matching on augmented training data.On SST-2 and QQP, the distilled model outperforms the best reported ELMo model; on MNLI, it trails ELMo by a few points.
- Model Quality: The distilled model is competitive with two prior ELMo BiLSTM implementations, while nondistilled BiLSTM baselines do not outperform BERT’s results.The base implementation obtains scores comparable with BiLSTMs from previous work.
- Model Quality: 4–7 points worse on average than deep transformer models remain the distilled model’s quality gap.The model nevertheless uses far fewer parameters and has better efficiency.
- Inference Efficiency: 98 and 349 times fewer parameters than ELMo and BERTLARGE, respectively, are used by the single-sentence model.The same analysis reports that it is 15 and 434 times faster, using inference on SST-2 with batch size 512 on one NVIDIA V100 GPU.
- Inference Efficiency: The 2.2-million-parameter 300-dimensional-LSTM variant is twice as large as the smaller variant but remains substantially smaller than ELMo.For sentence-pair tasks, the siamese counterpart uses no pairwise word interactions, so runtime scales linearly with sentence length.
6 Conclusion and Future Work
The paper concludes that BERT distillation can make simple BiLSTM models competitive with ELMo while using fewer parameters and less inference time, and proposes exploring simpler or slightly richer architectures.
- Conclusion: Distilling BERT knowledge into a simple BiLSTM produces ELMo-comparable results with much fewer parameters and less inference time.The conclusion presents shallow BiLSTMs as more expressive for natural language tasks than previously thought.
- Future Work: Future work includes testing extremely simple architectures such as convolutional networks, support vector machines, and logistic regression.Another direction is slightly more complex architectures using pairwise word interaction and attention.