PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ibrahim Alabdulmohsin, Xiaohua Zhai, Alexander Kolesnikov, Lucas Beyer

arXiv:2305.13035v5cs.CVcs.LG

TL;DR

For a fixed compute budget, choosing model size alone may miss the role of architecture shape in vision transformers. The paper derives a joint shape-scaling strategy and applies it to ViTs; its smaller SoViT model achieves strong results at compute comparable to a much larger model.

  • Problem

    Compute-optimal model sizing for a fixed budget has received attention, but vision research indicates that parameter count alone may not capture the performance relevance of architecture shape.

    Source: §1 · pp. 1–2

  • Method

    The paper develops a recipe to optimize neural-network shape dimensions jointly, deriving scaling strategies for ViT width, depth, and MLP size with fewer experiments.

    Source: §1 · p. 2

  • Results

    90.3% fine-tuning accuracy on ILSRCV2012 is achieved by SoViT-400m/14, which is 2.5× smaller than the compute-equivalent ViT-g/14; smaller shape-optimized vision models can match larger ones at equal compute.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Smaller models trained at an optimal architecture shape for the available compute can match much larger models, as demonstrated for vision transformers.

    Source: §6 · p. 10

    Dense segmentation can be a limitation of the proposed optimal model shape, and a different shape may be derived for that domain.

    Source: §5.4 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Compute-optimal scaling can require choosing model size for a fixed budget, while vision performance may also depend on architecture shape. This work develops a joint shape-scaling recipe and finds that shape-optimized smaller vision transformers can match larger models at equal compute.

1 Introduction

  • For a fixed compute budget, scaling model size outside the compute-optimal frontier can deteriorate performance, so the best choice may not train to convergence.The passage cites Chinchilla outperforming Gopher despite being 4× smaller.Source: §1 · p. 1
  • Earlier scaling analyses equated model size with parameter count, while vision research suggests that scaling dimensions jointly or favoring width can matter.The paper motivates shape optimization from findings that parameter count poorly predicts vision performance.Source: §1 · pp. 1–2
  • The recipe jointly optimizes width, depth, and MLP size, identifying a ViT scaling strategy in 115 experiments versus over 400 to optimize parameter count alone.The predicted frontier scales MLP dimension faster than depth, and depth faster than width.Source: §1 · p. 2, Figure 1 · p. 2
  • 90.3% fine-tuning accuracy on ILSRCV2012 is achieved by SoViT-400m/14, a 2.5× smaller model optimized for ViT-g/14-equivalent compute.The paper also reports 82.2% zero-shot accuracy in the LiT setup and evaluates captioning, VQA, and panoptic segmentation.Source: §1 · p. 2, Figure 2 · p. 3
  • The evaluation spans image classification, captioning, VQA, zero-shot classification, and panoptic segmentation, and identifies both gains and limitations.The paper also reports qualitative insights into scaling individual dimensions and differences in optimal shapes across domains.Source: §1 · p. 2

2 Related Work

Prior work has reported competing views on compute-optimal training and explored scaling vision transformers, but evidence on optimal ViT shape remained limited. This work distinguishes its approach from compound scaling and emphasizes downstream evaluation and compute-proxy caveats.

2 Related Work

  • Prior work disagrees on whether compute-optimal language models should scale parameters faster than data or scale them uniformly, with some instead training models longer.The paper says its ViT analysis partially agrees with the uniform-scaling view.Source: §2 · p. 3
  • 90.9% top-1 accuracy on ImageNet is reported for a 4-billion-parameter ViT, while a 22-billion-parameter model shows strong alignment to human visual perception on shape and texture bias.The passage also reports 90.4% top-1 accuracy for a 2-billion-parameter model.Source: §2 · p. 3
  • Earlier ViT shape studies recommend either increasing depth or prioritizing width, but their reported evidence is limited by small models, increased FLOPs, or no observed improvement.One study considered ViT-S and ViT-B and increased FLOPs by up to 4×; another reported no improvement from its width-first strategy.Source: §2 · p. 3
  • Unlike EfficientNet’s exhaustive search followed by uniform dimension scaling, this work fits scaling laws beyond the efficient frontier to derive separate optimal exponents for each dimension.The approach draws inspiration from compound scaling in MobileNet and EfficientNet.Source: §2 · p. 3
  • The analysis uses downstream metrics rather than upstream performance, and treats GFLOPs as a hardware-agnostic compute proxy that may not predict core hours for every model or hardware type.The paper focuses on architecture shape rather than training-protocol improvements.Source: §2 · p. 3

3 Scaling Strategy

