Source-linked AI summary
Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers
Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, Joseph E. Gonzalez
TL;DR
Compute constraints make it important to identify how Transformer size affects training and inference efficiency. The paper evaluates model scaling across self-supervised pretraining and machine translation, showing that large models can train efficiently and remain advantageous after compression. It concludes that training large models briefly and then compressing them can outperform smaller alternatives under resource constraints.
Problem
Compute-limited Transformer training must balance accuracy with constrained training time, hardware, and inference resources.
Method
The paper varies Transformer model size across self-supervised pretraining and machine translation, measuring convergence with wall-clock time and evaluating compression for inference.
Results
Large Transformers converge faster in gradient steps and wall-clock time, are more robust to compression, and can outperform small models at comparable inference costs.
Takeaways & Limitations
For resource-constrained training, the supported strategy is to train large models and then heavily compress them.
Takeaways & Limitations
The training strategy assumes initial hyperparameters and model settings are available; the cost of searching architectures and hyperparameters remains for future study.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper reframes compute-efficient Transformer training: larger models can converge faster in wall-clock time when stopped early, then be compressed for inference. This strategy addresses both training-resource limits and the inference costs of large models.
- Motivation: Maximizing accuracy under fixed hardware and training time requires rethinking the assumption that models must train until convergence.The paper frames compute efficiency as achieving the highest model accuracy given constrained hardware and training time.
- Training efficiency: Larger Transformer models converge in fewer gradient updates, and this acceleration typically exceeds their additional computational overhead.The experiments cover self-supervised pretraining and WMT14 English→French machine translation.
- Training efficiency: Convergence acceleration is largely determined by parameter count and is only weakly influenced by model width, depth, and batch size.
- Inference efficiency: Large models are more robust to compression, allowing heavily compressed models to outperform lightly compressed small models at comparable inference costs.The paper discusses compression in the context of increased inference computation and memory requirements for larger models.
- Scaling conditions: The optimal model size is closely linked to dataset size, with large models performing favorably when overfitting is limited.The paper attributes greater compressibility to reduced weight error under quantization or sparsity as model size increases.
2. Experimental Setup
The experiments evaluate Transformer scaling for self-supervised pretraining and high-resource machine translation using task-specific datasets and standard metrics. Training efficiency is measured directly with wall-clock time rather than only gradient steps or FLOPs.
- Tasks: The study evaluates model-size effects on masked-language-modeling pretraining and high-resource machine translation.The pretraining task uses ROBERTA-style models, while translation uses the standard Transformer on WMT14 English→French.
- Pretraining setup: ROBERTA pretraining uses BookCorpus and Wikipedia totaling roughly 3.4 billion words, with 0.5% held out for validation.Validation performance is reported as masked language modeling perplexity.
- Translation setup: The machine-translation setup uses 36M WMT14 English→French training sentences, with newstest2013 for validation and newstest2014 for testing.The reported metric is tokenized case-sensitive BLEU with compound splitting.
- Evaluation: Wall-clock time is the main training metric because gradient steps omit batch and model costs, while FLOPs overlook hardware differences between parallel and sequential operations.Timing is benchmarked on one NVIDIA 16GB V100 GPU using stabilized averages over 100 steps.
3. Larger Models Train Faster
Larger Transformer models reach comparable performance in fewer gradient steps, and their faster convergence can outweigh their higher per-step computational cost. Model size is generally more effective than batch-size scaling after the critical batch-size region, though the largest models show diminishing wall-clock returns.
- Convergence: Larger Transformer models reach the same performance using fewer gradient steps, and their faster convergence can outweigh the computational overhead of increased model size.This makes larger models faster to train after adjusting for wall-clock time.
- Convergence: Deeper RoBERTA models converge faster than shallow models with respect to gradient steps.
- Scaling choices: For masked language modeling, convergence depends more on total parameter count than model shape, while machine translation generally favors increasing width.Increasing either width or depth accelerates masked-language-model training, whereas wider translation models usually achieve better final performance than deeper ones.
- Scaling choices: Beyond a certain point, larger batch sizes provide only marginal wall-clock improvements, so training should shift toward increasing model size after reaching the critical batch-size region.Larger batches reduce gradient steps, but their wall-clock benefit saturates; batch-size scaling also requires learning-rate tuning.
- Convergence: Wider models converge faster than narrower models with respect to both gradient steps and wall-clock time.
- Limits: The largest models eventually reach diminishing wall-clock returns because per-step convergence gains decrease while computational overhead rises.For machine translation, the 6L, 2048H model converges more slowly than the 6L, 1024H model with respect to wall-clock time.
4. Larger Models Compress Better
Large Transformer models can be trained efficiently and then compressed to achieve stronger accuracy-efficiency trade-offs than small models. Quantization and pruning make large models more suitable for inference under comparable memory or parameter budgets.
- Evaluation setup: The compression evaluation trains ROBERTA models for the same wall-clock budget, finetunes them on MNLI or SST-2, and then applies quantization or pruning.Larger models are trained for fewer pretraining steps under this setup.
- Compression methodology: Quantization is applied post-hoc across parameter precisions from 4 to 32 bits, while pruning iteratively removes small-magnitude weights and finetunes to recover accuracy.The pruning procedure targets specified sparsity levels and typically requires little additional training.
- Quantization: Large models are more robust to quantization, so heavily quantized large models achieve the highest accuracy across almost all memory budgets.The evaluation compares accuracy at different quantization levels as a function of total memory usage.
- Pruning: Larger models can be pruned more aggressively without significantly hurting accuracy, producing the best accuracy-efficiency trade-off.Pruning results are evaluated by total nonzero parameters.
- Combined compression: Pruning and quantization are complementary, with large heavily compressed models remaining near the Pareto frontier; pruning 30–40% followed by 6–8-bit quantization is particularly strong.
5. When and Why Are Larger Models Better?
Larger Transformers train faster because they converge more quickly, but their compression robustness—not pretraining convergence—explains why they can also meet inference constraints. These advantages weaken when data is scarce and overfitting becomes important.
- Why larger models compress better: Large models remain the best-performing option at each test-time budget when pruning and quantization are combined.The two compression gains are complementary.
- Why larger models train faster: Larger models converge faster because they minimize training error quickly and face small generalization gaps on very large datasets.In large pretraining corpora, fitting rather than overfitting is the main challenge.
- Dataset-size boundary: When overfitting matters, the largest models can lose their convergence advantage: with 1% of the data, they have worse perplexity than smaller models.
- Why larger models train faster: Larger models use available parallel hardware more fully, making their per-iteration slowdown sublinear; gradient accumulation can trade memory for compute.
- Why larger models compress better: Quantization robustness depends mainly on model size rather than pretraining convergence: smaller models suffer much larger MNLI accuracy drops at 4-bit precision.
- Why larger models compress better: Quantization error and pruning error have lower mean and variance in deeper models, indicating weights that are easier to approximate with low-precision or sparse matrices.
6. Related Work
Related work improves training through optimizers, losses, architectures, storage, schedules, and distributed systems. This paper instead selects model size for a fixed small hardware budget, while noting open questions about larger budgets and tuning costs.
- Training efficiency: Prior work accelerates training through improved optimizers, loss functions, model structures, sparsity, storage requirements, and learning-rate schedules.
- Scaling model training: Distributed-training research scales computation across large hardware systems, whereas this work chooses model size for a fixed small hardware budget.
- Hyperparameter tuning and AutoML: The optimal training strategy may need to account for the cost of experimenting with architectures and hyperparameters, which this work does not optimize.
- Training efficiency: Concurrent work similarly finds that large, undertrained Transformer models can outperform small, well-trained models, while this paper additionally studies machine translation, downstream finetuning, and compression.
7. Conclusion and Future Work
The paper finds that increasing Transformer width and depth accelerates convergence in gradient steps and wall-clock time, while larger models remain viable for inference after heavier compression. It concludes that resource-constrained training should train large models and then compress them, while leaving broader-domain validation and explanatory questions for future work.
- Conclusion: Increasing Transformer width and depth accelerates convergence in both gradient steps and wall-clock time.
- Conclusion: Although large models are less efficient during inference, their greater compression robustness supports training large models and then heavily compressing them.
- Future work: Future work will test whether these conclusions extend to domains such as computer vision.
- Future work: Future work will examine why larger Transformers train quickly and compress well, and how model size affects overfitting and hyperparameter tuning.
A.3. The Impact of Dataset Size
This appendix section identifies learning curves for models trained with 5% and 1% of the original training data.
- Dataset-size experiments: Learning curves are shown for models trained using 5% and 1% of the training data.
B. Finetuning Models of Different Sizes
Larger ROBERTA models are not harder to finetune: when pretrained to comparable perplexity, they achieve comparable or often better downstream accuracy.
- Models with more parameters can match or outperform smaller models after finetuning.
- Comparable downstream accuracies show that larger ROBERTA models are not harder to finetune.The comparison covers finetuning on MNLI and SST-2 after models are stopped at roughly the same pretraining perplexity.
C. Negative Results: Layer Sharing
Layer sharing reduces memory use but slows convergence and worsens final performance, making it unsuitable for compute-efficient Transformer training or inference.
- Sharing layers can reduce memory movement and inter-device communication, slightly accelerating execution.These execution benefits accompany reduced memory consumption but do not offset the convergence penalty in the reported setting.
- 41% lower total memory usage results when all ROBERTA layers are shared at batch size 32.The memory reduction is especially pronounced for small batch sizes.
- Both all-layer and attention-only sharing slow training convergence and worsen performance in resource-constrained settings.Figure 11 attributes slower convergence and worse final accuracy to sharing attention layers.
- The authors do not recommend layer sharing for compute-efficient Transformer training or inference.
D. Compression Results for SST-2
Compression results on SST-2 favor models trained large and then heavily compressed, although dataset size and batch-size scaling affect the efficiency comparison.
- SST-2 finetuning requires tuning batch size, learning rate, seed, and dropout because the dataset is smaller than MNLI.
- Bigger models achieve better results than smaller models at the same number of floating-point operations.The comparison is reported in exaFLOPs rather than wall-clock time.
- For most compression budgets, the highest-accuracy SST-2 models are trained large and then heavily compressed.The result holds across quantization, pruning, and their combination.
- Beyond batch size 2048, increasing batch size provides only marginal wall-clock improvements.Larger batches accelerate convergence in gradient steps, but gradient accumulation costs are included in wall-clock time.
- With 5% of the data, the biggest models do not improve speed over smaller models, while at 1% they have worse perplexity due to overfitting.On the full dataset, the largest trained models are faster in wall-clock time.