PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Vision with Sparse Mixture of Experts

Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, Neil Houlsby

arXiv:2106.05974v1cs.CVcs.LGstat.ML

TL;DR

Vision models are usually dense, making scale expensive, so the paper asks whether sparse conditional computation can scale vision effectively. It introduces V-MoE, a sparse Vision Transformer with routed expert MLPs and adaptive routing, and shows competitive or better performance with reduced inference cost, including a 15B-parameter model. The approach also supports changing inference sparsity after training, while its demonstrated scope leaves further scaling and transfer improvements open.

  • Problem

    Dense vision networks process every input with every parameter, making large-scale models expensive to train and serve.

  • Method

    V-MoE replaces selected ViT feedforward layers with sparse routed expert MLPs and adds Batch Prioritized Routing for adaptive patch-level computation.

  • Results

    V-MoE reaches 90.35% ImageNet test accuracy at 15B parameters and can match the largest dense model using as little as half the inference compute.

  • Takeaways & Limitations

    Trained V-MoE models can be reused with adjusted sparsity at inference, trading performance against compute without further model adaptation.

  • Takeaways & Limitations

    The paper identifies further work in scaling expert counts, reducing data dependence, and improving transfer of sparse representations.

Abstract

from arXiv · show

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.

1 Introduction

V-MoE brings sparse conditional computation to vision by replacing selected ViT feedforward layers with routed expert MLPs. It scales to 15B parameters while matching or exceeding dense-model performance at reduced or adjustable inference cost.

  • Vision models at scale: V-MoE replaces selected dense ViT feedforward layers with sparse MoE layers that route each image patch to a subset of expert MLPs.The architecture distributes experts across devices while replicating routers, attention layers, and normal MLP blocks.
  • Batch Prioritized Routing: Batch Prioritized Routing discards the least useful patches, reducing computation on uninformative image regions and enabling adaptive inference cost.The paper reports matching dense-model performance while saving 20% of training FLOPs.
  • Vision models at scale: 15B parameters is the largest reported V-MoE scale, with models stably trained, transferred, and fine-tuned using as few as 1 000 datapoints.The largest model achieves 90.35% test accuracy on ImageNet after fine-tuning.
  • Vision models at scale: 90.35% test accuracy on ImageNet is achieved by the largest fine-tuned V-MoE model.
  • Performance and inference: V-MoE models can match the largest dense model while using as little as half the inference compute, or outperform it at the same cost.

2 The Vision Mixture of Experts

The Vision Mixture of Experts applies sparse expert routing to ViT patch representations, activating only a small subset of experts per token. Its routing and capacity mechanisms address load imbalance while trading model computation against capacity and quality.

  • 2.1 Conditional Computation with MoEs: Conditional computation activates different network subsets for different inputs, while MoEs assign different experts to different regions of the input space.
  • 2.1 Conditional Computation with MoEs: Sparse routing restricts each input to k ≪ E non-zero expert weights, so unused experts need not be computed.This permits model-parameter scaling beyond inference and training compute scaling.
  • 2.2 Vision Mixture of Experts: V-MoE replaces selected ViT MLPs with same-architecture expert MLPs having different weights, while ViT represents images as sequences of patches.
  • 2.3 Routing: TOPk(softmax(Wx + ϵ)) routes each image-token representation to the k highest-weight experts, with k = 1 or k = 2 in practice.Applying TOPk after softmax enables training with k = 1 and improves performance for k > 1.
  • 2.4 Capacity and load balancing: Expert collapse causes statistical and computational inefficiency by reducing effective capacity and producing imbalanced hardware utilization.The implementation fixes expert buffer capacities and uses auxiliary load-balancing losses.
  • 2.4 Capacity and load balancing: Capacity ratio C controls expert buffers: C > 1 adds slack for routing imbalance, whereas C < 1 forces some assignments to be ignored and can save inference compute.Tokens exceeding capacity retain information through residual connections rather than being entirely discarded.

3 Transfer Learning