The section develops a compute-optimal strategy for choosing vision-transformer shape dimensions, rather than treating parameter count as the only design choice. It combines a loss model and staged experiments to estimate how dimensions should scale with compute.

  • For small models, the compute-optimal shape for ImageNet few-shot evaluation is not necessarily optimal for captioning or VQA.Source: Figure 3 · p. 4

Single Dimension

  • At fixed compute, the goal is to find a shape minimizing downstream loss; the single-dimension model assumes other dimensions are non-bottlenecks and data is unlimited.Source: §3 · pp. 3–4
  • At fixed compute, loss varies non-monotonically with an individual shape dimension, motivating optimization along a compute-optimal frontier.Source: §3 · pp. 4–5
  • The compute-optimal dimension scales as x*_k = O(t^s_k), and the resulting frontier loss is a sum of power-law terms that fits and extrapolates empirical measurements.Source: §3 · p. 5
  • The procedure reduces the large-scale experiments needed to identify compute-optimal architectures by an order of magnitude compared with prior work.Source: §3 · p. 5

Multiple Dimensions

  • A star sweep varies one dimension at a time around a large-model center, instead of using an O(2^D) grid to estimate scaling parameters.Source: §3 · p. 5
  • A small-model grid sweep locates a Pareto-optimal starting architecture, since a suboptimal starting point can skew the estimates; it absorbs leading coefficients.Source: §3 · p. 5
  • When compute increases by a factor τ, the method scales each dimension by τ^w_k, using w_k = 1/D while noting more sophisticated strategies may do better.Source: §3 · p. 5

4 Shape-optimized ViT

The paper extends compute-optimal scaling to jointly choose vision-transformer shape dimensions, then evaluates the resulting scaling strategy in image classification and multitask decoding. Its findings show that dimension-specific scaling can produce smaller competitive models, with scaling exponents that are similar across the two domains.

4 Shape-optimized ViT

  • The evaluation spans ImageNet linear 10-shot classification and LiT-decoding, whose metric averages perplexity across COCO captioning, OCR, VQAv2, and GQA.The implementation fixes patch size at 14 × 14 while evaluating these two domains.Source: §4 · pp. 5–6

4.1 Image Classification

  • The estimated exponents were 0.45 for depth, 0.22 for width, and 0.6 for MLP size; these were robust to evaluation-metric choice.The data-scaling exponent was approximately 0.65 ± .06, independently of the chosen shape dimension.Source: §4.1 · p. 6, Figure 5 · p. 6
  • The resulting scaling order is MLP dimension faster than depth, and depth faster than width; each 10-fold compute increase raises optimized parameter count by approximately 2.5-fold.Small ViTs can match much larger models when shape and training duration are jointly optimized for available compute.Source: §4.1 · pp. 6–7
  • 78.5% ImageNet 10-shot performance for SoViT-150m/14 exceeded ViT-B/14's 76.6% at equal compute, and deviations from the predicted shape could notably reduce performance.The comparison also found that matching SoViT-150m/14's size without optimizing its shape did not achieve the same improvement.Source: Figure 6 · p. 7, §4.1 · p. 7
  • 90.3% top-1 accuracy on ImageNet followed compute-equivalent shape optimization, yielding SoViT-400m/14, significantly smaller and faster yet equally competitive with ViT-g/14.The comparison uses the compute equivalent of ViT-g/14 pretrained on 16 billion JFT-3B examples.Source: §4.1 · p. 7

4.2 Multitask Decoder

  • For multitask decoding, the derived width, depth, and MLP-size exponents were approximately 0.25, 0.49, and 0.62, respectively, nearly identical to the classification scaling exponents.Optimal shapes in small architectures could differ between image classification and multitask decoding despite the similar exponents.Source: §4.2 · p. 8

5 Evaluations

The evaluation compares SoViT-400m/14 with ViT-L/16 and ViT-g/14 across classification, zero-shot transfer, multitask decoding, panoptic segmentation, and flexified patch sizes. Results are broadly competitive with ViT-g/14, while panoptic segmentation is identified as a limitation of the proposed shape.

5.1 Image Classification

  • SoViT-400m/14 matches ViT-g/14 in fine-tuning and is generally on par in linear-probe and 10-shot classification.The linear-probe and 10-shot results are summarized in Tables 2 and 3.Source: §5.1 · pp. 8–9, Table 2 · p. 8, Table 3 · p. 9

5.2 Contrastive image-text tuning

  • SoViT-400m/14 is competitive with ViT-g/14 and substantially better than ViT-L/16 in zero-shot classification.The comparison follows locked-image text tuning on WebLI.Source: §5.2 · p. 9, Table 4 · p. 10

