SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ibrahim Alabdulmohsin*, Xiaohua Zhai*, Alexander Kolesnikov, Lucas Beyer*
{ibomohsin,xzhai,akolesnikov,lbeyer}@google.com
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.
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, it has been observed that the benefit of scale often follows a predictable power law in which the performance  $f(x)$  (e.g. error rate or log-perplexity) satisfies  $f(x) \sim \beta x^{-c} + \varepsilon_{\infty}$  for some  $\beta, c > 0$  as one varies the scaling dimension  $x$  (e.g. data or model size), if the remaining dimensions are not bottlenecks [34, 39, 27, 26, 3, 1]. Here,  $\varepsilon_{\infty}$  is the irreducible loss.
However, the simple power-law relation becomes more complicated when compute is considered. In this case, power laws are observed only along the compute-optimal frontier. Otherwise, scaling up the model size for a 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 to 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  $4 \times$  smaller in size.
Unfortunately, in both [39] and [35] among others, the "size" of a model is equated with its parameter count, with no special consideration for model "shape dimensions", such as "depth" or "width". The rationale behind this choice follows from the surprising observation that the transformer shape had little impact on its scaling behavior in language modeling (LM) when performance is measured upstream (e.g. using log-perplexity) [39, 32, 33]. Nevertheless, follow-up analysis suggests that shape plays a pivotal role in other domains, such as in machine translation [47] and also in language modeling for downstream performance [66], with recent works even advocating for extreme aspect ratios, such as a single wide attention layer [12].
* Significant technical contributions.
Figure 1: Predicted efficiency frontier (depth, width, MLP dimension, and parameter count) in SoViT. In large models, optimal shapes follow a 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 1T to 100T GFLOPS.
In vision, in particular, much earlier works using convolutional neural networks (CNNs) pointed out that the parameter count is indeed a 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 performance of much larger models despite being pre-trained with equal compute. It is derived from a recipe we introduce for optimizing the shape of neural architectures, such as their depth and width. A 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 experiments, we identify a scaling strategy in ViT for all three dimensions: width (internal representation), depth, and MLP size. For comparison, [35] requires over 400 experiments to optimize a single dimension (the parameter count) alone.
One major finding is 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-parameter GPT3 [13] — as well as LLaMA with its 13B parameter variant outperforming GPT3 on most benchmarks [69]. By introducing SoViT, we establish this phenomenon in vision as well.
Figure 1 summarizes how the various shape dimensions are scaled in SoViT (see Section 3 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 a ViT for the compute-equivalent of ViT-g/14 [80] pretrained on 16B JFT images [63]. We call this  $2.5 \times$  smaller model SoViT-400m/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/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.
Figure 2: Optimizing for the compute-equivalent of ViT-g/14 results in the  $2.5 \times$  smaller SoViT-400m/14 model achieves equivalent results across a wide range of benchmarks. Our model performs exceptionally well on the competitive ImageNet (ILSRCV2012) benchmark in comparison with significantly larger models from the recent literature [61, 78, 49, 80].
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, 55]. 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 were 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 on ImageNet (ILSRCV2012) with 2 billion parameters [80] and  $90.9\%$  top-1 accuracy with 4 billion parameters [15]. When scaled to 22 billion parameters, ViT exhibits state-of-the-art alignment to 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 ViT-B models and the reported accuracy improvement comes with an increase in FLOPs of up to  $\times 4$ , making it difficult to draw conclusions about the suggested shape'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 a fixed increase in compute (e.g.  $\times 2$ ). Afterwards, each dimension is scaled up by the same ratio with every subsequent increase in compute. In contrast, we expand scaling laws to simultaneously account for model size and compute beyond the efficient frontier and leverage them to derive the optimal scaling exponents for each dimension separately, as outlined in Section 3.
Throughout our analysis, we use downstream metrics, e.g. ImageNet 10-shot error, when measuring performance instead of upstream metrics. This follows recent reports arguing that upstream performance may not reflect downstream performance in language and vision [65, 80].
We use GFLOPs as a proxy for compute since it is hardware-agnostic and correlates well with actual wall-clock core-hours (see Figure 4). However, GFLOPs can have limitations [5, 20] and may not be a perfect predictor for the metric of interest (e.g. core hours) in all model and hardware types. Note that we focus on scaling the shape of the architecture, not on improving its training protocol, which can be similarly beneficial [5, 67, 62, 68].
3 Scaling Strategy
Notation. We begin with a formal description of the problem. We represent a neural architecture as a tuple  $\mathbf{x} = (\mathbf{x}_1, \mathbf{x}_2, \dots, \mathbf{x}_D) \in \mathbb{N}^D$  containing  $D$  shape dimensions, such as width, depth and MLP size. We denote compute such as GFLOPs by  $\mathbf{t}$ . We designate  $f: \mathbb{N}^D \times \mathbb{R}^+ \to \mathbb{R}$  a performance metric of interest, such as downstream ImageNet 10-shot error rate. Specifically,  $f(\mathbf{x}, \mathbf{t})$  results from (pre)-training an architecture  $\mathbf{x}$  for a fixed compute budget  $\mathbf{t}$ . We always assume that  $f$  corresponds to a loss, meaning lower values are better.
Figure 3: A 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-1k 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/VQA), respectively. For captioning/VQA, we average log-perplexity scores (see Section 4.2). In the leftmost three figures, usingImagenet-1k 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-text 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  $\mathbf{t}$  is to identify  $\mathbf{x}^{\star}$  (depending on  $\mathbf{t}$ ) such that:
$$
f \left(\mathbf {x} ^ {\star}, \mathbf {t}\right) - \inf  _ {x \in \mathbb {N} ^ {D}} f (x, \mathbf {t}) \leq \epsilon , \tag {1}
$$
for some small tolerance  $\epsilon > 0$ . Due to modeling assumptions, approximations, and the finite possible number of experiments conducted, we cannot hope for  $\epsilon = 0$  and have to tolerate a small excess loss.
Single Dimension. As demonstrated in Figure 3, the shape of a pretrained vision transformer has an impact on its downstream performance. To determine an optimal shape scaling strategy, we begin by considering both compute  $\mathbf{t}$  and a single shape dimension  $\mathbf{x}_k$  for  $k\in [D]$ , such as depth. In prior works, optimizing a single dimension  $\mathbf{x}_k$  for compute involves running a large number of experiments in order to identify the Pareto optimal frontier, from which power laws on  $\mathbf{x}_k$  or  $\mathbf{t}$  are derived [39, 35]. Since this is expensive, we propose the following joint functional form instead:
$$
f _ {k} \left(\mathbf {x} _ {k}, \mathbf {t}\right) \sim \alpha_ {k} \mathbf {x} _ {k} ^ {- a _ {k}} + \left(\beta_ {k} \mathbf {x} _ {k} ^ {b _ {k}} + \xi_ {k}\right) \mathbf {t} ^ {- c} + \varepsilon_ {k}, \tag {2}
$$
where  $\alpha_{k}, a_{k}, \beta_{k}, b_{k}, c, \xi_{k}, \varepsilon_{k} > 0$ . Here,  $f_{k}$  focuses on the dimension  $k$  alone and assumes that all other shape dimensions  $j \neq 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_{k}$  are scaling exponents when varying the corresponding shape dimension in the compute-unbounded regime,  $c$  is the data scaling exponent, while  $b_{k}$  relates to the impact of the model shape on compute.
Our argument for this particular functional form is six-fold:
I. If compute is unbounded, we recover the familiar power law relation on model size  $f_{k}(\mathbf{x}_{k}) \sim \alpha_{k}\mathbf{x}_{k}^{-a_{k}} + \varepsilon_{k}$  [34, 2, 38, 39]. In addition, increasing the model size  $x_{k}$  while keep the data size fixed does not imply that  $f_{k}(\mathbf{x}_{k},\mathbf{t}) \to \varepsilon_{k}$  because  $\mathbf{x}_k^b$  can increase faster than  $\mathbf{t}^c$  in (2).
II. For any fixed model size, the relation above reduces to the power law  $f_{k}(\mathbf{t}) \sim A\mathbf{t}^{-c} + B$ , where  $A = \beta_k\mathbf{x}_k^{b_k} + \xi_k$  and  $B = \alpha_k\mathbf{x}_k^{-a_k} + \varepsilon_k$ . Since the model size is fixed,  $\mathbf{t}$  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.  $\mathbf{x}_k$  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 = \Theta(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 (\mathbf {x}, \mathbf {t}) = \sum_ {k} \alpha_ {k} \mathbf {x} _ {k} ^ {- a _ {k}} + \sum_ {k} \beta_ {k} \mathbf {x} _ {k} ^ {b _ {k}} \mathbf {t} ^ {- c} + \xi \mathbf {t} ^ {- c} + \varepsilon_ {\infty}, \tag {3}
$$
for some constants  $\xi, \varepsilon_{\infty} > 0$
V. When optimizing the shape dimension  $\mathbf{x}_k$  for fixed compute  $\mathbf{t}$ , the optimal value  $\mathbf{x}_k^\star$  is:
$$
\mathbf {x} _ {k} ^ {\star} = \left(\frac {\alpha_ {k} a _ {k} \mathbf {t} ^ {c}}{\beta_ {k} b _ {k}}\right) ^ {\frac {1}{b _ {k} + a _ {k}}} = O \left(\mathbf {t} ^ {s _ {k}}\right), \quad \text {w h e r e :} s _ {k} = \frac {c}{b _ {k} + a _ {k}}. \tag {4}
$$
Recall that the scaling exponent  $s_k$  in (4) is positive because  $a_k, b_k, c > 0$ . Using the relation (4), we rearrange the terms in Eq. (2), and obtain the scaling law for model performance along the compute-optimal frontier (Appendix A):
$$
f _ {k} \left(\mathbf {x} _ {k}, t\right) = F \mathbf {x} _ {k} ^ {- a _ {k}} + G \mathbf {t} ^ {- c} + \varepsilon_ {k}, \quad \text {(i n t h e c o m p u t e - o p t i m a l f r o n t i e r)} \tag {5}
$$
for some constants  $F$  and  $G$ , which is a 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 a 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 to prior work [35].
Star Sweep – Conducting a brute-force grid search to estimate scaling parameters across all dimensions is expensive, since it requires  $O(2^{D})$  experiments to cover the search space. Instead, we demonstrate that a "star sweep" is sufficient: (1) starting from a large model  $\mathbf{x}^{(c)}$  (the star center), we vary a single dimension  $k \in [D]$  at a time in an exponentially-spaced grid, such that all values are much smaller than  $\mathbf{x}_k^{(c)}$ . In our experiments, for instance, we optimize three shape parameters: width, depth, and MLP dim (see Section 4 for a brief definition of each dimension). Our star center is  $\mathbf{x}^{(c)} = (1968, 40, 6144)$ ; i.e. has width 1968, depth 40, and MLP dim 6144. When varying MLP dim in the star sweep, we use the grid (1088, 1360, 1728, 2160, 2592, 3072), corresponding 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 a bottleneck when estimating the parameters in (2). This gives us the scaling exponents  $s_k$  in (4).
Grid Sweep - The second stage is a 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 a single architecture  $\mathbf{x}^{(0)}$  that lies in the Pareto optimal frontier for small compute as illustrated in Figure 3. This is important since a suboptimal  $\mathbf{x}^{(0)}$  can significantly skew results [5]. Our grid sweep identifies  $\mathbf{x}^{(0)}$  to be (608, 10, 928), the blue star in Figure 3. The advantage of this step is to absorb the leading coefficients in  $\mathbf{x}_k^* = O(\mathbf{t}^{s_k})$  in (4) so that the star sweep focuses on estimating the exponents  $s_k$  alone. We demonstrate in Figure 5 that the scaling exponents  $s_k$  are 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 architecture  $\mathbf{x}^{(0)}$  and the amount of compute  $\mathbf{t}^{(0)}$  it is optimal for, suppose we increase compute by a factor  $\tau > 1$  (i.e. the new compute is  $\tau \mathbf{t}^{(0)}$ ). By treating this increment  $\tau$  as a sequence of  $D$  smaller increments of size  $\tau^{w_k}$  each with  $\sum_{k} w_k = 1$ , an increase in compute by a factor of  $\tau$  is accompanied by an increase in every shape dimension  $k$  by a factor of  $\tau^{w_k}$ , respectively. In this work, the adopt the simplest strategy of setting  $w_k = 1 / D$ , but acknowledge that more sophisticated approaches might lead to better results.
4 Shape-optimized ViT
We implement the scaling strategy in Section 3 in vision transformers [24] pretrained on JFT-3B, a 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 hyper-parameters used in the
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  $\mathbf{x}^{(c)}$  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. RIGHT: GFLOPs is well-correlated with actual TPU core hours across models (correlation coefficient  $\sim 0.99$ ).
Figure 5: A plot of the estimated values of the exponents in (2) for different evaluation metrics  $f$ . The scaling exponent  $s_k$  tends to be less sensitive to the choice of metric than other exponents. Moreover, the data scaling exponent  $c$  is approximately  $c \approx 0.65 \pm .06$ , independently of the choice of 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 \times 14$ , but study "flexifying" to arbitrary sequence lengths following [7] in Section 5.5.
As an evaluation metric  $f$ , we consider two domains: (1) image classification, with ImageNet linear 10-shot error rate as the metric, and (2) image-to-text LiT-decoding following [8]. In the latter case, the evaluation metric  $f$  is an 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 to 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  $\mathbf{x}^{(c)} = (1968, 40, 6144)$  as our starting point. To estimate the scaling exponents  $s_k$  in (4) for each dimension separately, we 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 use an exponential spacing with all values being much smaller than in the star center  $\mathbf{x}^{(c)}$ . Following [24], we evaluate quality using few-shot linear transfer by using pre-trained models to extract features and fitting a linear regression head mapping them to the one-hot encoding of the target labels.
The individual scaling exponents we find are  $s_{\mathrm{depth}} \approx 0.45$ ,  $s_{\mathrm{width}} \approx 0.22$ , and  $s_{\mathrm{MLP}} \approx 0.6$ . Importantly, these exponents are 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  $a_k, b_k, c_k$  in (2) but the scaling exponent  $s_k$  is relatively unchanged, since it is formed as a ratio over other exponents. In addition, the data scaling exponent  $c$  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  $s_k$  point to the following picture:
I. MLP dimension should be scaled faster than depth, and depth faster than width.