Source-linked AI summary
Scaling Vision Transformers
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, Lucas Beyer
TL;DR
The paper asks how Vision Transformers scale with model size, data, and compute, a question left open by language-model scaling studies. It scales ViT models and datasets across broad ranges while refining architecture and training, achieving 90.45% top-1 ImageNet accuracy and 84.86% with 10 examples per class. The results characterize ViT scaling behavior and demonstrate an efficient recipe for training large, high-performing models, though the conclusions may not generalize beyond the studied scale or ViT family.
Problem
How Vision Transformers scale with model size, data, and compute remains unclear despite careful study of scaling laws for language Transformers.
Method
The study varies ViT architecture size, training images, and duration across public and private datasets, evaluates transfer quality, and refines architecture, optimizer, and training schedules.
Results
90.45% top-1 accuracy is achieved on ImageNet, while ViT-G/14 reaches 84.86% top-1 accuracy with 10 examples per class.
Takeaways & Limitations
ViT performance-compute scaling roughly follows a saturating power law, and staying on the frontier requires scaling compute and model size together.
Takeaways & Limitations
The scaling-law study uses the proprietary JFT-3B dataset, although the authors verify that the laws also apply on public ImageNet-21k.
Abstract
from arXiv · showhide
Attention-based neural networks such as the Vision Transformer (ViT) have recently attained state-of-the-art results on many computer vision benchmarks. Scale is a primary ingredient in attaining excellent results, therefore, understanding a model's scaling properties is a key to designing future generations effectively. While the laws for scaling Transformer language models have been studied, it is unknown how Vision Transformers scale. To address this, we scale ViT models and data, both up and down, and characterize the relationships between error rate, data, and compute. Along the way, we refine the architecture and training of ViT, reducing memory consumption and increasing accuracy of the resulting models. As a result, we successfully train a ViT model with two billion parameters, which attains a new state-of-the-art on ImageNet of 90.45% top-1 accuracy. The model also performs well for few-shot transfer, for example, reaching 84.86% top-1 accuracy on ImageNet with only 10 examples per class.
1. Introduction
The paper studies how Vision Transformer performance scales with model size, data, and compute, while refining training to improve transfer and enable larger models. It characterizes scaling behavior and reports strong few-shot results from the resulting ViT-G model.
- Vision Transformer scaling remains unclear despite established scaling laws for language Transformers.
- The study varies ViT models from five million to two billion parameters, datasets from one million to three billion images, and compute from below one TPUv3 core-day to beyond 10 000 core-days.
- The paper characterizes the performance-compute frontier for ViT models across two datasets.
- An improved training recipe uses strong L2 regularization on the final linear layer to strengthen few-shot transfer.
- 84.86% top-1 accuracy is achieved on ImageNet with 10 examples per class using the ViT-G model.
- The study also analyzes representation quality through ImageNet fine-tuning and linear 10-shot error as functions of compute, model size, and dataset size.
2. Core Results
Scaling ViT model size, data, and compute together improves transfer representation quality, while performance is constrained by model or dataset capacity and eventually saturates. Larger models are more sample efficient, and the two-billion-parameter ViT-G/14 achieves strong few-shot and benchmark results.
- Scaling up compute, model and data together: Scaling compute, model size, and data together improves representation quality, with the largest configuration attaining the lowest error rate.
- Scaling up compute, model and data together: Small models are bottlenecked by limited capacity and cannot benefit fully from the largest datasets or compute budgets.The Ti/16 model maintains a high error rate even after training on many images.
- Scaling up compute, model and data together: Large models benefit from additional data beyond 1B images, whereas increasing dataset size does not help small models such as Ti/16 or B/32.For L/16, error reaches 19% on larger datasets versus 27% on the 30M dataset; the largest models improve from 1B to 3B images.
- Double-saturating power law: Across over two orders of magnitude of compute, performance follows a power law but saturates at both low and high compute.The high-compute saturation is represented by an additive constant in E = aC^-b + c, indicating a non-zero performance ceiling.
- Double-saturating power law: Larger models are more sample efficient: matching L/16 requires nearly 100 times more images for Ti/16 in 10-shot transfer and about 20 times more during fine-tuning.The comparison uses images seen during pre-training across few-shot transfer and fine-tuning evaluations.
- Do scaling laws still apply on fewer images?: On ImageNet-21k, the double-saturation power law persists, and the best representation quality comes from scaling compute, model size, and data together.Representation quality is measured by ImageNet linear 10-shot error rate as a function of total training compute.
- ViT-G/14 results: ViT-G/14, with nearly two billion parameters, reaches 84.86% ImageNet accuracy with 10 examples per class and 90.45% top-1 ImageNet accuracy.It is evaluated on multiple downstream benchmarks, including ImageNet-v2, ReaL, ObjectNet, and VTAB.
- ViT-G/14 results: ViT-G/14 improves over prior methods on several benchmarks, while ImageNet classification appears to be approaching saturation and ObjectNet performance depends on evaluation differences.It is about 2% behind CLIP on ObjectNet, whose direct evaluation without ImageNet fine-tuning likely improves robustness.
3. Method details
The paper improves ViT architecture and training to reduce memory use and strengthen transfer performance. Key changes include decoupled head regularization, memory-efficient pooling and optimizers, and schedules that support varied training durations.
- System efficiency: The resulting improvements allow ViT-G/14 to be trained with data parallelism alone while fitting the entire model on one TPUv3 core.The approach also fits thin ViT models with up to 100 encoder blocks.
- Decoupled weight decay: Decoupling weight decay between the prediction head and body improves few-shot transfer performance.High head decay can improve transfer while decreasing upstream pre-training performance, and the authors do not fully explain this behavior.
- Memory-efficient heads: GAP and MAP replace the [class] token as more memory-efficient aggregation heads, with similar performance across head types.The authors choose MAP because it is more expressive and produces a uniform architecture.
- Scaling data: Switching from JFT-300M to JFT-3B improves both small and large models without further model scaling.The improvement appears in linear few-shot evaluation and transfers to full-dataset evaluation.
- Memory-efficient optimization: The improved optimizer reduces memory overhead to 50% above the space required for model parameters while performing on par with or slightly better than Adam.Half-precision momentum reduces Adam overhead, while modified Adafactor uses rank-1 factorization for the second momentum.
- Learning-rate schedules: Cooldown-based learning-rate schedules enable evaluating multiple training durations without restarting every run from scratch.The schedules combine warmup, a slowly decaying or constant middle phase, and linear annealing toward zero.
4. Related Work
The paper extends scaling-law research from language and generative vision models to discriminative image modeling with ViTs. It relates this work to smaller ViTs, CNN scaling strategies, and memory-efficient large-model training.
- Smaller Vision Transformers: Earlier ViT work focused on medium and large models, while later variants extend ViT architectures down to approximately five million parameters.The related work discusses smaller ViT sizes and cheaper architectural variants.
- Memory-efficient training: The Shapefinder simulation compares original ViT, the paper’s improvements with half-precision Adam, and modified Adafactor across memory and training speed.White regions indicate configurations that ran out of memory, while dot brightness represents relative training speed.
- Scaling laws: This paper presents the first study of scaling laws for discriminative image modeling.Prior work established empirical power laws for language models and other generative modalities.
- Scaling-up vision models: CNN research developed scaling strategies that balance depth, width, and resolution to improve visual recognition.The paper places ViT scaling alongside EfficientNet, ResNet, and other large CNN efforts.
5. Discussion
The discussion acknowledges the proprietary JFT-3B dataset as a limitation while reporting that the scaling-law insights also hold on public ImageNet-21k. It also considers the energy cost of large-scale experiments and possible amortization through reuse and transfer.
- Limitations: The scaling-law study relies on proprietary JFT-3B data, limiting the direct generality of its evidence.The authors address this scope boundary by verifying that the scaling laws also apply to public ImageNet-21k.
- Societal impact: Training the largest ViT-G model carries an energy cost, although the authors suggest that one-time scaling studies and downstream transfer can amortize it.The paper contrasts the cost of pre-training with the lower cost of transferring pretrained weights.
6. Conclusion
The paper concludes that ViT performance follows a saturating power law when sufficient data are available, and that compute and model size must scale together to remain efficient. Larger models are more sample efficient and support strong few-shot learning, while the proposed training recipe enables efficient large-model training.
- Scaling conclusions: ViT performance-compute follows a roughly saturating power law when training data are sufficient.The conclusion frames this as the paper’s central scaling result.
- Scaling conclusions: Keeping model size fixed as additional compute becomes available is suboptimal for staying on the performance-compute frontier.The authors state that compute and model size must be scaled simultaneously.
- Transfer conclusions: Larger ViT models are more sample efficient and are strong few-shot learners.This conclusion follows the paper’s scaling analysis and transfer evaluations.
- Training conclusions: The new training recipe enables efficient training of large, high-performing ViT models.The paper presents this as a concluding contribution alongside its scaling findings.
- Scope: The conclusions may not generalize beyond the studied scale or beyond the ViT model family.This is the explicit scope boundary stated in the conclusion.
A. More few-shot transfer results
Scaling laws similar to those observed elsewhere also appear on Oxford IIIT Pets, CIFAR-100, and Caltech-UCSD Birds.
- Similar scaling laws are observed on Oxford IIIT Pets, CIFAR-100, and Caltech-UCSD Birds.
B. Pre-training details
The models use Adafactor-based pre-training with large batches for the largest ViT variants and different weight decay for the head and body.
- All ViT models are pre-trained with Adafactor using half-precision momentum.
- Models smaller than ViT-g use batch size 4096, while ViT-g and ViT-G use batches up to 32,768 across 2,048 TPUv3 chips.
- Weight decay is set to 3.0 for the head and 0.03 for the body.
C. Configuration file for pre-training ViT-g
The ViT-g/14 pre-training configuration follows big_vision conventions and specifies data processing, optimization, batch size, and evaluation settings. Experiments also examine resolution and patch-size scaling while keeping token count fixed.
- Configuration: The full ViT-g/14 pre-training configuration follows big_vision codebase conventions.
- Evaluation: Few-shot linear regression solves an L2-regularized problem using frozen embeddings from 224 × 224 images.
- Evaluation: Fine-tuning uses SGD with momentum, batch size 512, gradient clipping, no weight decay, and higher input resolutions for larger models.
- Resolution and patch size: Scaling resolution and patch size linearly together produces subtle quality differences when the number of tokens is fixed.
- Scaling-law evaluation: Figure 9 measures representation quality as few-shot error rate across four datasets, although the law does not always fit perfectly.
F. Full table of few-shot results
The few-shot results cover 5-shot and 10-shot learning on four datasets, including the largest ViT models and smaller models summarized across multiple tables.
- 5-shot and 10-shot learning results are reported on the four datasets from Figure 9.
- ViT-g/14 and ViT-G/14 results are summarized in Table 4, using batch sizes up to 32,768.
- Table 4 normalizes training steps as Images Seen/4096 to make comparisons readable across batch sizes.
- Results for smaller ViT models are distributed across Tables 5 through 13, with some missing rows that do not affect the scaling-law trend.
G. Full table of finetune results
This section compiles few-shot and finetune results across ViT model variants. Finetune evaluations cover ImageNet, ImageNet V2, and ImageNet ReaL, with training-step and learning-rate details reported alongside the results.
- Training details: The finetuning appendix reports total steps, cooldown steps, and the best learning rate selected using a held-out 2% ImageNet training split.A few table rows are missing, but the text states that they do not affect the scaling-laws trend.