Scaling Vision with Sparse Mixture of Experts

In plain words

Most image models are dense: every part of the network works on every image. Making them bigger therefore makes them slower too. This paper takes a Vision Transformer, which cuts an image into small square patches and processes them as a sequence, and gives it a bank of specialist sub-networks, called experts. A small router looks at each patch and sends it to only a couple of experts. The network can hold far more parameters than it uses for any one patch.

The paper also adds a rule for what to do when an expert is overloaded. Instead of letting patches arrive in a fixed order and dropping whatever comes last, the model first handles the patches with the strongest router weights. Under a tight compute budget it can then skip the lowest-priority patches and spend the saved work on the rest. This rule can be applied to a model that was trained without it, so the compute budget can be changed after training.

TL;DR

V-MoE replaces every other MLP block of a ViT with a sparse mixture of 32 MLP experts, routing each patch to its top-2 experts §2.2–2.3 §3.1. The authors report accuracy on par with the largest dense ViTs at as little as half the inference compute §1, and a 14.7B model at 90.35% ImageNet fine-tuned accuracy §3.6 tab. 2. Batch Prioritized Routing lets a trained model discard patches at test time to trade accuracy for compute §4.1–4.2. Both results are author-reported.

Problem

Dense ViTs spend the same compute on every patch of every image, so scaling capacity scales cost §1. Sparse MoE layers decouple parameter count from per-token compute, but they have known failure modes in vision: routing can collapse onto a few experts, and fixed buffers can overflow §2.4. The paper asks whether sparse experts can be trained at vision scale, transferred, and used at a chosen compute budget §1.

Method

Architecture. A V-MoE is a ViT with L blocks. In some blocks the MLP is replaced by a sparse MoE layer whose experts are MLPs with the same shape but separate weights §2.2. The non-expert parts (attention, routers, the remaining MLPs) are shared across devices fig. 1. Experts are placed on every other block ("Every-2") by default, following GShard §3.1. An alternative, "Last-n", places experts only on the last n even blocks §3.1.

Routing. For a patch representation x, the router computes

g(\mathbf{x}) = \mathrm{TOP}_k\big(\mathrm{softmax}(\mathbf{W}\mathbf{x} + \boldsymbol{\epsilon})\big)

with ε drawn independently per entry from N(0, 1/E²) §2.3. The top-k operation is applied after the softmax, not before as in Shazeer et al. This is needed to train with k=1, where the other order gives zero gradient almost everywhere, and it also performed better at k=2 §2.3 app. A.1. In practice k is 1 or 2 §2.3; the default is k=2 §3.1.

Capacity. Each expert processes a fixed number of tokens, its buffer size

B_e = \mathrm{round}\left(\frac{k\,N\,P\,C}{E}\right)

where N is the number of images per batch, P the tokens per image, E the number of experts and C the capacity ratio §2.4. Tokens that overflow an expert are not processed by that expert, but the residual connection carries them forward §2.4. Underfilled buffers are zero-padded §2.4. C>1 gives slack; C<1 forces the router to drop assignments §2.4. Upstream training uses C=1.05 §3.1.

Load balancing. Training adds an auxiliary loss. It combines an importance loss, the squared coefficient of variation of the summed routing weights per expert, and a load loss, the same statistic applied to a smooth estimate of how many tokens each expert receives app. A.2. The two are averaged and added to the classification loss with weight λ=0.01 app. A.2. The load loss is needed because balanced average weights can still hide an expert that is never selected app. A.2.

Upstream and transfer. The models are pre-trained on JFT-300M, deduplicated against the validation and test sets §3.2. They are evaluated by JFT precision@1, which ignores the label hierarchy app. B.6. Transfer uses:

• linear few-shot: a linear regression from frozen features to one-hot labels, on ImageNet 1/5/10-shot §3.2 §3.4;
• full fine-tuning: the standard ViT protocol with expert dropout 0.1 and half the schedule length for the other datasets §3.5 app. B.4;
• VTAB: 19 tasks with 1,000 examples each §3.2 tab. 1.

