Source-linked AI summary
Reducing Transformer Depth on Demand with Structured Dropout
Angela Fan, Edouard Grave, Armand Joulin
TL;DR
Large Transformer models deliver strong NLP performance but impose substantial computation and memory costs. The paper introduces LayerDrop, which trains models with randomly dropped layer structures so sub-networks of arbitrary depth can be extracted without finetuning. Across multiple generation and language-understanding tasks, LayerDrop stabilizes deep-model training and supports strong smaller models, with a recommended higher drop rate for very small inference-time models.
Problem
Over-parameterized Transformers are computationally demanding, motivating automatic extraction of smaller sub-networks under memory or latency constraints without finetuning while maintaining good performance.
Method
LayerDrop randomly drops groups of weights, particularly entire layers, during training to make one over-parameterized network robust to structured pruning at inference time.
Results
LayerDrop achieves state-of-the-art or strong performance across diverse text-generation and pre-training benchmarks while enabling extraction of models at various depths and outperforming smaller models trained from scratch or reduced with BERT-specific strategies.
Takeaways & Limitations
A single large Transformer can provide multiple shallow, efficient models on demand without finetuning, while LayerDrop also stabilizes training of substantially deeper networks.
Takeaways & Limitations
The authors use a LayerDrop rate of p = 0.2 in experiments but recommend p = 0.5 when targeting very small inference-time models.
Abstract
from arXiv · showhide
Overparameterized transformer networks have obtained state of the art results in various natural language processing tasks, such as machine translation, language modeling, and question answering. These models contain hundreds of millions of parameters, necessitating a large amount of computation and making them prone to overfitting. In this work, we explore LayerDrop, a form of structured dropout, which has a regularization effect during training and allows for efficient pruning at inference time. In particular, we show that it is possible to select sub-networks of any depth from one large network without having to finetune them and with limited impact on performance. We demonstrate the effectiveness of our approach by improving the state of the art on machine translation, language modeling, summarization, question answering, and language understanding benchmarks. Moreover, we show that our approach leads to small BERT-like models of higher quality compared to training from scratch or using distillation.
1 INTRODUCTION
LayerDrop trains over-parameterized Transformers to remain robust when layers are dropped, enabling shallow sub-networks of different depths to be extracted from one model. The approach regularizes deep networks and supports competitive performance without finetuning extracted models.
- Transformer layers contain millions of parameters, creating substantial memory and latency demands during training and inference.
- LayerDrop randomly drops layers during training so the network becomes robust to pruning at inference time.Dropping entire layers targets efficient shallow models while preserving the original network structure.
- One trained network can yield shallow models of any desired depth on demand, without finetuning.This avoids training a separate model from scratch for each model size.
- The method achieves competitive performance for smaller extracted models and outperforms training smaller models from scratch and dedicated BERT reduction strategies.
- LayerDrop regularizes very deep Transformers and stabilizes training, contributing to state-of-the-art performance across varied benchmarks.
2 RELATED WORK
The paper situates LayerDrop within structured dropout and structured pruning, distinguishing its goal of learning networks that tolerate pruning from approaches that directly learn pruning decisions. It also differs from related compression and adaptive-memory methods by training Transformers for pruning resilience.
- LayerDrop applies dropout to groups of weights, especially entire layers, to induce redundancy for pruning shallow efficient models.
- Unlike Stochastic Depth, which targets faster training of deep ResNets, LayerDrop targets pruning and extraction of shallow Transformer sub-networks.
- Structured pruning removes coherent weight groups while preserving network structure, distinguishing it from unstructured weight pruning.
- Unlike distillation and related compression methods, LayerDrop trains models to be resilient to pruning, reducing performance loss from test-time pruning.
- Trainable mechanisms for controlling memory footprint are complementary to LayerDrop and may benefit from combination with it.
3 METHOD
The method trains Transformers with structured dropout, especially by dropping layers, so subnetworks can be pruned efficiently at inference without further finetuning.
- Structured dropout: LayerDrop targets efficient pruning because dropping coherent groups can reduce inference time, unlike unstructured DropConnect, which produces smaller but not necessarily efficient models.Attention heads may be computed in parallel, so dropping them does not necessarily reduce runtime.
- Structured dropout: Structured dropout constrains the mask to be constant over predefined groups of weights, including heads, feedforward matrices, or layers.This general formulation can target overlapping structures while preserving coherent network components.
- Structured dropout: LayerDrop drops entire layers during training to make Transformer subnetworks robust to subsequent structured pruning.It applies dropout-like random masking to groups of weights, focusing on layers because some structures provide greater inference-time efficiency than others.
- Inference-time pruning: Every Other, Search on Valid, and Data Driven Pruning are considered strategies for selecting layers to prune after LayerDrop training.Every Other is intuitive and balanced; Search on Valid is computationally intensive, while Data Driven Pruning learns individual layer drop rates.
- Inference-time pruning: Every Other works surprisingly well across many tasks, while Search on Valid and Data Driven Pruning provide only marginal gains, without further finetuning of pruned networks.This supports extracting shallower models directly from the trained network.
- Drop-rate selection: For N groups and fixed drop ratio p, training uses N(1 −p) groups on average; targeting r groups determines the optimal drop rate.The experiments use p = 0.2, while p = 0.5 is recommended for very small inference-time models.
4 EXPERIMENTAL SETUP
The experiments evaluate LayerDrop across machine translation, language modeling, summarization, long-form question answering, and sentence representation using established datasets and Transformer or BERT architectures.
- Tasks and implementation: The evaluation covers neural machine translation, language modeling, summarization, long-form question answering, and natural language understanding.Models are implemented in PyTorch using fairseq-py.
- Machine translation: WMT English-German uses 4.5M sentence pairs for training, newstest2013 for validation, and newstest2014 for testing with a Transformer Big architecture.The LayerDrop rate is p = 0.2, and decoding follows checkpoint averaging, length-penalty, and beam-search settings.
- Language modeling: Wikitext-103 language modeling uses 100M tokens, a 260K vocabulary, and a 16-layer Transformer, reporting test perplexity.The LayerDrop rate is p = 0.2, with standard dropout tuned on validation data.
- Summarization: CNN-Dailymail summarization uses over 280K news articles paired with multi-sentence summaries and evaluates generated text with ROUGE.The setup uses a Transformer base architecture, p = 0.2 LayerDrop, generation-length tuning, and 3-gram blocking.
- Long-form question answering: ELI5 contains 272K question-answer pairs with supporting web documents and evaluates long answers with ROUGE.The experiment uses Transformer Big, beam size 5, and 3-gram blocking.
- Sentence representation: BERT base and large models are pretrained on Bookscorpus plus Wiki or a larger corpus combination, then evaluated on multiple language-understanding tasks.The reported evaluation includes accuracy on MRPC and QNLI among the listed benchmarks.
5 RESULTS
LayerDrop improves Transformer performance across language modeling, sequence-to-sequence tasks, and sentence representation, while enabling strong shallower models from a single trained network.
- Language Modeling: 0.6 perplexity points: a 40-layer Transformer with LayerDrop improves the state of the art on WikiText-103.A 16-layer Transformer also improves by 0.4 perplexity and matches Transformer-XL.
- Sequence to sequence modeling: 30.2 BLEU: a 12-encoder-layer Transformer with LayerDrop improves the state of the art on WMT14 English-German.A standard Transformer without LayerDrop diverges at 12 encoder layers.
- Sequence to sequence modeling: LayerDrop boosts performance on machine translation, summarization, and long-form question answering tasks.The experiments apply LayerDrop to state-of-the-art Transformer architectures.
- Bi-Directional Pre-training: LayerDrop improves RoBERTa performance on several natural language understanding tasks and supports models double the size for stronger performance.The comparisons include training on additional data and longer pre-training.
- Pruning: One LayerDrop-trained network can be pruned to a range of depths without fine-tuning and outperforms small models trained from scratch.Pruning a network trained without LayerDrop performs poorly because it was not trained to tolerate missing layers.
- Pruning BERT-like Models: LayerDrop-pruned BERT-like models outperform BERT and RoBERTa models trained from scratch, with stronger results when additional data is used.The comparison includes DistilBERT and uses every-other-layer pruning followed by downstream-task fine-tuning.
6 ABLATION STUDIES
Ablations show that dropping entire layers is effective, every-other-layer pruning is difficult to beat, input and output layers are especially important, and larger training-time LayerDrop benefits deep pruning.
- Comparison of Structured Dropout: Dropping entire layers has strong performance, while dropping attention heads is worse and does not improve runtime because heads are computed in parallel.Dropping sub-layers and layers shows no large difference in the tested relatively shallow networks.
- Comparison of Various Pruning Strategies: Every-other-layer pruning is tough to beat; validation-set layer search and learned dropping provide only marginal gains.Removing consecutive chunks, especially the first or last half, is harmful.
- Choosing which Layers to Prune: The input and output layers are the most important layers in Wikitext-103 pruning experiments.Importance is estimated from perplexity after removing each layer across random 8-layer selections.
- Relationship between LayerDrop at Training Time and Pruning at Inference Time: Larger training-time LayerDrop benefits inference-time pruning when substantial depth reduction is desired.The training and test settings become more closely matched under larger LayerDrop.
7 CONCLUSION
LayerDrop regularizes Transformers for structured pruning, enabling efficient subnetworks of varied depths while supporting deeper-network training and strong performance across text tasks.
- 7 CONCLUSION: LayerDrop regularizes Transformers to withstand layer pruning at inference time and extract shallow, efficient models of varied depths.The paper focuses on layers as the prunable structures.
- 7 CONCLUSION: Across text generation and pre-training tasks, LayerDrop enables and stabilizes substantially deeper networks while preserving strong performance in extracted subnetworks.The conclusion connects training deeper models with extracting models at different depths.
A.1.1 NEURAL MACHINE TRANSLATION
The neural machine translation setup uses task-specific vocabulary and optimization choices, while pruning removes every other decoder layer without fine-tuning.
- NEURAL MACHINE TRANSLATION: WMT English-German uses a 32K joint byte-pair encoding, cosine learning-rate scheduling, and 0.1 label smoothing.The schedule follows Wu et al. (2019).
- NEURAL MACHINE TRANSLATION: IWSLT German-English uses 160K lowercased training pairs, a 10K joint BPE vocabulary, beam size 4, and a 6-by-6 Transformer base model.The smaller dataset setting uses FFN size 1024, hidden dimension 512, and four attention heads.
- NEURAL MACHINE TRANSLATION: Inference pruning applies the every-other-layer strategy to the decoder without fine-tuning.This is the specified pruning procedure for the translation setup.
A.1.2 LANGUAGE MODELING
The experiments combine LayerDrop training with structured layer pruning across language-modeling and related benchmark settings. Results indicate faster training and pruning to shallower models, while additional finetuning often offers limited gains.
- Pruning setup: Pruning applies the Every Other Layer strategy, with some configurations evaluated without finetuning.The strategy is used for language-modeling and decoder pruning, while RoBERTa pruning uses task-specific finetuning for smaller models.
- Training setup: LayerDrop experiments use varied Transformer and RoBERTa configurations, including 12-, 24-, and 48-layer architectures.The 48-layer model uses LayerDrop 0.5, leaving 24 layers active on average during each forward pass.
- Training-time efficiency: Almost 2x training speed is achieved when LayerDrop removes half the layers from a fixed-size 16-layer model.The measurement uses words per second on 8 V100 GPUs for Wikitext-103 training.
- Inference-time pruning: Higher LayerDrop rates permit more aggressive inference-time pruning across MNLI, SST2, and QNLI, although downstream finetuning makes the relationship less straightforward.The comparison is reported in Figure 8 for natural language understanding tasks.
- Finetuning: Additional finetuning produces only marginal improvements after pruning a 16-layer language model to 8 layers.The finetuning parameters depend on the model depth at test time and are difficult to optimize.