[Paper Review] Reducing Transformer Depth on Demand with Structured Dropout

Angela Fan, Édouard Grave+1 more|arXiv (Cornell University)|Sep 25, 2019

Topic Modeling56 references273 citations

TL;DR

LayerDrop trains a single over-parameterized Transformer so that any sub-network depth can be extracted at inference without fine-tuning, enabling efficient on-demand models while maintaining strong performance.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for memory- and compute-efficient Transformer models in NLP tasks.
• Introduce a training regularization that makes sub-networks of varying depths robust without finetuning.
• Show that pruning to smaller depths yields competitive or state-of-the-art performance across benchmarks.

Proposed method

• Apply random structured dropout by dropping groups of weights that align with model structures (e.g., layers).
• Focus on dropping entire Transformer layers (LayerDrop) to enable on-demand depth at inference time.
• Describe pruning strategies (Every Other, Search on Valid, Data Driven) and prefer Every Other for simplicity and effectiveness.
• Provide a relationship for optimal pruning drop rate p* = 1 - r/N to target a desired depth.
• Train large Transformer models once; extract shallower sub-networks at test time without finetuning.

Experimental results

Research questions

• RQ1 Can LayerDrop regularize Transformers to be robust to layer-wise pruning at inference time?
• RQ2 How does on-demand depth via LayerDrop compare to training separate smaller models or distillation across NLP tasks?
• RQ3 What pruning strategies are effective for selecting which layers to keep when pruning?
• RQ4 Does LayerDrop enable state-of-the-art results across translation, language modeling, summarization, QA, and NLU benchmarks?

Key findings

• LayerDrop regularizes very deep Transformers, stabilizing training and achieving strong results on multiple NLP benchmarks.
• From one large pre-trained model, small, efficient sub-networks of any depth can be extracted at test time without finetuning.
• LayerDrop-enabled pruning often outperforms training small models from scratch and standard pruning without LayerDrop, across generation and pre-training tasks.
• Dropping entire layers is effective, with Every Other layer being a strong, simple pruning strategy across tasks.
• Pruning RoBERTa-like models with LayerDrop yields better results than BERT/RoBERTa trained from scratch or distillation in several settings, especially with more data.
• Training with larger LayerDrop improves performance when significant depth reduction is desired, aligning train-time and test-time conditions.