Paper: Navigating Scaling Laws: Accelerating Vision Transformer's Training via Adaptive Strategies (ICLR-2024)

Reviewer comment (Novelty):
Novelty: The concept of adaptive models that can modify their shape during training seems very similar to Neural Architecture Search (NAS) which isn't covered at all in the related work or in the comparisons. See [1] for a comprehensive overview of NAS.

Author response:
We thank the reviewer for pointing out this connection. We believe however that there is a misunderstanding and we apologise for the confusion in advance. Compute-optimality, as considered in our work, as well as in the relevant literature [1, 2, 3, 4, 5, 6], concerns finding a model **within** a specified family of models (e.g. the family of Vision Transformers), which has optimal test error for a pre-specified level of **compute** (e.g. number of FLOPs/training time etc.). We are not concerned with finding **new** optimal models but rather with how we can **reach them more efficiently**, i.e. by training with **fewer** FLOPs.
Neural architecture search on the other hand is primarily concerned with finding a new optimal model architecture for a given task by changing architectural components such as the activation function, normalisation layers, the computational graph etc. Such an approach always focuses on finding a **single, fixed model** while we simply propose to reach such an optimal model through **smarter training**. We have expanded the related works section and made this point clearer.

Reviewer comment (Experiments):
Complexity: While adaptive models sound promising, they might introduce additional complexities in terms of training dynamics and hyperparameter tuning. I appreciate that the authors have mentioned some of these issues in the "limitations" section, but I feel a more comprehensive treatment is necessary especially related to how the optimal scheduling will change for other tasks / architectures.

Author response:
We agree that our method introduces more complexity since it involves training a series of models. We remark however that we **did not tune hyper-parameters** for the different models but used the **same set throughout training** (e.g. no change in learning rate, momentum etc.). One could argue that the transitioning times enhance complexity, but those are determined by the scaling law, which is precisely why we think our framework offers a powerful toolbox for model training. We observe our approach to further be stable across different modalities (see point 3 in global response), thus we believe that the introduced additional complexity is minimal.

Reviewer comment (Evaluation):
Overfitting concerns: Altering a model's shape during training could lead to overfitting issues. I think the paper should address this concern more explicitly eg. what do the train / test gaps look like for the adaptive model.

Author response:
Thank you for this question. We actually believe that the opposite is true. (1) When starting to train adaptively, the initial models are usually the lower capacity ones, e.g. we start with smaller width, smaller depth (see 2 in global response) and larger patch size, thus leading to less (potential overfitting). (2) Changing the model shape acts as a regulariser, as transitioning momentarily decreases test and training accuracy as a portion of the network is randomly initialised (at least for width and depth adaptation). (3) We are pre-training on a large corpus of images (around 15 million images) with data augmentation, thus overfitting in general is less of an issue in such a setting. 
We further back these arguments with empirical evidence in Figure 12, where we plot (upstream) training and test accuracies against training steps. We can indeed see that overfitting is not an issue with our approach.

Reviewer comment (Theory):
How do NAS approaches relate to neural scaling laws and the method proposed in the paper?

Author response:
We thank the reviewer for pointing out this connection. We believe however that there is a misunderstanding and we apologise for the confusion in advance. Compute-optimality, as considered in our work, as well as in the relevant literature [1, 2, 3, 4, 5, 6], concerns finding a model **within** a specified family of models (e.g. the family of Vision Transformers), which has optimal test error for a pre-specified level of **compute** (e.g. number of FLOPs/training time etc.). We are not concerned with finding **new** optimal models but rather with how we can **reach them more efficiently**, i.e. by training with **fewer** FLOPs.
Neural architecture search on the other hand is primarily concerned with finding a new optimal model architecture for a given task by changing architectural components such as the activation function, normalisation layers, the computational graph etc. Such an approach always focuses on finding a **single, fixed model** while we simply propose to reach such an optimal model through **smarter training**. We have expanded the related works section and made this point clearer.

Reviewer comment (Writing):
The authors should provide a more in-depth overview of the complexities introduced by adaptive models, potentially with guidelines to address these.

Author response:
We agree that our method introduces more complexity since it involves training a series of models. We remark however that we **did not tune hyper-parameters** for the different models but used the **same set throughout training** (e.g. no change in learning rate, momentum etc.). One could argue that the transitioning times enhance complexity, but those are determined by the scaling law, which is precisely why we think our framework offers a powerful toolbox for model training. We observe our approach to further be stable across different modalities (see point 3 in global response), thus we believe that the introduced additional complexity is minimal.

Reviewer comment (Evaluation):
Is overfitting an issue? What does the train / test gap look like for the models?

