Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design

Summary

Scaling laws in LLMs have typically been used to derive compute optimal model sizes. In fact one of the initial scaling laws papers in language modeling has indicated that as long as the model size is kept constant, the model shape (corresponding to embedding dim, mlp ratio, number of heads, depth) are not as important for performance. This paper on the contrary discovers that in the case of a Vision Transformer, given an equivalent amount of compute it is indeed feasible to design parameter and inference cost optimal models. These models are competitive with the larger models or outperform them on different tasks. Based on these observations the paper presents qualitative insights for scaling individual shape dimensions of a vision transformer across domains. Further the models derived based on these insights are evaluated on several computer vision tasks.

Strengths

The paper empirically investigates the effect of optimizing shapes of vision transformers to yield competitive models at a fraction of parameter size. Given the parameter and inference cost efficiency this finding is quite impactful towards designing hardware aware vision transformers. Furthermore, though the empirical investigation is expensive it is at a cost of significantly fewer experiments in comparison to previous scaling laws studies. The paper is very well written and clear in most parts.

Weaknesses

The insights derived in the paper are very important and impactful. However since this is more of an empirical investigation to derive scaling behaviours I find the paper lacks originality. Also the scaling methodology and behaviour would very likely change for different vision applications which also use transformers eg: superresolution, self-supervised learning. Scaling behaviors might also change for different transformer types (swin, deit etc). Deriving scaling laws for every application would mean repeating the analysis, which would incur significant computational cost. Furthermore since the JFT dataset and the code of the paper is not open-sourced the analysis and observations derived are not reproducible. For Zero shot transfer, Linear probes only and imgnet finetuning ViT-G/14 still seems to dominate in most cases, hence I am not very convinced if the representations learned by the SoViT-400m are indeed comparably or more effective than the ones learnt by ViT-G/14.

Questions

