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

By Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Daniel Klein, Joseph E. Gonzalez. 2020

157 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Large Language Models, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Large Language Models

In Efficiency & Compression

Citations in the corpus

Pretrained Transformers Improve Out-of-Distribution Robustness (2020, methodology): “This highlights that testing model compression methods for BERT (Shen et al., 2020; Ganesh et al., 2020; Li et al., 2020) on only in-distribution examples gives”
Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers (2021, background): “Given that every moment spent on hardware accelerators is a cost incurring activity, we believe that research in distilling practical scaling insights and recommendations to be highly crucial (Li et al., 2020; Kaplan et al., 2020; Bello et al., 2021).”
On the Predictability of Pruning Across Scales (2020, positive): “This analytic result parallels and extends the findings of Li et al. (2020) on NLP tasks.”
Feature Learning in Infinite-Width Neural Networks (2020, background): “(Li et al., 2020a) proved a complexity separation between NTK and finite-width networks by showing the latter approximates a sort of infinite-width feature learning network.”
Pre-training via Paraphrasing (2020, background): “Recent work showed that large models train more efﬁciently [Li et al., 2020, Kaplan et al., 2020].”
Movement Pruning: Adaptive Sparsity by Fine-Tuning (2020, background): “As shown in previous works [Li et al., 2020, Han et al., 2016] quantization and pruning are complimentary and can be combined to further improve the performance/size ratio.”
General Purpose Text Embeddings from Pre-trained Language Models for Scalable Inference (2020, background): “Recent trends show that scaling models up (Liu et al., 2019b; Lan et al., 2019; Raf-fel et al., 2019; Li et al., 2020) in terms of computation still improves end task performance, raising ∗ Equal contribution. questions about whether and how the most accurate models can be applied in real-world…”
Vokenization: Improving Language Understanding via Contextualized, Visually-Grounded Supervision (2020, result): “Since some smaller tasks in GLUE are reported as unstable (Dodge et al., 2020), recent papers (e.g., Li et al. (2020b)) only report on selected tasks.”
Block Pruning For Faster Transformers (2021, result): “This result is in line with Li et al. (2020): the larger the model, the more pruning is effective.”
Reformulating Unsupervised Style Transfer as Paraphrase Generation (2020, methodology): “The models are ﬁnetuned using a small learning rate of 5e-5 and converge to a good solu-tion fairly quickly as noticed by recent work (Li et al., 2020; Kaplan et al., 2020).”
Data Movement Is All You Need: A Case Study on Optimizing Transformers (2020, methodology): “provided empirical recommendations on model size, batch size, and so on [23], [73], [74].”