Scaling to 15B. The largest model, V-MoE-15B, has 48 blocks, Every-2 placement, 32 experts and k=2, for 14.7B parameters §3.6 tab. 2. It is trained on JFT-3B, with the changes listed in app. B.3 §3.6.

Batch Prioritized Routing (BPR). BPR is applied at inference to models trained with vanilla routing, with unchanged parameters, and it trades compute for accuracy by changing the capacity ratio §4.2 app. C.2. Vanilla routing, as in GShard and Switch, goes through the rows of the batch in a fixed order and assigns each token to its top-i expert while the buffer has room. Priority is set by row order, and within an image the token order is fixed §4.1 app. C.1. BPR instead gives each token a priority score and sorts all tokens across the batch before allocation. The score is the token's largest router weight. The sum of the top-k weights worked equally well §4.1 app. C.1. Every top-i choice is still assigned before any top-j choice, for i < j §4.1.

# Batch Prioritized Routing allocation  [§4.1] [app. C.1]
# Input:  gates (N*P, E) router weights; k; capacity per expert
# Output: list of (token, expert, weight) assignments

def batch_prioritized_routing(gates, k, capacity):
    # Priority score per token: its largest router weight  [§4.1]
    scores = max(gates, axis=expert)             # (N*P,)
    # Visit tokens from highest to lowest score  [§4.1]
    order = argsort(scores, descending=True)     # (N*P,)
    # Fill buffers choice by choice: every top-1 before any top-2  [§4.1]
    buffer_fill = zeros(E)                       # (E,)
    assignments = []
    for i in range(k):
        for token in order:
            # i-th highest expert for this token (0-based index)  [app. C.1]
            expert = top_i_expert(gates, token, i)
            if buffer_fill[expert] < capacity:
                buffer_fill[expert] += 1
                weight = gates[token, expert]
                assignments.append((token, expert, weight))
            # Full buffer: drop this choice, keep the residual  [§2.4]
    return assignments

Skip-patch (a fixed fraction S of patches is kept, the rest is dropped) is also described, but the paper presents it only as a related algorithm it tried app. C.1. Priority routing with a capacity ratio C is the main method §4.2.

Hyperparameters

Upstream model and routing settings (default V-MoE configuration):

Upstream training on JFT-300M app. B.2:

Notes for V-MoE-15B app. B.2 app. B.3:

• weight decay 0.03 on kernels, with biases not regularised, and 3.0 on the head kernel app. B.3;
• gradient moving average stored in bfloat16 app. B.3;
• classification head is a self-attention head with an extra MLP, not the token head app. B.3;
• the paper says ViT-S settings do not match DeiT-Small app. B.2.

Fine-tuning app. B.4:

V-MoE-15B ImageNet fine-tune: 560×560 input, 30,000 steps (about 6 epochs), base LR 0.006, debiased Polyak averaging with momentum 0.999999 app. B.4.

Results

Author-reported. The upstream set is JFT-300M; ImageNet numbers are few-shot or fine-tuned as labelled. Compute columns are as printed in the table tab. 2.

VTAB (19 tasks, 1,000 examples each; mean with 95% CI) tab. 1:

Inference with BPR, FLOPs-matched to dense ViT-H/14 app. C (matching-cost table; ViT-H/14 row: JFT 56.68, IN/5-shot 76.95):

Text claims for BPR §4.2: a V-MoE-H/14 trained with vanilla routing can beat the dense model while using about half the FLOPs and a little over half the runtime; the authors also state that BPR stays close to the dense model while processing roughly 15–30% of the tokens fig. 6. The low-capacity points that support these claims are in fig. 5 and app. C; the table in app. C for vanilla routing does not show them.

Training with BPR (the claims from app. C.3, not tabulated in the main text): matches dense upstream performance with about 80% of training FLOPs for V-MoE-S/32 and V-MoE-B/32, and with 70–75% for V-MoE-L/32 (k=1 and k=2) app. C.3.

