Source-linked AI summary
Scaling Vision Transformers to 22 Billion Parameters
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetić, Dustin Tran, Thomas Kipf, Mario Lučić, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, Neil Houlsby
TL;DR
Vision Transformers have not reached the scale of large language models, motivating a 22B-parameter model. The paper presents ViT-22B and a recipe for efficient, stable training, then evaluates it broadly, finding increasing downstream performance with scale and benefits in robustness, fairness tradeoffs, and human visual alignment.
Problem
Vision models have lagged far behind language models in parameter scale, with the largest dense ViT at 4B parameters while language models reach hundreds of billions.
Method
The paper develops ViT-22B, a 22B-parameter ViT trained with architectural changes for efficient, stable scaling and evaluated across diverse downstream tasks.
Results
ViT-22B shows increasing performance with scale across evaluations, including frozen-feature, zero-shot, dense-prediction, robustness, fairness, and human-alignment results.
Takeaways & Limitations
ViT-22B demonstrates the potential for LLM-like scaling in vision, with improved robustness, fairness tradeoffs, and alignment to human shape/texture perception.
Takeaways & Limitations
Fairness experiments use CelebA with binary gender and attractive-or-smiling targets only to verify technical claims, not to endorse those tasks.
Abstract
from arXiv · showhide
The scaling of Transformers has driven breakthrough capabilities for language models. At present, the largest large language models (LLMs) contain upwards of 100B parameters. Vision Transformers (ViT) have introduced the same architecture to image and video modelling, but these have not yet been successfully scaled to nearly the same degree; the largest dense ViT contains 4B parameters (Chen et al., 2022). We present a recipe for highly efficient and stable training of a 22B-parameter ViT (ViT-22B) and perform a wide variety of experiments on the resulting model. When evaluated on downstream tasks (often with a lightweight linear model on frozen features), ViT-22B demonstrates increasing performance with scale. We further observe other interesting benefits of scale, including an improved tradeoff between fairness and performance, state-of-the-art alignment to human visual perception in terms of shape/texture bias, and improved robustness. ViT-22B demonstrates the potential for "LLM-like" scaling in vision, and provides key steps towards getting there.
1 Introduction
ViT-22B addresses the large scaling gap between vision and language models by combining architectural and training changes with broad evaluation. At 22B parameters, it improves performance across tasks and shows benefits in robustness, fairness tradeoffs, and human visual alignment.
- Vision models have lagged behind language models in scale, with the largest dense ViT at 4B parameters versus language models reaching 540B.
- ViT-22B is the largest dense ViT and uses architectural changes to address scaling instabilities while enabling efficient model-parallel training.
- 89.5% accuracy on ImageNet is achieved using ViT-22B as a frozen visual feature extractor.
- 85.9% zero-shot ImageNet accuracy and 88.6% student accuracy are obtained using a text tower and distillation, respectively.
- 87% shape bias accompanies improved out-of-distribution behavior, reliability, uncertainty estimation, and fairness tradeoffs.
2 Model Architecture
ViT-22B modifies the Vision Transformer with parallel layers, query/key normalization, and omitted biases to improve efficiency and training stability at scale.
- ViT-22B incorporates parallel layers, query/key normalization, and omitted biases as its three main architectural modifications.
- Parallel layers: Parallel attention and MLP blocks enable fused linear projections and additional parallelization.
- QK Normalization: Query/key normalization applies LayerNorm before dot-product attention to prevent divergence from uncontrolled attention-logit growth.
- Omitting biases on QKV projections and LayerNorms: Removing biases from QKV projections and LayerNorms improves accelerator utilization by 3% without quality degradation.
- The encoder uses multi-head attention pooling to aggregate per-token representations in the head.
3 Training Infrastructure and Efficiency
ViT-22B combines explicit model/data sharding with asynchronous communication and computation to train efficiently across TPUv4 devices. These techniques achieve 54.9% MFU while processing 1.15k tokens per second per core.
- Parallelization: A 2D logical mesh of size t × k combines data parallelism across t groups with model parallelism across k devices.Each device stores 1/k of activations and computes 1/k of every linear layer’s output.
- Asynchronous operations: Asynchronous parallel linear operations overlap communication with matrix multiplication while preserving equivalence to the unsharded computation.The design aims to minimize communication and keep matrix-multiply units busy.
- Asynchronous operations: Row-sharding communicates input blocks, whereas column-sharding computes partial output vectors that are scatter-reduced across devices.The two layouts communicate (k −1)(n/k) and (k −1)(m/k) floats, respectively, under the stated block arrangement.
- Parameter sharding: Parameter sharding fits larger models and batch sizes by gathering weights for forward computation and scattering gradients asynchronously.While one layer computes, the device can begin communicating the next layer’s weights.
- Efficiency results: 1.15k tokens per second per core and 54.9% MFU were achieved during TPUv4 training.The reported MFU was 8.7 percentage points above PaLM’s 46.2% and 10.9 points above ViT-e’s 44.0% on the same hardware.
4 Experiments
Experiments evaluate ViT-22B as a frozen or lightly adapted backbone across classification, transfer, dense prediction, video, robustness, fairness, calibration, and human alignment. Performance generally improves with scale, while the largest model also improves several robustness, fairness, calibration, and perceptual-alignment measures.
- Linear probing: ViT-22B linear probing shows notable gains with scale and can approach or exceed smaller models’ full fine-tuning at higher resolution.The comparison uses ImageNet-1k linear evaluation with models pretrained on large datasets.
- Zero-shot transfer: ViT-22B achieves comparable or better zero-shot results across ImageNet test sets and sets a new state of the art on ObjectNet.ObjectNet zero-shot performance is highly correlated with ViT model size.
- Out-of-distribution: Scaling continues to improve out-of-distribution performance, although fine-tuning increases in-domain accuracy more than OOD accuracy and therefore reduces effective robustness.The trend holds for models trained only on JFT images and for models fine-tuned on ImageNet.
- Semantic segmentation: 44.7 mIoU with only 1,200 ADE20K images exceeds DeiT-III Large by 8.6 mIoU and ViT-G by 2.3 mIoU.With more transfer data, ViT-G and ViT-22B performance converges.
- Monocular depth estimation: Frozen ViT-22B features achieve the best monocular-depth performance across metrics, with gains attributable to both larger architecture and larger pretraining data.The gap between DPT and linear decoding indicates that the features retain geometric information but do not expose all of it to a trivial readout.
- Fairness and reliability: ViT-22B improves fairness tradeoffs, benefits all subgroups, reduces subgroup performance disparities, and improves accuracy-calibration tradeoffs.Scale alone does not change demographic parity, while the accuracy-calibration improvement is measured using expected calibration error.
- Human alignment: ViT-22B improves alignment with human visual recognition, with ViT-22B-384 reaching 87% shape bias and 13% texture bias.Different resolutions lead on different toolbox metrics: OOD robustness, human accuracy alignment, and error consistency.
5 Conclusion
The conclusion presents ViT-22B as a 22-billion-parameter vision transformer enabled by architectural and systems changes for stable, efficient scaling. It reports strong transfer and zero-shot behavior, including unusual generated images, while identifying training instability as a scaling challenge addressed by QK normalization.
- Conclusion: ViT-22B reaches 22 billion parameters through parallel layers, QK normalization, omitted biases, and efficient model-parallel training.These changes target training stability and hardware utilization at scale.
- Conclusion: Frozen ViT-22B embeddings support strong downstream performance when thin layers are trained on top.The conclusion highlights frozen-feature transfer as a computationally efficient use of the model.
- Zero-shot examples: ViT-22B classifies generated images containing unusual objects and scenes absent from its training distribution.The examples come from images generated by Parti and Imagen.
- Training stability: Without normalization, attention logits exceed 50,000 in magnitude, producing nearly one-hot attention and unstable losses and gradients.QK normalization keeps the 1e-3 learning rate stable at larger model scales.
C Model Card
The model card documents ViT-22B’s 22-billion-parameter, randomly initialized model, its datasets and broad evaluation coverage, alongside research-use and ethical-use boundaries.
- Model Characteristics: ViT-22B is a static model trained from random initialization on a JFT variant containing around 4B images.
- Model Characteristics: 22 billion parameters define the model’s scale.
- Evaluation Results: Evaluation spans transfer, dense prediction, depth estimation, video classification, fairness, perceptual alignment, robustness, reliability, and calibration.
- Evaluation Results: The model card identifies datasets including CelebA, ImageNet, Waymo Open, Kinetics 400, Moments in Time, and BAPPS.
- Sensitive Use: Applications involving surveillance, law enforcement, healthcare, hiring, employment, or self-driving cars require measures to mitigate ethical risks.
- Known Limitations: ViT-22B is designed for research and has not been tested outside research settings that can affect performance.
D Transfer to image classification: More results and addition details
This section compares frozen-feature transfer methods and protocols across image, few-shot, video, and segmentation tasks, emphasizing efficiency, representation choice, and backbone freezing.
- Linear Probing: L-BFGS provides a convergent alternative to SGD for linear probing on ImageNet-1k embeddings.
- Transfer Trade-offs: Fine-tuning becomes more resource-intensive as model size increases, whereas linear probing is cheaper but often performs worse.
- Head2Toe: Head2Toe achieves 7% better results on VTAB-1k than linear probing but remains 6% below full fine-tuning.
- Head2Toe: Head2Toe uses concatenated intermediate features from ViT blocks, positional embeddings, pooling outputs, pre-logits, and logits, producing a 349081-dimensional vector.
- Few-shot Transfer: Few-shot transfer across 25 tasks uses 1, 5, 10, and 25 shots; scaling improves many tasks, but not all.
- Dense Prediction Transfer: For semantic segmentation, fine-tuning substantially helps a linear decoder, while the gap is much smaller for UperNet, which has ∼870 times more parameters.
E.2.2 Decoder Architectures
The decoder experiments describe DPT and linear architectures for dense prediction, their feature-processing choices, and the frozen-backbone video classification setup.
- DPT Decoder: DPT uses four reassemble-and-fusion blocks operating at 4×4, 8×8, 16×16, and 32×32 spatial resolutions.
- DPT Decoder: Reusing one 16×16 feature map at every DPT stage produced similar scores, while directly feeding pre-norm feature maps caused instabilities.
- Linear Decoder: The linear decoder transforms 16 × 16 × 6144 ViT-22B features through transpose convolutions and a 1 × 1 convolution into a resized depth map.
- Linear Decoder: ViT-22B’s larger decoder capacity was controlled by down-projecting features or approximately matching decoder parameter counts.
- Depth Evaluation: Depth estimation uses log(1 + depth) targets, MSE training, and AbsRel and δ metrics measured after undoing the log transformation.
- Video Classification: Video classification extracts one frozen pre-logit embedding per frame and processes frame-token sequences with a shallow transformer classifier.
G Fairness
The fairness section evaluates demographic parity alongside accuracy, calibration, and uncertainty-related metrics, using subgroup and debiasing analyses across ViT variants.
- Fairness Evaluation: Demographic parity measures fairness by comparing ViT-22B with smaller ViT variants on CelebA.
- Metrics: The evaluation reports classification accuracy, expected calibration error, and Oracle Collaborative AUC.
- Debiasing: Performance is plotted after debiasing each ViT variant to prescribed bias levels, using accuracy, ECE, and OC-AUC.
- Subgroup Performance: Before bias mitigation, ViT-22B performs better overall and within each CelebA subgroup across all three metrics.
- Subgroup Fairness: ViT-22B shows smaller absolute performance differences between female and male subgroups across all three metrics.
- Calibration: Temperature scaling uses 20% of ImageNet validation data for calibration and reports accuracy and ECE on the remaining 80%.
I.4 Results of Plex-22B and challenges
At 22B scale, Plex variants underperform vanilla ViT-22B across nearly all reported metrics, while the authors identify fine-tuning, hyperparameter, and numerical-stability challenges as possible explanations.
- ViT-22B and Plex-22B models are fine-tuned at resolution 384 in the reported comparison.
- Across all metrics except label uncertainty on ImageNet-ReaL-H, Plex-22B variants perform worse than vanilla ViT-22B.
- Pre-training vs. fine-tuning: The authors attribute the gap partly to adding BatchEnsemble and heteroscedastic layers only during fine-tuning, where they may not be properly learned.
- Hyperparameter tuning: A finer-grained search over fine-tuning duration, learning rate, and temperature may be required to close the performance gap.
- Numerical stability: Plex components may require scale-specific stabilization techniques analogous to those used to stabilize ViT-22B training.
J Error Consistency & Human Alignment
Additional model-vs-human benchmarking results are provided for ViT-22B fine-tuned on ImageNet across four reported result panels.
- Figure 19 reports additional model-vs-human benchmark results for ViT-22B fine-tuned on ImageNet.
K Perceptual similarity
ViT-22B improves low-resolution ImageNet accuracy but remains on an unfavorable accuracy–perceptual-similarity tradeoff, with perceptual scores varying by training and distance-function choices.
- ViT-22B’s perceptual similarity is evaluated against prior ImageNet-trained models to examine whether diverse large-scale classification breaks the observed tradeoff.
- 84.2 accuracy on ImageNet 64 × 64 is 16% better than the best models trained directly on ImageNet.
- Adapting ViT-22B to 64 × 64 images requires resizing the embedding layer from patch-size 14 to patch-size 8.
- An untrained ViT-22B obtains a Perceptual Score of 62.3, while the model with the Mean Pool distance function reaches 66.2.
- ViT-e falls in the same ballpark as ViT-22B, with slightly lower accuracy and Perceptual Scores.
- ViT-22B lies at the bottom-right of the accuracy-perceptual-similarity tradeoff, combining the best validation accuracy with the worst perceptual scores.
L Feature attribution analysis
The feature-attribution analysis compares ViT-22B saliency maps before and after cooldown using Integrated Gradients on example inputs.
- Integrated Gradients is applied to three example datapoints before and after ViT-22B cooldown to inspect how the checkpoints process inputs.
- A gray baseline of 0.5 with 1024 steps yields the qualitatively best saliency results.
- Figure 21 presents saliency maps before and after model cooldown.