SoViT

Analysis

Why This Paper Matters

In the race to build ever-larger vision models, the prevailing wisdom has been to simply scale up width, depth, and parameters. SoViT challenges this brute-force approach by showing that the shape of a model—how its width, depth, and MLP dimensions are balanced—matters as much as its size. For AI practitioners, this means achieving state-of-the-art performance with models that are 2.5× smaller, saving compute and memory without sacrificing accuracy. The paper provides a principled framework for finding compute-optimal architectures, making it a must-read for anyone designing or deploying vision transformers.

How It Works

The authors formalize the problem: given a fixed compute budget (in GFLOPs), find the optimal shape dimensions (width, depth, MLP size) that maximize performance. They propose a functional form fk that relates each dimension xk to performance f, motivated by observations from scaling laws and data efficiency.

To optimize a single dimension, they fit this function to data from grid sweeps. For multiple dimensions, they introduce a star sweep: start from a Pareto-optimal architecture at small compute, then vary one dimension at a time while keeping others fixed. This reduces the number of experiments needed.

Using JFT-3B (3 billion images, 30,000 classes), they estimate scaling exponents for width, depth, and MLP size. Key findings: MLP dimension should scale fastest, then depth, then width. This leads to SoViT-400m/14 (width 1152, depth 27, MLP 4304) and SoViT-150m/14.

Results

SoViT-400m/14 achieves 90.3% top-1 accuracy on ImageNet, matching ViT-g/14 (which has 2.5× more parameters). SoViT-150m/14 improves over ViT-B/14 from 76.6% to 78.5%—a significant gain for the same compute.

In multitask decoding (COCO, OCR, VQAv2, GQA), SoViT-400m/14 performs comparably to ViT-g/14, showing the shape-optimized model generalizes well across tasks.

Significance

SoViT demonstrates that compute-optimal shape design can dramatically improve efficiency. For practitioners, this means: before scaling up, optimize your model's proportions. The star sweep method is practical and reproducible. The paper also provides concrete scaling exponents that can guide future architecture search. While focused on ViTs, the principles likely apply to other transformer families. This work is a step toward more sustainable AI—getting more performance per FLOP.