5.3 Multitask Decoding

  • SoViT-400M performs on par with ViT-g/14 in multitask decoding evaluations of captioning, OCR, and question answering.The decoder is fixed to two layers; reported evaluations include COCO CIDEr, OCR, VQAv2, and GQA.Source: §5.3 · p. 9, Table 4 · p. 10

5.4 Panoptic Segmentation

  • On COCO panoptic segmentation at 512 × 512 input resolution, SoViT-400m/14 achieves 43.7 PQ, versus 44.8 PQ for ViT-g/14.The authors identify dense segmentation as a possible limitation of the proposed optimal shape and leave domain-specific shape investigation for future work.Source: §5.4 · p. 9

5.5 Flexifying SoViT-400M

  • Flexified SoViT-400m/14 maintains a clear advantage over reference ViT models when patch size varies at inference time.This does not establish that other patch sizes could not be compute-optimal.Source: §5.5 · pp. 9–10, Figure 7 · p. 9

6 Conclusion

The paper introduces an efficient method for optimizing neural-network shape and applies it to vision transformers. Its analysis finds that smaller models with compute-optimal shapes can match much larger models.

  • The paper introduces an efficient method for optimizing the shape of neural architectures.Source: §6 · p. 10
  • The method is successfully applied to vision transformers.Source: §6 · p. 10
  • Smaller models trained with an optimal architecture shape for the right amount of compute can match much larger models.Source: §6 · p. 10

ArXiv Version History

The version history records successive layout, citation, content, and typo corrections. Later revisions added model-shape and experimental details and corrected the SoViT-400m pretraining duration.

A Scaling Laws Analysis

The appendix presents proofs of two claims and details the experimental sweeps used in the analysis. The grid-sweep guidance emphasizes interior optima, reliable compute ranges, and sufficiently large training runs.

A Scaling Laws Analysis

A.1 Quasiconvexity Proof

  • For positive a_k and b_k, the proof establishes quasiconvexity by showing the function decreases before its stationary point and increases after it.Source: §A.1 · p. 16

A.2 Derivation of (5)

  • The derivation of Equation (5) rearranges Equation (4), combines it with Equation (2), and simplifies using constants F, G ≥ 0.Source: §A.2 · p. 16

B.1 Hyper-parameters

  • The hyperparameter table covers both sweep types; a batch size of 128 lets multiple models train in parallel on small hardware topologies.Source: §B.1 · p. 17, Table 5 · p. 17

B.2 Star Sweep

  • The star sweep varies width, depth, and MLP dimension across training steps while fixing patch size at 14 × 14 and attention heads at 16.Source: §B.2 · p. 17

B.3 Grid Sweep

  • The grid sweep should place the selected optimum inside the grid, verify it across a compute range, use tightly spaced grid increments, and train on enough examples for reliable power laws.Source: §B.3 · p. 17

C Multitask Decoding Setup

The multitask decoding setup fixes the decoder at two layers, a choice reported to generally perform well [8].

C Multitask Decoding Setup

D LiT Training Setup

The section gives hyperparameter settings for LiT training and ImageNet fine-tuning, distinguishing explored choices from good defaults.

D LiT Training Setup

E.1 Full model fine-tuning

  • For ImageNet fine-tuning, only learning-rate decay, random augmentation, and mixup strengths vary across resolutions.Table 8 identifies resolution-varying settings, explored settings, and unexplored good defaults.Source: §E.1 · p. 20, Table 8 · p. 20
  • The authors explored learning rates, mostly shorter training durations, and Polyak averaging; the tabulated setting performed best across the board, while other settings were unexplored good defaults.Source: §E.1 · p. 20, Table 8 · p. 20

E.2 Linear probe on frozen encoder

The linear-probe procedure extracts pre-logits representations and selects settings through a per-model grid search on minival.

E.2 Linear probe on frozen encoder

  • The probe uses the 1152-dimensional pre-logits vector output by the MAP head, which feeds the linear classification layer.Source: §E.2 · p. 20
  • For ViT-L/16, SoViT-400m/14, and ViT-g/14, a grid search varies augmentation, learning rate, epochs, and weight decay, selecting the best model on minival (2% of train).The search tests two augmentation choices, learning rates 0.001, 0.0003, and 0.0001, epochs 1, 3, and 10, and weight decay 0.0001 or none.Source: §E.2 · p. 20, Table 9 · p. 21
  • Other probe settings are kept at known-good defaults based on prior explorations with similar models, such as plain ViTs.Table 9 distinguishes grid-searched settings from unexplored good defaults.Source: §E.2 · p. 20, Table 9 · p. 21
Loading 2305.13035v5…