Scaling Vision with Sparse Mixture of Experts
Introduces sparse MoE into Vision Transformer, achieving 90.35% on ImageNet with 15B parameters, halving inference compute compared to dense models.
Key Findings
Methodology
This work integrates sparse gating mixture of experts (MoE) into the Vision Transformer (ViT), replacing select MLP layers with MoE layers across multiple layers. The routing uses TOPk (k=1 or 2) softmax-based algorithms with noise augmentation, combined with capacity control via a ratio C and a batch prioritization strategy. The model is pretrained on large datasets (JFT-300M, JFT-3B), with auxiliary losses ensuring load balancing among experts. During inference, the model dynamically adjusts sparsity to balance performance and compute, enabling efficient scaling up to 15B parameters.
Key Results
- The 15B-parameter V-MoE model achieves 90.35% top-1 accuracy on ImageNet, surpassing many dense counterparts, while inference FLOPs are reduced by approximately 50%. Using Batch Prioritized Routing (BPR), training FLOPs are reduced by 20%, with minimal performance loss, demonstrating effective sparse routing. The model also exhibits strong transfer learning capabilities, maintaining high accuracy with minimal fine-tuning on various datasets.
- Pretraining on JFT-300M shows that V-MoE models outperform dense models in accuracy and training stability, with scalable architecture validated across different sizes (L, H, 15B). The models' ability to balance performance and efficiency via adjustable sparsity levels is confirmed through extensive experiments.
- The BPR algorithm enables the model to intelligently skip less informative image patches, reducing computational load during inference without sacrificing accuracy. This adaptive compute strategy allows the model to operate efficiently under resource constraints, making it suitable for deployment in real-world scenarios.
Significance
This research pushes the boundary of large-scale vision models by successfully applying sparse MoE mechanisms, traditionally successful in NLP, to vision tasks. It addresses the critical challenge of scaling models without prohibitive computational costs, opening avenues for deploying ultra-large models in practical applications. The ability to dynamically trade performance and compute at inference time enhances model flexibility, making it adaptable to diverse hardware environments and real-time requirements. The approach also provides insights into expert specialization and routing interpretability, fostering further research into sparse neural architectures.
Technical Contribution
The core technical innovation lies in adapting TOPk routing with noise and capacity control to vision tasks, enabling stable training of models with up to 15B parameters. The integration of Batch Prioritized Routing allows selective skipping of image patches, optimizing resource utilization. The work demonstrates that sparse MoE layers can be seamlessly incorporated into existing ViT architectures, maintaining or improving accuracy while significantly reducing computational costs. These advancements extend the theoretical understanding of sparse routing and open new possibilities for scalable vision models.
Novelty
This is the first large-scale application of sparse MoE to vision transformers, achieving unprecedented parameter sizes (15B) with competitive accuracy. The introduction of Batch Prioritized Routing for adaptive patch skipping is a novel contribution, enabling efficient inference. Unlike prior dense models, this approach offers a flexible trade-off between performance and efficiency, setting a new standard for scalable vision architectures.
Limitations
- The model's performance degrades under extreme sparsity (C<0.3), risking information loss. Routing stability depends on careful hyperparameter tuning, which can be complex. Large-scale training still demands substantial computational resources, limiting accessibility. The interpretability of expert specialization remains limited, and generalization to diverse tasks needs further validation.
Future Work
Future directions include developing more robust routing algorithms, enhancing expert specialization and interpretability, and exploring self-supervised pretraining to further improve transferability. Efforts will also focus on reducing training costs via more efficient hardware utilization, extending the approach to other vision tasks like detection and segmentation, and integrating with edge devices for real-time applications.
AI Executive Summary
Deep learning has revolutionized computer vision, with Vision Transformers (ViT) achieving remarkable accuracy on large-scale datasets. However, their dense architecture results in high computational costs, limiting scalability and deployment. To address this, the current work introduces a sparse mixture of experts (MoE) mechanism into ViT, creating the Vision MoE (V-MoE). This approach replaces select MLP layers with multiple experts, each responsible for different input regions, and employs a TOPk routing algorithm with noise and capacity controls to ensure load balancing. The key innovation is the Batch Prioritized Routing (BPR) strategy, which dynamically skips less informative image patches, significantly reducing inference cost without sacrificing accuracy. The models are pretrained on massive datasets like JFT-300M and JFT-3B, with parameters reaching 15 billion, and fine-tuned on ImageNet, achieving 90.35% top-1 accuracy. Extensive experiments demonstrate that V-MoE models outperform dense counterparts across various benchmarks, especially in resource-constrained settings. The ability to dynamically adjust sparsity during inference offers a flexible trade-off between performance and efficiency, making these models suitable for real-world applications such as autonomous driving and medical imaging. This work marks a significant step toward scalable, efficient, and adaptable vision models, opening new avenues for deploying large neural networks in practical scenarios. Future research will focus on improving routing robustness, expert specialization, and extending the framework to other vision tasks, promising a new era of large-scale, resource-efficient visual intelligence.
Deep Analysis
Background
近年来，深度学习在视觉任务中取得巨大突破，特别是Transformer架构的引入带来了性能飞跃。早期以ResNet、EfficientNet等卷积网络为代表，逐步发展到Vision Transformer（ViT），在大规模预训练和迁移学习中表现优异。然而，随着模型参数不断扩大，训练和推理成本急剧上升，成为实际应用的瓶颈。稀疏模型和条件计算逐渐成为研究热点，尤其在自然语言处理（NLP）领域，稀疏门控混合专家（MoE）模型（如Switch Transformer、GShard）已实现参数的指数级扩展，显著降低了资源消耗。将此机制引入视觉Transformer，旨在突破模型规模与效率的限制，推动大规模视觉模型的实际落地。
Core Problem
当前大规模视觉Transformer多为密集结构，导致参数冗余和计算瓶颈。虽然模型规模不断扩大，但实际应用中推理速度慢、资源消耗大，限制了其普及。如何在保证性能的同时，降低模型的计算成本，成为核心难题。此外，模型在不同任务和数据规模下的适应性不足，限制了推广。引入稀疏机制虽有潜力，但在视觉领域应用尚不成熟，面临路由算法不稳定、专家负载不均等挑战。
Innovation
本研究的创新点包括：1）将稀疏门控MoE机制引入视觉Transformer架构，突破模型规模极限；2）设计了TOPk路由算法，结合容量调节和批次优先策略，有效解决专家负载不均问题；3）提出动态调节稀疏度的机制，实现性能与计算成本的平滑折衷；4）在大规模数据集（JFT-300M、JFT-3B）上训练出参数达15B的模型，验证其扩展性和实用性。这些创新极大提升了模型效率和适应性，为未来视觉模型的规模化提供了技术基础。
Methodology
- �� 构建基于ViT的架构，将部分MLP层替换为MoE层，每层E个专家，采用TOPk路由机制，输入经过线性变换后，计算专家权重。• 路由算法在softmax后选择前k个最大值，结合噪声增强鲁棒性，确保梯度流通。• 采用容量调节参数C，控制每个专家的最大处理能力，避免负载不均。• 引入Batch Prioritized Routing（BPR），根据每个图像区域的重要性优先分配计算资源，跳过无关区域。• 训练过程中结合辅助损失，确保专家负载均衡，模型在JFT-300M预训练，支持迁移和微调。• 在不同模型规模（如L、H、15B）上验证性能，调节k值和容量比以优化性能与效率。• 采用大规模分布式训练，利用TPU集群，确保模型训练的稳定性和效率。
Experiments
- �� 在JFT-300M和JFT-3B数据集上预训练，评估模型在精度和计算效率上的表现。• 使用ImageNet-1K进行微调，验证模型在实际分类任务中的效果。• 设计少-shot和全数据微调实验，比较稠密与稀疏模型的性能差异。• 通过不同规模（L、H、15B）模型验证参数规模与性能关系。• 测试BPR在不同稀疏度下的推理速度和准确率，验证动态调节策略的有效性。
Results
- �� 15B参数的V-MoE模型在ImageNet达90.35%准确率，超越多数密集模型，推理FLOPs减半。• BPR策略实现训练FLOPs减少20%，推理时智能跳过无关区域，极大提升效率。• 在JFT-300M预训练中，模型表现优异，迁移能力强，支持少量微调。• 多模型实验显示，稀疏MoE在不同任务中均表现出良好的扩展性和性能优势，验证了稀疏机制的有效性。
Applications
- �� 适用于大规模图像识别、自动驾驶、医疗影像分析等场景，尤其在计算资源有限的边缘设备上。• 支持动态资源调度，满足不同应用对速度和精度的需求。• 未来可结合边缘计算，推动智能监控、无人驾驶等行业的技术升级。
Limitations & Outlook
- �� 在极端稀疏（C<0.3）时，模型性能可能下降，信息丢失风险增加。• 路由算法依赖大量调参，训练不稳定性较高。• 大规模训练仍需大量计算资源，限制了模型的普及。• 模型在某些特定任务中的泛化能力尚待验证，未来需优化鲁棒性。
Plain Language Accessible to non-experts
想象你在一家大型工厂工作，工厂里有很多不同的车间，每个车间专门生产不同的零件。平时，所有车间都在同时工作，效率低且浪费能源。现在，工厂决定只让最需要的车间工作，比如你要做汽车时，只让发动机和车身车间忙碌，其他的暂停。这个想法就像论文中的稀疏专家模型，模型会根据不同图片的内容，智能选择只用一部分“专家”来处理，既保证效果，又节省时间和能源。就像你只找最擅长某个零件的车间一样，模型也会根据图片不同部分，调动最合适的专家，既快又准。
ELI14 Explained like you're 14
想象你在学校里，有很多老师，每个老师都擅长不同科目。平时，所有老师都在上课，虽然都很厉害，但浪费时间。现在，老师们决定只在需要的时候出现，比如你问数学问题时，只请数学老师来讲课。这样既快又省力。模型也是这样，它可以根据图片的内容，选择只用一部分“专家”来处理，不用每个都动，节省了很多计算时间。这样，模型既能保持高准确率，又能快很多，特别适合在资源有限的设备上使用。
Abstract
Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are "dense", that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance of state-of-the-art networks, while requiring as little as half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains 90.35% on ImageNet.