Author response:
Thank you for this question. We actually believe that the opposite is true. (1) When starting to train adaptively, the initial models are usually the lower capacity ones, e.g. we start with smaller width, smaller depth (see 2 in global response) and larger patch size, thus leading to less (potential overfitting). (2) Changing the model shape acts as a regulariser, as transitioning momentarily decreases test and training accuracy as a portion of the network is randomly initialised (at least for width and depth adaptation). (3) We are pre-training on a large corpus of images (around 15 million images) with data augmentation, thus overfitting in general is less of an issue in such a setting. 
We further back these arguments with empirical evidence in Figure 12, where we plot (upstream) training and test accuracies against training steps. We can indeed see that overfitting is not an issue with our approach.

Reviewer comment (Experiments):
While the proposed method is interesting and makes intuitive sense, don't such shape parameters exist for transformer models in other modalities. An experiment to show that the scheduler generalizes to text or audio or other modalities would make this work far stronger.

Author response:
1. **More modalities:** Thank you for your comment, which prompted us to verify our framework further on text-based tasks. For a detailed description, we point to the global response to all reviewers, especially point 3 where we address it as well as the new section Appendix E in our work.

Reviewer comment (Experiments):
The work explores the axis being tuned one at a time (patch size or model width). It would be nice to show that the result holds when two shape parameters are scheduled together. Do the authors expect the trend to hold?

Author response:
2. **Simultaneous shape adaptation:** This is a very good point! We verified that our framework can indeed handle such simultaneous adaptions by varying both the patch size and the model width. For more details, we point the reviewer to the global response to all reviewers, especially point 1.

Reviewer comment (Experiments):
Considering patch size and width of the transformer blocks are reasonable first choices to tune flops per training example. It would be interesting to consider other choices and their effects.

Author response:
We thank the reviewer for acknowledging the simplicity and effectiveness of our approach.

**More shape parameters:** We agree that using more “shape” parameters in our study would strengthen the results. To this end, we performed experiments on adapting the depth of the Vision Transformer. For more details, we point the reviewer to the global response to all reviewers, especially point 2.

We hope that we could address the questions of the reviewer and are happy to further clarify remaining concerns. If not, we would be grateful if the reviewer would consider raising the score.

Reviewer comment (Experiments):
I would like to see more ablations on other choices of the transformer architecture on their effect on optimal architecture.

Author response:
We thank the reviewer for acknowledging the simplicity and effectiveness of our approach.

**More shape parameters:** We agree that using more “shape” parameters in our study would strengthen the results. To this end, we performed experiments on adapting the depth of the Vision Transformer. For more details, we point the reviewer to the global response to all reviewers, especially point 2.

We hope that we could address the questions of the reviewer and are happy to further clarify remaining concerns. If not, we would be grateful if the reviewer would consider raising the score.

Reviewer comment (Novelty):
The authors have mentioned 3 contributions in the paper. However, these contributions could be summarized as one point---an adaptive strategy to change the patch size and the network width to adapt training.

Author response:
We thank the reviewer for the points raised. 

1. **Summarise contributions:** We would like to highlight that the devised framework applies to any “shape” parameter and not just width or patch size. We thus do believe that this is a contribution of its own due to its generality. We then verify our framework by specialising it to adaptive patch size as well as width. We agree that the latter two points can also be summarised as one. We have re-written the main text accordingly.

Reviewer comment (Experiments):
Many arguments in the paper lack theoretical connections, and only some empirical results were shown in the paper. I think if the paper focuses on empirical results, more experiments such as different network property changes, and different changes with multiple kinds of properties should be discussed and compared. If the authors keep current experiments, more applications or tasks could be included to further validate the proposed method.

Author response:
We thank the reviewer for the points raised. 

2. **More empirical evidence:** We agree that more experimental results would strengthen our work. To address this point, we have performed more experiments on (1) adapting the depth of the network, (2) adapting patch size and width simultaneously, as well as (3) adapting depth on text-based tasks. For more details, we point the reviewer to the global response to all reviewers.

Reviewer comment (Experiments):
Following the above discussion, one example is: when adapting patch size, there are no width changes, and when adapting network width, there are no patch size changes. I wondered if the authors could do further experiments on various changes including cross adapting to make the argument stronger.

Author response:
We thank the reviewer for the points raised. 

3. **Simultaneous shape adaptation:** This is a very nice suggestion. We verified that our framework can indeed handle such simultaneous adaptions by varying both the patch size and the model width. For more details, we point the reviewer to the global response to all reviewers, especially point 1.

Reviewer comment (Presentation):
There are also some ambiguities in the figures/paragraphs. For example, in Figure 3, what do those arrows mean?

Author response:
We thank the reviewer for the points raised. 

4. **Arrows in Figure 3:** We apologise for the confusion. For the left-most and right-most figure, the colorful arrows denote which partial derivative is maximal for a given level of error E. The black arrows on the other hand correspond to transition points where the shape of the model is changed. This figure purely serves as an illustration of our framework. We have updated the caption in the main text accordingly to clarify this.