Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design
Authors: Ibrahim Alabdulmohsin
Xiaohua Zhai
Alexander Kolesnikov
Lucas Beyer
Venue: NeurIPS 2023 poster (2023)
Keywords: Vision transformer
scaling laws
compute-optimal model design
vision
TL;DR: We introduce new scaling laws for deriving compute-optimal model shapes and implement this in ViT. We demonstrate that scaled-down architectures trained at their optimal shapes for the right amount of compute, are comparable to fully-scaled models.

Abstract: Scaling laws have been recently employed to derive compute-optimal model size (number of parameters) for a given compute duration. We advance and refine such methods to infer compute-optimal model shapes, such as width and depth, and successfully implement this in vision transformers. Our shape-optimized vision transformer, SoViT, achieves results competitive with models that exceed twice its size, despite being pre-trained with an equivalent amount of  compute. For example, SoViT-400m/14 achieves 90.3% fine-tuning accuracy on ILSRCV2012, surpassing the much larger ViT-g/14 and approaching ViT-G/14 under identical settings, with also less than half the inference cost. We conduct a thorough evaluation across multiple tasks, such as image classification, captioning, VQA and zero-shot transfer, demonstrating the effectiveness of our model across a broad range of domains and identifying limitations. Overall, our findings challenge the prevailing approach of blindly scaling up vision models and pave a path for a more informed scaling.

Getting ViT in Shape:
Scaling Laws for Compute-Optimal Model Design
Ibrahim Alabdulmobsin', Xiaohua Zhai", Alexander Kolesnikov, Lucas Beyer*
Google DeepMind
Zürich, Switzerland
{ibomohsin, xzhai, akolesnikov, horg://r
Abstract
Scaling laws have been recently employed to derive compute-optimal model size (number of parameters) for given compute duration. We advance and refine such methods to infer compute-optimal model shapes, such as width and depth, and successfully implement this in vision transformers. Our shape-optimized vision transformer, So VIT, achieves results competitive with models that exceed twice its size, despite being pre pre-trained with equivalent amount of compute. For example, SoViT-400m/ 14 achieves 90.3% line-tuning accuracy on ILSRCV2012, surpassing the much larger ViT-g/14 and approaching ViT-G/14 under identical settings, with also less than half the inference cost. We conduct a thorough evaluation across multiple tasks, such image classification, captioning, VQA and zero-shot transfer, demonstrating the effectiveness of our model across a broad range of domains and identifying limitations. Overall, our findings challenge the prevailing approach of blindly scaling vision models and pave path for a more informed scaling.
Introduction
The de-facto approach for improving performance of vision and language models today is scale:
large models are trained on more data for longer [64, 43, 24, 19, 80, 23, 13, 16]. Empirically, has been observed that the benefit of scale often follows predictable power law in which the performance f(x) (e.g. error rate or log-perplexity) satisfies f(x) ~ Bx- + &oo for some B, e > O as one varies the scaling dimension x (e.g. data or model size), the remaining dimensions are not bottlenecks [34, 39, 27, 26, 3, I]. Here, Soo is the irreducible loss.
However, the simple pawer-law relation becomes more complicated when compute is considered. In this case, power laws are observed only along the compute-optinme frontier. Otherwise, scaling up the model size for fixed compute budget can deteriorate performance (see [39, 35] and Figure 4). Since one often has a fixed compute budget in mind (e.g. available hardware and time), one should pick the model size that maximizes performance subject the compute budget constraint, which may imply not training until convergence. Indeed, this approach was used successfully in the recent Chinchilla [35] that outperformed its predecessor Gopher [55] despite being 4x smaller in size.
Unfortunately, in both [39] and [35] among others, the "size" a model is equated with its parameter count, with no special consideration for model "shape dimensions", such "depth" or "width". The rationale behind this choice follows from the surprising observation that the transformer shape had little impact its scaling behavior in language modeling (LM) when performance is measured upstream (e.g. using log pog-eerllxitt) [39, 32, 33]. Nevertheless, follow-up analysis suggests that shape plays pivotal role in other domains, such in machine translation [47] and also in language modeling for downstream performance [66], with recent works even advocating for extreme aspect ratios, such single wide attention layer [12].
"Significant technical contributions.
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

Scaling law
B/16
L/16
H/16
g/14
G/14
c/14
2000
60
Width
Depth
15k
MLP dimension
Parameters
1500
40-
x2.0
IB
X6.
10k
20
100M
500
1T
10T
100T
1T 10T
100T
1T
10T
100T
10T
100T
Pre training compute in GFLOPs
Figure 1: Predicted efficiency frontier (depth, width, MLP dimension, and parameter count) in SoViT. In large models, optimal shapes follow similar trajectory in both image classification and multimodal tasks (see Section 4) although they can be different in small models (see Figure 3). We provide on the right (in blue) the amount of increase when compute goes from IT to 100T GFLOPS.
In vision, in particular, much earlier works using convolutional neural networks (CNNs) pointed out that the parameter count is indeed poor predictor of performance. For example, scaling all dimensions [64, 43, 5] in ResNets [29] is more effective than scaling a single dimension such as depth alone. In addition, scaling width [79] is often more effective than depth, especially for small models [36, 58, 75]. Hence, optimizing the "shape" of transformers seems worthwhile.
In this work, we present SoViT: a shape-optimized vision transformer [24] that matches the per- formance of much larger models despite being pre-trained with equal compute. derived from recipe we introduce for optimizing the shape of neural architectures, such their depth and width. principled approach for scaling multiple dimensions is advantageous because although one can scale dimensions via brute-force search, this requires extensive computation and often remains sub-optimal [64]. Our recipe allows us to extrapolate without having to conduct an extensive set of experiments. For example, after only 115 experments, we identify scaling strategy in VIT for all three dimensions: width (internal representation), depth, and MLP size. For comparison, [35] requires 400 experiments optimize single dimension (the parameter count) alone.
One major finding that small vision models can perform on par with larger ones with the same compute if we optimize their shape. In language, recent works have demonstrated the value of scaled-down architectures, such as the Chinchilla model [35] discussed earlier a 70B parameter model that outperforms the 280B-parameter Gopher [55] and 175B-paramete GPT3 [13] - as well as LLaMA with its 13B parameter variant outperforming GPT3 on most benchmarks [69]. By introducing VIT establish this phenomenon in vision well.
Figure/ 1 summarizes how the various shape dimensions are scaled SoViT (see Section for derivation). The MLP dimension is scaled faster than depth, which in turn is scaled faster than width. When summarized by their parameter count (rightmost plot), compute-optimal ViTs are smaller than was previously used. With this scaling strategy, we find the shape of ViT for the compute- equivalent of ViT-g/14 [80] pretrained on 16B JFT images [63]. We call this 2.5x smaller model SoVIT-400m/14 14. It achieves 90.3% fine-tuning accuracy on ILSRCV2012 [22] and 82.2% zero-shot accuracy in the locked-image text tuning (LiT) setup [81]. We further evaluate SoViT-400m/1 14 on captioning, VQA and panoptic segmentation and highlight some results in Figure 2.
Statement of Contribution. In summary, our contribution is to:
Introduce a new method for optimizing the shape of neural networks, such as their depth and width. Our technique expands and improves previous methods by optimizing multiple shape dimensions jointly while requiring significantly fewer experiments.
Demonstrate the effectiveness of scaled-down architectures in vision. We optimize ViT for the compute-equivalent of ViT-g/14, leading to a smaller, faster model of equal quality.
Present new qualitative insights for scaling vision transformers, such as on how to scale individual shape dimensions and how optimal ViT shapes vary across domains.
Conduct extensive evaluation across tasks like image classification, image captioning, VQA, zero-shot classification and panoptic segmentation, identifying both gains and limitations.

---

INet-G-lot
OljesNa
COOO-CIDE:
VEL
Paraem: (M)
90.3
90
90.1
128
GHLOOSs poroooss.iin
203
VIE
bags form /s
S0AD WIP Cudesall Vi2g. SwaV2. SAVT
VESG
VIELL SAYT ViT
V.12 SANT Mix
VILL SNVT ViT
1000
Figure 2: Optimizing for the compute-equivalent of ViT-g/14 results in the 2.5x smaller SoViT- 400m/14 model achieves equivalent results across wide range of benchmarks. Our model performs exceptionally well the competitive ImageNet (ILSRCV2012) benchmark in comparison with significantly larger models from the recent literature [61, 78, 49, 80].
Related Work
Optimizing training for compute has received a significant amount of attention in recent years, partly due to the financial and environmental costs of training large models [52, 55J. However, conflicting results are sometimes reported. For example, in language modeling. [39] argues that the model size should be scaled faster than the data size, implying it is compute optimal to "undertrain" large models. Similar conclusions are found in [47]. On the other hand, [35] argues that the model size should be scaled uniformly with the data size, and highlights that transformers
not trained long enough, leading to some recent efforts [69] "overtraining" their models instead. Our analysis for ViT in Section 4 agrees partially with the latter result.
Scaling the size of vision transformers has led to remarkable results achieving, for instance, 90.4% top-1 accuracy ImageNet (ILSRC22012) with 2 billion parameters [80] and 90.9% top-1 accuracy with billion parameters [15]. When scaled to 22 billion parameters, ViT exhibits state-of-the-art alignment human visual perception in terms of shape/texture bias, among other findings [21].
Despite the clear benefit of scale, there has been little investigation into optimally scaling the shape of ViTs. [66] suggest preferentially increasing depth before scaling other dimensions uniformly. For ViT, however, they only consider small ViT-S and VYT-B models and the reported accuracy improvement comes with an increase in FLOPs of up to x4, making it difficult to draw conclusions about the suggested shape's s quality. In contrast [12] recommend scaling width over depth, but the authors do not observe any improvement when applying their strategy to VIT.
Our analysis draws inspiration from "compound scaling" in MobileNet [36] and EfficientNet [64], while differing in significant ways. EfficientNet uses an exhaustive grid search to determine the optimal architecture for fixed increase compute (e.g. x2). Afterwards, each dimension scaled up by the same ratio with every subsequent increase in compute. In contrast, we expand scaling laws simultancously account for model size and compute beyond the efficient frontier and leverage them to derive the optimal scaling exponents for each dimension separately, outlined in Section 3.
Throughout our analysis, ing performance instead of upstream metrics. This follows recent reports arguing that upstream performance may not reflect downstream performance in language and vision [65, 80].
downstream metrics, e.g. ImageNet 10-shot error, when measur-
We use GFLOPs as proxy for compute since wall-clock core-hours (see Figure 4). However, GFLOPs have limitations [5, 20] and may not be a perfect predictor for the metric of interest (c.g. core hours) in all model and hardware types. Note that we focus on scaling the shape of the architecture, not improving its training protocol, which can be similarly beneficial [5, 67, 62, 68].
aandstie-agnostic and correlates well with actual
3 Scaling Strategy
Notation. We begin with a formal description of the problem. We represent neural architecture as xuple = Xx,,x xD)) E N containing D shape dimensions, such as width, depth and MLP size. We denote compute such as GFLOPs by We designate / ND x R+ - a performance metric of interest, such downstream ImageNet 10-shot error rate. Specifically, f(x, t.) results from (pre)-training an architecture x for fixed compute budget We always assume that f corresponds to loss, meaning lower values are better.
3

