Model growth and recursion improve transformer training efficiency
How Model Growth, Recursion, and Boundary Operators Influence Scaling Exponents
Machine Learning
Summary
Scaling laws explain how adding more computing power usually makes AI models better at tasks. This paper shows that changing the model’s design, like looping parts of it or growing it differently, can make these improvements happen much faster than expected. The authors found that using looping (making the model process data multiple times) or special connections between layers can lead to big gains in how efficiently the model learns. In some cases, a smaller model can do as well as a much larger one but with far less computing. These tricks help models use their ‘depth’ better to get more done with the same money spent on computing.
What this means in practice
- •For machine learning engineers: Train large transformer models more efficiently by adjusting recursion and model growth to reduce computing costs significantly.
- •For cloud infrastructure teams: Optimize resource allocation for AI training jobs by adopting architectures with looping and boundary operators to improve compute utilization.
Authors
Zixi Chen, Akshay Vegesna, Samip Dahal, Andrew Gordon Wilson
Abstract
Scaling laws predict how loss decreases with increases in computation. We show, contrary to conventional wisdom, that architectural interventions can modify scaling exponents in pre-training, leading to exponential improvements in performance with increases in computation. As an anchoring point, we consider the architectural formulation of looped transformers. Although not typically used in this way, looping, also known as recursive depth, provides a mechanism for model growth, by increasing the number of loops during training. Model growth, with and without shared weights, provides the biggest changes to the scaling exponents. In particular, a 7.4B model growth architecture matches GPT-3 13B on CORE with roughly $20\times$ less compute, and has compute efficiency gains that increase with scale. Moreover, simply using a boundary operator in a vanilla transformer, which normalizes and injects an earlier block, also provides increasing compute-efficiency gains, although to a lesser extent. In the data-constrained, multi-epoch setting, standard looping has a useful regularizing effect, where we find it is compute-optimal to increase the number of loops with scale. These results can be understood through the lens of computational depth: for a given computational budget, we wish to increase the usable depth of the transformer, which can lead to efficiency gains that increase with scale.