1. For given parameter size, does optimizing shape i.e trading-off dimensions while parameter size is held constant help? This was not found to not be helpful in [Kaplan et-al](https://arxiv.org/abs/2001.08361) and it would be interesting to know the observations in this study. 2. How closely are the laws derived tied to the size of the pre-training dataset, would the observations be similar if the laws were studied on imagenet for example? 3. Could the authors report the total TPU-hours/cost of analysis? 4. Minor: Should scaling “vision transformations” be scaling “vision transformers”? Line 70 5. How does this analysis compare to NAS methods (eg: [AutoFormer ](https://openaccess.thecvf.com/content/ICCV2021/papers/Chen_AutoFormer_Searching_Transformers_for_Visual_Recognition_ICCV_2021_paper.pdf)) which automatically derive optimal shapes? What are the advantages and disadvantages between these? 7. How hyper-parameters are set for different shapes is unclear to me. Could you please clarify this? 8. What do the question marks in Table 1 indicate? 9. ViT-g/14 seems to work better for multitask decoding in table-3. Do the authors have an intuition of why that could be? If all my concerns and questions are addressed appropriately I am willing to increase my score for the paper.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The limitations of the work are not adequately discussed: 1. The compute cost of the analysis? 2. Reproducibility of the experiments? 3. Are there any implicit assumptions which may effect the laws derived empirically in the paper? I encourage the authors to discuss these limitations in the paper.

Summary

The authors study the recent empirical insight that test performance follows a predictable power-law structure in terms of (optimally-allocated) compute and extend this notion to take into account the “shape” parameters of underlying model such as width, depth etc. They demonstrate that power-law behaviour can indeed be leveraged to design a strategy to explore the shape space, enabling a significant decrease in the amount of computation needed, compared to a naive grid search. The authors focus on the class of vision transformers and show that their discovered shape-optimal models match or outperform larger models (trained with same compute) while at the same time offering more efficient inference.

Strengths

1. The idea to leverage the (predictable) powerlaw behaviour of test performance to determine the ideal shape of a network is very well founded, and it’s very surprising that nothing similar has been done before. The results show that is indeed worthwile to optimize for the shape of vision transformers, as one can apparently achieve similar performance with a significantly smaller model. 2. An important quality of this work which in my opinion goes beyond most previous work on scaling behaviour is that the predictability is actively leveraged to reduce the search space of compute optimal models. By devising the so-called star-sweep strategy, the amount of compute needed is strongly reduced (albeit still very large). 3. The experimental setup is very extensive and the discovered architecture is evaluated on a broad range of tasks, ensuring that its optimality is not just an artifact of optimizing for ImageNet downstream accuracy.

Weaknesses

The definitions of some terms in this paper are sometimes unclear or never really stated. The most prominent and important one is compute **t**. Its definition seems to switch from context to context, which makes it difficult to follow. As far as I understood, compute **t** refers to the total compute, i.e. how many samples/tokens, training epochs and FLOPs per forward pass are needed. In other parts of the text compute seems to be referred as the number of samples solely, e.g. in line 127, compute unbounded seems to mean infinite sample size. Similarly in line 133, how can compute be fixed if the model size can be scaled arbitrarily (and compute is a function of model size)? In general I find it confusing that equation (2) depends both on compute **t** and shape parameter x_k, while **t** iself is actually also dependent on $x_k$, i.e. $t(x_k)$. Wouldn’t it be cleaner to simply replace compute **t** by the number of samples? I would appreciate if the authors could elaborate on this. I don’t think this has any major implications for the results but a clarification would certainly enhance readability. Similarly, a quick definition of width, MLP dimension and depth could be helpful for readers less familiar with the details of vision transformers.

Questions

1. How does each shape parameter (i.e. width, MLP dim and depth) contribute to the total number of FLOPs for say, a single forward pass? Is it for instance “cheaper” to make a ViT wider compared to adding another layer? What about wall-clock time and memory? It would be nice to at least have an approximate understanding of how the FLOP count is affected by these shape dimensions. 2. Something that is maybe clear and I missed it, but were the ViTs that you compare against, i.e. VIT-G/14 and ViT-g/14, considered “compute-optimal” before this work, and thus a “valid baseline”? 3. How does flexifying the architecture rule out that other patch sizes could be shape-optimal? The corresponding section is very short unfortunately (for space reasons I assume) but it would be great if the authors could expand on this.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have addressed the limitations of their work.

Summary

This paper proposes a novel and empirical take on the design of large vision transformers (ViTs), in the continuity of a previous paper aiming at optimizing the training of transformers. Whereas the previous paper was aiming to optimize a single parameter (optimal model size) given a fixed training budget, this paper goes one step further and attempts to discover the optimal ViT architecture (that is 3 parameters: token dimension, depth, MLP size) given a fixed training budget. The paper empirically solves this question via numerous experiments, while answering related questions and distilling interesting insights on the way.

Strengths

### Writing - The introduction is well written and clear - quite complex and unclear at times ### Method - The optimization method is novel and an improvement w.r.t. the "optimal model size" paper, as it can deal with multiple hyper-parameters while requiring much less experiments. - The insights about the choice of the joint functional form (eq. (2)) are interesting and well-grounded - Definitely a valuable paper for the community, even though some aspects could be improved ### Experiments - multiple experiments on several benchmarks support the initial claim that a smaller-but-optimal architecture achieves as well as vanilla SotA architectures - Experimental findings are useful and valuable (Section 4.1) - it is nice to see that other tasks than image classification are being experimented with, and that the findings hold for a wide variety of downstream tasks

Weaknesses

### Exposition - line 57: "Figure 1: The MLP dimension is scaled faster than depth, which in turn is scaled faster than width." --> I do not find this to be clear, looking at Figure 1 ### Method - The scaling parameters defined in eq (4) seems central in the analysis. Unfortunately, I'm not sure to understand exactly what it is and where it comes from. In particular, I don't understand why it should be invariant to the choice of the shape dimension (see Figure 5) - About the difference between small and large models - it is claimed that "in small models, an optimal shape in one domain is not necessarily optimal in others." (Figure 3) - How are "small" and "large" models exactly defined? This seems like a very convenient subjective definition. - Why does the model size matter so much when it comes to the optimal architecture across application domains? ### Experiments - Figure 6: - "while keeping compute fixed." --> how exactly? Does this mean other model hyper-parameters are decreased to keep the same model complexity? Or does this mean the training is short/longer? - the supposedly optimal model is clearly not optimal, since increasing the depth or MLP size results in (slightly) better performance. Since there experiments are on a relatively "small" model, according to the paper, does this mean the claims do not hold anymore? (see above)

Questions

I'm ready to upgrade my rating if the authors clarify some of the points mentioned above.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

yes

Summary

This paper introduces an efficient approach to investigate the scaling laws for compute-optimal model shapes, such as model width and depth. It proposes a shape-optimized vision transformer called SoViT. A comprehensive evaluation across various tasks highlights the effectiveness of the proposed architecture. SoViT-400m/14 achieves 90.3% fine-tuning accuracy on ILSRCV2012, surpassing ViT-g/14 with a larger model size. This study makes a valuable contribution to the design of vision transformers and is expected to have a certain degree of impact in this era.

Strengths

1. The paper is well-written, with well-defined formulas and sufficient supporting materials. 2. The three shape parameters are analyzed well. The paper proposes star sweep and grid sweep strategies to investigate the scaling laws avoiding the expensive search cost. 3. Extensive experiments are implemented to validate the method, including image classification, multitask decoding, and segmentation tasks.

Weaknesses

1. Figure 3. is challenging to understand, and it may benefit from more elaborate annotations and explanations to clarify the significance of certain data points within the figure. 2. For the experiments, SoViT-400m/14 can surpass ViT-g/14 for the image classification task. However, SoViT-400m/14 does not show a significant advantage over ViT-g/14 and ViT-L/16 in other tasks, such as the OCR and VQA. The detail of metrics such as Log-PPL and CIDEr in Table 4. should be explained in the paper. 3. A few typos in the paper, such as x_{k} in A.1 Quasiconvexity Proof.

Questions

1. Figure 6. indicates that deviating from the optimal depth/MLP configuration does not lead to performance degradation. Furthermore, the evidence supporting the superiority of the optimal shape is limited, and it seems that only three experiments (33%<, 33%>, 200%) were conducted, as shown in the Figure. 2. The compute-optimal model shape is different among tasks. How to search for an optimal model which can be applied to various downstream tasks?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

As above.

acknolegment

I have read the rebuttal, and I must say I am relatively disappointed by the lack of willingness from the authors to improve their paper based on my and other reviewer suggestions. Overall, I see no promise to improve the manuscript except for fixing typos and adding references. I know what is a plot and an axis, and I know how to read them, thank you. My role as a reviewer is not just to accept or reject a paper, it is also to help improve the paper quality and readability, for the sake of readers. Arguably, Figure 1 is a good example of my point: it is cited in the paper line 56, followed by the explanation "The MLP dimension is scaled faster than depth, which in turn is scaled faster than width", but this does not show at all at first sight when someone looks at Figure 1! And not on 2nd sight either. It actually requires some calculations based on the axes' tick labels to realize than, indeed, MLP dimension is scaled faster than depth, etc. The figure is completely counter-intuitive in that regard. Is it too much to ask for improving it? Same comments for the scaling exponents. Section 3 is really technical, and it would't hurt to explain things a bit more (as pointed out by other reviewers too). Same comments for other things the reviewer found difficult, unclear or counter-intuitive. Also, I am not satisfied by the answer regarding the difference between large and small models. I know that asymptotically, shape will not matter when the model grow sufficiently large. What would be useful is to be able to characterize when this happens, because so far the paper does not answer this question at all. By the way, Figure 3 is, again, not super clear nor intuitive (and what does all the gray circle exactly denote?)

Clarification

Dear reviewer, We thank you again for the insightful feedback and we acknowledge the areas where clarification is required. We would like to clarify that in places where we have not detailed how certain feedback will be incorporated, it's mainly because we are still deliberating on the best way to incorporate those suggestions. We are definitely taking all comments into account when revising the paper. This includes, for example, improving the clarity of Figure 1, Figure 3, and Section 3. In addition, we plan to include further details about the CIDEr and log-perplexity metrics, what “equivalent-compute” means, a brief description of the shape dimensions, the meaning of the exponents b and c, a link to the code, and adding further discussions in Section 5.5 to highlight the role of the sequence length. We will address all of these points. Thank you again for the constructive feedback, and for your suggestions to enhance the quality and readability of the paper.

Clarification

Dear reviewer, We thank you again for the insightful feedback and acknowledge the areas where clarification is required. We would like to clarify that in places where we have not detailed how certain feedback will be incorporated, it's mainly because we are still deliberating on the best way to incorporate those suggestions. We are definitely taking all comments into account when revising the paper. This includes, for example, improving the clarity of Figure 1, Figure 3, and Section 3. In addition, we plan to include further details about the CIDEr and log-perplexity metrics, what “equivalent-compute” means, a brief description of the shape dimensions, the meaning of the exponents b and c, a link to the code, and adding further discussions in Section 5.5 to highlight the role of the sequence length. We will address all of these points. Thank you again for the constructive feedback, and for your suggestions to enhance the quality and readability of the paper.

Clarification

Dear reviewer, We thank you again for the insightful feedback and acknowledge the areas where clarification is required. We would like to clarify that in places where we have not detailed how certain feedback will be incorporated, it's mainly because we are still deliberating on the best way to incorporate those suggestions. We are definitely taking all comments into account when revising the paper. This includes, for example, improving the clarity of Figure 1, Figure 3, and Section 3. In addition, we plan to include further details about the CIDEr and log-perplexity metrics, what “equivalent-compute” means, a brief description of the shape dimensions, the meaning of the exponents b and c, a link to the code, and adding further discussions in Section 5.5 to highlight the role of the sequence length. We will address all of these points. Thank you again for the constructive feedback, and for your suggestions to enhance the quality and readability of the paper.