The paper evaluates V-MoE transfer through fixed-representation few-shot learning, full fine-tuning, and low-data adaptation, while varying expert placement, selection, and capacity. V-MoE preserves upstream gains on ImageNet few-shot evaluation, performs competitively or better during fine-tuning, and scales to a 14.7B-parameter model.

  • Transfer Learning: Transfer experiments use linear few-shot learning on fixed representations, full fine-tuning, and VTAB adaptation with 1,000 data points per task.The few-shot setup uses 1, 5, or 10 examples per ImageNet class; full fine-tuning uses approximately 1M images.
  • Transfer Learning: V-MoE variants are built from ViT-S, ViT-B, ViT-L, and ViT-H, with expert layers placed every other block or in the last n even blocks.The study also varies the number of selected experts and buffer capacity; the default settings are k = 2, E = 32, and upstream C = 1.05.
  • Linear few-shot results: Upstream gains are preserved on ImageNet 5-shot evaluation, where V-MoE representations outperform ViT models for a new task under compute and time comparisons.JFT performance is measured with precision@1, while ImageNet transfer uses fixed representations and linear regression.
  • Low-data regime: VTAB results show significant gains from experts at the ViT-L/16 level, while V-MoE-H/14 performance is similar in the low-data regime.These results indicate that the models can be fine-tuned with small amounts of data without further tricks.
  • Scaling up V-MoE: The 14.7B-parameter V-MoE-15B achieves 82.78% ImageNet 5-shot accuracy and 90.35% when fully fine-tuned.It is trained with 16.8k TPUv3-core-days on the larger JFT-3B dataset.

4 Skipping Tokens with Batch Prioritized Routing

Batch Prioritized Routing (BPR) ranks tokens across the batch before expert allocation, allowing V-MoE to discard less useful patches and reduce inference or training compute. It improves the performance–compute trade-off over vanilla routing, including at low capacity.

  • Inference: BPR can make V-MoE competitive with dense models while processing only 15-30% of the tokens.The comparison is reported for a V-MoE-H/14 model with k = 2 and varying inference capacity ratios.
  • Routing algorithm: BPR prioritizes tokens using their maximum routing weight before allocating them to experts.The sum of TOP-k weights performs equally well, while directly learning the priority function performed worse.
  • Routing algorithm: BPR sorts all batch tokens by priority, so reducing expert capacity discards lower-priority patches rather than following input order.Vanilla routing processes rows sequentially, giving priority based on batch and within-image ordering.
  • Inference: BPR outperforms vanilla routing at inference, with the largest difference when capacity C ≤0.5 and tokens are fully dropped.Figure 5 evaluates the performance–inference-FLOPs trade-off for large V-MoE models using BPR with C ∈{0.6,0.7,0.8} and k ∈{1,2}.
  • Training: Max-weight BPR matches dense performance while saving around 20% of total training FLOPs.It also strongly outperforms vanilla routing at a similar FLOP budget.

5 Model Analysis

The model analysis examines expert specialization, router usefulness, routing-weight distributions, and robustness to changing k at inference. Experts specialize more clearly in later layers, while routing decisions and configurations show different degrees of flexibility.

  • Overview: The analysis is exploratory and is used to guide the design of new routing algorithms for complex sparse models.
  • Specialized experts: Experts in later MoE layers specialize in discriminating between small sets of ImageNet classes, whereas earlier layers focus on features shared across classes.The reported shared features include background, basic shapes, and colours.
  • The value of routers: Router interventions show that models are robust to early routing changes but more sensitive to decisions in the last layers.The analysis replaces routers one at a time with uniformly random routers after training.
  • Routing weights distributions: Selected routing-weight distributions vary widely across different mixture-of-experts layers.
  • Changing k at inference time: Sparse models remain fairly robust when inference uses k′ ≠ k from training, enabling control over FLOPs per input.The paper evaluates changing the number of selected experts at inference time to decrease or increase computation in production systems.

6 Related work

Related work frames V-MoE within conditional computation and mixture-of-experts research, especially sparse expert models in language. Earlier vision MoEs generally operated at considerably smaller scale or used different convolutional and routing designs.

  • Conditional Computation: Conditional computation activates input-dependent portions of a model to increase parameter count without proportionally increasing computational cost.Prior mechanisms include stochastic neurons, reinforcement learning, and decision-tree-like selection.
  • Mixture of Experts: Mixture-of-experts models combine sub-model outputs through an input-dependent router, using either all experts or a sparse subset.
  • MoEs for Language: Language MoEs scaled this approach to very large models, including transformer systems with over one trillion parameters and sparse expert selection.The cited work also reported faster pre-training than a dense baseline with transfer and distillation gains.
  • MoEs for Vision: Earlier vision MoEs were considerably smaller than language models and V-MoE, including convolutional approaches with channel, kernel, or shallow-router selection.

