Navigating Scaling Laws: Accelerating Vision Transformer's Training via Adaptive Strategies

The main argument of the paper is that, by being guided by scaling laws, one can develop compute-optimal adaptive models that surpass the performance of their static counterparts. The paper then proposes an "adaptive" ViT model that can alter its patch size during training. The authors claim the model optimally navigates the underlying scaling laws, thereby reducing the computational resources required to attain a specified performance level.

• The work is motivated by the following observation that the shape of the model remains fixed throughout the training process. Shape in this context refers to hyperparameters of the model that do not degrade performance when changed (such as width, depth and patch size).
• In turn, the paper proposes an adaptive training methodology to achieve equivalent performance for a specific model while using fewer computational resources (termed as FLOPs).
• The algorithm exploits the fact that the scaling laws proposed till now are in the form of a generalized power law. Thus, they compute the inverse function (as the fitted scaling law is bijective) and then compute the set where the shape parameter being tuned achieves the highest gradient. This set provides a scheduler for the shape parameter being tuned across training time.
• The paper leverages existing flexible transformer architectures like FlexVIT to tune shape-parameters like patch size and model width.

In this paper, authors introduce a new approach to training deep learning models that we call “adaptive training.” Adaptive training allows the model to change its shape during the course of training, which can lead to significant reductions in the required compute to reach a given target performance.

This work focuses on vision tasks and the family of Vision Transformers (ViTs), where the patch size as well as the width naturally serve as adaptive shape parameters. Authors demonstrate that, guided by scaling laws, one can design compute-optimal adaptive models that outperform their “static” counterparts.

They then propose a simple and effective strategy to traverse scaling laws, opting for the one that leads to the fastest descent, i.e. maximum performance gain for the same amount of compute. The work showcases the efficiency of our approach by optimally scheduling the patch size of a Vision Transformer, leading to significant reductions in the required amount of compute to reach optimal performance. They further confirm the validity of the approach by optimally scheduling the model width of a Vision Transformer.

The work demonstrates that adaptive training is a promising new approach to training deep learning models that can significantly reduce the computational resources required to achieve state-of-the-art performance.

The paper proposes a “compute-optimal” model, where the characteristics of the network, such as width, depth, and patch size, could be changed/adapted during training. The authors argue that such an “adaptive” model can utilize the scaling property of the neural network, and achieve “optimal” computation status---in a way, it balances the tradeoff between the computation and the performance. The modification of the network mainly focus on ViT, and the authors schedule the width and patch size to reduce training computation resource. Specifically, for both patch size and the network width, the authors apply adapting, scaling, and scheduling strategies. Experiments on ViT of different patch sizes/widths show promising computation reduction.

We thank the reviewers for the valuable feedback that we truly believe has helped us improve our work. We have updated the paper, including more experiments that we hope further bolster the claims made throughout our work. We refer the reviewers to the newly created sections C, D, E, and F in the appendix.

1. Simultaneous shape adaptation: We thank the reviewers for this very nice suggestion! Our framework can indeed handle multiple shape parameters at once. We verified this by adapting the patch size and the model width of a Vision Transformer simultaneously and indeed observed that such a strategy leads to even more computational savings, leading to a reduction in FLOPs of 65% (Appendix C). We refer the reviewers to Figure 19.

2. More shape parameters: In order to demonstrate the validity of our framework further, we included the depth of the network as a “shape” parameter in our experimental study (Appendix D). We use the same experimental setup as for patch size and model width and grow the network over the course of training from a starting depth of 6 to the final depth 18, using the set of values {6, 12, 18} as allowed depths during training. We show in Figure 20 how performance matches the predicted scaling. We would like to highlight further that our method works even more generally, allowing one to switch between different training objectives (e.g. self-supervised to supervised) or even datasets (pre-training to target) or adjusting hyperparameters such as batch size during training. We leave such explorations to future work.

3. More modalities: We agree with the reviewers that verifying our framework on more modalities would strengthen the results. As a step towards this, we also evaluated our shape-adaptive strategy on text-based data when growing the model’s depth (Appendix E). We perform auto-regressive learning on the “Code Parrot” dataset, following the evaluation protocol of [1]. We use upstream evaluation loss as the metric of interest. We display the corresponding results in Figure 21. Varying context size (the natural “dual” to patch size in vision) or varying the context window of the attention operations are also interesting directions that we leave for future work.

If no further concerns remain, we kindly ask the reviewers to reconsider their scores.

[1] He, Bobby, and Thomas Hofmann. "Simplifying Transformer Blocks." arXiv preprint arXiv:2311.01906 (2023).

The discussion period is ending soon, and we still haven’t heard back from any of the reviewers. The reviewers provided questions which we addressed in our rebuttal, and it would be great to have a discussion in case some points have remained unclear. We would really appreciate it if you provide feedback at your earliest convenience.

Best wishes, The authors

In this paper, the authors introduce a new training of deep neural networks that relies on adaptively changing the shape and size of the model during the course of training. The goal of the training algorithm is to trade off compute versus predictive performance during the course of training. The authors demonstrate the efficacy of their results by training a series of vision transformer models for image classification tasks. The reviewers commented positively on the question being addressed by the authors and the underlying hypothesis of the work. The reviewers also commented negatively on the limited scope of the experiments in terms of identifying more shape parameters that may be adapted during training, as well as concerns about the complexity of the method, and concerns about overfitting. The authors' rebuttals were not responded to by the authors, so the AC examined the concerns raised by the reviewers to see if this would pass the bar of acceptance. At face value, the AC agrees with the reviewers that additional experiments are needed to justify the complexity of this method and make for stronger empirical results. One item of concern is how the authors could provide evidence that their method generalized beyond the narrow parameters of vision transformers and image classification. Given the uniformly low scores and the review by the AC, this paper does not warrant acceptance by the conference. The authors however are encouraged to incorporate their follow up experiments and do their best to address some of these concerns and resubmit to a future workshop or conference.

Reject