[Paper Review] Navigating Scaling Laws: Compute Optimality in Adaptive Model Training

Sotiris Anagnostidis, Gregor Bachmann+2 more|arXiv (Cornell University)|Nov 6, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes adaptive training strategies for Vision Transformers that dynamically adjust model shape—specifically patch size and width—during training to achieve compute-optimal performance. By scheduling these parameters based on neural scaling laws, the method reduces required FLOPs by over 50% compared to static models while maintaining or improving accuracy.

ABSTRACT

In recent years, the state-of-the-art in deep learning has been dominated by very large models that have been pre-trained on vast amounts of data. The paradigm is very simple: investing more computational resources (optimally) leads to better performance, and even predictably so; neural scaling laws have been derived that accurately forecast the performance of a network for a desired level of compute. This leads to the notion of a `compute-optimal' model, i.e. a model that allocates a given level of compute during training optimally to maximize performance. In this work, we extend the concept of optimality by allowing for an `adaptive' model, i.e. a model that can change its shape during training. By doing so, we can design adaptive models that optimally traverse between the underlying scaling laws and outpace their `static' counterparts, leading to a significant reduction in the required compute to reach a given target performance. We show that our approach generalizes across modalities and different shape parameters.

Motivation & Objective

• To address the inefficiency of static model training, where fixed architecture shapes lead to suboptimal compute utilization despite known scaling laws.
• To investigate whether dynamically adapting model shape during training can lead to faster convergence and reduced computational cost for a given performance target.
• To develop a principled, scalable strategy for scheduling shape parameters (patch size and width) in Vision Transformers based on empirical scaling laws.
• To demonstrate that adaptive models can outperform static models in terms of compute efficiency while maintaining or exceeding performance targets.

Proposed method

• The method uses neural scaling laws to predict the optimal model configuration (patch size or width) at each compute level during training.
• For patch size adaptation, the model is trained with a schedule that transitions between different patch sizes based on the predicted performance gain per FLOP.
• For width adaptation, the model is expanded from a smaller initial width (e.g., 192) to larger widths (e.g., 384, 768) using a simple expansion mechanism that preserves the original weights and randomly initializes new parameters.
• The expansion process is applied to all layers, including attention weights, MLP weights, and normalization parameters, ensuring compatibility with the Transformer architecture.
• The scheduler selects the next shape parameter (patch size or width) based on which scaling law offers the steepest performance gain for the next increment of compute.
• Performance is monitored throughout training, and the method is validated by comparing adaptive models against static models of fixed size.

(a) Illustration of the effect of patch size on a given image.

Experimental results

Research questions

• RQ1 Can dynamically adapting model shape (patch size or width) during training lead to faster convergence and reduced FLOP consumption compared to static models?
• RQ2 How does the performance of an adaptive model compare to static models when trained under the same compute budget, particularly in terms of accuracy and efficiency?
• RQ3 To what extent can the principles of neural scaling laws be used to guide the scheduling of shape parameters to achieve compute-optimality?
• RQ4 Does the performance drop during shape adaptation (e.g., after width expansion) recover quickly, and is the model still optimal in the long term?
• RQ5 Can the proposed adaptive strategy be generalized to other shape parameters beyond patch size and width, such as depth or sparsity?

Key findings

• Adaptive scheduling of patch size in Vision Transformers reduces the required training FLOPs by more than 50% to reach a target ImageNet-1k error rate, significantly outperforming static models.
• The adaptive model based on the computed scheduler recovers quickly from performance drops after shape transitions and remains aligned with the predicted scaling law.
• Scheduling model width also leads to improved compute efficiency, though the gains are slightly less pronounced than with patch size adaptation.
• The method achieves compute-optimality by continuously selecting the shape configuration (patch size or width) that maximizes performance gain per FLOP at each training stage.
• The simple expansion mechanism—preserving original weights and randomly initializing new ones—proves sufficient for recovery and performance improvement, outperforming more complex function-preserving schemes.
• The results confirm that dynamic adaptation allows models to optimally traverse multiple scaling laws, leveraging the best performance characteristics of each shape configuration at different compute levels.

(b) ViT-B FLOPs.