---

[Net-1k 5 shel Hrror
1Net-lk 10-shed Hrror
1Net 1k 25 abot Hiroe
Average Perpicxity
0.3
0.00
0.60
245
0:53
0.51
253
0.50
0.00
0.30
645
25,0
0.50
0.75
Compute (Giflaps)
1:10
Compute (Killaps)
1:10
Compute (Killops)
Le10
Compute (Gillos)
lel0
Figure grid sweep over multiple ViT shapes pretrained on 600M JFT examples highlights the important role of shape. Each dot corresponds to a model architecture pretrained on 600M examples and evaluated on a downstream metric, e.g. Imagenet- 5-shot in the leftmost plot. The two architectures marked in blue and red - identical in all four figures are compute-optimal for classification and image-to-text tasks (captioning/VO4) //Anon)in///),,,, respectively. For captioning/VQA, we average re-perplexity scores (see Section 4.2). In the leftmost three figures, using Imagenet- Ik few-shot evaluation, the compute-optimal model highlighted in blue is compute-optimal in all three cases, but it is not compute-optimal for image-to-texs tasks as shown in the rightmost figure. So, in small models, an optimal shape in one domain is not necessarily optimal in others.
The goal of optimizing shape for fixed compute t is to identify x" (depending on t) such that:
f(x
inf f(z t) S
=END
for some small tolerance 0. Due to modeling assumptions, approximations, and the finite possible number of experiments conducted, we cannot hope for c = 0 and have to tolerate a small excess loss.
Single Dimension. As demonstrated in Figure the shape pretrained vision transformer has an impact on its downstream performance. To determine an optimal shape scaling strategy, we begin by considering both compute and a single shape dimension x for [D]. such as depth. In prior works, optimizing a single dimension X3: for compute involves running a large number of experiments in order to identify the Pareto optimal frontier, from which power laws on X or t are derived [39, 35]. Since this is expensive, we propose the following joint functional form instead:
fu(XX, t) ~ xx + (Bxx + + +
where Ork, Ogg Bg, bx, C, Ea, ck O. Here, f& focuses the dimension k alone and assumes that all other shape dimensions j * k are sufficiently large such that they do not constitute a bottleneck. We also assume that data is unlimited so that there is no risk of overfitting. We estimate the parameters in (2) by minimizing the relative error. In (2), aş are scaling exponents when varying the corresponding shape dimension in the compute-unbounder regime, is the data scaling exponent, while by relates to the impact of the model shape on compute.
Our argument for this particular functional form is six-fold:
I. If compute unbounded, we recover the familiar power law relation on model size fi (xx:) ~ ay d2 + E& [34, 2, 38, 39]. In addition, increasing the model size 2y while keep the data size fixed does not imply that (x t) - E because x can increase faster than t in (2). II. For any fixed model size, the relation above reduces to the power law fx(()) ~ At + B, where A A= 3x by + and B agx; 44 +EA- Since the model size is fixed, is proportional to the size of the data. Such data scaling laws have been demonstrated extensively in various domains [1-3, 27, 34, 39, 59, 80].
III. For fixed compute, the relation w.r.t. X is non-monotone, quasiconvex (see Appendix A), in agreement with empirical measurements [39, 35]. See IsoFlop curves in Figure 4.
IV. Arguments for power law behavior using space partitioning suggest that the exponent c is independent of the shape dimension. In particular, c = @(1/d), where d is the intrinsic dimension of the data manifold [2, 38, 59]. From this, we conclude that assuming the functional form in (2) for every shape dimension separately cannot lead to any contradictions since this assumption is satisfied by the decomposable loss:
f(x, t) Σ C
+
Bex t to + =o + 6oo,
(3)
for some constants &0 0.

---

V. When optimizing the shape dimension X3: for fixed compute t, the optimal value xj is:
t
TTE
x
= =0 -t*),
where: 8k by: ak
(4)
B bk
Recall that the scaling exponent 85 in (4) is positive because ag. by, O. Using the rela- tion (4), we rearrange the terms in Eq. (2), and obtain the scaling law for model performance along the compute-optimal frontier (Appendix A):
fx((x, 1) Fxx aR Gt
- + Ek.
(in the compute-optimal frontier)
for some constants F and G, which is sum of power law terms involving the model size and compute. Indeed, this decomposition has been demonstrated to hold within the compute-optimal frontier by [39] and [35].
VI. Eq. (2) fits empirical measurements and extrapolates accurately as well, see Figure 4.
Multiple Dimensions. Next, we expand upon the previous approach by incorporating multiple dimensions. To reiterate, our method involves both functional form (2) and a novel procedure. Our procedure significantly decreases the number of large-scale experiments required to identify compute-optimal architectures, by an order of magnitude compared prior work [35].
Stur Sweep- Conducting a brute-force grid search to estimate scaling parameters across all dimensions expensive, since it requires (22)) experiments cover the search space. Instead, demonstrate that a "star sweep" is sufficient: (1) starting from large model x(e) (the star center), vary single dimension E [D] at a time an exponentially-spaced grid, such that all values are much smaller than x In experiments, for instance, optimize three shape parameters: width, depth, and MLP dim (see Section 4 for a brief definition of each dimension). Our star center is x = (1968, 40, 6144); i.e. has width 1968, depth 40, and MLP dim 6144. When varying MLP dim in the star sweep, to about 20% increase in each step, while fixing width to 1968 and depth to 40. We do this to ensure that other dimensions do not form bottleneck when estimating the parameters in (2). This gives us the scaling exponents 8% in (4).
the grid (1088, 1360, 1728, 2160, 2592, 3072), corresponding
Grid Sweep The second stage is grid sweep for small models trained for short compute. Depending on the number of shape dimensions involved, the cost of running this grid sweep can be negligible. Its goal is to identify single architecture (0) that lies the Pareto optimal frontier for small compute as illustrated in Figure 3. This is important since a suboptimal x")) can significantly skew results [5]. Our grid sweep identifies (0) to be (608, 10, 928), the blue star in Figure 3. The advantage of this step is to absorb the leading coefficients in x = (t**) ) in (4) so that the star sweep focuses on estimating the exponents 8& alone. We demonstrate in Figure 5 that the scaling exponents 8y robust to the choice of the evaluation metric f. In Appendix B.3, we discuss important considerations that were taken into account during this analysis.
Scaling. Finally, we scale all dimensions jointly. Starting from the small compute-optimal archi- tecture (0) and the amount of compute (()) it is optimal for, suppose we increase compute by a factor (i.c. the new compute is r (O)) By treating this increment T as a sequence D smaller an increase in compute by factor of is accompanied
increments of size " each with Σx Wk by an increase in every shape dimension l: by a factor of 7**,, respectively. In this work, the adopt the simplest strategy of setting Wik = 1/D, but acknowledge that more sophisticated approaches might lead to better results.
Shape-optimized ViT
We implement the scaling strategy in Section 3 in vision transformers [24] pretrained on JFT-3B, proprietary dataset with about 30k classes and around 3 billion examples [80], using the Adam optimizer [41]. As mentioned in Section 3, we focus on optimizing three shape dimensions: width (size of internal representation), depth (number of encoder blocks) and MLP dim (hidden dimension). Following [43, 24, 80], we remove near-duplicate examples between upstream JFT-3B data and all the downstream train and test sets. Appendix B contains the full set of hapereparameters used in the
5

---

INet-Ik 10-shot Accuracy
1soFLOP Curves
Compute Comparison
0.8
[Ner- Ik 10-shot Error
0.64
10x to
TPU Core Hours
10t
0.60
Actual
20x
0.55
0.50
10'
0.2
0.45
0.2
0.5 Predicted
0.8
500
1000
1500
2000
10
10
width
GFLOPs
Figure 4: LEFT: Comparison between ILSRCV2012 (denoted INet- 1k) 10-shot error rate predicted by Eq. (2) and actual. The value marked in violet corresponds to the star center x(e) that is never used when estimating scaling parameters. Eq. (2) is consistent with empirical measurements and extrapolates accurately. MIDDLE: IsoFlop curves in ViT as one varies the width dimension. RIGITT GFLOPs well-correlated with actual TPU core hours across models (correlation coefficient 0.99).
depch
2.0
09
14
exponent
S-bet
10-dui
25-bbt
10-
25-bbot
5-bott
10-bre
Figure A plot of the estimated values the exponents in (2) for different evaluation metrics f. The scaling exponent 8 tends to be less sensitive the choice of metric than other exponents. Moreover, the data scaling exponent c is approximately c = 0.65 + .66, independently of the choice the shape dimension, in agreement with what would be expected using space partitioning arguments [2, 38, 59].
experiments, including full details about the star and grid sweeps described in Section 3. We fix the patch size in our analysis to 14 14, but study "flexifying" to arbitrary sequence lengths following [7] in Section 5.5.
As evaluation metric /, we consider domains: (1) image classification, with ImageNet linear 10-shot error rate as the metric, and (2) image-to-text LiTT-decoding following [8]. In the latter case, the evaluation metric is average of four perplexity scores: COCO captioning, optical character recognition (OCR), and question answering (VQAv2 and GQA). Refer to [8] for details about the LiT-decoder setup. By considering such distinct domains, our goal is identify similarities and differences (if any) in how to optimally scale the shape of vision transformers (ViT).
4.1 Image Classification
We use the aforementioned star center ( = (1968, 40, 6144) as our starting point. To esti- mate the scaling exponents 8& in (4) for each dimension separately,
vary width in the grid (608, 768, 928, 1088, 1328, 1648), depth in the grid (8, 10, 12, 16, 20, 24), and MLP dim in the grid (1088, 1360, 1728, 2160, 2592, 3072). As discussed in Section 3, we ing with all values being much smaller than in the star center (()). Following [24], evaluate quality using few-shot linear transfer by using pre-trained models to extract features and fitting linear regression head mapping them to the one-hot encoding of the target labels.
exponential spac-
The individual scaling exponents we find are Sdepth RJ 0.45, Swidth R 0.22, and SMLP 0.6. Impor- tantly, these exponents quite robust to the choice of the metric. As shown in Figure 5, changing the metric from ImageNet 10-shot to either 5-shot or 25-shot can change the best-fit estimate of the other exponents ag, 6k, Cj in (2) but the scaling exponent šę is relatively unchanged, since it formed as ratio over other exponents. In addition, the data scaling exponent e appears to be independent of the choice of the shape dimension. As mentioned earlier, this is consistent with space partitioning arguments for power law scaling [2, 38, 59].
The estimated scaling exponents point to the following picture:
I. MLP dimension should be scaled faster than depth, and depth faster than width.

---

1Net 10 (0-dot Accuris
widhb
80
depti
MLP
78.9
78.5
78.5
78.5
78.5
78.6
78
77.9
78-
76-
77.1
76.7 76.5
76.4
Optimizing VIT shape for the compute-equivalent of ViT-B/14 results in SoViT- 150m/14, which improves performance significantly. See Section 4. 1. CENTER & RIGHT: Impact of deviating from the optimal shape in SoViT-150m// (in green) while keeping compute fixed by changing the training duration such that the total FLOPs is the same in all models.
Figure 6: LEFT:
II. The size of ViT, as quantified by its parameter count, is scaled more slowly than the allocated compute. More precisely, for every increment compute by factor of 10, the parameter count of the optimized model shape increases by factor of x 2.5.
III. As demonstrated in Figure 1, small ViT models can match the performance of much larger ones when their shape and training duration jointly optimized for the available compute.
We validate these predictions by optimizing the shape of ViT for the cumpute-equivaleee of ViT-g/14 when the latter is pretrained 16 billion JFT-3B examples as done in [80]. The resulting model, SoViT-400m/ 4, is significantly smaller and faster, yet equally competitive. It has a width 1152, depth 27, and MLP dim 4304. Fine-tuning it ImageNet results in 90.3% top-1 accuracy, Figure 2. Section 5 presents various other evaluations.
In Figure 6, also optimize the shape of VIT for the compute-equivalent of ViT-B/14 pretrained on billion examples of JFT-3B using Imagenet 10-shot error rate as an evaluation metric, resulting in SoViT-150m/1l 14. It has width of 880, depth 18, and MLP dim 2320. As shown in Figure optimizing the shape of ViT leads to a significant improvement in performance, from 76.6% in ViT-B/14 to 78.5% in So VOVV--150m 14 when both are trained for the same amount of compute. We also vary the optimized shape by decreasing/increasing one dimension at a time and retraining the corresponding model while keeping the total compute fixed. As shown in Figure small deviations from the predicted optimal shape can lead to a notable drop in performance, especially for width since it has the smallest scaling exponent (see Figure 5). We also include in Figure 6 (LEFT) comparison with a model, denoted B-150m, which has the same shape as ViT-B/14 4 but the same size SoViT-1 14. This confirms that while optimizing the model size improves performance, optimizing the shape improves it even further.
Importantly, the model shapes in Figure 6 bear resemblance those observed during the star or grid sweeps. To recall, the star sweep is centered around architecture x(e) whose shape dimensions are significantly larger than in ViT-B/14, whereas the grid sweep pretrains models that are substantially smaller and for only 600M examples. The ability of our strategy accurately identify near-optimal model shape within this context underscores its robust extrapolation capability.
4.2 Multitask Decoder
Besides image classification, there has been a significant interest in multimodal applications, mostly fueled by the convergence across language and vision the transformer architecture [72, 24]. In particular, an encoder-decoder transformer with an autoregressive decoder is a popular choice because allows reusing pretrained image encoders. We repeat the analysis conducted in Section 4.1 to optimize the shape of the image encoder, while fixing the decoder architecture to two layers as was used in [8]. Further details are provided in Appendix C. As evaluation metric f. we use the average of four perplexity scores: COCO captioning [48, 14], OCR [50]. VQAv2 [28] and GQA [37], without normalization since they share similar scale. For the learning rate and weight decay we conduct sweep where we vary the learning rate in {10-3 3 3 10 44 10 10=4) and the weight decay in {3 x 10-4 10-1, 1 x 10-33. We pick the largest learning rate and the corresponding weight decay that result in a stable training run (i.e. smooth training loss curve and gradient norms) for both the largest and smallest image encoder architectures. From this, a learning rate of 3 x 10-1 and weight decay of 10-1 are selected.
7

---

Table ImageNet fine-tuning. The top shows models trained in the same controlled setting. and the bottom a representative set of large well-performing models. SoViT compares favorably. Contrary to common practice, we use held-out 2% Train select hyper-parameters. Selecting them on Val would increase all scores. FL.PSs according to XIA; PyTorch reports MACs.
Size
ImageNet variant
Model
Pretraining
Input
Params
FLOPs
Val [57]
ReaL [6]
v2 [56]
SoViT-400m/14
JFT-3B
224
428 M
221 G
88.9
90.3
80.7
VIT-L/16 [80]
JFT-3B
384
303
383 G
88.5
90.4
80.4
SoViT-400m/14
JFT-3B
384
428 M
672 G
90.0
90.9
83.2
ViT-g/14 [80]
JFT-3B
5182
1011
3208 G
90.2
90.9
SoViT-400m/14
JFT-3B
518²
428 M
1374 G
90.3
91.0
83.4
VIT-G/14 [80]
JFT-3B
5182
M
5668 G
90.4
90.8
83.3
SwinV2-G [49]
IN-21K + 70M
640
3000 M
90.2
84.0
CoAtNet-6 [19]
JFT-3B
5122
1470
1521
90.4
MAE-WSS [61]
IG-3B
518
1890 M
5679
89.7
90.9
83.0
CoCa [77]
JFT-3B ALIGN-1.8B
5762
2100M
91.0
Using this analysis, the derived scaling exponents approximately 0.25, 0.49 and 0.62 for width, depth and MLP size, respectively. Hence, whereas the optimal shape dimensions in small architec- tures be quite different between image classification and multitask decoding, shown in Figure 3, the scaling exponents are nearly identical, so the same scaling recipe is used in both domains.
5 Evaluations
Overview. We now evaluate SoViT-400M in various contexts to verify whether it broadly matches ViT-g/14's performance, only in the ILSRCV2012 10-shot metric was optimized for. The settings we cover are few-shot, frozen linear probes on ImageNet, zero-shot transfer, multitasking including captioning. OCR, and question answering, as well as panoptic segmentation. In each of these settings, we compare SoViT-400m/14 to ViT-L/16 and a ViT-g/14, all trained on the
Compute. Experiments are executed on Tensor Processing Units (TPU). SoViT-400m/14 pre- trained on 40 billion examples, which amounts to 9T GFL.Pss and 230K T core-hours. ViT-g/14 was pretrained on 16 billion examples, corresponding to 9T GFLOPs and 210K TPUv3 core-hours.
5.1 Image Classification
We verify classification performance in three common and widely useful setups: full fine-tuning, linear probes on the frozen model, and few-shot linear classification.
Fine-tuning on ImageNet. Pre- Pre-trained image encoders most commonly [18] evaluated by fine- tuning them on the ILSVRC2012 classification task. The detailed fine-tuning settings are provided in Appendix E. One important aspect is increase image resolution [70] as way of further increasing the capacity of the pre-trained model during fine-tuning [43]. Table shows the performance of So ViTTT-400m/1 in comparison with ViT-L/16, ViT-g/14 fine-tuned at various resolutions, along with few more representative models from the literature. The results confirm that SoViT-400m/ 14 achieves the goal of matching ViT-g/14 while being significantly smaller.
Linear probing on ImageNet. The quality of the pre- trained representation learned by the model is often more directly assessed by performing linear probes, meaning learning a linear classifier on top of unmodified, frozen output features from the model. We present results of this evaluation on the full ImageNet-1k [57] dataset in Table 2, including robustness evaluations of the learned
Table 2: Linear ILSVRC2012 probes.
Val ReaL v2
-R
Obj
L/16
86.7 90.0 78.5
88,9
67.8
63.5
SoViT
88.2
90.3
80.6
89.0
76.4
68.7
g/14
88.4 90.2 80.8
90.3
76.6
67.7

---

Table 3: SoViT-400m/14 perforss competitively with ViT-g/14 in 10-shot classification.
INet
CIFARI00 Pets
Birds
Caltech [25]
Cars [45]
Colorectal
DTD
UC
[22]
[46]
[51]
[74]
[40]
[17]
[76]
VIT-L/16
81.5
82.2
97.0
97.1
89.9
93,8
79.4
72.0
96.3
SoViT-400m/14
84.1
86.7
97.6
88.8
91.3
93.6
81.5
72.5
97.7
ViT-g/14
84.0
87.2
97.4
88.5
89.3
93.9
78.9
74.1
98.2
probe according ReaL [6], ImageNNt-v2 [56], ImageNet- Renditions [30], ImageNet-AAdversariall [31], and Object-
Net [4] testsets. SoViT-400m/14 is generally on par with viT-g/14 despite its smaller output width.
Broad few-shot linear transfer. We follow [24, 80] and evaluate closed-form linear regression probe for 10-shot classification across a wide range of classification tasks in Table 3. Again, So 400m/14 performs on-par with ViT-g/14 I/T-g/14 across the board.
5.2 Contrastive image-text tuning
Next, follow the locked-image text tuning (LiT) recipe [81] on the WebLI dataset [15] add zero-shot classification abilities to the pre-trained ViT-L/16, SoViT-400m/14 and ViT-g/14 image encoders. In this setup, new text encoder is trained using the contrastive image-text matching objective [54]. See Appendix D for details. Table 4 (second column) shows that SoViT-400m/1k is competitive with ViT-g/14, and substantially better than ViT-L/16.
5.3 Multitask Decoding
We also evaluate the three pretrained ViT models in multitask decoding as described in Section 4.2, where we follow the setup studied in [8]. We fix the decoder architecture to two layers since was found to perform well [8]. For evaluation, we report COCO CIDEr [48, 14, 731, OCR [50], VQAv2 [28] and GQA [37] accuracy and log-perplexity. In brief, the CIDEr score measures the similarity between a generated caption and reference captions, considering n-gram statistics, OCR question-answeriig evaluations. Results are summarized in Table 4. SoViT-400M performs on par with ViT-g/14.
evaluates optical character recognition, whereas both VQAv2 and GQA
5.4 Panoptic Segmentation
Additionally, evaluate SoViT-400m/14 14 panoptic segmentation [42], which is challenging dense scene understating task by closely following the setup in UViM [44]. At a high level, UVIM panoptic segmentation model consists of a visual image encoder and a decoder which maps the image representation to an intermediate code. The code later decoded to the panoptic segmentation mask using fixed VQVAE [71] model, which was pretrained panoptic masks [44]. In our experiments we initialize UViM's image encoder with ViT-L/16, SoVIT-400m/14 and ViT-g/14.
Following [44], we train the UVIM model using the COCO panoptic dataset (with 512 x 512 input resolution) and report the PQ metric. We achieve 43.5, 43.7 and 44.8 PQ points for ViT-L/16. SoViT-400m/14 and ViT-g/14 respectively. Our results indicate that dense segmentation tasks can be a limitation of the proposed optimal model shape, and different model shape might be derived in this domain. We leave this investigation for future work.
ImageNet top-1 1 Aaaaa
90%
So/14 224px
1/16@384px
B/16@384px
85%
B/32@ 384px
5.5 Flexifying So VIT-400M
Finally, since we do not include the patch size (sequence length) as part of the shape optimization, we verify that this is not limitation by flexifying [7] SoViT-400m/14 on I for 300 epochs. The performance of the resulting RixiiSoViT-4000 is shown in Fig 7 as green
HP 1214111172224 374
Sequence length (at inference)
Figure Flexification of SoViT-
400m/14 14 (abbr. So/14). See Section 5.5.

---

Table Summary of multitask decoding and zero-shot transfer results, see Sections 5.2 5.3.
ImgNet
OCR-VQA [50]
GQA [37]
VQAv2 [28]
COCO Capt. [14]
Model
Zero-shot
Acc [%6
Log-PPL
Acc [%c]
Log-PPL
Acc [%)
Log-PPL
CIDEr
Log-PPL
ViT-L/16
79.9
48.3
17.9
55.3
24.9
66.4
20.9
120
28.7
SoViT-400M
82.2
52.9
15.3
56.0
23.9
67.7
20.9
125
28.1
ViT-g/14
82.4
52.5
15.9
58.0
22.5
68.8
21.5
126
27.9
curve when varying the patch-size at inference time. A few reference VIT models from Table and [80] are added, confirming that SoViT-400m maintains a clear advantage. It is worth noting that flexifying does not rule out that other patch sizes could be compute-optimal. It merely demonstrates that So ViT-400M continues to perform quite well for other patch sizes when it is flexified.
Conclusion
In conclusion, we introduce efficient method for optimizing the shape of neural architectures and successfully apply it to vision transformers. Our analysis demonstrates that smaller models, trained at their optimal architecture shape for the right amount of compute, can match much larger models.
Acknowledgments and Disclosure of Funding
We thank Mostafa Dehghani, Andreas Steiner, Daniel Keysers, Neil Houlsby, Sam Smith, David Schneider-Joseph, Rodolphe Jenatton and the anonymous reviewers for their valuable feedback and discussions. We also thank the Google DeepMind unit at large for providing supportive research environment. use the big_ vision codebase [10, 9] for conducting experiments this project.
References
[1] Alabdulmohsin, I., Neyshabur, B., and Zhai, X. (2022). Revisiting neural scaling laws in language and vision. In Advances in neural information processing systems (NeurIPS). 1,
[2] Bahri, Y., Dyer, E., Kaplan, J., Lee, J., and Sharma, U. (2021). Explaining neural scaling laws. arXiv preprint arXiv:2102.06701. 4.
[3] Bansal, Y., Ghorbani, B., Garg, A., Zhang, B., Krikun, M., Cherry, C., Neyshabur, B., and Firat, O. (2022). Data scaling laws in NMT: The effect of noise and architecture. arXiv preprint arXiv:220..01994, 1,4
[4] Barbu, A., Mayo, D., Alverio, J. Luo, W., Wang, C., Gutfreund, D., Tenenbaum, J., and Katz, B. (2019). Objectnet: large-scale bias-controlled dataset for pushing the limits of object recognition models. Advances in neural information processing systems, 32. 9
[5] Bello, I., Fedus, W., Du, X., Cubuk, E. D., Srinivas, A., Lin, T.-Y,, Shlens, and Zoph, B. (2021). Revisiting resnets: Improved training and scaling strategies. Advances in neural information processing systems (NeurIPS). 3.
[6] Beyer, L., Hénaff, O. J., Kolesnikov, A., Zhai, X., and van den Oord, A. (2020). Are we done with imagenet? CoRR, abs/2006.07159. 8, 9, 19,
[7] Beyer, L., Izmailov, P., Kolesnikov, A., Caron, M., Kornblith, S., Zhai, X., Minderer, M., Tschannen, M., Alabdulmohsin. I., and Pavetic, F. (2023a). Flexivit: One model for all patch sizes. CVPR. 6,
[8] Beyer, I.., Wan, B., Madan, G., Pavetic, F., Steiner, A., Kolesnikov, A., Pinto, A. S., Bugliarello, E., Wang, X., Yu, Q. Chen, I.-C.,, and Zhai, X. (2023b). study of autoregressive decoders for multi-tasking in computer vision. 6, 9, 17
10

---

[9] Beyer, L., Zhai, X., and Kolesnikov, A. (2022a). Better plain vit baselines for imagenet- lk. 10
[10] Beyer, L.,, Zhai, X., and Kolesnikov, A. (2022b). Big vision. https://githubb.oo com/
10
[11] Boyd, S. P. and Vandenberghe, L. (2004). Convex optimization.. Cambridge university press. 15
[12] Brown, J. R., Zhao, Y. Shumailov, L,, and Mullins, R. D. (2022). Wide attention is the way forward for transformers. arXiv preprint arXh::2210.00644 1,3
[13] Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. Advances neural information processing systems (NeurIPS). 1,2
[14] Chen, X., Fang, H., Lin, T.-Y., Vedantam, R., Gupta, S., Dollár, P., and Zitnick, C. L. (2015). Microsoft coco captions: Data collection and evaluation server. arXiv preprint 5r.iv:1504.00325. 7,9. 7,9, 10
[15] Chen, X., Wang, X., Changpinyo, S., Piergiovanni, A., Padlewski, P., Salz, D., Goodman, S., Grycner, A., Mustafa, B., Beyer, L.., Kolesnikov, A., Puigcerver, J., Ding, N., Rong. K., Akbari,
H., Mishra, G., Xue, L.. Thapliyal, A., Bradbury, J., Kuo, W., Seyedhosseini, M., Jia, C., Ayan, B. K., Riquelme, C., Steiner, A., Angelova, A., Zhai, X., Houlsby, N., and Soricut, R. (2022). Pali:
A jointly-scaled multilingual language-image model. 3,
[16] Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung,
H. W., Sutton, C., Gehrmann, S., al. (2022). Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311.
[17] Cimpoi, M. Maji. S., Kokkinos, L. Mohamed, S., and Vedaldi, A. (2014). Describing textures in the wild. In Proceedings of the IEEE Conf. Computer Vision and Pattern Recognition (CVPR).
9
[18] Code, P. W. (2023). Papers With Code: ImageNet Benchmark. https: https://paperrsithccdd.e com/sota/ astta/iagge-aa celassification www.iiiiitttiioooiieieee imageenet.. [Online; accessed 16-May-2023). 8
[19] Dai, Z., Liu, H., Le, Q. V., and Tan, M. (2021). Coatnet: Marrying convolution and attention for
all data sizes. Advances in neural information processing systems (NeurIPS). 1, 8
[20] Dehghani, M., Amab, A., Beyer, L., Vaswani, A., and Tay, Y. (2022). The efficiency misnomer. ICLR.
[21] Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A., Caron,
M., Geirhos, R., Alabdulmohsin, I., Jenatton, R., Beyer, I.., Tschannen, M., Arnab, A., Wang. X., Riquelme, C., Minderer, M., Puigcerver, J., Evci, U., Kumar, M., van Steenkiste, S., Elsayed, G. F., Mahendran, A., Yu, F., Oliver, A., Huot, F., Bastings, J., Collier, M. P., Gritsenko, A., Birodkar, V., Vasconcelos, C., Tay, Y., Mensink, T., Kolesnikov, A., Pavetić, F., Tran, D., Kipf, T., Lučić, M., Zhai, X., Keysers, D., Harmsen, J., and Houlsby, N. (2023). Scaling vision transformers to 22
billion parameters. 3
[22] Deng, J., Dong, W., Socher, R., L.-.,,, K., and Fei-Fei, L. (2009). Imagenet: A large-scale hierarchical image database. In Conference Computer Vision and Pattern Recognition (CVPR).
2, 9
[23] Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint w0t.:1880.0805. 1
[24] Dosovitskiy, A., Beyer, L., Kolesnikov. A., Weissenborn, D., Zhai X., Unterthiner, T., De- hghani, M., Minderer, M. Heigold, G., Gelly, S., et al. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Representation Learning (ICLR). 1, 6, 7,9
[25] Fei-Fei, L., Fergus, R., and Perona, P. (2004). Learning generative visual models from few training examples: incremental bayesian approach tested on 101 object categories. Conference Computer Vision and Pattern Recognition (CVPR) Workshops.
11

---

[26] Ghorbani, B.. Firat, O., Freitag, M., Bapna, A., Krikun, M., Garcia, X., Chelba, C., and Cherry, C. (2021). Scaling laws for neural machine translation. arXiv preprint arXiv:2109.07740.
[27] Gordon, M. A., Duh, K., and Kaplan, J. (2021). Data and parameter scaling laws for neural machine translation. In Conference on Empirical Methods in Natural Language Processing. 1,
[28] Goyal, Y., Khot, T., Summers-Stay, D., Batra, D., and Parikh, D. (2017). Making the in vqa matter: Elevating the role of image understanding in visual question answering. In CVPR. 7,9, 10
[29] He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR). 2
[30] Hendryeks, D., Basart, S., Mu, N., Kadavath, S., Wang, F., Dorundo, E., Desai, R., Zhu,
T., Parajuli, S., Guo, M., Song, D., Steinhardt, J., and Gilmer, J. (2021a). The many faces of robustness: A critical analysis of out-of-distribution generalization. ICCV. 9
[31] Hendrycks. D., Zhao, K., Basart, S., Steinhardt, J., and Song, D. (2021b). Natural adversarial examples. CVPR. 9
[32] Henighan, T., Kaplan, J., Katz, M., Chen, M., Hesse, C., Jackson, J., Jun, H., Brown, T. B., Dhariwal, P., Gray, S., et al. (2020). Scaling laws for autoregressive generative modeling, arXiv preprint arXiv:2010.14700
[33] Hernandez, D., Kaplan, J., Henighan, T., and McCandlish, S. (2021). Scaling laws for transfer. arXiv preprint arXiv:2102.01229)
[34] Hestness, J., Narang, S., Ardalani, N., Diamos, G., Jun, H., Kianinejad, H., Patwary, M., Ali, M., Yang, Y., and Zhou, Y. (2017). Deep learning scaling is predictable, empirically, arXiv preprint arXiv: 004099 1,
[35] Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E.,, Cai, T., Rutherford, E., Casas, D. d. L.,, Hendricks, L. A., Welbl, J., Clark, A., et al. (2022). Training compute-optimg large language models. In Advances in neural information processing systems (NeurlPS). 1, 3, 4,5
[36] Howard, A. G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andrectto, M., and Adam, H. (2017). MobileNets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint 170.:v:7704008880 2,
[37] Hudson, D. A. and Manning, C. D. (2019). GQA: dataset for compositional question answering over real-world images. arXiv preprint arXiv:1902.09506. 7, 9, 10
[38] Hutter, M. (2021). Learning curve theory. arXiv preprint arXiv:2102.04000 4,
[39] Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. (2020). Scaling laws for neural language models. arXiv preprint arXiv:2001.08361. 1, 4,
[40] Kather, J. N., Weis, C.-A., Bianconi, F., Melchers, S. M., Schad, L. R., Gaiser, T., Marx, A., and Z'ollner, F. G. (2016). Multi-class texture analysis in colorectal cancer histology. Scientific reports, 6:27988. 9
[41] Kingma, D. and Ba, J. (2014). Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980. 5
[42] Kirillov, A., He. K., Girshick, R., Rother, C., and Dollar, P. (2019). Panoptic segmentation. In Conference on Computer Vision and Pattern Recognition (CVPR). 9
[43] Kolesnikov, A., Beyer, L.., Zhai, X., Puigcerver, J., Yung, J., Gelly, S., and Houlsby, N. (2020). Big transfer (BiT): General visual representation learning. In European Conference on Computer Vision (ECCV). 2,
[44] Kolesnikov, A., Susano Pinto, A., Beyer, L., Zhai, X., Harmsen, J., and Houlsby, N. (2022). UViM: unified modeling approach for vision with learned guiding codes. Advances in neural information processing systems (NeurIPS).
12

---

[45] Krause, J., Stark, M.. Deng, J., and Fei-Fei, L. (2013). 3d object representations for fine-graing categorization. In 4th Interaatinaal IEEE Workshop on 3D Representation and Recognition (3dRR-13), Sydney, Australia. 9
[46] Krizhevsky, A. (2009). Learning multiple layers of features from tiny images. Technical report.
[47] Li, Z.,, Wallace, E., Shen, S., Lin, K., Keutzer, K., Klein, D., and Gonzalez, J. (2020). Train big, then compress: Rethinking model size for efficient training and inference of transformers. In International Conference on Machine Learning (ICML). 1,3
[48] Lin, T.-Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L. (2014). Microsoft coco: Common objects context. In ECCV. 7.9
[49] Liu, Z.,, Hu, H., Lin, Y., Yao, Z.,, Xie, Z., Wei, Y., Ning. J., Cao, Y., Zhang, Z.., Dong, I.., et al. (2022). Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12009-12019. 3, 8
[50] Mishra, A., Shekhar, S., Singh, A. K., and Chakraborty, A. (2019). OCR-VQA: Visual question answering by reading text in images. In ICDAR. 7,9, 10
[51] Parkhi, O. M., Vedaldi, A., Zisserman, A., and Jawahar, C. V. (2012). Cats and dogs. In IEEE Conference on Computer Vision and Pattern Recognition. 9
[52] Patterson, D., Gonzalez, J., Le, Q., Liang, C., Munguia, L.-M., Rothchild, D., So, D., Texier, M., and Dean, J. (2021). Carbon emissions and large neural network training. arXiv preprint arXiv:2 104. 10350.
[53] Pham, H., Dai, Z., Ghiasi, G., Kawaguchi, K., Liu, H., Yu, A. W. Yu, Chen, Y-T., Luong, M.-T., Wu, Y., et al. (2021). Combined scaling for zero-shot transfer learning. arXiv preprint arXiv:2111.10051 18
[54] Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. (2021). Learning transferable visual models from natural language supervision. In ICML. 9
[55] Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J., Song, F., Aslanides, J., Henderson, S. Ring, R., Young. S., et al. (2021). Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:21 1446. 1, 3,3
[56] Recht, B., Roclofs, R., Schmidt, L., and Shankar, V. (2019). Do imagenet classifiers generalize to imagenet? CoRR, abs/1902.108) 10811. 8.
[57] Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M. S., Berg, A. C., and Fei-Fei, L. (2014). Imagenet large scale visual recognition challenge. CoRR, abs/1 40s/1500
[58]
Sandler, M., Howard, A., Zhu, M., Zhmoginov, A., and Chen, L.-C. (2018). MobileNetV2: In- verted residuals and linear bottlenecks. In Conference on Computer Vision and Pattern Recognition (CVPR). 2
[59] Sharma, U. and Kaplan, J. (2022). Scaling laws from the data manifold dimension. Journal of Machine Learning Research, 23. 4, 6
[60] Shazeer, N. and Stern, M. (2018). Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning (ICML). 16, 17, 18
[61] Singh, M., Duval, Q., Alwala, K. V., Fan, H., Aggarwal, V., Adcock, A., Joulin, A., Dollár, P.. Feichtenhofer, C., Girshick, R., et al. (2023). The effectiveness of mae pre-pretraining for billion-scal pretraining. arXiv preprint arXiv:2303. 22223333334433 13496. 3, 8
[62] Steiner, A. P., Kolesnikov, A., Zhai, X., Wightman, R., Uszkoreit, J., and Beyer, L. (2022). How train your vit? data, augmentation, and regularization vision transformers. Transactions on Machine Learning Research. 3
13

---

[63] Sun, C., Shrivastava, A.. Singh, S., and Gupta, A. (2017). Revisiting unreasonable effectiveness of data in deep learning era. International Conference on Computer Vision (ICCV). 2
[64] Tan, M. and Le, Q. (2019). EfficientNet: Rethinking model scaling for convolutional neural networks. In International Conference on Machine Learning (ICML). 2,
[65] Tay, Y., Dehghani, M., Abnar, S., Chung, H. W., Fedus, W., Rao, J., Narang, S., Tran, V. Q.. Yogatama, D., and Metzler, D. (2022a). Scaling laws vs model architectures: How does inductive bias influence scaling? arXiv preprint arXiv:2207.1053
3
[66] Tay. Y., Dehghani, M., Rao. J., Fedus, W., Abnar, S., Chung, H. W., Narang, S., Yogatama, D., Vaswani, A., and Metzler, D. (2022b). Scale efficiently: Insights from pre-training and fine-tuning transformers. In International Conference on Representation Learning (ICIR). 1,
[67] Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. (2021). Training data-efficient image transformers distillation through attention. In International Conference on Machine Learning (ICML). 3
[68] Touvron, H., Cord, M., and Jégou, H. (2022). DeÎT III: Revenge of the VIT. In Computer Vision- ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXIV, pages 516-533. Springer. 3
[69] Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. (2023). Open and efficient foundation language models. 2,
Llama:
[70] Touvron, H., Vedaldi, A., Douze, M., and Jégou, H. (2019). Fixing the train-test resolution discrepancy. Advances neural information processing systems, 32. 8
[71] Van Den Oord, A., Vinyals, O., et al. (2017). Neural discrete representation learning. Advances in neural information processing systems (NeurIPS). 9
[72] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, E., and Polosukhin, I. (2017). Attention is you need. Advances in neural information processing systems (NeurIPS). 7
[73] Vedantam, R., Lawrence Zitnick, C., and Parikh, D. (2015). Cider: Consensus-based image description evaluation. In CVPR. 9
[74] Welinder, P., Branson, S., Mita, T., Wah, C., Schroff, F., Belongie, S., and Perona, (2010). Caltech-UCSD Birds 200. Technical Report CNS-TR-2010-001, California Institute of Technology.
[75] Wu, Z., Shen, C., and Van Den Hengel, A. (2019). Wider or deeper: Revisiting the resnet model for visual recognition. Puttern Recognition. 2
[76] Yang, Y. and Newsam, S. (2010). Bag-of-visual-wordd and spatial extensions for land-use classi- fication. In ACM SIGSPATIAL International Conference on Advances in Geographic Information Systems (ACM GIS).
[77] Yu, J., Wang, 7., Vasudevan, V., Yeung, I.., Seyedhosseini, M., and Wu, Y. (2022). CoCa: Contrastive captioners are image-text foundation models. arXiv preprint arXiv:2205.00 000000...1......
[78) Yuan, L., Chen, D., Chen, Y.-L., Codella, N., Dai, X., Gao, J., Hu, H., Huang, X., Li, B., Li, C., et al. (2021). Florence: A new foundation model for computer vision. arXiv preprint
arXiv:2111.11432.
[79] Zagoruyko. S. and Komodakis, (2016). Wide residual networks. arXiv preprint 1rXiv:160..07146. 2
[80] Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. (2022a). Scaling vision transformers, In Conference on Computer Vision and Pattern Recognition (CVPR). 1, 2, 3, 4, 5, 7, 8, 9, 10
[81] Zhai, X., Wang, X., Mustafa, B., Steiner, A., Keysers, D., Kolesnikov, A., and Beyer, L. (2022b). LiT: shot transfer with locked-imen -im-ee text tuning. In CVPR, pages 18123-18133, 2,
14

---

A Scaling Laws Analysis
In this appendix, we present proofs of two claims in the paper. First, we show that (2) is quasiconvex on its first argument X5. Second, we derive (5).
A.1 Quasiconvexity Proof
We assume throughout the proof that ag. by strictly positive, otherwise (xg, t) is monotone function on its first argument and the statement holds trivially.
To establish the quasiconvexity of fx(XE,, t.) in (2), we observe that:
Ofk
ox x (1+0L) + Babat "x -Ax (I+aa) + Bx
Setting the derivative to gives the unique solution in R+:
A
B
At the limit X2 - oc, the term involving X vanishes and have the asymptotic relation:
fa(x t) ~ Bat= 3xt=xx
which is an increasing function since by O. Since à is the only point in RT where 0fz/Ox; we conclude that ['x, t) is monotone increasing for all Xx 2 L.
Similarly, when X o have:
Ju () Or.Xx ak
which is monotone decreasing. Therefore, f'(xx t) < O for all X < &. Combining both results implies that (a t) monotone decreasing the domain E (0, &) and is monotone increasing in the domain E (#, oo).
function f(y) is said to be quasi-convex if for any yn and y/2 in its domain and any À [o, 1). one has [11]:
f(xy1 + (1 A)12) max (f()), f(z2)-
Suppose for the purpose of obtaining a contradiction that fi(Xx, t) is not quasiconvex on its first argument. Then, there exists two points y/, 3/2 R+ and [0, such that the above condition is violated. Let ú Ay11 + (1 A)82. But, then, by the mean-value theorem, there must exist two points Cj E [y1. and C2 E (g. y2l where:
fl(11) f(ÿ) ((V)
0
1 91
fl((2) f(12] - (()
S <0
1/2 ÿ
with 02 > 61 This implies that e 2 - and 02 S I. which a contradiction. Therefore, fx (x t)
quasi-convex on its first argument.
A.2
Derivation of (5)
Rearranging the expression in (4), we have:
Babe
(xt))atay to
agaa
From this and (2), we obtain:
fx t) an (x+) -as + Ba Bx(x))
C4 Bgbb (xz
Eut -o Ekr
where we plugged in the last expression. Simplifying yields (5) for some constants F,G 0.
15

---

B Shape Optimization
B.1 Hyper-parameters
Table Common hyper-paramelers settings for both star and grid sweeps.
Image Resolution
224 224
Batch size
128
Preprocessing
Rescale(-1, 1)
Augmentation
InceptionCrop, Left- Right Flip
Optimizer
AdaFactor [60]
Gradient Clipping
1.0
Learning Rate
8c-4
Label Smoothing
0
Weight Decay
0.03 x 8e-4
Schedule
Reverse SQRT, 10K Warmup steps, 50K Cooldown steps
Table 5 provides the set of hyperparameters used in the star and grid sweeps. We small batch size of 128 here in order to train multiple models parallel on small hardware topologies.
B.2 Star Sweep
In the star sweep, we use the center x(6) = (1968, 40, 6144) as our starting point. To esti- mate the scaling exponents 8; (4) for each dimension separately, we vary width the grid (608, 768, 928, 1088, 1328, 1648), depth in the grid (8, 10, 12, 16, 20, 24), and MLP dim in the grid (1088, 1360, 1728, 2160, 2592, 3072). We train each model on 500K, 1M, and 2M steps. We always fix the patch size to 14 x 14 and the number of attention heads to 16.
B.3 Grid Sweep
In the grid sweep, we pretrain each architecture on 600M examples. We use the cross-product of:
1. width: 416, 512, 608, 768
2. depth: 6, 8, 10, 12
3.
MLP Size: 768, 928, 1088, 1360
Some important considerations to be taken into account include:
When designing the grid sweep, made sure that the compute-optimal model selected lies strictly in the interior of the grid, not on its boundary. This is because if it lies at the boundary (e.g. its depth the maximum depth used in the grid), one cannot determine if is compute-optimal or if increasing that dimension will yield even better models. This can be an iterative process, which additional grid points are added to the sweep if necessary.
When identifying the model, we ensured that it is compute-optimal for a good range of compute (not only at some isolated point). Since the model is now compute-optimal for range of compute budgets, select starting point in recipe the least compute is optimal for. For example, if model is compute-ope optimal for computes ranging from TFLOPs to 2 TFLOPs, use 1 TFLOPS in recipe. In other words, eir on the side caution, giving preference to larger models as we scale up the vision transformer (ViT).
Generally, the grid sweep should be tightly packed; e.g. with increments of 20% only in each dimension. By contrast, increments in the star sweep should be large in order to identify the scaling exponents reliably.
Even though the goal in the grid sweep is to identify "small" architecture that is compute- optimal for "small" amount of compute, the amount of compute used the analysis should be large enough for results to be reliable and for power laws take effect. That is why in our experiments, we use 100M training examples in the grid sweep as opposed, for instance, to using only few million examples.
16

---

C Multitask Decoding Setup
Table 6: Multi-task decoding Hyperparameter Settings.
Image Resolution
Batch size
512
Preprocessing
Rescale(-1, 1), ResizeSmall(256) CentralCrop(224)
Augmentation
InceptionCrrop(222) Left-Right Flip
Optimizer
AdaFactor [60]
Epochs
50
Gradient Clipping
Label Smoothing
0.1
Learning Rate
3c-4
Weight Decay
le-4
Schedule
Cosine, 10% Warmup period
Vocabulary Size
32k
Encoder Dropout Rate
Decoder Dropout Rate
0.1
Table 6 summarizes the hyperparameter settings for the multitask decoding setup in Section 4.2 and Section 5.3. We always fix the decoder to layers since it generally performs well [8].
17

---

D LiT Training Setup
Table Locked -mmgee text tuning (LiT) Hyperparameter Settings.
Image Resolution
224 x 224
Batch size
32K
Preprocessing
Rescale(-1, 1)
Augmentation
None
Optimizer
AdaFactor [60]
Total Examples
900M
Gradient Clipping
1.0
Learning Rate
le-3
Weight Decay
lc-4
Schedule
Cosine, 20% Warmup period
Vocabulary Size
32k
Bias Init
-10
Temperature Init
10
Internal Representation
1,152
Table 7 summarizes the hyperparameter settings for the locked-image text turning (LiT) setup, which is used to report zero-shot classification accuracy in Table 4. We use large batch size of 32K in this setup because it improves the performance of contrastive training [53].
18

---

Table ImageNet fine-tuning settings. Settings in the first section vary with resolution, settings in the middle section were explored, and settings in the last section are unexplored good defaults.
Full model fine-tuning
224
384px
518
Learning rate decay
0.85
0.9
0.9
Random augment
2,10
2,10
Mixup
0.2
0.2
Training duration
50k steps (20 epochs)
Learning rate
0.03
Polyak averaging (EMA)
Optimizer
SGD with 0.9 Momentum
Gradient clipping
1.0
Weight decay
Batch size
512
Learning rate schedule
Cosine with 500 steps linear warmup inception, cepppioon_copppppppp (RandomResize)
Image crop
Random flip
Horizontal
Loss
Sigmoid cross-entropy [6]
Head init
kernel=O, bias=-6.9
Train and minival splits
train[:98%] and train [98%:]
E Transfer to ImageNet-1k
E.1
Full model fine-tuning
Table lists the settings for the ImageNet-1K fine-tuning results presented in Table in the main paper. The only three settings which differ across resolutions are learningrate decay, random augment and mixup strenghts. We did explore various learningrates, training durations (mostly shorter) well as Polyak averaging, although the same setting shown the table appears to best across the board. Finally, we list various other settings which did not explore. We simply used good default values from experience.
E.2 Linear probe on frozen encoder
We take the image representation the pre-logits, i.c. the 1152-dimensional vector that comes out of the MAP-head and feeds right into the linear classification layer. For each of ViT-L/16, So ViT-400mm and ViT-g/14, perform grid-search over the following settings, and select the best-performing model on minival (2% of train) to be reported in Table 2: Augmentation: resize (256) random, pp (224) vs. inception_ cereppiooocooppppp (224). learning rate: 001, 0. 0003, 0. 0.0011, epochs: 1, 3, 10, weight decay: 0. 0001, None. should be noted that we keep various other settings to "known good defaults" based on prior explorations with similar models (i.e. plain ViTs). Table 9 summarizes key settings.
19

---

Table 9: ImageNet linear probing settings. Settings in the first section were grid-searched for each model, settings in the last section are unexplored good defaults.
Linear probe 224 px
VIT-I/16
SoViT-400m/14
ViT-g/14
Learning rate
0.001
0.0003
0.001
Weight decay
0.0001
Training duration
24.7 k steps (10 epochs)
Image crop
resize (256) raaddm_ crop (224)
Random augment
Mixup
0.1
Learning rate decay
Polyak averaging (EMA)
Optimizer
SGD with 0.9 Momentum
Gradient clipping
Batch size
512
Learning rate schedule
Cosine with 10% linear warmup Horizontal
Random flip
Loss
Sigmoid cross-entropy [6]
Head init
kernel=0, bias=-6.9
Train and minival splits
train[:99K] and train [99%:]
20


Reviews:
[{"review_id": "WB3uD5Kxjl", "reviewer": "Reviewer_T7K3", "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. \nA 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.", "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.\n2. The compute-optimal model shape is different among tasks. How to search for an optimal model which can be applied to various downstream tasks?", "limitations": "As above.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written, with well-defined formulas and sufficient supporting materials.\n2. 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.\n3. Extensive experiments are implemented to validate the method, including image classification,  multitask decoding, and segmentation tasks.\n", "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.\n2. 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.\n3. A few typos in the paper, such as x_{k} in A.1 Quasiconvexity Proof.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and careful review. Please see our response below:\n\n- In Figure 3, each dot corresponds to a model architecture pretrained on 600M examples and evaluated on one downstream metric.\nThe metrics  from left to right are: 5-shot, 10-shot, and 25-shot (all in ImageNet). In the rightmost figure, the metric is an average log-perplexity score across a mixture of four tasks including VQA and captioning, as we describe in Section 4.2. \nIn the first three figures, the downstream metrics are for classification and we observe that the compute-optimal model highlighted in blue is compute-optimal in all three cases because it lies in the efficient frontiers in all three cases. But, it is not compute-optimal for image-to-text tasks as shown in the rightmost figure. For this reason, a compute-optimal shape in one domain may not be optimal in others. We will clarify this in the paper.\n\n- In Table 4, we show that SoViT-400M is comparable to ViT-g/14 in multimodal tasks. As you mention, it does not perform strictly better but performing equally well is itself a significant gain, because SoViT-400M is much smaller and less costly (GFLOPs, Images/Core/s) than ViT-g/14, as shown in Figure 2.\n\n- We did not explain log-perplexity and CIDEr because they are standard metrics in the literature (e.g. [1] and Section 9.3.2 in [2]). We will include references for their definitions in the revised version of the paper.\n\n- Thanks for spotting the typo in A.1. We’ll fix it.\n\n- Regarding the optimality of the shape and Figure 6, we clarify this in Lines 115-116. Due to modeling assumptions, approximations, and the finite possible number of experiments, we can only approximate the optimal shape. In Figure 6, we see that deviating from the predicted optimal width, for example, does degrade performance. For other dimensions, the performance does not change significantly when they are increased but it does degrade when they are decreased. We do believe that this highlights our argument that one can identify a near-optimal shape using the recipe we propose.\n\n- The reason we only report <33%, >33%, and >200% in Figure 6 is because we do observe a drop in all cases when using <33% so it is a sufficient demonstration. On the other hand, increasing the dimension by >33% or >200% both give similar results, so we expect similar results when using, for example, >100%.\n\n- It is true that the compute-optimal shape is different across domains when the model is small. The scaling exponents, however, are the same in both domains: (1) image classification and (2) image-to-text, as discussed in Section 4.2 . This means that in terms of order-of-magnitude, the optimal shape is similar in both domains when the model is sufficiently large. We do provide evidence to support this in our experiments by evaluating SoViT-400M in multimodal tasks, such as zero-shot classification, captioning, and VQA.  \n\n[1] Vedantam, R, et al. \"Cider: Consensus-based image description evaluation.\" CVPR, 2015.\n\n[2] Zhang, A., Lipton, Z. C., Li, M., & Smola, A. J. Dive into deep learning. 2021.\n"}, {"review_id": "Q9JqZt3ENC", "reviewer": "Reviewer_DDss", "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.\nWhereas 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. \nThe paper empirically solves this question via numerous experiments, while answering related questions and distilling interesting insights on the way. ", "questions": "I'm ready to upgrade my rating if the authors clarify some of the points mentioned above.", "limitations": "yes", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "### Writing\n- The introduction is well written and clear\n- quite complex and unclear at times\n\n### Method \n- 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.\n- The insights about the choice of the joint functional form (eq. (2)) are interesting and well-grounded\n- Definitely a valuable paper for the community, even though some aspects could be improved\n\n### Experiments\n- multiple experiments on several benchmarks support the initial claim that a smaller-but-optimal architecture achieves as well as vanilla SotA architectures\n- Experimental findings are useful and valuable (Section 4.1)\n- 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\n- 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\n\n### Method\n- 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)\n- About the difference between small and large models\n  - it is claimed that \"in small models, an optimal shape in one domain is not necessarily optimal in others.\" (Figure 3)\n  - How are \"small\" and \"large\" models exactly defined? This seems like a very convenient subjective definition.\n  - Why does the model size matter so much when it comes to the optimal architecture across application domains?\n\n### Experiments\n- Figure 6: \n  - \"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?\n  - 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)\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and careful review. Please see our response below:\n\n- In Figure 1, we use different axes for different dimensions. This may make it difficult to see how the MLP dimension is scaled faster than the others. Note, for example, that going from 1T to 100T GFLOPs corresponds to an increase in MLP by around a factor of $\\times 3$. For depth, on the other hand, the increase is around a factor of $\\times 2$. Generally, the scaling exponent for the MLP dimension is $\\approx 0.6$, which is larger than depth or width.\n\n- Equation 4 is obtained by setting the derivative in (3) to zero and solving for $x$. The scaling exponents need not be invariant to the choice of the dimension as you mention. However, based on prior theoretical works that explain scaling laws via a space partitioning argument (Bahri, et al. 2021; Hutter, 2021; Sharma and Kaplan, 2022), we would expect the exponent $c$ to be roughly similar across all dimensions, which seems to be indeed the case in Figure 5. So, only $c$ is expected to be invariant. \n\n- Regarding the model size, the terminology “small” and “large” is meant to convey in simpler terms the following results. In Figure 3, we observe that the compute-optimal model for classification highlighted in blue is not compute-optimal for image-to-text tasks as shown in the rightmost figure. For this reason, a compute-optimal shape in one domain may not be optimal in others. However, we also show that the scaling exponents are similar. Hence, in terms of order-of-magnitude, the optimal shape is similar in both domains when the model is sufficiently large. We do provide evidence to support this in our experiments by evaluating SoViT-400M in multimodal tasks, such as zero-shot classification, captioning, and VQA.  \n\n- We keep compute fixed by changing the training duration, making it longer for smaller models and shorter for larger models, such that the total FLOPs is the same.\n\n- Regarding the optimality of the shape and Figure 6, we clarify this in Lines 115-116. Due to modeling assumptions, approximations, and the finite possible number of experiments, we can only approximate the optimal shape. In Figure 6, we see that deviating from the predicted optimal width, for example, does degrade performance. For other dimensions, the performance does not change significantly when they are increased but it does degrade when they are decreased. We do believe that this highlights our argument that one can identify a near-optimal shape using the recipe we propose.\n"}, {"review_id": "bstHJAA3G5", "reviewer": "Reviewer_s4n2", "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.\n", "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.\n2. 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”?\n3. 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. ", "limitations": "The authors have addressed the limitations of their work. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. 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).\n3. 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.\n", "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. \n\nSimilarly, a quick definition of width, MLP dimension and depth could be helpful for readers less familiar with the details of vision transformers.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and careful review. Please see our response below:\n\n- Compute is defined in terms of FLOPs throughout the paper. However, when the architecture is fixed, compute becomes proportional to the number of seen examples. That’s why in Line 127, we refer to infinite data as compute-unbounded.\nThere is indeed a dependence between the architecture and its compute as you mention. However, we can always treat them separately by first fixing the architecture and, then, training for the chosen amount of compute (in FLOPs). This is why we can have both $t$ and $x$ in Equation 2. The reason we do not use the sample size is because large models are more sample-efficient (Zhai, et al. 2022) so if the goal is to minimize compute by minimizing the size of the training data, the solution is to (trivially) scale up the size of the model indefinitely.\n\n- We will add a brief description of the shape dimensions in the revised version of the paper.\n\n- The impact of the model shape on compute is captured by the exponent $b$, which are shown in Figure 5. These are automatically taken into account when optimizing the shape for compute in Eq 4. Note that in Eq 4, increasing the value of $b$ would decrease the scaling exponent $s$.\n\n- The baselines ViT-g/14 and ViT-G/14 are not compute-optimal. This is what we demonstrate in this work. We use them because they are widely used in the literature, and hence serve as valid baselines.\n\n- Flexifying does not rule out that other patch sizes could be compute-optimal. Our intent is merely to demonstrate that SoViT-400M continues to perform quite well for other patch sizes when it is flexified. \n"}, {"review_id": "E8MToWeJId", "reviewer": "Reviewer_hmtB", "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. \n", "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. \n2. 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?\n3. Could the authors report the total TPU-hours/cost of analysis?\n4. Minor:  Should scaling “vision transformations” be scaling “vision transformers”? Line 70\n5. 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? \n7. How hyper-parameters are set for different shapes is unclear to me. Could you please clarify this?\n8. What do the question marks in Table 1 indicate?\n9. ViT-g/14 seems to work better for multitask decoding in table-3. Do the authors have an intuition of why that could be? \n\nIf all my concerns and questions are addressed appropriately I am willing to increase my score for the paper. \n", "limitations": "The limitations of the work are not adequately discussed:\n1. The compute cost of the analysis?\n2. Reproducibility of the experiments?\n3. Are there any implicit assumptions which may effect the laws derived empirically in the paper?\n\nI encourage the authors to discuss these limitations in the paper. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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. \nFurthermore, 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. \n", "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. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and careful review. Please see our response below:\n\n- We disagree that our work lacks originality. We provide an approach based on scaling laws for inferring compute-optimal model shapes, which has not been done in the literature before. In addition, we introduce the functional form in Eq 3, which helps in reducing the number of experiments significantly. We also demonstrate that a combination of a star and a grid sweep is sufficient for inferring compute-optimal shapes, thereby reducing the cost of the analysis further.\n\n- We agree that having different architectures on different domains may require repeating a similar analysis. However, the fact that different domains may result in different optimal shapes is not a limitation of our work per se, since our goal is to discover those optimal shapes, not alter them. We would like to emphasize, however, that the scaling exponents are similar across the two domains we studied. Hence, in terms of order-of-magnitude, the optimal shape is similar in both domains when the model is sufficiently large. We provide evidence to support this in our experiments by evaluating SoViT-400M in multimodal tasks, such as zero-shot classification, captioning, and VQA.  \n\n- Regarding open-sourcing the code, we use a publicly available codebase that was removed to preserve anonymity and will be included in the final version of the paper. In the Appendix, we do provide the full training configuration for our experiments and provide a full description of the star and grid sweeps. \n\n- Optimizing the shape while keeping the parameter count fixed indeed helps. We demonstrate this in Figure 6 (leftmost figure). Here, both SoViT-150m and the baseline (denoted B-150m) have the same size, yet SoViT-150m performs better.\n\n- In our analysis, we assume that no overfitting occurs. This is crucial for Eq 2 to be valid, otherwise increasing the compute t can increase the loss at some point. To ensure that there is no overfitting, the size of the pretraining dataset should be quite large. For this reason, we do our analysis on JFT-3B.\n\n- The relation between TPU hours and GFLOPs is reported in Figure 4. \n\n- Thank you for spotting the typo in Line 70. We’ll fix it.\n\n- The same hyperparameters are used for all architectures. These are provided in the Appendix (e.g. see Appendix B.1 for supervised pretraining and Appendix C for multitask decoding).\n\n- We will fix the question marks in Table 1 and replace them with dashes. Thanks for pointing this out.\n\n- The difference between SoViT-400m and ViT-g/14 in Table 4 for multitask decoding is negligible in our opinion; compare it for instance with the difference between SoViT-400m with ViT-L/16 that is similar to it in size. \n"}]

Author rebuttal:

Meta-review: This paper explores scaling laws and optimization strategies for ViT shapes / architecture hyper-parameters, introducing a new shape optimized variant of ViTs, evaluated on multiple downstream tasks. The paper received all positive reviews from 4 expert reviewers, who appreciated a well written paper, extensive experiments and analyses, 

Minor weaknesses raised were on clarity, in particular some Figures, and metrics. The paper was extensively discussed between reviewers and authors, which led to improvements of the presentation of an already initially well received paper, thus a positive outcome of the peer reviewing process.

The AC recommends acceptance.
Decision: Accept (poster)