SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 6: LEFT: 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/14 (in green) while keeping compute fixed by changing the training duration such that the total FLOPs is the same in all models.
II. The size of ViT, as quantified by its parameter count, is scaled more slowly than the allocated compute. More precisely, for every increment in compute by a factor of 10, the parameter count of the optimized model shape increases by a factor of  $\approx 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 are jointly optimized for the available compute.
We validate these predictions by optimizing the shape of ViT for the compute-equivalent of ViT-g/14 when the latter is pretrained on 16 billion JFT-3B examples as done in [80]. The resulting model, SoViT-400m/14, is significantly smaller and faster, yet equally competitive. It has a width of 1152, depth 27, and MLP dim 4304. Fine-tuning it on ImageNet results in a  $90.3\%$  top-1 accuracy, see Figure 2. Section 5 presents various other evaluations.
In Figure 6, we also optimize the shape of ViT for the compute-equivalent of ViT-B/14 pretrained on 4 billion examples of JFT-3B using Imagenet 10-shot error rate as an evaluation metric, resulting in SoViT-150m/14. It has a width of 880, depth 18, and MLP dim 2320. As shown in Figure 6, optimizing the shape of ViT leads to a significant improvement in performance, from  $76.6\%$  in ViT-B/14 to  $78.5\%$  in SoViT-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 6, 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) a comparison with a model, denoted B-150m, which has the same shape as ViT-B/14 but the same size as SoViT-150m/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 no resemblance to those observed during the star or grid sweeps. To recall, the star sweep is centered around an architecture  $\mathbf{x}^{(c)}$  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 to accurately identify a 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 on the transformer architecture [72, 24]. In particular, an encoder-decoder transformer with an autoregressive decoder is a popular choice because it 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 an 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 a similar scale. For the learning rate and weight decay hyper-parameters, we conduct a sweep where we vary the learning rate in  $\{10^{-3}, 3 \times 10^{-4}, 10^{-4}\}$  and the weight decay in  $\{3 \times 10^{-4}, 10^{-4}, 3 \times 10^{-5}\}$ . 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 \times 10^{-4}$  and a weight decay of  $10^{-4}$  are selected.
Table 1: 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 a held-out  $2\%$  of Train to select hyper-parameters. Selecting them on Val would increase all scores. FLOPs according to XLA; PyTorch reports MACs.
Model	Pretraining	Size	ImageNet variant
Input	Params	FLOPs	Val [57]	RealL [6]	v2 [56]
SoViT-400m/14	JFT-3B	2242	428 M	221 G	88.9	90.3	80.7
ViT-L/16 [80]	JFT-3B	3842	303 M	383 G	88.5	90.4	80.4
SoViT-400m/14	JFT-3B	3842	428 M	672 G	90.0	90.9	83.2
ViT-g/14 [80]	JFT-3B	5182	1011 M	3208 G	90.2	90.9	-
SoViT-400m/14	JFT-3B	5182	428 M	1374 G	90.3	91.0	83.4
ViT-G/14 [80]	JFT-3B	5182	1882 M	5668 G	90.4	90.8	83.3
SwinV2-G [49]	IN-21k + 70M	6402	3000 M	-	90.2	-	84.0
CoAtNet-6 [19]	JFT-3B	5122	1470 M	1521 G	90.4	-	-
MAE→WSP [61]	IG-3B	5182	1890 M	5679 G	89.7	90.9	83.0
CoCa [77]	JFT-3B + ALIGN-1.8B	5762	2100 M	-	91.0	-	-
Using this analysis, the derived scaling exponents are approximately 0.25, 0.49 and 0.62 for width, depth and MLP size, respectively. Hence, whereas the optimal shape dimensions in small architectures can be quite different between image classification and multitask decoding, as 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, or only in the ILSRCV2012 10-shot metric it was optimized for. The settings we cover are few-shot, frozen linear probes on ImageNet, zero-shot transfer, image-language 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 is pretrained on 40 billion examples, which amounts to 9T GFLOPs and 230K TPUv3 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-trained image encoders are 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 to increase image resolution [70] as a way of further increasing the capacity of the pre-trained model during fine-tuning [43]. Table 1 shows the performance of SoViT-400m/14 in comparison with ViT-L/16, ViT-g/14 fine-tuned at various resolutions, along with a 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 pretrained 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	-A	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 performs competitively with ViT-g/14 in 10-shot classification.
INet [22]	CIFAR100 Pets [46]	Pnts [51]	Birds [74]	Caltech [25]	Cars [45]	Colorectal [40]	DTD [17]	UC [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 to RealL [6], ImageNet-v2 [56], ImageNet
Renditions [30], ImageNet-Adversarial [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 a closed-form linear regression probe for 10-shot classification across a wide range of classification tasks in Table 3. Again, SoViT-400m/14 performs on-par with ViT-g/14 across the board.
5.2 Contrastive image-text tuning
Next, we follow the locked-image text tuning (LiT) recipe [81] on the WebLI dataset [15] to add zero-shot classification abilities to the pre-trained ViT-L/16, SoViT-400m/14 and ViT-g/14 image encoders. In this setup, a 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/14 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 it was found to perform well [8]. For evaluation, we report COCO CIDEr [48, 14, 73], 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 evaluates optical character recognition, whereas both VQAv2 and GQA are question-answering evaluations. Results are summarized in Table 4. SoViT-400M performs on par with ViT-g/14.
5.4 Panoptic Segmentation
Additionally, we evaluate SoViT-400m/14 on panoptic segmentation [42], which is a 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 is later decoded to the panoptic segmentation mask using a fixed VQVAE [71] model, which was pretrained on 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 \times 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 a different model shape might be derived in this domain. We leave this investigation for future work.
5.5 Flexifying SoViT-400M
Finally, since we do not include the patch size (sequence length) as part of the shape optimization, we verify that this is not a limitation by flexifying [7] SoViT-400m/14 on ILSVRC2012 for 300 epochs. The performance of the resulting FlexiSoViT-400m is shown in Fig 7 as green
Figure 7: Flexification of SoViT-400m/14 (abbr. So/14). See Section 5.5.
Table 4: Summary of multitask decoding and zero-shot transfer results, see Sections 5.2 & 5.3.
Model	ImgNet	OCR-VQA [50]	GQA [37]	VQAv2 [28]	COCO Capt. [14]
Zero-shot	Acc [%]	Log-PPL	Acc [%]	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 1 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 SoViT-400M continues to perform quite well for other patch sizes when it is flexified.
In conclusion, we introduce an 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.
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 a supportive research environment. We use the big_vision codebase [10, 9] for conducting experiments in this project.
[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, 4
[2] Bahri, Y., Dyer, E., Kaplan, J., Lee, J., and Sharma, U. (2021). Explaining neural scaling laws. arXiv preprint arXiv:2102.06701. 4, 6
[3] Bansal, Y., Ghorbani, B., Garg, A., Zhang, B., Krikun, M., Cherry, C., Neyshabur, B., and First, O. (2022). Data scaling laws in NMT: The effect of noise and architecture. arXiv preprint arXiv:2202.01994. 1, 4
[4] Barbu, A., Mayo, D., Alverio, J., Luo, W., Wang, C., Gutfreund, D., Tenenbaum, J., and Katz, B. (2019). Objectnet: A 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, J., and Zoph, B. (2021). Revisiting resnets: Improved training and scaling strategies. Advances in neural information processing systems (NeurIPS). 2, 3, 5
[6] Beyer, L., Henaff, O. J., Kolesnikov, A., Zhai, X., and van den Oord, A. (2020). Are we done with imagenet? CoRR, abs/2006.07159. 8, 9, 19, 20
[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. In CVPR. 6, 9
[8] Beyer, L., Wan, B., Madan, G., Pavetic, F., Steiner, A., Kolesnikov, A., Pinto, A. S., Bugliarello, E., Wang, X., Yu, Q., Chen, L.-C., and Zhai, X. (2023b). A study of autoregressive decoders for multi-tasking in computer vision. 6, 7, 9, 17
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  $\mathbf{x}_k$ . Second, we derive (5).
A.1 Quasiconvexity Proof
We assume throughout the proof that  $a_{k}, b_{k}$  are strictly positive, otherwise  $f_{k}(\mathbf{x}_{k}, \mathbf{t})$  is a monotone function on its first argument and the statement holds trivially.
To establish the quasiconvexity of  $f_{k}(\mathbf{x}_{k},\mathbf{t})$  in (2), we observe that:
$$
\frac {\partial f _ {k}}{\partial \mathbf {x} _ {k}} = - \alpha_ {k} a _ {k} \mathbf {x} _ {k} ^ {- (1 + a _ {k})} + \beta_ {k} b _ {k} \mathbf {t} ^ {- c} \mathbf {x} _ {k} ^ {b _ {k} - 1} \doteq - A \mathbf {x} _ {k} ^ {- (1 + a _ {k})} + B \mathbf {x} ^ {b _ {k} - 1}.
$$
Setting the derivative to zero gives the unique solution in  $\mathbb{R}^+$ :
$$
\hat {\mathbf {x}} = \left(\frac {A}{B}\right) ^ {\frac {1}{a _ {k} + b _ {k}}}.
$$
At the limit  $\mathbf{x}_k\to \infty$  , the term involving  $\mathbf{x}_k^{-a_k}$  vanishes and we have the asymptotic relation:
$$
f _ {k} (\mathbf {x} _ {k}, \mathbf {t}) \sim \beta_ {k} \mathbf {t} ^ {- c} \mathbf {x} _ {k} ^ {b _ {k}},
$$
which is an increasing function since  $b_{k} > 0$ . Since  $\hat{x}$  is the only point in  $\mathbb{R}^{+}$  where  $\partial f_{k} / \partial \mathbf{x}_{k} = 0$ , we conclude that  $f(\mathbf{x}_k, \mathbf{t})$  is monotone increasing for all  $\mathbf{x}_k \geq \hat{x}$ .
Similarly, when  $\mathbf{x}_k\to 0^+$  , we have:
$$
f _ {k} (\mathbf {x} _ {k}, \mathbf {t}) \sim \alpha_ {k} \mathbf {x} _ {k} ^ {- a _ {k}},
$$
which is monotone decreasing. Therefore,  $f^{\prime}(\mathbf{x}_k,\mathbf{t})\leq 0$  for all  $\mathbf{x}_k\leq \hat{x}$ . Combining both results implies that  $f_{k}(x,\mathbf{t})$  is monotone decreasing in the domain  $x\in (0,\hat{x})$  and is monotone increasing in the domain  $x\in (\hat{x},\infty)$ .
A function  $f(y)$  is said to be quasi-convex if for any  $y_{1}$  and  $y_{2}$  in its domain and any  $\lambda \in [0,1]$ , one has [11]:
$$
f \left(\lambda y _ {1} + (1 - \lambda) y _ {2}\right) \leq \max  \left\{f \left(y _ {1}\right), f \left(y _ {2}\right) \right\}. \tag {6}
$$
Suppose for the purpose of obtaining a contradiction that  $f_{k}(\mathbf{x}_{k},\mathbf{t})$  is not quasiconvex on its first argument. Then, there exists two points  $y_{1},y_{2}\in \mathbb{R}^{+}$  and  $\lambda \in [0,1]$  such that the above condition is violated. Let  $\hat{y} = \lambda y_1 + (1 - \lambda)y_2$ . But, then, by the mean-value theorem, there must exist two points  $c_{1}\in [y_{1},\hat{y} ]$  and  $c_{2}\in [\hat{y},y_{2}]$  where:
$$
f _ {k} ^ {\prime} (c _ {1}) = \frac {f (\hat {y}) - f (y _ {1})}{\hat {y} - y _ {1}} \geq 0
$$
$$
f _ {k} ^ {\prime} (c _ {2}) = \frac {f (y _ {2}) - f (\hat {y})}{y _ {2} - \hat {y}} \leq 0,
$$
with  $c_{2} > c_{1}$ . This implies that  $c_{1} \geq \hat{x}$  and  $c_{2} \leq \hat{x}$ , which is a contradiction. Therefore,  $f_{k}(\mathbf{x}_{k},\mathbf{t})$  is quasi-convex on its first argument.
A.2 Derivation of (5)
Rearranging the expression in (4), we have:
$$
\left(\frac {\beta_ {k} b _ {k}}{\alpha_ {k} a _ {k}}\right) \left(\mathbf {x} _ {k} ^ {\star}\right) ^ {b _ {k} + a _ {k}} = \mathbf {t} ^ {c}
$$
From this and (2), we obtain:
$$
f _ {k} (\mathbf {x} _ {k} ^ {\star}, \mathbf {t}) = \alpha_ {k} (\mathbf {x} _ {k} ^ {\star}) ^ {- a _ {k}} + \beta_ {k} (\mathbf {x} _ {k} ^ {\star}) ^ {b _ {k}} \left(\frac {\alpha_ {k} a _ {k}}{\beta_ {k} b _ {k} (\mathbf {x} _ {k} ^ {\star}) ^ {b _ {k} + a _ {k}}}\right) + \xi_ {k} \mathbf {t} ^ {- c} + \varepsilon_ {k},
$$
where we plugged in the last expression. Simplifying yields (5) for some constants  $F, G \geq 0$ .
B Shape Optimization
B.1 Hyper-parameters
Table 5: Common hyper-parameters 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	8e-4
Label Smoothing	0
Weight Decay	0.03 × 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 use a small batch size of 128 here in order to train multiple models in parallel on small hardware topologies.
B.2 Star Sweep
In the star sweep, we use the 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). We train each model on 500K, 1M, and 2M steps. We always fix the patch size to  $14 \times 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, we 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 is the maximum depth used in the grid), one cannot determine if it is compute-optimal or if increasing that dimension will yield even better models. This can be an iterative process, in 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 a range of compute budgets, we select as a starting point in our recipe the least compute it is optimal for. For example, if a model is compute-optimal for computes ranging from 1 TFLOPs to 2 TFLOPs, we use 1 TFLOPS in our recipe. In other words, we err on the side of 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 a "small" architecture that is computationally optimal for a "small" amount of compute, the amount of compute used in the analysis should be large enough for results to be reliable and for power laws to take effect. That is why in our experiments, we use  $>100\mathrm{M}$  training examples in the grid sweep as opposed, for instance, to using only a few million examples.
C Multitask Decoding Setup
Table 6: Multi-task decoding Hyperparameter Settings.
Image Resolution	224 × 224
Batch size	512
Preprocessing	Rescale(-1, 1), ResizeSmall(256), CentralCrop(224)
Augmentation	InceptionCrop(224), Left-Right Flip
Optimizer	AdaFactor [60]
Epochs	50
Gradient Clipping	1.0
Label Smoothing	0.1
Learning Rate	3e-4
Weight Decay	1e-4
Schedule	Cosine, 10% Warmup period
Vocabulary Size	32k
Encoder Dropout Rate	0
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 2 layers since it generally performs well [8].
D LiT Training Setup
Table 7: Locked-image text tuning (LiT) Hyperparameter Settings.
Image Resolution	224 × 224
Batch size	32K
Preprocessing	Rescale(-1, 1)
Augmentation	None
Optimizer	AdaFactor [60]
Total Examples	900M
Gradient Clipping	1.0
Learning Rate	1e-3
Weight Decay	1e-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 a large batch size of  $32\mathrm{K}$  in this setup because it improves the performance of contrastive training [53].
Table 8: 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 px	384 px	518 px
Learning rate decay	0.85	0.9	0.9
Random augment	-	2,10	2,10
Mixup	-	0.2	0.2
Training duration	50 k 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
Image crop	inception.crop (RandomResize)
Random flip	Horizontal
Loss	Sigmoid cross-entropy [6]
Head init	kernel=0, bias=-6.9
Train and minival splits	train[:98%] and train[98%:]
E Transfer to ImageNet-1k
E.1 Full model fine-tuning
Table 8 lists the settings for the ImageNet-1k fine-tuning results presented in Table 1 in the main paper. The only three settings which differ across resolutions are learningrate decay, random augment and mixup strengths. We did explore various learningrates, training durations (mostly shorter) as well as Polyak averaging, although the same setting shown in the table appears to be best across the board. Finally, we list various other settings which we did not explore. We simply used good default values from experience.
E.2 Linear probe on frozen encoder
We take the image representation at the pre-logits, i.e. 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, SoViT-400m/14 and ViT-g/14, we perform a 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.crop(224) vs. inception.crop(224), learning rate: 0.001, 0.0003, 0.0001, epochs: 1, 3, 10, weight decay: 0.0001, None. It 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.