7 Conclusions

The paper concludes that sparse conditional computation can scale vision models and support input sparsity through Batch Prioritized Routing. It presents this as an early stage whose extensions include broader expert scaling, reduced data dependence, and improved transfer.

  • Conclusions: Sparse conditional computation produced some of the largest vision models to date, with improvements in representation learning and transfer learning.
  • Conclusions: Batch Prioritized Routing repurposes model sparsity to introduce input sparsity without further adapting the model.This allows trained models to be reused with sparse conditional computation.
  • Conclusions: Sparse models can be made more efficient at inference without further model adaptation, a point connected to inference-dominated CO2 footprints.The passage states that 90% of the footprint in cited recent NLP work stems from inference costs.
  • Future directions: The work identifies scaling expert counts, reducing dependence on data, and improving transfer as directions for future development.It also mentions heterogeneous expert architectures and conditional variable-length routes.

A.1 Ablation on the modification of the routing function

The routing formulation reverses softmax and top-k to preserve trainability for k = 1, while auxiliary losses encourage balanced expert usage. The training setup combines importance and load regularization with standard classification loss.

  • Routing function: Reversing softmax and top-k avoids the zero-gradient problem that prevents easy training when k = 1.The authors also report better performance than the original ordering for k > 1.
  • Importance loss: Importance loss encourages similar normalized routing weights across experts for each batch.The importance distribution is regularized using its squared coefficient of variation.
  • Load loss: Importance balancing alone can leave a small subset of experts receiving most assignments despite balanced aggregate weights.This motivates an additional differentiable proxy for assignment balancing.
  • Load loss: The load loss estimates expert-selection probabilities using the k-th maximum noisy routing score as a threshold.The threshold is evaluated under resampled router noise to obtain a differentiable assignment-balance proxy.
  • Combined objective: The final auxiliary loss averages importance and load losses, and the overall objective adds it to classification loss.All experiments use λ = 0.01, which the authors found robust and insensitive.
  • Experimental setup: The experiments span JFT pretraining, few-shot and full-data ImageNet, VTAB, and other natural, specialized, and structured datasets.The models use sparse and dense ViT variants with documented architectural and fine-tuning settings.

B.5 Results and details for all models

The appendix reports architectural, training-cost, and performance details for dense and sparse models across upstream, few-shot, and downstream evaluation. It also defines how multi-label JFT precision is computed.

  • Model results: Table 8 provides architecture, training cost, upstream, few-shot, and downstream results for dense and sparse models.V-MoE models use E = 32 experts and report selected experts per token, model size, and evaluation metrics.
  • Compute accounting: The appendix includes total training time on one TPUv3 core and total training compute in exaFLOPs.These costs accompany the reported model performance comparisons.
  • Model results: Figure 8 compares JFT-300M, ImageNet 5-shot, and full-ImageNet fine-tuning against total training time.Colors identify ViT variants, markers distinguish dense ViT from V-MoE, and lines show their Pareto frontiers.
  • Evaluation: JFT Precision@1 ignores the class hierarchy and checks whether the highest-probability prediction is among the image’s true labels.This defines the reported metric for the multi-label upstream dataset.

B.7 Training data deduplication

Deduplicating ImageNet-near images from training has no consistent significant effect on V-MoE-S/32 few-shot results. The appendix also documents routing algorithms and compute–performance evaluation figures.

  • Training data deduplication: Deduplication produces no consistent and significant effect on V-MoE-S/32 few-shot evaluation.Across seeds, deduplicated models can outperform or underperform the original, except that 1-shot worsened.
  • Compute comparisons: Figures 9–13 compare sparse and dense models using upstream, few-shot, and training-compute or training-time axes.The figures cover ImageNet 1-shot, 5-shot, and 10-shot settings alongside upstream performance.
  • Inference routing: Figure 14 compares performance with inference FLOPs and runtime for vanilla routing, Batch Prioritized Routing, and ViT models.The sparse-routing comparisons vary C over {0.6,0.7,0.8} and k over {1,2}.
  • Batch Prioritized Routing: Batch Prioritized Routing orders patches by routing scores before assigning them to expert buffers, allowing capacity-based dropping.The explored scoring function uses each patch’s maximum routing weight.
  • Skip-patch routing: Skip-patch directly discards the lowest-scoring fraction 1 − S of patches before routing the remaining M = SN patches.S and C are distinct parameters and may be adjusted to avoid excessive FLOPs waste.

