Scaling Vision Transformers to 22 Billion Parameters
Proposes ViT-22B with parallel layers, QK normalization, and bias removal, achieving stable training at 22B parameters, surpassing prior models.
Key Findings
Methodology
This work introduces a novel architecture combining parallel attention and MLP layers, applying QK normalization to prevent logits explosion, and removing biases for efficiency. Using JAX/FLAX, it employs model and data parallelism with asynchronous communication and parameter sharding. The training on a 4B-image JFT dataset utilizes cosine learning rate scheduling over 177k steps with batch size 65k, optimizing stability and scalability. Evaluation includes image classification, transfer learning, segmentation, depth estimation, and video tasks, demonstrating broad applicability and performance gains.
Key Results
- ViT-22B achieves 89.5% accuracy on ImageNet zero-shot classification, outperforming previous dense models. It shows improved robustness on out-of-distribution datasets like ObjectNet, indicating better generalization. In linear probing, it surpasses prior models with a 4% increase over ViT-e/14, confirming the benefits of scale. The model also excels in fine-grained classification and depth estimation, validating its detailed feature extraction capabilities.
- The model’s hardware utilization reaches 54.9% MFU, a 15% increase over prior models, enabling efficient large-scale training. Its multi-task performance demonstrates significant improvements across classification, segmentation, and video understanding, with strong transferability. These results confirm that architectural innovations and training strategies effectively support massive parameter scaling.
- Additional experiments show the model’s superior out-of-distribution robustness, with notable gains on datasets like ObjectNet and ImageNet-A. The model’s multi-modal transferability, including video classification with frozen backbones, indicates broad potential for real-world applications. The combination of high accuracy and efficiency marks a new state-of-the-art in vision model scaling.
Significance
This research pushes the boundary of vision model scaling, demonstrating that with architectural and training innovations, models can reach 22B parameters while maintaining stability and efficiency. It bridges the gap between vision and language models, opening avenues for more capable, robust, and generalizable AI systems. The ability to train such large models efficiently paves the way for future breakthroughs in computer vision, enabling applications in autonomous driving, medical imaging, and multimedia understanding. It also provides a blueprint for scaling other neural architectures, emphasizing the importance of normalization, parallelization, and hardware-aware design.
Technical Contribution
The paper introduces a parallel layer architecture, QK normalization, and bias removal, which collectively address the instability issues in scaling large transformers. It leverages advanced distributed training techniques, including asynchronous communication and parameter sharding, to optimize hardware utilization. These innovations enable training of a 22B parameter model with high efficiency, setting new standards for large-scale vision transformers. The approach also offers insights into stability mechanisms and hardware-aware design, contributing to the theoretical understanding of large neural network training.
Novelty
This is the first work to successfully train a 22B dense vision transformer, employing a combination of parallel layers, QK normalization, and bias-free design. Unlike prior models limited to 4B parameters, this work demonstrates that architectural modifications can stabilize training at unprecedented scales. The integration of these techniques with distributed training strategies represents a significant leap forward in vision model scaling, establishing new benchmarks and opening new research directions.
Limitations
- High computational and hardware costs limit accessibility, requiring extensive resources for training and deployment. Despite stability improvements, training still demands sophisticated infrastructure and tuning, restricting rapid experimentation. The model's robustness on highly out-of-distribution data, while improved, still faces challenges, especially in real-world noisy environments. Future work should focus on reducing costs, enhancing robustness, and simplifying training procedures.
Future Work
未来将探索模型压缩和剪枝技术以降低部署成本，同时结合多模态数据，扩展模型在视频、文本等领域的应用。还需研究更高效的训练算法和自适应架构，以实现更大规模模型的稳定训练。推动硬件优化和算法创新，提升模型的普及率和实用性，最终实现更智能、更鲁棒的视觉AI系统。
AI Executive Summary
The rapid evolution of deep learning has seen Transformer architectures revolutionize natural language processing, with models scaling to hundreds of billions of parameters, exhibiting emergent capabilities and remarkable generalization. In contrast, vision models have lagged, with the largest dense ViT reaching only 4 billion parameters, limiting their potential. This paper introduces ViT-22B, a landmark achievement in scaling vision transformers, combining architectural innovations and training strategies to overcome stability barriers at unprecedented size.
The core technical advances include a parallel layer design that enables efficient computation, the application of QK normalization to control attention logits, and the removal of biases to streamline the architecture. These modifications, coupled with distributed training on JFT-4B data, allow the model to be trained stably and efficiently, achieving a hardware utilization of 54.9% MFU and a 15% efficiency boost. The training process employs a cosine learning rate schedule, extensive sharding, and asynchronous communication, ensuring scalability.
Empirical results demonstrate that ViT-22B significantly outperforms previous models across multiple benchmarks. On ImageNet zero-shot classification, it reaches 89.5% accuracy, surpassing prior dense models. Its robustness on out-of-distribution datasets like ObjectNet and ImageNet-A is notably improved, indicating better real-world applicability. Transfer learning experiments show superior performance in segmentation, depth estimation, and video classification, validating its broad utility.
This work marks a pivotal step toward 'LLM-like' scaling in vision, opening new horizons for large-scale, robust, and versatile AI systems. Despite high costs and complexity, the architectural insights and training techniques developed here set a new benchmark for future research. Continued efforts in efficiency, robustness, and multimodal integration will further advance the field, bringing us closer to truly general-purpose visual AI.
Deep Analysis
Background
近年来，Transformer架构在自然语言处理中的成功激发了视觉模型的研究热潮。最早的ViT模型（Dosovitskiy et al., 2021）开启了视觉Transformer的时代，但其最大参数规模仅为4B，难以匹敌语言模型的百亿级参数。随着硬件能力提升，学界尝试扩展模型规模，但训练不稳定性和效率瓶颈成为主要障碍。此前工作如ViT-G（Zhai et al., 2022a）和ViT-e（Chen et al., 2022）在参数规模上取得一定突破，但仍未突破百亿大关。大规模模型的潜力在于其更强的表达能力和迁移能力，但同时带来了训练难题，包括梯度爆炸、收敛缓慢等问题，尤其在参数超过8B时表现尤为明显。解决这些问题，成为推动视觉Transformer发展的关键。
Core Problem
现有视觉Transformer模型在参数规模扩展中遇到训练不稳定和效率瓶颈，限制了模型性能的进一步提升。大模型易出现梯度爆炸、收敛缓慢等问题，尤其在参数超过8B时表现尤为明显。此外，硬件资源的限制也制约了模型的训练和应用。如何在保证训练稳定的前提下，实现参数规模的突破，是当前研究的核心难题。解决这一问题，将极大推动视觉模型的性能极限，为复杂视觉任务提供更强的基础。
Innovation
本文提出三项核心创新：第一，采用平行层架构，将Attention和MLP块并行处理，提升并行效率，减少训练时间；第二，引入QK归一化（LayerNorm应用于查询和键），有效控制attention logits的数值范围，避免梯度爆炸；第三，取消偏置参数，简化模型结构，提升硬件利用率。这些设计共同解决了大模型训练中的不稳定问题，显著提升了训练效率和模型性能。结合异步通信和参数切片技术，实现了22B参数模型的高效训练，为视觉Transformer的规模极限提供了技术支撑。
Methodology
- �� 架构设计：采用平行层结构，将Attention和MLP块在同一层中并行计算。• QK归一化：在计算attention前，对查询和键进行LayerNorm，控制logits范围。• 无偏置：移除偏置参数，简化模型结构。• 分布式训练：利用JAX/FLAX框架，结合模型和数据并行，采用异步通信和参数切片技术。• 训练策略：在JFT-4B数据集上预训练，使用余弦学习率调度，逐步缩放模型参数。• 评估方法：在多任务、多数据集上测试，包括分类、检测、分割和视频任务，验证模型泛化能力。
Experiments
采用JFT-4B图像数据集，训练177k步，批量65k，模型参数达22B。评估指标包括ImageNet准确率、零样本分类、偏离分布数据表现等。对比不同架构变体，验证QK归一化和无偏置的效果。进行多任务迁移学习，验证模型在语义分割和深度估计中的表现。还测试了视频分类和多模态任务，确保模型的广泛适用性。
Results
模型在ImageNet零样本分类中达89.5%准确率，超越以往最大模型。在ObjectNet等偏离分布数据集上表现优异，提升鲁棒性。线性探测任务中，ViT-22B在ImageNet达89.5%，比前一最大模型提升约4%。在细粒度分类和深度估计中也优于较小模型，验证了大规模参数带来的性能提升。硬件利用率达54.9%，训练效率提升15%，验证了架构设计的有效性。
Applications
该模型可作为基础视觉特征提取器，支持高精度分类、检测、分割等任务，适用于自动驾驶、安防监控、医疗影像等行业。模型的迁移能力强，可用于多模态融合和视频理解，推动智能视觉系统的发展。未来，结合模型压缩和优化算法，有望实现边缘设备上的部署。
Limitations & Outlook
训练成本高昂，依赖大规模硬件资源，限制普及。架构复杂，调优难度大，影响推广。模型在某些偏离训练分布的数据上仍存在鲁棒性不足的问题。未来需优化训练流程，降低成本，提升模型的实用性和鲁棒性。
Plain Language Accessible to non-experts
想象你在一家大型工厂里，生产各种商品。每个工人都负责不同的任务，比如组装、包装、检验。以前，工人们都按顺序工作，一个接一个，效率不高。现在，工厂引入了一种新方法，把一些任务同时进行，工人们可以同时合作，效率大大提高。这个工厂还用了一种新工具，可以让工人们更好地理解彼此的工作流程，避免误会和错误。这样一来，生产速度快了，商品质量也更稳定。类似的，研究中的大模型也是这样，通过优化架构和训练方法，让计算机“同时”学习更多信息，变得更聪明、更可靠。
ELI14 Explained like you're 14
想象你在学校里，有很多不同的老师教你不同的科目。以前，老师们都按顺序讲课，你得等前面的人讲完才能开始。现在，有一种新方法，老师们可以同时讲不同的内容，这样你学到东西就快多了。科学家们也是这样，让计算机学习很多很多信息，但要确保它们不会出错。为了做到这一点，他们设计了特别的“工具”和“规则”，让大模型可以同时处理大量信息，还能保持稳定。结果，这个大模型变得比以前更聪明，能更好地理解图片、视频，甚至能在新场景中表现得更棒。虽然需要很多计算资源，但未来它会帮助我们解决很多复杂的问题，比如自动驾驶、医疗诊断等。
Abstract
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.