Fine-tuning with BPR on V-MoE-S/32 (k=2) app. C.4: accuracy falls from 77.4% to 68.6% when capacity drops to 0.1; BPR in both pre-training and fine-tuning gives 71.8% at that capacity; keeping high capacity during fine-tuning matters more than during pre-training (74.1% vs 71.7%).

Model analysis §5 app. E: experts in deeper MoE layers specialize to small sets of ImageNet classes; earlier layers do not show this fig. 7 §5. Replacing one router at a time with a random router hurts little in early layers and more in the last layers §5 app. E.

Paper vs code

No official code. The paper's LaTeX source contains no code links or project page [raw/src/2106.05974/main.tex]. The arXiv abstract page lists no code repository. A GitHub repository search for the title returned one result, owned by an account that is not an author, so it was not used as official code. A search for the method name ("vmoe sparse mixture of experts vision") returned 0 results.

Pitfalls

• Router noise and inference. The paper says noise is added to the router logits during the forward pass and that it "improves performance" §2.3. It does not state whether noise is used at test time. Not stated.
• Top-k weights are not renormalised, as far as the text says. The paper applies top-k after the softmax §2.3, but does not state whether the selected weights are rescaled to sum to one. Not stated.
• Capacity and rounding. B_e uses round(); the tie-breaking and the rounding of very small buffers are not stated §2.4.
• Overflow order. Under vanilla routing, which tokens overflow depends on the fixed order within images and the random order of images in a batch §4.1. Under BPR, the token ranking is global across the batch, so each image's compute depends on the other images in the batch §4.1. The batch size used at test time is not stated, so reproducing the BPR numbers requires matching batch composition.
• Every-2 indexing. The paper says MoEs are "on every other layer" §3.1 and that Last-n uses "the last n even blocks" §3.1. Whether blocks are counted from 0 or 1 is not stated.
• Inconsistent upstream numbers. V-MoE-L/16 Every-2 has JFT prec@1 57.65 in tab. 2 but 57.64 in app. C. The ImageNet 5-shot values are the same (77.10) in both. The paper does not explain the difference.
• FLOPs and time are not the same. The paper repeatedly says dense models need less runtime for the same FLOPs because of data transfer §4.2. In the matching-cost table, several V-MoE rows are below the dense model in FLOPs but well above it in time (for example 155.83% time at 100.03% FLOPs) app. C.
• Hardware. All timing figures are on TPUv3 cores tab. 2 fig. 2. Other hardware is not covered.
• Layer choice. The authors state that fewer MoE layers (Last-n) have "typically little impact on quality" and are faster, because of less communication overhead §3.1 app. E. This is for their setup only.
• Fine-tuning. Full fine-tuning uses the auxiliary loss, though the authors say it is often not needed at that step §3.5. They also say the fine-tuning protocol "could be further improved" for sparse models §3.5.
• Scale-dependent tricks. The 15B model needs bfloat16 moving averages, rsqrt LR decay, changed weight decay and a new head app. B.3. These are described as changes "based on" another paper app. B.3; their effect on the 15B result is not isolated.
• Dropout and repeats. The ImageNet fine-tuning curve is averaged over three runs §3.5. Few-shot numbers for the training-data deduplication study vary a lot across seeds app. B.7.
• Priority scores are a heuristic. The authors note that router outputs mainly reflect token–expert fit rather than how much a token matters for the final label §4.1. Learned scoring functions were tried and not better app. C.1.
• Skip-patch is not the main method. Enforced dropping of a fraction S of patches is presented only as a related experiment app. C.1. Do not implement it as the BPR of this paper.

Lineage

• Builds on: Vision Transformer (ViT) §2.2 §3.1; sparsely gated MoE for deep learning (Shazeer et al. 2017), changing the top-k order and the load-balancing losses §2.3 app. A.2; GShard-style placement of MoEs on every other layer §3.1; Switch-style expert dropout app. B.4 §3.5.
• Compared against: dense ViT variants (S/32 to H/14), NFNet-F4+ and Meta Pseudo-Labels on ImageNet fine-tuning tab. 2.
• Superseded by: not stated.