C.2 Applied during Inference

Batch Prioritized Routing can be applied at inference to pretrained V-MoEs, reducing processed patches as capacity decreases. It yields a smooth performance–compute trade-off without further training.

  • Inference procedure: Batch Prioritized Routing is applied to models trained with vanilla routing without changing their parameters or performing additional learning.The method changes routing only during inference.
  • Capacity trade-off: Performance decreases slowly and smoothly as expert capacity C is constrained for every-2 models with k = 2.Figure 15 reports this relationship across the main model capacities.
  • Routing comparison: For V-MoE-H/14, L/16, B/16, and S/32, Batch Prioritized Routing shows a clear advantage over vanilla routing in the low-capacity regime.The most interesting regime is C < 0.5, where the comparison shows a large performance gap.
  • Practical implication: The method provides a smooth performance–compute trade-off after training, without further training or adjustment.This flexibility allows deployment resources and constraints to determine the operating point.

C.3 Applied during Training

Batch Prioritized Routing improves training efficiency across V-MoE-S/32, V-MoE-B/32, and V-MoE-L/32, while routing behavior varies across network depth and inference settings.

  • C.3 Applied during Training: Training with Batch Prioritized Routing improves results across V-MoE-S/32, V-MoE-B/32, and V-MoE-L/32, especially at reduced capacity.With full capacity, the two routing algorithms are expected to behave similarly because no tokens need dropping.
  • C.3 Applied during Training: 20% fewer training FLOPs suffice for V-MoE-S/32 to match dense upstream performance and preserve few-shot performance.The upstream match requires around 80% of training FLOPs; few-shot matching requires around 85% or 80%, depending on k.
  • C.3 Applied during Training: 10% fewer training FLOPs suffice for V-MoE-B/32 to match dense upstream performance while keeping or improving few-shot representation quality.The upstream match requires at most 80% of the dense model’s training FLOPs.
  • C.3 Applied during Training: 70–75% of training FLOPs suffice for V-MoE-L/32 to match dense upstream performance, while few-shot performance remains identical with over 20% savings.At k = 2 and C = 0.1, the expert model already outperforms dense upstream precision.
  • E.1 Routing analysis: Routing is most important in the last two MoE layers, while most intermediate layers tolerate mis-routing except layer 9.Replacing consecutive layers with random routers causes performance drops, indicating compounded mis-routing is harmful.
  • E.2 Specialized experts: Experts specialize by class in later layers but appear to process shared patch properties such as backgrounds and basic shapes in earlier layers.Layer-wise routing weights and expert selection can therefore support different computational roles across depth.
  • E.4 Changing k at inference: Increasing k at inference can improve performance, but excessive k eventually hurts when the model is unprepared for the altered routing-weight distribution.Models trained with larger k generally improve when k increases during fine-tuning, whereas k = 1 downstream can erase benefits from larger upstream k.

E.6 Pre-training with less data

V-MoE can remain competitive with dense models as pre-training data decreases, but its advantage weakens at very small data scales and depends on regularization choices.

  • Training on JFT300M with less data: V-MoE-L/32 is initially robust to reduced JFT-300M data, but a dense model becomes slightly preferable at 9M pre-training samples.The 9M subset is about 3% of JFT-300M; the remaining evaluated data sizes begin at 30M, about 10%.
  • Training on ImageNet21k: On ImageNet-21k, sparse scaling generally harms performance relative to dense counterparts except for V-MoE-S.The authors observed overfitting in both pre-training validation accuracy and transfer performance as training continued.
  • Training on ImageNet21k: RandAugment with N = 2 and M = 10 helps expert models while harming dense models, yielding an expert model that outperforms the dense baseline for each architecture.This intervention is presented as an initial response to overfitting on ImageNet-21k.
  • Discussion: The reduced-data study is explicitly incomplete and points to data augmentation and regularization as important considerations for improving V-MoE with less pre-training data.The authors expect methods developed for dense and data-efficient vision transformers to be useful for V-MoE.
Loading 2106.05974v1…