[Paper Review] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers

Zhuohan Li, Eric Wallace+5 more|arXiv (Cornell University)|Feb 26, 2020

Topic Modeling65 references51 citations

TL;DR

The paper shows that training very large Transformer models and stopping early can be more compute-efficient than training small models to convergence, and that such large models are more compressible, enabling better accuracy at similar inference costs after quantization/pruning.

ABSTRACT

Since hardware resources are limited, the objective of training deep learning models is typically to maximize accuracy subject to the time and memory constraints of training and inference. We study the impact of model size in this setting, focusing on Transformer models for NLP tasks that are limited by compute: self-supervised pretraining and high-resource machine translation. We first show that even though smaller Transformer models execute faster per iteration, wider and deeper models converge in significantly fewer steps. Moreover, this acceleration in convergence typically outpaces the additional computational overhead of using larger models. Therefore, the most compute-efficient training strategy is to counterintuitively train extremely large models but stop after a small number of iterations. This leads to an apparent trade-off between the training efficiency of large Transformer models and the inference efficiency of small Transformer models. However, we show that large models are more robust to compression techniques such as quantization and pruning than small models. Consequently, one can get the best of both worlds: heavily compressed, large models achieve higher accuracy than lightly compressed, small models.

Motivation & Objective

• Investigate how Transformer model size affects compute efficiency under fixed hardware and time constraints.
• Evaluate whether larger models trained briefly can outperform smaller models trained to convergence in downstream tasks.
• Examine robustness of large models to compression techniques like quantization and pruning.
• Analyze when and why larger models offer better trade-offs between training speed and inference efficiency.

Proposed method

• Systematically vary Transformer width and depth for RoBERTa-style MLM pretraining and WMT14 En→Fr translation.
• Measure wall-clock time and gradient steps to convergence across model sizes.
• Apply post-hoc compression (quantization and pruning) to finetuned large models and assess accuracy vs memory.
• Quantize weights to various bitwidths and prune weights iteratively to different sparsity levels and evaluate impact on downstream tasks.
• Compare compression impact across model sizes to determine robustness of large models.

Experimental results

Research questions

• RQ1 Does increasing Transformer size accelerate convergence in terms of wall-clock time and gradient steps across MLM pretraining and MT tasks?
• RQ2 Are larger models more compressible via quantization and pruning than smaller models when finetuned on downstream tasks?
• RQ3 What is the trade-off between training larger models and subsequent compression in terms of accuracy under fixed inference budgets?
• RQ4 How does dataset size and potential overfitting influence the benefits of large versus small models?
• RQ5 Why do larger models exhibit better compressibility and how can this be exploited in practice?

Key findings

• Larger models converge to lower validation error in fewer gradient updates and faster in wall-clock time after accounting for overhead.
• For MLM and MT, bigger models provide faster convergence per wall-clock time, with gains exceeding the overhead of larger size.
• Larger models are more robust to both quantization and pruning than smaller models, achieving higher accuracy at similar memory budgets after compression.
• Combining pruning and quantization on large models yields Pareto-optimal accuracy-efficiency trade-offs, with particular strength in pruning around moderate sparsity followed by low-bit quantization.
• The optimal model size is linked to dataset size, with large models performing well in big data settings where overfitting is less of a concern.