From Sparse to Soft Mixture of Experts

ArXiv: 2308.00951

🎯 Pitch

Soft MoE replaces discrete token routing with differentiable slot-based mixtures, slashing inference cost while scaling to 40× the parameters of a ViT Huge—adding only 2% latency. It matches or exceeds dense models on few-shot and fine-tuning while running 5.7× faster at inference, eliminating the token-dropping and instability that plague sparse MoEs.

1. Executive Summary

This paper introduces Soft MoE, a fully-differentiable sparse Transformer layer that replaces the discrete, non-differentiable token-to-expert assignment in conventional Sparse Mixture of Experts with an implicit soft assignment — each expert processes a learned weighted combination of all input tokens (dispatch weights), and outputs are reconstructed via another learned combination (combine weights). In image classification experiments on JFT-4B with Vision Transformers at Small through Huge scales, Soft MoE dominates both dense ViTs and popular Sparse MoEs (Tokens Choice and Experts Choice routing) on training cost/performance Pareto frontiers, while enabling models with over 40× more parameters than ViT Huge/14 to run at only 2% increased inference time. The paper also demonstrates that Soft MoE B/16 matches or outperforms ViT H/14 on few-shot and finetuning metrics while being 5.7× faster at inference, establishing that the benefits of test-time sparsity can be preserved without the training instability, token dropping, and expert imbalance that limit conventional Sparse MoEs — though only in encoder architectures where per-token determinism and full-sequence mixing are permissible.

2. Context and Motivation

The Core Problem: Scaling Model Capacity Without Paying Full Computational Cost

The fundamental tension this paper addresses is one of the central challenges in modern deep learning: how do you make models bigger and more capable without making them proportionally slower and more expensive to run? For Transformer models, this tension is particularly acute because their cost scales directly with model size — a Transformer with $N$ parameters costs roughly $O(N)$ to train and run at inference time. The empirical evidence from scaling laws (Kaplan et al., 2020; Hoffmann et al., 2022; Zhai et al., 2022a) tells us that, to optimally use a given training compute budget, we need to scale model size and training data together. This creates a practical bottleneck: we can design architectures with enormous capacity, but we cannot afford to activate every parameter on every input.

The Mixture of Experts (MoE) paradigm offers a conceptual solution to this dilemma. Rather than activating all model parameters for every input, an MoE model selectively routes different inputs through different subsets of parameters, called experts. The key insight is that, if the routing is done efficiently, the model can have many more total parameters than it uses for any single forward pass, breaking the linear relationship between total capacity and per-input cost.

Where Existing Sparse MoE Approaches Fall Short

The paper identifies several specific, interrelated failure modes in the dominant Sparse MoE approaches that motivated the development of Soft MoE:

The discrete assignment problem is fundamentally awkward to optimize. All conventional Sparse MoE routing algorithms face a core difficulty: they must solve a discrete matching problem between tokens and experts at every layer, during every forward pass, while simultaneously training the entire network end-to-end with gradient descent. This is inherently problematic because the routing decision is non-differentiable — you cannot take a meaningful gradient through a discrete choice of which expert processes which token. The paper catalogs the variety of techniques that have been proposed to work around this, each with their own complexity and limitations:

"Many techniques have been devised to find good token-to-expert matches: linear programs (Lewis et al., 2021), reinforcement learning (Bengio et al., 2015), deterministic fixed rules (Roller et al., 2021), optimal transport (Liu et al., 2022), greedy top-k experts per token (Shazeer et al., 2017), or greedy top-k tokens per expert (Zhou et al., 2022)."

Each of these represents a different compromise. Linear programming and optimal transport approaches are principled but computationally expensive. Reinforcement learning introduces its own optimization difficulties. Greedy top-k approaches are fast but can produce load-imbalanced assignments that require additional heuristic auxiliary losses to correct.

Token dropping silently degrades performance. In capacity-constrained Sparse MoEs where experts can only process a fixed number of tokens, some tokens inevitably go unprocessed. This token dropping appears in both major routing paradigms:

• Tokens Choice (Shazeer et al., 2017; Lepikhin et al., 2020): Each token selects the top-$K$ experts with the highest routing scores, but experts have limited capacity. When an expert's buffer fills up, tokens that selected it are dropped. A token is fully dropped when none of its $K$ selected experts have capacity remaining.

• Experts Choice (Zhou et al., 2022): Each expert selects the top-capacity tokens. Since experts make independent selections, some tokens may be selected by many experts (receiving more computation than necessary) while others are selected by none (getting zero computation).

The paper quantifies this in Appendix B (Figures 8-10): as the number of experts increases while capacity remains tight ($C=1$), dropping rates quickly grow to 15% or more. At extreme scales (hundreds to thousands of experts), Experts Choice can drop 40-50% of tokens in some layers — meaning nearly half the input receives no expert processing at all. This is not a minor implementation detail; it directly sacrifices model quality for the sake of routing efficiency.

Training instability from load imbalance. Even when tokens are not completely dropped, uneven distribution of tokens across experts creates training difficulties. Some experts receive many more tokens than others, leading to uneven gradient updates. The standard remedy — auxiliary load-balancing losses — adds another heuristic term to the training objective that must be carefully tuned and can conflict with the primary task loss. The paper notes that these "challenges can be greater in out-of-distribution settings: small inference batch sizes, novel inputs, or in transfer learning," since routing patterns that were balanced on the training distribution may collapse when the input distribution shifts.

Inability to scale the number of experts effectively. A particularly striking finding in the paper (Section 3.5, Figure 6) is that existing Sparse MoEs hit a ceiling on the number of experts they can productively use. In the experiments varying expert count:

"For the two Sparse MoEs, there is a point at which training difficulties outweigh the benefits of additional capacity, resulting in the a modest optimum number of experts."

For Tokens Choice routing, performance actually degraded beyond 128-256 experts as training instability and token dropping overwhelmed any capacity benefit. For Experts Choice, the optimum was higher but still finite. This is a critical limitation because, in principle, more experts should always be better — more total parameters with the same per-input FLOPs is the entire promise of MoE architectures. The fact that existing routers cannot realize this promise at scale is a fundamental barrier.

Hardware-inefficient routing operations. Sparse MoE routing requires sorting tokens by routing scores, applying top-k selections, and managing expert buffers. These operations are "typically not well suited for hardware accelerators" (Section 2.2), particularly on TPUs and GPUs which are optimized for dense matrix multiplications, not branching logic and sparse data movement. As Figure 6 (bottom) shows, the training throughput of Sparse MoEs drops dramatically as the number of experts increases — Tokens Choice throughput falls to less than half when going from 8 to 4096 experts — even though the theoretical FLOPs remain constant. The routing overhead itself becomes the bottleneck.

Per-sequence non-determinism. In capacity-constrained Sparse MoEs, tokens from different sequences within a batch compete for fixed-capacity expert buffers. This means the routing for a given input sequence depends on what other sequences happen to be in the same batch:

"the model is no longer deterministic at the sequence-level, but only at the batch-level."

This is problematic both for reproducibility (the same input can produce different outputs depending on batch composition) and for deployment scenarios where batch sizes may be small or inputs arrive one at a time. Under small batch sizes, the capacity constraints become more severe because tokens have fewer opportunities to find available expert slots, potentially amplifying the token dropping problem.

The Broader Importance of Solving These Problems

The combined effect of these failure modes is that Sparse MoEs, despite their theoretical appeal, are fragile in practice. They require careful tuning of routing hyperparameters (number of experts, capacity multiplier $C$, number of selected experts $K$, auxiliary loss weight), are sensitive to batch size and input distribution, and cannot trivially scale to very large numbers of experts without diminishing returns or training collapse.

This fragility matters enormously for the field's trajectory. As models continue to scale — and as the economics of training and deployment increasingly favor architectures that can amortize parameter cost across inputs — the inability to reliably build large MoE models is a bottleneck on practical progress. The promise of MoEs is models with hundreds of billions or trillions of parameters that cost only as much to run as models 10-100× smaller. If the routing mechanism is unreliable, that promise goes unrealized.

Furthermore, the specific challenges around token dropping and load balance have real consequences for fairness and robustness. A model that silently drops tokens is, in effect, allocating less computation to those inputs — and if those inputs come disproportionately from certain classes, data sources, or user populations, the dropping introduces a form of computational bias that is difficult to detect or correct.

How Soft MoE Positions Itself Relative to Prior Work

The paper positions Soft MoE not as an incremental improvement to existing routing algorithms, but as a fundamentally different approach that sidesteps the discrete assignment problem entirely. Rather than trying to solve the hard matching problem better (with more sophisticated optimization, better auxiliary losses, or clever heuristics), Soft MoE asks: what if we don't do discrete assignment at all?

The key conceptual move is to replace hard, one-hot routing (token $i$ goes to expert $j$) with soft, continuous routing (token $i$ contributes fractionally to every expert's input, according to learned weights). This is made precise through the dispatch and combine weight matrices in Equations 1 and 2:

• Every expert's input is a learned convex combination of all input tokens — no token is ever dropped, because every token contributes to every expert.
• Every output token is a learned convex combination of all expert outputs — every expert contributes to every token's final representation.

This design directly addresses each of the identified failure modes:

• Fully differentiable: All operations are continuous softmax-weighted averages, so gradients flow through the routing weights. No reinforcement learning, no discrete optimization, no auxiliary losses.

• No token dropping: Every token participates in every expert's input (to varying degrees), so no computation is wasted.

• No expert imbalance: Every expert processes exactly $p$ slots, each of which is a weighted average of all tokens. The computation is perfectly balanced by construction.

• Scales to many experts: Since the cost depends on the total number of slots ($n \cdot p$) rather than the number of experts per se, and since each expert can have as few as one slot, the number of experts can grow without increasing FLOPs. The paper demonstrates this up to 4096 experts (Figure 23).

• Hardware-efficient: The routing involves only dense matrix multiplications and softmax operations — no sorting, no top-k, no buffer management. As Figure 6 shows, Soft MoE throughput barely changes when scaling from 8 to 4096 experts.

• Per-sequence deterministic: Because slots are computed as weighted averages of tokens within a single sequence, the computation for each input is independent of batch composition.

The paper also explicitly distinguishes Soft MoE from two related but distinct ideas:

Token merging / sequence reduction methods (Jaegle et al., 2021; Ryoo et al., 2021; Renggli et al., 2022; Wang et al., 2022) also compute weighted averages of input tokens, but their goal is to reduce sequence length to mitigate the quadratic cost of self-attention. Soft MoE's goal is different: it preserves the original sequence length (through the combine weights that reconstruct output tokens) and uses the weighted averaging as a routing mechanism to distribute computation across expert parameters.

Parameter averaging / soft merging of experts (Yang et al., 2019; Tian et al., 2020; Muqeeth et al., 2023) are also fully differentiable, but they merge expert parameters rather than routing inputs. These approaches can be more expensive because (a) averaging large expert parameter matrices is costly in memory and time, and (b) every input uses a different weighted combination of parameters, preventing efficient batched computation. Soft MoE instead routes inputs to experts, where each expert is applied to multiple slots through standard vectorized operations.

The Paper's Explicit Scope Boundaries

The paper is careful to acknowledge that Soft MoE does not solve all MoE problems. It explicitly identifies one major limitation in Section 6: auto-regressive decoding. Because Soft MoE computes each slot as a weighted average of all input tokens, it breaks causality if applied naively to decoder architectures where future tokens must not influence past predictions. The paper frames this as "a promising research avenue that we leave for future work," acknowledging that the current contribution is specific to encoder (or encoder-like) architectures where full bidirectional context is available.

This is a significant scope constraint: it means Soft MoE, as presented, applies to vision encoders, BERT-style language encoders, and other non-autoregressive contexts, but not to GPT-style language models where the MoE paradigm has seen some of its most prominent deployments (Fedus et al., 2022; Lepikhin et al., 2020). The contrastive learning experiments in Section 4 extend the evaluation to image-text alignment but still use Soft MoE only in the image encoder tower.

In summary, the paper addresses a clear, well-documented set of practical problems with Sparse MoEs that have limited their adoption and scalability, and proposes a solution that — within its encoder-only scope — avoids all of these problems through a fundamental redesign of the routing mechanism rather than incremental improvements to the existing discrete-assignment paradigm.

3. Technical Approach

3.1 Reader Orientation

We are building a drop-in replacement layer for the MLP blocks in a Vision Transformer that allows the model to have vastly more parameters (many parallel "expert" MLPs) while keeping the per-input computational cost roughly constant. The problem is that conventional Sparse MoE layers require solving a discrete token-to-expert assignment problem that introduces training instability, token dropping, expert imbalance, and scaling bottlenecks. The solution is to eliminate discrete assignment entirely: instead of sending each token to one or a few specific experts, every expert processes a learned weighted average of all tokens, and every output token is a learned weighted average of all expert outputs. This makes the entire layer fully differentiable, perfectly load-balanced by construction, and free of token dropping — at the cost of one extra matrix multiplication per routing direction.

3.2 Big-Picture Architecture (Diagram in Words)

A Soft MoE layer sits in place of a standard MLP block inside a Transformer encoder. It has the following major components:

1. Input tokens $X \in \mathbb{R}^{m \times d}$ — the $m$ token representations (e.g., 196 patches for a 224×224 image at patch size 16), each of dimension $d$, arriving from the previous attention block after layer normalization.

2. Slot parameters $\Phi \in \mathbb{R}^{d \times (n \cdot p)}$ — a learned matrix with one $d$-dimensional column per slot, where $n$ is the number of experts and $p$ is the number of slots per expert. This matrix is the only additional learned parameter beyond the expert MLP weights, and it controls the soft routing.

3. Dispatch weights computation — a softmax over columns of $X\Phi$ that produces a matrix $D \in \mathbb{R}^{m \times (n\cdot p)}$ where $D_{ij}$ is the fractional contribution of token $i$ to slot $j$. Each column of $D$ sums to 1 (convex combination over tokens).

4. Input slots $\tilde{X} \in \mathbb{R}^{(n\cdot p) \times d}$ — computed as $\tilde{X} = D^\top X$, meaning each slot is a different weighted average of all input tokens. The number of slots $(n\cdot p)$ can be much smaller than $m$, controlling the computational cost.

5. Expert functions $\{f_1, \ldots, f_n\}$ — typically standard MLPs with the same architecture as the dense block they replace. Expert $\lfloor i/p \rfloor$ processes the $i$-th slot; with one slot per expert ($p=1$), each expert processes exactly one slot.

6. Output slots $\tilde{Y}$ — the result of applying each expert to its assigned slots. These are still in "slot space" (one vector per slot) and need to be mapped back to "token space."

7. Combine weights computation — a softmax over rows of the same logit matrix $X\Phi$ that produces a matrix $C \in \mathbb{R}^{m \times (n\cdot p)}$ where $C_{ij}$ is the fractional contribution of slot $j$ to output token $i$. Each row of $C$ sums to 1.

8. Output tokens $Y \in \mathbb{R}^{m \times d}$ — computed as $Y = C\tilde{Y}$, meaning each output token is a different weighted average of all output slots. The output sequence has the same shape as the input.

9. Residual connection — the output $Y$ is added back to the original input tokens (standard Transformer residual) and passed to the next layer.

Information flow: Input tokens X → compute routing logits XΦ → normalize columns to get dispatch weights D → compute input slots as D^⊤X → apply experts to produce output slots Ỹ → normalize rows of same logits to get combine weights C → compute output tokens as CỸ → add residual.

The same logit matrix $X\Phi$ is used for both dispatch and combine, just softmax-normalized along different axes. This ensures the routing is consistent in both directions while being parameter-efficient — only one learned matrix controls the entire soft assignment.

3.3 Roadmap for the Deep Dive

• First, the dispatch mechanism (Equation 1): how we compute soft weighted averages of input tokens to create the slots that experts process. This is the core innovation that replaces discrete routing.
• Second, the combine mechanism (Equation 2): how we reconstruct output tokens from expert outputs using another soft weighted average. This completes the round-trip through "slot space."
• Third, the time complexity analysis: why the cost depends on slots, not experts, and what this implies for scaling.
• Fourth, the L2 normalization modification: a critical stability fix that prevents softmax collapse when model dimension is large.
• Fifth, the full algorithm in code: a walkthrough of Algorithm 1 to cement understanding, including the JAX einsum operations.
• Sixth, model placement and configuration: how Soft MoE layers are distributed in the Transformer, how many experts and slots to use, and why these choices matter.
• Seventh, connections to standard MoEs: what properties Soft MoE shares with Sparse MoEs and what distinguishes it from dense models.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a novel architectural component paper whose core idea is that soft, fully-differentiable token mixing can replace hard discrete expert assignment in MoE layers, preserving the key benefit of MoEs (large parameter count with controlled compute cost) while eliminating the main sources of fragility (non-differentiability, token dropping, load imbalance, scaling bottlenecks).

The Dispatch Mechanism: Computing Expert Input Slots

The dispatch step transforms $m$ input tokens into $(n\cdot p)$ slots that will be processed by the experts. Each slot is a different convex combination of all input tokens, with the combination weights learned through parameters $\Phi$ and conditioned on the token representations $X$.

The dispatch weights $D$ are defined in Equation 1:

$D_{ij} = \frac{\exp((X\Phi)_{ij})}{\sum_{i'=1}^{m} \exp((X\Phi)_{i'j})}$

where $X \in \mathbb{R}^{m \times d}$ is the matrix of $m$ input tokens, each of dimension $d$; $\Phi \in \mathbb{R}^{d \times (n\cdot p)}$ is the learned slot parameter matrix with one $d$-dimensional column per slot; $n$ is the number of experts; and $p$ is the number of slots per expert, so there are $(n\cdot p)$ slots total.

What it computes: first, the raw routing logits are computed as the matrix product $X\Phi$, giving an $m \times (n\cdot p)$ matrix where entry $(i, j)$ is a scalar score representing how relevant token $i$ is for slot $j$. Then a softmax is applied over the token axis (rows of $X\Phi$, indexed by $i'$ in the denominator) for each slot independently. This means each column $j$ of $D$ — representing the dispatch weights of all tokens to slot $j$ — sums to 1, forming a convex combination. Slot $j$ is therefore a weighted average of all $m$ tokens, with the weights given by column $j$ of $D$.

The input slots $\tilde{X}$ are then computed as:

$\tilde{X} = D^\top X$

where $\tilde{X} \in \mathbb{R}^{(n\cdot p) \times d}$. Each row of $\tilde{X}$ is one slot vector of dimension $d$, computed as $\tilde{X}_j = \sum_{i=1}^{m} D_{ij} X_i$ — the weighted average of all input token representations, with weights specific to slot $j$.

Why this form: the softmax over columns normalizes dispatch weights per-slot, meaning every slot receives a total weight of 1 distributed across tokens. This guarantees that no token is ever dropped — every token contributes to every slot, though the contribution may be very small for some token-slot pairs if the corresponding logit is low. The convex combination structure ensures that slot representations stay in the same seminumerical range as the input tokens (rather than growing with sequence length), which is important for training stability. The same $X\Phi$ logits will be reused for the combine weights (but with a different normalization axis), creating a parameter-efficient and consistent bidirectional routing.

Crucially, this design means the model learns which tokens to pay attention to for each expert implicitly, through gradient descent on the task loss — no separate assignment optimization, no auxiliary losses, no reinforcement learning. The softmax temperature is fixed at 1 (no learned temperature parameter), which provides a natural balance between hard and soft assignment: the model can still produce near-one-hot dispatch weights when beneficial by making the logit for one token much larger than others.

The Combine Mechanism: Reconstructing Output Tokens

After experts process the slots to produce output slots $\tilde{Y} \in \mathbb{R}^{(n\cdot p) \times d}$, we need to map these back to the original $m$ token positions. The combine weights $C$ perform this mapping, again as learned convex combinations:

$C_{ij} = \frac{\exp((X\Phi)_{ij})}{\sum_{j'=1}^{n\cdot p} \exp((X\Phi)_{ij'})}$

What it computes: using the exact same logit matrix $X\Phi$ as the dispatch step, a softmax is now applied over the slot axis (columns of $X\Phi$, indexed by $j'$ in the denominator) for each token independently. This means each row $i$ of $C$ — representing the combine weights of all slots to output token $i$ — sums to 1. Output token $i$ is therefore a weighted average of all $(n\cdot p)$ output slots, with weights determined by how relevant each slot's expert output is for reconstructing token $i$.

The final output tokens are:

$Y = C\tilde{Y}$

where $Y \in \mathbb{R}^{m \times d}$ preserves the original sequence shape. Each output token $Y_i = \sum_{j=1}^{n\cdot p} C_{ij} \tilde{Y}_j$ is a different convex combination of all expert-processed slots.

Why this form: the softmax over rows normalizes combine weights per-token, meaning every output token receives a total weight of 1 distributed across slots. This ensures the output sequence has exactly $m$ tokens (matching the input for the residual connection) and that each token's representation is a proper convex combination, staying in a well-behaved range. Using the same logit matrix $X\Phi$ for both dispatch and combine is a deliberate design choice that forces the routing to be consistent: the same token-slot affinities govern how tokens flow into slots and how expert outputs flow back to tokens. This is more parameter-efficient than learning separate matrices and provides a natural symmetry — if token $i$ contributes strongly to slot $j$ in dispatch, it will also tend to draw strongly from slot $j$'s expert output in combine, though the different normalization axes (slot-wise vs. token-wise) prevent the two weight matrices from being simple transposes.

The combine step is what distinguishes Soft MoE from token-pooling or sequence-reduction methods. Those methods keep the sequence compressed after the weighted averaging, whereas Soft MoE expands back to the original sequence length, making it a drop-in replacement for any Transformer MLP block without needing to modify subsequent attention layers.

Time Complexity and the Slot-Expert Decoupling

The computational cost of a Soft MoE layer breaks down into two components:

$\text{Cost} = O(m \cdot (n\cdot p) \cdot d) + O((n\cdot p) \cdot k)$

where $m$ is the number of tokens, $(n\cdot p)$ is the total number of slots, $d$ is the token/slot dimension, and $O(k)$ is the cost of applying one expert function to one slot.

The first term $O(m \cdot (n\cdot p) \cdot d)$ is the cost of computing $X\Phi$ (the routing logits) and the two matrix multiplications $D^\top X$ and $C\tilde{Y}$. This is a straightforward dense matrix multiplication — exactly the kind of operation that GPUs and TPUs are optimized for.

The second term $O((n\cdot p) \cdot k)$ is the cost of applying the expert MLPs. Since each slot is processed by exactly one expert (the $\lfloor i/p \rfloor$-th expert for slot $i$), this cost is proportional to the number of slots, not the number of experts.

The critical scaling insight: by choosing $p = O(m/n)$ — that is, making the number of slots per expert inversely proportional to the number of experts such that the total slot count $(n\cdot p)$ stays constant — the time complexity reduces to:

$O(m^2 d + mk)$

The number of experts $n$ disappears from the complexity entirely. This means you can increase the number of experts (and thus the total parameter count) arbitrarily without changing the FLOPs, as long as you decrease the slots per expert correspondingly. In the extreme case of $p=1$ (one slot per expert), the total number of slots equals the number of experts $n$, and the cost depends on $n$ rather than on the number of tokens $m$. When $n$ is chosen to be similar to $m$ (e.g., 256 experts and 256 tokens), the cost is dominated by $O(mk)$, which is the same as applying a single expert to all tokens.

Furthermore, the $m^2 d$ term is the same asymptotic complexity as single-headed self-attention, which Transformers already compute in every block. Therefore, the routing overhead of Soft MoE does not create a new computational bottleneck beyond what the architecture already handles.

What this enables in practice: the paper demonstrates scaling from 8 to 4096 experts with essentially flat throughput (Figure 6, bottom row), while Sparse MoEs see throughput drop by more than half over the same range due to routing overhead. With Soft MoE, the number of experts becomes a free parameter for increasing model capacity — limited only by memory, not by compute.

L2 Normalization: Preventing Softmax Collapse at Scale

A critical implementation detail that the paper identifies and solves is the tendency of the softmax routing to collapse to near-one-hot vectors as model dimension $d$ increases. This is not a minor numerical issue — it threatens to turn Soft MoE back into a hard-assignment (and thus non-differentiable) system at large scale, defeating its entire purpose.

The theoretical problem (Appendix E.1): when softmax is applied to layer-normalized inputs and the dimension $d$ grows, the inputs to the softmax scale as $O(\sqrt{d})$ because layer normalization produces roughly unit-norm vectors. The softmax operation $\text{softmax}(z)_i \propto \exp(z_i)$ amplifies differences exponentially. With $d$ large, even small relative differences in the $\sqrt{d}$-scaled logits become enormous after exponentiation, causing the softmax to saturate toward a one-hot vector.

The paper formalizes this in Equation 9 of Appendix E.1: as $d \to \infty$, the softmax output tends to a vector that is nonzero only for indices achieving the maximum logit value, and uniform across those maximizers. When the maximum is unique (which is typical with random initialization and different expert parameters), the softmax collapses to a one-hot vector.

The empirical consequence (Appendix E.2): Figures 13 and 14 demonstrate that without normalization, the average maximum dispatch weight per slot grows toward 1.0 as $d$ increases from 384 to 1664, and the ImageNet 10-shot accuracy degrades significantly. The collapse also makes the model sensitive to learning rate choice — higher learning rates accelerate the collapse (Figure 14).

The solution (Section 2.3, Algorithm 2): instead of directly computing $X\Phi$, normalize both inputs:

$\text{logits} = \text{scale} \cdot (X_{\text{normalized}})(\Phi_{\text{normalized}})^\top$

where:

• X_normalized = l2_normalize(X, axis=1) — each token representation is normalized to unit L2 norm along its feature dimension.
• Phi_normalized = l2_normalize(Phi, axis=0) — each slot parameter vector is normalized to unit L2 norm along its feature dimension.
• scale is a trainable scalar parameter that controls the overall magnitude of logits.

The L2 normalization ensures that the inner products \langle \text{X_normalized}_i, \text{Phi_normalized}_j \rangle are bounded in $[-1, 1]$ regardless of dimension $d$, since both vectors have unit norm. The softmax inputs no longer scale with $\sqrt{d}$, and the softmax retains a meaningful, non-saturated distribution. The trainable scale parameter allows the model to learn the appropriate sharpness of the routing distribution — a small scale produces nearly uniform weights (all tokens contribute equally), while a large scale produces sharper, more selective weights.

Why this matters beyond stability: the normalization also improves training at smaller $d$ (as Figures 13 and 14 show, the benefits appear even at $d=384$), and it makes the model less sensitive to the peak learning rate. This is because the normalization decouples the softmax temperature from the model dimension and the parameter initialization scale, making the routing behavior more predictable and controllable.

The Full Algorithm in Code

Algorithm 1 in the paper provides a complete JAX implementation. Walking through it line by line clarifies the exact operations and their shapes:

def soft_moe_layer(X, Phi, experts):
    # Compute the dispatch and combine weights.
    logits = jnp.einsum('md,dnp->mnp', X, Phi)
    D = jax.nn.softmax(logits, axis=(0,))
    C = jax.nn.softmax(logits, axis=(1, 2))
    # The input slots are a weighted average of all the input tokens,
    # given by the dispatch weights.
    Xs = jnp.einsum('md,mnp->npd', X, D)
    # Apply the corresponding expert function to each input slot.
    Ys = jnp.stack([
        f_i(Xs[i, :, :]) for i, f_i in enumerate(experts)],
        axis=0)
    # The output tokens are a weighted average of all the output slots,
    # given by the combine weights.
    Y = jnp.einsum('npd,mnp->md', Ys, C)
    return Y

Line 3 — logits computation: the einsum 'md,dnp->mnp' computes the routing logits for every pair of token and slot. Here m is the number of tokens, d is the feature dimension, n is the number of experts, and p is slots per expert. The resulting tensor has shape (m, n, p) — for each of the $m$ tokens, there is a score for each of the $n \times p$ slots. In the actual implementation, the paper flattens the (n, p) dimensions into a single slot dimension (n*p,) for simplicity, but keeping the expert-slot structure explicit is useful for applying experts by index.

Line 4 — dispatch weights: softmax(logits, axis=(0,)) normalizes over the token axis (axis 0) independently for each slot (axes 1 and 2). This gives $D \in \mathbb{R}^{m \times n \times p}$ where $\sum_{i=1}^{m} D_{ijk} = 1$ for every expert $j$ and slot-within-expert $k$.

Line 5 — combine weights: softmax(logits, axis=(1, 2)) normalizes over all slot axes (expert and slot-within-expert) independently for each token. This gives $C \in \mathbb{R}^{m \times n \times p}$ where $\sum_{j=1}^{n} \sum_{k=1}^{p} C_{ijk} = 1$ for every token $i$.

Line 8 — input slots: einsum('md,mnp->npd', X, D) computes $\tilde{X} = D^\top X$. For each slot (indexed by expert $j$ and slot-within-expert $k$), this computes $\tilde{X}_{jk} = \sum_{i=1}^{m} D_{ijk} X_i$ — a weighted average of all tokens. The output shape is (n, p, d).

Line 10-12 — expert application: the jnp.stack with list comprehension applies each expert $f_i$ to the $p$ slots assigned to it. Expert $i$ receives Xs[i, :, :] which is a $p \times d$ matrix (all $p$ of its slots), processes them (typically with independent MLP forward passes for each slot, which can be batched), and produces Ys[i, :, :] of the same shape $p \times d$. The resulting Ys has shape (n, p, d).

Line 15 — output tokens: einsum('npd,mnp->md', Ys, C) computes $Y = C\tilde{Y}$. For each token $i$, this computes $Y_i = \sum_{j=1}^{n} \sum_{k=1}^{p} C_{ijk} \tilde{Y}_{jk}$. The output shape is (m, d), exactly matching the input token shape.

A subtle implementation note: the paper mentions using l2_normalize(X, axis=1) and scale * l2_normalize(Phi, axis=0) as inputs to line 3 in the actual implementation (Algorithm 2 defines the normalization). This is not shown in the simplified Algorithm 1 but is critical for the stability reasons discussed above.

Model Placement and Configuration Strategy

The paper systematically studies where and how to place Soft MoE layers in the Transformer architecture.

Placement: following standard MoE practice, Soft MoE blocks replace the MLP blocks in a subset of Transformer layers, not every layer. The paper's default configuration replaces the second half of the blocks. For example, in a 12-layer ViT, layers 6-11 would use Soft MoE while layers 0-5 remain dense. The rationale is that later layers benefit more from increased capacity because they process more abstract features, while early layers work with low-level features where a few dense MLPs suffice.

Appendix D provides an expert placement ablation using a Soft MoE S/16 with 12 layers and 512 total experts, distributed in various ways (all in one layer, split across 2, 4, 6, or 8 layers). The best configurations use 4-8 MoE layers (e.g., layers 8-11 with 128 experts each achieves 72.8% ImageNet 10-shot vs. 70.0% with all 512 experts in layer 11). Spreading experts across more layers consistently outperforms concentrating them, suggesting that having multiple routing decisions throughout the network is more beneficial than a single, large routing step.

Number of experts and slots: the per-layer number of experts and slots per expert are the central hyperparameters. The paper's experiments in Section 3.3 sweep from 16 to 4096 experts and 1 to 16 slots per expert. The key findings from Section 3.5 and Figures 6 and 12:

• One slot per expert is optimal. Across all backbone sizes, the model with the most experts (and thus one slot per expert) achieves the best performance for a given total slot count. Increasing slots per expert provides only marginal gains at significant computational cost. For instance, with 32 experts, going from 1 to 32 slots per expert improves ImageNet 10-shot from ~72% to ~74% but increases step time by over 2× (Figure 12, right).

• The paper hypothesizes (Section 6, "Lazy experts" and Appendix I) that multiple slots per expert tend to align — the learned slot parameters for the same expert become highly correlated, so the slots compute similar weighted averages and the expert processes redundant information. Figures 28-30 in Appendix I show this empirically: with one slot per expert (Figure 28), the inner products between slot parameter vectors show no clear structure beyond the diagonal. With 4 or 16 slots per expert (Figures 29-30), consecutive slots from the same expert show very high inner products (bright off-diagonal blocks), confirming alignment. This suggests that experts lack the flexibility to handle very different input distributions across their slots, and the additional slots add computational cost without commensurate representational benefit.

• For Sparse MoEs, there is an optimal number of experts. Beyond a point (typically 128-256 for Tokens Choice, higher for Experts Choice), training difficulties from increased token dropping and routing overhead outweigh the benefits of additional capacity. Soft MoE shows no such ceiling — performance improves monotonically with expert count up to at least 4096 experts (Figure 23, Appendix F.3).

Total slot count relative to sequence length: when replacing the second half of blocks, the typical setup matches the total number of slots to the input sequence length to keep FLOPs comparable to a dense model. For a ViT with 196 tokens (14×14 patches at 224×224), using 128 experts with 1 slot each totals 128 slots — slightly fewer than the number of tokens, making the Soft MoE layer compute cheaper than the equivalent dense MLP block. This is the source of Soft MoE's inference speed advantage: it processes fewer slot-MLP calls than a dense model processes token-MLP calls, while having many more expert parameters to draw from.

What Soft MoE Is and Is Not: Connections to Sparse and Dense MoEs

The paper is explicit about where Soft MoE sits in the spectrum between dense and sparse architectures.

Soft MoE is not technically sparse. In a conventional Sparse MoE, each expert's parameters are applied to only a subset of input tokens — this is the source of the sparsity and the computational savings. In Soft MoE, every slot is a weighted average of all input tokens with generally nonzero weights. Every input token therefore fractionally activates all model parameters. From a strict sparsity perspective, there is no sparse computation — the full dispatch and combine matrices are dense.

Soft MoE is not a Dense MoE either. In a hypothetical Dense MoE, every expert would process every input token, resulting in FLOPs that scale as $O(n \cdot m \cdot k)$ — proportional to both the number of experts and the number of tokens. Soft MoE's cost scales as $O((n\cdot p) \cdot (k + md))$, where $(n\cdot p)$ (the number of slots) can be much smaller than $m \cdot n$ when $p$ is small and $n$ is moderate. The slots act as a computational bottleneck that limits the total expert computation.

The paper describes Soft MoE as having "features of both sparse and dense" (Section 2.2): it achieves the parameter-count scaling benefits of Sparse MoEs (many expert parameters with controlled per-input FLOPs) through a mechanism that is continuous and dense (weighted averages), avoiding the discrete optimization problems entirely. The sparsity-like property comes from the fact that each expert processes only $p$ slots rather than $m$ tokens — if $p \ll m$, the expert computation is substantially less than a dense model would require, even though the routing is continuous.

This hybrid nature becomes particularly clear in the per-sequence determinism property. In Sparse MoEs operating on batches of multiple sequences, tokens from different sequences compete for expert capacity, making the routing for a given input depend on other inputs in the batch. Soft MoE avoids this because each sequence's dispatch and combine weights are computed solely from that sequence's tokens — the batch dimension is never mixed in the routing. Each sequence is processed independently and deterministically, regardless of batch composition or size. This is a significant practical advantage for deployment, where batch sizes may vary or inputs arrive one at a time.

Parameter Count and Scaling

The relationship between Soft MoE parameters and the standard ViT architecture provides a clear picture of how capacity scaling works. In a dense ViT block, there is one MLP (with typically two linear layers, e.g., expanding from $d$ to $4d$ and back to $d$). The total MLP parameters per block are approximately $8d^2$ (counting both weight matrices).

In a Soft MoE block with $n$ experts, each expert is a full MLP with the same architecture. The total expert parameters are therefore approximately $n \times 8d^2$ — $n$ times the dense block's MLP parameters. The additional routing parameters $\Phi$ are $d \times (n\cdot p)$, which is usually negligible compared to the expert parameters when $d$ is large.

For example, a ViT H/14 has $d=1280$ and roughly 669M total parameters. A Soft MoE H/14 with 128 experts and 1 slot per expert in 16 MoE layers has $128 \times 8 \times 1280^2 \times 16$ additional parameters from the expert MLPs, yielding approximately 54B parameters — over 80× the dense model's parameter count. Yet the inference FLOPs increase by only about 2% (Table 1) because the number of slots (128 per layer × 16 layers = 2048 total slots processed per forward pass) is similar to the number of tokens that the dense MLP blocks would process.

This decoupling of parameter count from computational cost is the fundamental value proposition of MoEs, and Soft MoE achieves it without any of the discrete routing problems that limit Sparse MoE scaling. The key enabler is the $O(m^2 d + mk)$ time complexity that is independent of the number of experts $n$, allowing $n$ to grow arbitrarily within memory constraints.

Training Details and Hyperparameter Choices

While the paper does not exhaustively list every training hyperparameter for every experiment in the main text, several configurations are specified:

For the Pareto frontier experiments (Section 3.3): 300k training steps, batch size 4096, resolution 224×224, reciprocal square root learning rate schedule. Models from S/32 to H/14 are trained, with Soft MoE variants sweeping expert counts and slot configurations.

For the long training durations (Section 3.4): up to 4M steps (H/14 at 2M for cost reasons), following a similar setup to Zhai et al. (2022a). The last half of blocks are replaced with Soft MoE layers with 128 experts and one slot per expert. The paper notes that longer cooldowns (linear learning rate decay) work better for Soft MoE, increasing from 50k to 500k steps for the longest runs.

For the normalization fix: the scale parameter multiplying the L2-normalized routing logits is trainable, initialized to a value that produces reasonable softmax temperatures at the start of training. The L2 normalization itself uses an epsilon of $10^{-6}$ for numerical stability (Algorithm 2).

Distributed model training: for models with very large expert counts, standard expert-sharding techniques from Lepikhin et al. (2020) and Riquelme et al. (2021) are used to distribute expert parameters across devices. The paper accounts for this distributed overhead by reporting both FLOPs and wall-clock time (TPUv3-chip-hours) in all comparisons, ensuring that communication costs are reflected in the efficiency analysis.

4. Key Insights and Innovations

Innovation 1: MoE Routing as Continuous Token Mixing Rather Than Discrete Assignment

The dominant conceptual framing in the MoE literature treats routing as a matching problem: tokens on one side, experts on the other, and the router's job is to find good pairings subject to capacity constraints. This framing traces back to the original Sparsely-Gated MoE (Shazeer et al., 2017) and has shaped virtually all subsequent work — whether the matching is done via top-k selection, optimal transport, reinforcement learning, or linear programming, the core mental model is "which token goes to which expert?"

Soft MoE makes a fundamental conceptual break with this framing. It asks: what if we don't match tokens to experts at all, but instead let every expert see a different learned summary of the entire input? The routing mechanism is not assignment but soft content-based pooling — each expert receives a weighted average of all tokens, with the weights learned end-to-end through the same task loss that trains the rest of the network.

This shift in framing carries several non-obvious intellectual consequences:

It reframes routing as representation learning. In the discrete assignment view, the router's job is a combinatorial optimization problem — allocating scarce expert capacity efficiently. The difficulty of this problem (non-differentiability, load imbalance, token dropping) is what generates the auxiliary losses, capacity constraints, and optimization tricks that make Sparse MoEs fragile. In Soft MoE's pooling view, the router's job is purely representational: learn attention-like weights that aggregate information usefully for each expert. There is no capacity allocation problem because every expert receives the same computational budget (one slot) by construction — the model learns what information to give each expert, not which tokens to route.

It eliminates the assumption that experts should specialize by token. Sparse MoEs implicitly assume that different experts should handle different tokens — that specialization at the token level is the right granularity for conditional computation. Soft MoE challenges this assumption by showing that experts can specialize by learned feature combination rather than by token identity. An expert that receives a weighted average emphasizing edges, or textures, or a particular spatial region, can develop expertise in processing those features regardless of which specific tokens carry them. This is a more flexible form of specialization because the same token can contribute useful information to multiple experts simultaneously (high-level semantic features to one, low-level texture features to another).

It redefines sparsity. The paper is careful to note that Soft MoE is "not technically sparse" because every token fractionally activates all parameters. Yet it achieves the same practical benefit as Sparse MoEs — large parameter count with controlled compute — through a different mechanism: the slot bottleneck. Rather than sparsifying the token-expert interaction matrix, Soft MoE compresses the sequence into fewer slot representations before applying expensive expert computation. This is a genuinely different way to achieve conditional computation efficiency, and it suggests a broader design principle: computational sparsity can be achieved through dimensionality reduction in a learned space rather than through hard gating.

Evidence for the importance of this framing shift comes from the ablation in Section 3.6 and Table 2. The "Uniform" baseline — where slots are simple uniform averages of all tokens with no learned routing — already outperforms a dense ViT (51.8% vs. 48.3% JFT precision-at-1), demonstrating that having multiple expert views of the same pooled information is beneficial even without learned routing. Learned routing (Soft MoE) pushes this further to 54.3%, showing that the content-dependent pooling adds significant value on top of the multi-expert architecture. This decomposition — multi-expert architecture providing capacity, learned pooling providing routing — is conceptually cleaner than the entangled capacity-allocation-plus-specialization framing of Sparse MoEs.

Compared to the prior work on token merging (Jaegle et al., 2021; Ryoo et al., 2021; Renggli et al., 2022), the conceptual distinction is also sharp. Those methods merge tokens to reduce sequence length for efficiency, then continue processing the compressed sequence. Soft MoE uses the slot bottleneck only within a single layer, expanding back to the original sequence length via combine weights before the next layer. The slot bottleneck is therefore a per-layer routing mechanism, not a permanent sequence compression. This allows the model to make different routing decisions at different depths, adapting which information each expert receives based on the layer's position in the network.

Innovation 2: Hard Assignment Is Not Necessary for the Core MoE Benefit

A tacit assumption running through the MoE literature is that conditional computation requires some form of hard gating — that the efficiency gains of MoEs come specifically from the fact that most parameters are zeroed out for most inputs. The Softmax gating in Shazeer et al. (2017), the top-k selection in Fedus et al. (2022), and the expert-capacity mechanism in Zhou et al. (2022) all share this assumption: sparsity of the token-expert interaction matrix is what saves computation.

Soft MoE presents a compelling counterexample. It achieves comparable or better efficiency than Sparse MoEs (Figure 5 shows Soft MoE B/16 matching ViT H/14 at 5.7× faster inference; Figure 6 shows Soft MoE throughput scales much better with expert count than Sparse MoEs) without any zeroing out of computation. Every token contributes to every slot, and every slot contributes to every output token. The efficiency comes entirely from the reduced number of expert forward passes (fewer slots than tokens) rather than from sparse connections.

This finding carries significant theoretical weight because it decouples two properties that the field had implicitly bundled together:

• Parameter scaling with sublinear compute: the ability to add parameters without proportionally adding FLOPs.
• Discrete routing: the use of hard, sparse connections to achieve that decoupling.

By demonstrating that the first property does not require the second, Soft MoE opens a design space that was previously overlooked. The slot bottleneck mechanism suggests that any form of learned dimensionality reduction that compresses the sequence before expensive computation can achieve the MoE benefit, not just hard gating. This is a more general principle, and it explains why "Uniform" routing (no learned routing at all, just averaging) already outperforms dense models: the multi-expert architecture itself provides benefit through increased parameter count, independent of routing quality.

This decoupling also explains why Soft MoE can scale to 4096 experts without the instability that plagues Sparse MoEs at scale (Figure 23). In Sparse MoEs, increasing expert count makes the discrete assignment problem harder — more experts competing for the same tokens, more opportunities for tokens to be dropped or experts to be starved. In Soft MoE, increasing expert count (while decreasing slots per expert to keep total slots constant) changes nothing about the optimization difficulty — the dispatch and combine weight computation is identical regardless of expert count. The number of experts becomes a free parameter for scaling capacity, decoupled from routing quality.

This insight reframes the research question for future MoE methods. Rather than asking "how can we solve the hard assignment problem better?", the question becomes "how can we design a routing mechanism that avoids hard assignment entirely while preserving efficiency?" The slot bottleneck is one answer, but the principle is broader.

Innovation 3: Slot Alignment as the Explanation for Diminishing Returns from Multiple Slots per Expert

A persistent ambiguity in MoE design is how to allocate total FLOPs between having more experts versus giving each expert more capacity (more tokens or slots). In Sparse MoEs, increasing capacity per expert reduces token dropping and can improve performance. In Soft MoE, the analogous choice is increasing slots per expert versus increasing the number of experts (with one slot each) for a fixed total slot budget.

The paper discovers a striking regularity: one slot per expert is consistently optimal, and additional slots per expert provide marginal gains at best despite increasing computational cost. Figure 6 (top) shows this clearly — in the Soft MoE heatmap, the best model at each row (slots-per-token) is always the one with the most experts (fewest slots per expert, ideally one). Figure 12 shows that going from 1 to 32 slots per expert with 32 total experts improves ImageNet 10-shot only modestly (~72% to ~74%) while more than doubling step time.

But the paper goes beyond this empirical observation to provide a mechanistic explanation, and this explanation is one of its most intellectually distinctive contributions. Appendix I (Figures 28-30) visualizes the inner products between slot parameter vectors for models with 1, 4, and 16 slots per expert.

With one slot per expert (Figure 28), there is no visible structure beyond the diagonal — each expert's single slot parameter is roughly orthogonal to others, suggesting each expert learns to look at different aspects of the input.

With 4 slots per expert (Figure 29), a block-diagonal pattern emerges: the 4 parameter vectors belonging to the same expert show high pairwise inner products (bright 4×4 blocks on the diagonal), indicating they are nearly aligned — they compute similar weighted averages. Slots from different experts remain decorrelated.

With 16 slots per expert (Figure 30), the block-diagonal pattern is even more pronounced, with the 16 slots per expert forming strongly correlated clusters.

This is not a trivial finding. It suggests a fundamental limitation in expert flexibility: a single MLP expert, when given multiple slots, tends to pull those slots toward similar input distributions. The slots "collapse" to redundancy rather than learning complementary views. The paper hypothesizes that "an expert may lack the flexibility to accommodate very different slot projections" (Section 2.3) — an MLP trained with one set of parameters cannot effectively process slot representations that differ substantially in their feature statistics or the token combinations they emphasize.

This insight has both theoretical and practical significance. Theoretically, it suggests that expert specialization operates at a coarser granularity than slot-level computation — an expert's parameters define a processing function, and that function imposes an inductive bias that constrains what kinds of inputs the expert can usefully handle. Multiple slots per expert force the same processing function to handle potentially different input types, which is less effective than having separate experts with separate parameters specialized to each input type.

Practically, this explains the paper's consistent design choice of one slot per expert and provides guidance for future MoE designs: invest in more experts rather than more capacity per expert. The memory cost of many experts is real (Section 6 acknowledges this), but the computational cost is the same, and the representational benefit of expert diversity outweighs the benefit of within-expert slot diversity.

This diagnostic contribution — identifying why more slots per expert fails — distinguishes the paper from work that merely sweeps hyperparameters and reports which settings work best. It provides a mechanistic hypothesis that can guide algorithmic improvements (could we regularize slots to be decorrelated? could we design experts that handle more diverse inputs?) and informs the broader understanding of how conditional computation interacts with neural network capacity.

Innovation 4: The Empirical Resolution of MoE Scaling Tensions

The paper resolves several tensions in the MoE scaling literature through careful, large-scale empirical comparison rather than algorithmic novelty. This is a contribution of consolidation and clarity rather than invention, but it is significant because prior work on these questions provided conflicting signals.

Inference cost vs. parameter count. Prior Sparse MoE work (Fedus et al., 2022; Riquelme et al., 2021) demonstrated that MoEs could achieve better quality per FLOP than dense models, but the scaling relationship was messy — increasing expert count often degraded throughput due to routing overhead, and token dropping created a soft ceiling on effective capacity. Soft MoE cleanly demonstrates that the relationship between parameter count and throughput can be made nearly flat (Figure 6, bottom): going from 8 to 4096 experts barely changes Soft MoE's training throughput, while Tokens Choice and Experts Choice throughput drops by >50%. This establishes a new empirical upper bound on MoE scaling efficiency — the routing overhead can be made negligible, and the only remaining constraint is memory, not compute. This is an existence proof that changes the conversation from "how much routing overhead can we tolerate?" to "how can we eliminate routing overhead entirely?"

Number of experts vs. performance. The observation that Sparse MoEs have an optimal number of experts beyond which performance degrades (Section 3.5) has been noted in practice but not systematically characterized or explained. Soft MoE's monotonic improvement with expert count up to 4096 (Figure 23) demonstrates that this ceiling is an artifact of routing failure, not an inherent property of expert architectures. This is a negative result about Sparse MoEs that is informative: the routing mechanism, not the expert design, is the bottleneck on MoE scaling. It redirects attention from "how many experts should we use?" (a hyperparameter question) to "how can we design routers that don't degrade with expert count?" (an algorithmic question).

The FLOPs-matched comparison. The paper's rigorous FLOPs- and wall-clock-time-matched comparisons (Figures 3, 4, 5) establish that Soft MoE is on the Pareto frontier across essentially all training budgets and backbone sizes. The dominance is not marginal — Figure 5 shows Soft MoE B/16 achieving the same ImageNet 10-shot accuracy as ViT H/14 while requiring 10.4× fewer inference FLOPs and 5.7× less inference time. This is an unusually clean result in the MoE literature, where comparisons are often complicated by differing architectures, training recipes, or hardware environments. By standardizing on ViT backbones and controlling for both FLOPs and wall-clock time, the paper provides one of the most convincing demonstrations of MoE efficiency in the vision domain.

The placement effect. Appendix D provides a systematic ablation of expert placement (total experts fixed, distributed across different numbers of layers) that had not been thoroughly characterized for MoE Transformers. The finding that 4-8 MoE layers outperform both fewer layers (1-2) and more layers (all 12) at matched total expert count suggests that there is an optimal granularity of conditional computation — enough layers to make multiple routing decisions throughout the network, but not so many that individual layers have too few experts to specialize. This provides practical guidance and also raises interesting theoretical questions about the interaction between routing depth and expert specialization that the paper does not fully explore but identifies for future work.

The contrastive learning transfer (Section 4). The demonstration that Soft MoE's benefits persist when the vision tower is frozen and paired with a trained text encoder (Table 3) addresses a concern specific to MoE architectures: that the routing might overfit to the pretraining task distribution, producing representations that don't transfer well. The consistent improvements on zero-shot ImageNet (+7% for S/16, +4% for B/16, +1.1% for L/16 over their ViT counterparts) suggest that the soft routing learns generally useful visual features, not just ones tuned to JFT classification. This is important for the practical deployment case where a pretrained vision encoder is used as a frozen feature extractor for downstream tasks.

Together, these empirical findings constitute a consolidation of MoE design principles — not inventing new concepts, but rigorously establishing which existing ideas work, which don't, and why — that the field had been lacking. The paper's comprehensive Pareto frontier analysis (over 100 models trained, Appendix J) sets a new standard for thorough MoE evaluation.

5. Experimental Analysis

Evaluation Methodology

• Dataset. All primary image classification experiments use JFT-4B (Zhai et al., 2022a), a proprietary dataset containing more than 4 billion images across 29k classes. The paper does not specify exact train/validation splits, but evaluation during pretraining uses upstream validation precision-at-1 on JFT-4B. For few-shot evaluation, models are frozen and evaluated on ImageNet 10-shot accuracy (10 labeled examples per class from ImageNet-1k, with a new classification head trained on those examples). For finetuning evaluation, models are finetuned on the full ImageNet-1k training set (1.3M images) at 384×384 resolution and evaluated on the ImageNet-1k validation set. The contrastive learning experiments in Section 4 use WebLI (Chen et al., 2022), a proprietary dataset of 10B image–alt-text pairs, and evaluate zero-shot on ImageNet, CIFAR-100, Oxford-IIIT Pet, and COCO retrieval.

• Base model(s). The vision backbone is the standard Vision Transformer (ViT) at scales ranging from Small (S/16: 33M parameters) through Base (B/16: 108M), Large (L/16: 333M), and Huge (H/14: 669M). These are chosen as representative dense baselines that span the practical range of compute budgets. Soft MoE models use the same ViT backbones but replace the MLP blocks in the second half of layers with Soft MoE layers. Soft MoE variants scale from 933M parameters (S/16 with 128 experts) to 54.1B parameters (H/14 with 256 experts). The paper states that PaLM-style autoregressive models are not used; the primary domain is vision encoders where full bidirectional context is available.

• Metrics. Three primary metrics are used throughout: (1) JFT-4B precision-at-1 — the fraction of validation examples where the model's top-1 prediction matches the ground-truth class, measured during pretraining; (2) ImageNet 10-shot accuracy — the accuracy on ImageNet-1k validation after training only a new classification head on 10 examples per class with frozen backbone weights; (3) ImageNet finetuning accuracy — the accuracy on ImageNet-1k validation after finetuning the full model on the 1.3M training images at 384×384 resolution. For contrastive learning, zero-shot accuracy on downstream datasets and COCO image-to-text / text-to-image retrieval scores are reported. All metrics are standard top-1 accuracy percentages. The paper reports both FLOPs (floating point operations, theoretical compute) and wall-clock time (TPUv3-chip-hours for training, TPUv3 ms/img for inference) to account for hardware efficiency differences that FLOPs alone miss.

• Baselines. The paper compares against: (1) Dense ViT — standard Vision Transformers without any MoE layers, at S/16, B/16, L/16, and H/14 scales; (2) Tokens Choice Sparse MoE (Shazeer et al., 2017; Riquelme et al., 2021) — each token selects the top-K experts with highest routing scores, using Batch Priority Routing (BPR) for K ∈ {1, 2} and capacity multiplier C ∈ {1.0, 1.125, 2.0}; (3) Experts Choice Sparse MoE (Zhou et al., 2022) — each expert selects the top-C tokens by routing score, with capacity multiplier C ∈ {0.5, 1.0, 1.125, 2.0}. For both Sparse MoE baselines, the paper sweeps expert counts (16–4096) and configurations. In total, over 100 models were trained for the Pareto frontier analysis (Appendix J, Table 9 lists all 107 model configurations).

• Generation budget / compute accounting. Compute is measured along two axes: training cost (total FLOPs for the full training run and total TPUv3-days of wall-clock time) and inference cost (FLOPs per image and milliseconds per image on TPUv3). Training cost comparisons use total exaFLOP consumed across all training steps; inference cost uses per-image metrics at batch size 1 to reflect realistic deployment. For Sparse MoEs, routing overhead (sorting, top-k, buffer management) is captured in wall-clock time but not in theoretical FLOP counts, which the paper notes as an important distinction — the FLOPs-only comparison favors Sparse MoEs by hiding their hardware-inefficient operations.

• Cross-validation / statistical protocol. The paper does not report standard cross-validation (training and evaluating multiple times with different random seeds). Instead, for the Pareto frontier analysis (Section 3.3), a large number of model configurations are trained once each, and the Pareto frontier is computed over all runs — a model lies on the frontier if no other model achieves both better performance and lower cost. This provides a robust comparison across the space of possible configurations without requiring multiple training runs per configuration. For the ablations (Section 3.6), a single training run per configuration is reported (300k steps on JFT-4B). The paper acknowledges in Section 6 that memory consumption of models with many experts is a practical limitation but does not treat this as a statistical concern.

Main Quantitative Results

Training Pareto Frontiers: Soft MoE Dominates at All Budgets

The headline result from Section 3.3 is that Soft MoE models lie on the training cost/performance Pareto frontier across all backbone sizes and training budgets, consistently outperforming both dense ViTs and Sparse MoEs (Tokens Choice and Experts Choice) at matched FLOPs or matched wall-clock time.

Figures 3a and 3b display JFT-4B precision-at-1 and ImageNet 10-shot accuracy against both total training FLOPs (bottom axis) and total training TPUv3-days (top axis). The key patterns:

• Soft MoE models (blue) consistently appear on the upper-left envelope of both plots, indicating they achieve the highest accuracy for any given training budget. No dense model (red) or Sparse MoE model (orange for Experts Choice, green for Tokens Choice) surpasses the Soft MoE frontier at any point.

• The advantage is largest in the mid-compute regime (roughly 10–100 TPUv3-days), where Soft MoE achieves approximately 4–6 percentage points higher JFT precision-at-1 and 5–10 points higher ImageNet 10-shot accuracy than comparably expensive dense models.

• The advantage persists but narrows at higher budgets, as all model classes begin to saturate on these metrics. At 100–1000 TPUv3-days, Soft MoE still leads but the gap over the best Sparse MoEs shrinks to 1–2 percentage points.

• Larger marker sizes (indicating larger backbone models) appear further right on the cost axis, and within each cost level, Soft MoE variants dominate. For example, Soft MoE S/16 configurations outperform dense ViT B/16 at similar training costs while being smaller backbone models.

Figure 19 in Appendix F.3 shows the full set of all 107 trained models (not just Pareto-optimal ones), confirming that the Pareto dominance is not an artifact of selective plotting — Soft MoE models cluster in the upper-left region, while Sparse MoEs show much wider variance with many configurations falling well below the frontier.

Table 9 in Appendix J provides the complete listing of all trained configurations with their exact metrics. Representative comparisons include:

• Ref 18: Soft MoE S/16 with 16 experts achieves 50.9% JFT P@1 and 68.1% IN/10shot at 12.4 exaFLOP / 14.2 TPUv3-days.
• Ref 20: Dense S/16 achieves 47.9% JFT P@1 and 60.8% IN/10shot at 17.0 exaFLOP / 15.3 TPUv3-days — Soft MoE uses fewer FLOPs and less time while substantially outperforming.
• Ref 52: Soft MoE B/16 with 128 experts achieves 55.3% JFT P@1 and 77.0% IN/10shot at 59.0 exaFLOP / 46.8 TPUv3-days.
• Ref 51: Dense B/16 achieves 52.0% JFT P@1 and 71.8% IN/10shot at 64.8 exaFLOP / 45.2 TPUv3-days — similar training cost, 3.3% higher JFT, 5.2% higher IN/10shot.

Figures 20-22 in Appendix F.3 break down the Pareto plots into pairwise comparisons: Soft MoE vs. Dense (Figure 20), Soft MoE vs. Experts Choice (Figure 21), and Soft MoE vs. Tokens Choice (Figure 22). In each pairwise comparison, Soft MoE models consistently dominate, though the margin over Experts Choice is smaller than over Tokens Choice or Dense.

Long Training Durations: Soft MoE Maintains Advantage with Extended Training

Section 3.4 reports results from training models for much longer durations — up to 4M steps for most models and up to 9M–10M steps for the smallest backbones — to study whether Soft MoE's advantages persist under "overtraining" regimes.

Figure 4 plots JFT P@1, ImageNet 10-shot, and ImageNet finetuning accuracy against total training FLOPs for matched backbone classes (S/16, B/16, L/16, H/14). The key finding is that both Soft MoE and ViT improve with extended training, and Soft MoE maintains its advantage at all points — the Soft MoE curves lie above the ViT curves at matched training cost for all three metrics.

Table 8 and Figure 16 in Appendix F provide the numerical results. Representative comparisons at the longest training durations:

• Soft MoE S/16 (128 experts, 10M steps, 437.7 TPUv3-days): 59.2% JFT P@1, 79.8% IN/10shot, 87.1% IN/finetune. Inference cost: 0.7 ms/img, 8.6 GFLOP/img.
• Soft MoE S/16 (10M steps with 500k cooldown): 60.9% JFT P@1, 80.7% IN/10shot, 87.7% IN/finetune.
• Soft MoE B/16 (128 experts, 4M steps, 449.5 TPUv3-days): 60.0% JFT P@1, 82.0% IN/10shot, 88.0% IN/finetune.
• ViT B/16 (4M steps, 410.1 TPUv3-days): 56.2% JFT P@1, 76.8% IN/10shot, 86.6% IN/finetune.
• ViT H/14 (2M steps, 2039.8 TPUv3-days): 59.7% JFT P@1, 83.3% IN/10shot, 88.9% IN/finetune. Inference cost: 8.6 ms/img, 334.2 GFLOP/img.
• Soft MoE H/14 (256 experts, 2M steps, 2553.7 TPUv3-days): 62.1% JFT P@1, 84.3% IN/10shot, 89.1% IN/finetune. Inference cost: 10.9 ms/img, 342.4 GFLOP/img.

Several patterns emerge from these long-training results:

• Soft MoE S/16 outperforms ViT B/16 across all metrics despite ViT B/16 being a 3.3× larger backbone and consuming roughly similar training FLOPs (529.8 vs. 864.1 exaFLOP for the shorter S/16 run). This demonstrates that Soft MoE's extra expert capacity more than compensates for a smaller backbone.

• Soft MoE B/16 approaches ViT L/16 performance (82.0% vs. 81.5% IN/10shot; 88.0% vs. 88.5% IN/finetune) while being 3× faster at inference (1.5 vs. 4.9 ms/img) and using nearly 4× fewer FLOPs (32.0 vs. 122.9 GFLOP/img).

• Soft MoE L/16 outperforms ViT H/14 (83.7% vs. 83.3% IN/10shot at 4M steps; 88.9% tied for finetuning) while being roughly 2× faster at inference (4.8 vs. 8.6 ms/img) and 3× faster at training (1355.4 vs. 2039.8 TPUv3-days for 4M and 2M steps respectively).

• The additional 500k-step cooldown for Soft MoE models (compared to 50k for ViT) was found to be beneficial. The paper notes that "longer cooldowns (linear learning rate decay) works well for Soft MoE" — the S/16 with 500k cooldown achieves 60.9% JFT P@1 vs. 59.2% with 50k cooldown at the same step count.

Inference-Time Optimized Models: Largest Gains at Small Backbone Sizes

Section 3.4 also reports models trained specifically for longer durations to optimize the accuracy-per-inference-cost tradeoff, with results in Figure 5 and Table 1.

Figure 5 is organized as a 2×3 grid: the top row shows performance vs. evaluation cost (GFLOP/img or TPUv3 ms/img); the bottom row shows the same metrics vs. training cost. The key takeaway from the evaluation-cost plots is:

• Soft MoE B/16 (9M steps, 500k cooldown) achieves 62.4% JFT P@1 and 82.9% IN/10shot at 1.5 ms/img and 32.0 GFLOP/img.
• ViT H/14 (2M steps) achieves 59.7% JFT P@1 and 83.3% IN/10shot at 8.6 ms/img and 334.2 GFLOP/img.
• This represents a 10.4× reduction in inference FLOPs and a 5.7× reduction in inference wall-clock time for slightly higher accuracy (62.4% vs. 59.7% JFT, 82.9% vs. 83.3% IN/10shot — essentially matching or beating the larger model).

Table 1 provides the complete numerical comparison. Additional highlights:

• Soft MoE S/14 with 256 experts (10M steps): 60.1% JFT P@1, 80.6% IN/10shot, 87.5% IN/finetune at 0.9 ms/img, 13.2 GFLOP/img. This model has 1.8B parameters (27× more than ViT H/14's 669M) but runs at 9.6× lower inference time and 25.3× lower inference FLOPs while achieving comparable JFT P@1 (60.1% vs. 59.7%).

• Soft MoE L/16 with 128 experts (4M steps, 500k cooldown): 63.0% JFT P@1, 84.3% IN/10shot, 89.2% IN/finetune at 4.8 ms/img, 111.1 GFLOP/img — the highest performance model across all configurations, exceeding ViT H/14 by 3.3% JFT P@1 and 1.0% IN/10shot while being 1.8× faster at inference.

• ViT H/14 (1M steps) vs. Soft MoE H/14 (256 experts, 2M steps): The Soft MoE variant reaches 62.1% JFT P@1 (vs. 58.8%) and 84.3% IN/10shot (vs. 82.7%) at similar inference cost (10.9 vs. 8.6 ms/img) but with 54.1B parameters vs. 669M.

The paper explicitly notes that "Soft MoE B/16 trained for 1k TPUv3 days matches or outperforms ViT H/14 trained on a similar budget" and that "Soft MoE B/16 matches the ViT H/14 model's performance when we double ViT-H/14's training budget (to 2k TPU-days)."

Scaling Expert Count: Soft MoE Uses More Experts More Effectively

Section 3.5 and Figure 6 present a systematic sweep of expert count and slots-per-expert across all three MoE methods.

Figure 6 (top) are heatmaps showing ImageNet 10-shot accuracy as a function of number of experts (columns, 8 to 4096) and slots-per-token (rows, 1 to 16, representing the average number of expert assignments per token). The three heatmaps correspond to Soft MoE, Experts Choice, and Tokens Choice.

For Soft MoE:

• At every row (fixed slots-per-token budget), accuracy increases monotonically with expert count. The best entry in each row is the rightmost column (most experts, fewest slots per expert, ideally one).
• The highest accuracy overall is in the bottom-right corner: 16 slots-per-token and 4096 experts achieves 75.4% IN/10shot, though this configuration is computationally expensive.
• At 1 slot-per-token (the cheapest expert compute setting), going from 8 to 4096 experts improves accuracy from 66.6% to 67.4% — a modest but consistent gain at zero additional FLOPs (since total slots remain constant).

For Experts Choice:

• There is a clear optimal expert count that depends on the slots-per-token budget. At 1 slot-per-token, accuracy peaks at 64 experts (69.4%) and then degrades to 65.8% at 4096 experts. At 16 slots-per-token, the optimum shifts right to 128 experts (72.5%) before declining to 71.0% at 4096.
• The degradation at high expert counts is attributed to increased token dropping and routing overhead, as documented in Appendix B.

For Tokens Choice:

• The degradation is even more severe. At 1 slot-per-token, accuracy peaks at 64–128 experts (67.3–67.5%) and plummets to 55.0% at 4096 experts — well below the performance with just 8 experts (64.7%). The paper's Appendix B shows that at 4096 experts with K=1 and C=1, approximately 25% of tokens are dropped entirely in some layers, which explains the catastrophic performance collapse.
• At 16 slots-per-token (K=2, C=2 configuration), the optimal is around 128–256 experts (72.1–72.2%), degrading more gracefully but still falling to 71.0% at 4096 experts.

Figure 6 (bottom) are heatmaps of training throughput (images/second) for the same sweep.

For Soft MoE, throughput is essentially flat across expert counts at each slots-per-token level — going from 8 to 4096 experts changes throughput by at most a few percent (e.g., 142 → 132 img/s at 1 slot-per-token). This confirms the theoretical analysis: the cost depends on total slots, not total experts.

For Experts Choice, throughput drops substantially at high expert counts due to routing overhead. At 1 slot-per-token, throughput falls from 374 img/s (8 experts) to 91 img/s (4096 experts) — a 4.1× slowdown despite identical theoretical FLOPs.

For Tokens Choice, the degradation is even worse: 366 → 45 img/s at the same settings, an 8.1× slowdown.

The paper summarizes: "Soft MoE's throughput is approximately constant when adding more experts. However, the Sparse MoEs' throughputs reduce dramatically from 1k experts."

Figure 23 (Appendix F.3) extends the sweep to a fixed 4096 total slots, varying expert count from 2 to 4096 while adjusting slots-per-expert to maintain the same total slot count. Soft MoE's JFT P@1 improves from 48.1% (2 experts) to 53.8% (4096 experts), while normalized training step time increases only from 0.90 to 1.18 — a 31% throughput reduction for a 5.7 percentage point accuracy gain and a 255× increase in expert parameters (38M to 9.7B parameters). This is the clearest demonstration that expert count is effectively a free parameter for Soft MoE, constrained only by memory, not compute.

Contrastive Learning Transfer: Soft MoE Representations Transfer Better

Section 4 tests whether the benefits of Soft MoE representations persist when the vision encoder is frozen and used for image-text contrastive learning. Following the LiT framework (Zhai et al., 2022b), a pretrained and frozen vision tower is paired with a randomly initialized text encoder and trained on WebLI for 18B image-text pairs (~5 epochs).

Table 3 reports zero-shot accuracy on ImageNet, CIFAR-100, and Oxford-IIIT Pet, plus COCO image-to-text and text-to-image retrieval scores.

Key results:

• Soft MoE-S/16 (128 experts) vs. ViT-S/16: 81.2% vs. 74.2% on ImageNet zero-shot (+7.0%), 67.2% vs. 56.6% on CIFAR-100 (+10.6%), 96.6% vs. 94.8% on Pet (+1.8%). COCO retrieval also improves (56.0% vs. 53.6% Img2Text). This is a striking transfer result — the Soft MoE vision encoder produces substantially better representations for a frozen feature extraction task despite being trained on the same JFT classification objective.

• Soft MoE-B/16 (128 experts) vs. ViT-B/16: 82.5% vs. 79.6% on ImageNet (+2.9%), 74.4% vs. 71.0% on CIFAR-100 (+3.4%). The gap narrows at larger backbone sizes, suggesting that larger dense models partially recover the representation quality, but Soft MoE still leads.

• Soft MoE-L/16 (128 experts) vs. ViT-L/16: 83.8% vs. 82.7% (+1.1%), with CIFAR-100 at 79.9% vs. 77.5% (+2.4%). A "souped" Soft MoE-L/16 (longer training or different configuration, details not fully specified) reaches 84.3% ImageNet and 81.3% CIFAR-100.

• Soft MoE-H/14 (256 experts) vs. ViT-H/14: 84.6% vs. 83.8% (+0.8% ImageNet), 86.3% vs. 84.7% CIFAR-100 (+1.6%). The COCO retrieval scores are slightly lower for Soft MoE-H/14 (61.0% vs. 62.7% Img2Text, 44.8% vs. 45.2% Text2Img), which the paper attributes to "the poor alignment between features learned on closed-vocabulary JFT and this open-vocabulary task."

The consistent improvement across backbone sizes — with margins largest at S/16 and shrinking at H/14 — suggests that Soft MoE's routing mechanism learns representations that are more transferable, and that this benefit is most pronounced when the base backbone has limited capacity. The paper does not deeply analyze why the representations transfer better, but the result is important for practitioners considering Soft MoE for downstream feature extraction tasks.

Appendix F.1 reports additional contrastive learning experiments on the publicly available LAION-400M dataset (Schuhmann et al., 2021), training both vision and text towers from scratch. Figure 15 shows that Soft MoE B/16 outperforms both ViT B/16 and Experts Choice B/16 (128 experts) on ImageNet 10-shot, ImageNet 0-shot, and Oxford-IIIT Pet 0-shot accuracy. Soft MoE also benefits from data augmentation ("Inception crop"), while ViT and Experts Choice do not — the paper hypothesizes this "is because Soft MoE can better utilize the expert parameters."

Ablation Studies and Robustness Checks

Expert placement in the network (Appendix D, Tables 4-6): Fixing 512 total experts across the network and distributing them across different numbers of Soft MoE layers in a 12-layer S/16 backbone:

• Concentrating all experts in one layer (layer 11): 70.0% IN/10shot, 51.5% JFT P@1.
• Splitting across 4 layers (layers 8-11): 72.8% IN/10shot, 53.2% JFT P@1 — the best configuration.
• Splitting across 8 layers (layers 4-11): 72.1% IN/10shot, 53.1% JFT P@1.
• Splitting across 8 layers with gaps (layers 1-4 and 8-11): 70.5% IN/10shot, 52.1% JFT P@1. The finding that 4-8 contiguous MoE layers in the second half of the network works best, and that interleaving MoE layers with gaps is worse, is consistent across all three routing methods (Soft MoE, Tokens Choice, Experts Choice — see Tables 5-6). This confirms the paper's default design of replacing the second half of blocks.

Number of slots per expert (Figure 12, Appendix C): With 32 experts in a Soft MoE S/16 trained for 300k steps, increasing slots per expert from 1 to 32:

• JFT P@1 improves from ~49.5% to ~52.5% (modest gain).
• IN/10shot improves from ~66% to ~73% (more substantial gain).
• Train step time increases from ~100 ms to ~250 ms (2.5× slowdown).

The paper notes this is inefficient compared to increasing expert count: "increasing the number of slots per expert only increases performance of Soft MoE a small amount, while increasing cost substantially." Figure 12 also compares against Experts Choice with group sizes of 1 and 8 images. Experts Choice benefits more from increased slots (catching up to Soft MoE at 8-image group size) but at much higher cost.

Routing mechanism components (Section 3.6, Table 2, Figure 7): Ablating the Soft MoE S/14 with 256 experts, 6 MoE layers, trained 300k steps:

Key findings: (1) Having experts and slots (the "Uniform" baseline) already provides substantial gains over a dense ViT (+3.5% JFT, +7.7% IN/10shot) — the multi-expert architecture alone is beneficial even without learned routing. (2) Learned dispatch weights contribute more than learned combine weights (+1.0% JFT and +2.8% IN/10shot when adding dispatch to Uniform, vs. +0.8% JFT and +1.8% IN/10shot when adding combine). (3) Full Soft MoE with both learned dispatch and combine achieves the best results. (4) All slot-based methods (Identity, Uniform) outperform the dense baseline, confirming that processing tokens through multiple expert views is beneficial regardless of routing quality.

L2 normalization for softmax stability (Appendix E, Figures 13-14): This ablation addresses the softmax collapse problem at large model dimensions. Figure 13 shows that without L2 normalization of inputs to the routing softmax, as model dimension d increases from 384 to 1664:

• The average maximum dispatch weight per slot grows toward 1.0 (indicating near-one-hot routing, i.e., the softmax is collapsing).
• The average maximum combine weight per token also grows but less dramatically.
• ImageNet 10-shot accuracy degrades significantly at d ≥ 1024 without normalization, while remaining stable with normalization. Figure 14 shows that the normalization also makes the model less sensitive to peak learning rate: with d=1664, the unnormalized model's maximum dispatch weight saturates near 1.0 at all learning rates, while the normalized model maintains non-saturated weights (0.2–0.4 range depending on learning rate) and achieves better accuracy. The paper states that "when d is big enough, the ImageNet 10shot accuracy is significantly worse than that achieved by properly normalizing the inputs."

Expert count scaling at fixed total slots (Figure 23): With 4096 total slots held constant, increasing expert count from 2 to 4096 (while decreasing slots-per-expert from 2048 to 1):

• JFT P@1 improves monotonically from ~48.1% to ~53.8%.
• IN/10shot improves from ~62% to ~75%.
• Normalized training step time increases only 31% (0.90 to 1.18 relative to the 32-expert baseline), with most of the increase attributed to communication overhead from distributing experts across more devices, not increased computation. This confirms that expert count can be treated as a free parameter for scaling capacity, bounded only by memory and communication constraints.

Batch Priority Routing for Tokens Choice (Appendix F.2, Table 7): For Tokens Choice Sparse MoE with S/16 backbone, BPR provides significant gains when K=1:

• 32 experts, K=1: 50.1% → 51.2% JFT P@1, 64.5% → 68.9% IN/10shot with BPR.
• 64 experts, K=1: 50.0% → 51.5% JFT P@1, 64.4% → 69.1% IN/10shot with BPR.
• For K=2, BPR provides minimal additional benefit (71.0% → 71.4% at 32 experts), since having K=2 already provides routing flexibility that partially mitigates token dropping.

Token dropping in Sparse MoEs (Appendix B, Figures 8-11): For both Experts Choice and Tokens Choice, increasing expert count while keeping capacity tight (C=1, K=1) leads to rapidly increasing token dropping rates:

• At 16–32 experts, dropping rates are ~15%.
• At 512–1024 experts, dropping rates exceed 40–50% for Experts Choice and ~25% for Tokens Choice.
• Slightly increasing buffer capacity (C=1.125) reduces dropping modestly (~5% absolute reduction) but doesn't solve the fundamental scaling problem (Figure 9-10).
• BPR for Tokens Choice reduces dropping and improves performance (Figure 11) but doesn't eliminate the expert-count ceiling.

Slot correlation analysis (Appendix I, Figures 28-30): In Soft MoE S/16 with different slots-per-expert configurations:

• 1 slot per expert (Figure 28): No clear correlation structure between slot parameter vectors beyond the diagonal — each slot learns a distinct projection.
• 4 slots per expert (Figure 29): Clear block-diagonal structure — the 4 slots from each expert show high pairwise inner products (bright 4×4 blocks), indicating alignment.
• 16 slots per expert (Figure 30): Even stronger block-diagonal patterns — the 16 parameter vectors per expert are highly correlated, suggesting near-redundancy.

This provides mechanistic evidence for why one slot per expert is optimal: multiple slots per expert don't learn complementary representations but instead collapse to similar weighted averages, adding computational cost without commensurate representational benefit.

Cumulative dispatch and combine weight distributions (Appendix H, Figures 26-27): Analyzing a finetuned Soft MoE H/14 on ImageNet:

• Dispatch weights: In earlier MoE layers (blocks 16–27), a few tens of tokens cover 80% of the weighted average mass for a given slot, and there is significant variation across slots (wide colored bands). In the final two layers (blocks 30–31), the dispatch weights are nearly uniform — all tokens contribute roughly equally, suggesting these layers pool information broadly.
• Combine weights: Even in the final layers where dispatch is uniform, combine weights are not uniform — some experts contribute much more than others to output tokens, indicating that expert specialization persists even when inputs are broadly pooled. The colored bands are wide, indicating significant slot-to-slot variation in importance.

Visual inspection of slot assignments (Appendix G, Figure 25): Visualizing dispatch weights for a cherry-picked set of 8 slots in a Soft MoE S/16 finetuned on ImageNet shows that different slots learn to focus on different image regions — some slots emphasize the central object, others emphasize background or edges, others are more uniform. This provides qualitative evidence that the soft routing learns meaningful, diverse feature pooling.

Per-token and per-expert contribution distributions (Appendix G, Figure 24):

• Left plot: The distribution of summed dispatch weights per token across slots shows that a small fraction of tokens (2–5%) contribute disproportionately (summed weight > 2), while 15–20% contribute very little (summed weight ≤ 0.25). This is essentially a "soft token dropping" — tokens with low total dispatch weight are effectively ignored by the expert processing, but unlike hard dropping in Sparse MoEs, their information still propagates through the residual connection.
• Middle plot: Expert importance (average combine weight, normalized) varies 3–14× across experts depending on the layer. Some experts are heavily used; others are lightly used. Unlike Sparse MoEs, this doesn't cause training problems because every expert receives exactly the same computational budget (one slot) regardless of its combine weight — unused experts still consume compute but don't slow down training.
• Right plot: For each slot, the number of tokens needed to reach 90% cumulative dispatch weight varies widely — from ~20–25 tokens (concentrated) to ~150+ tokens (distributed), indicating that different experts learn different pooling granularities.

Critical Assessment

The experiments in this paper are unusually thorough for an architectural contribution — over 100 models trained for the Pareto frontier analysis, systematic sweeps of expert count from 2 to 4096, multiple training durations from 300k to 10M steps, and evaluation across pretraining, few-shot, finetuning, and contrastive transfer. This breadth of empirical evidence is a genuine strength. However, several limitations and unexamined questions deserve attention.

Does Soft MoE genuinely dominate Sparse MoEs, or are the Sparse MoE baselines undertuned? The paper's claim that Soft MoE is on the Pareto frontier and dominates both Tokens Choice and Experts Choice is well-supported for the configurations tested — the 107-model sweep in Appendix J is extensive. However, Sparse MoEs are notoriously sensitive to auxiliary loss weighting, capacity factor, and batch size. The paper sweeps K (1-2), C (0.5-2.0), and expert count (16-4096), but does not report sweeping the auxiliary load-balancing loss coefficient or other training hyperparameters specific to Sparse MoEs. The Tokens Choice K=2, C=2 configurations at large expert counts show catastrophic performance drops (Figure 6), which might be partially addressable with better-tuned auxiliary losses or different capacity allocation strategies. The paper also uses Batch Priority Routing which improves Tokens Choice (Table 7), but other orthogonal improvements to Sparse MoE training (e.g., Z-loss from Zoph et al., 2022, not cited) are not explored. This doesn't invalidate the comparison — Soft MoE's advantage of not needing these heuristics is part of its value proposition — but it means the Sparse MoE baselines represent strong but not necessarily maximally optimized versions of those methods.

The evaluation is entirely on image classification, mostly on proprietary data. JFT-4B is not publicly available, making exact replication impossible. The paper partially addresses this with LAION-400M experiments (Appendix F.1), but those are at smaller scale (B/16 only, 275M images) and report only a few metrics. Whether Soft MoE's advantages extend to other vision tasks (object detection, segmentation, video) or to NLP encoder tasks (BERT-style pretraining) is untested. The contrastive learning experiments begin to address transfer but use only frozen vision encoders paired with text — a more complete picture would include end-to-end finetuning on diverse downstream tasks.

The auto-regressive decoding limitation has not been empirically probed. The paper identifies this as a limitation (Section 6) but provides no experiments on decoder architectures, even with causal masking modifications. A natural question is whether causal constraints on the dispatch/combine weights (preventing future tokens from influencing current expert inputs) would degrade the routing quality or whether the method could be adapted. Given the prominence of decoder-only MoEs in language modeling (Fedus et al., 2022; Lepikhin et al., 2020), this is a significant scope constraint that limits the paper's impact on the largest current application area for MoEs.

The memory consumption claim is acknowledged but not quantified. Section 6 notes that one slot per expert means many expert parameters, consuming memory. Table 1 shows Soft MoE H/14 with 256 experts at 54.1B parameters — this is an enormous model that requires significant device memory or model parallelism. The paper doesn't report the memory footprint, device count, or communication overhead for these models, making it difficult for practitioners to assess whether the approach is feasible for their hardware constraints. The training throughput numbers (Figures 6, 23) partially capture this, but inference-time memory requirements for deployment are not discussed beyond the ms/img timings.

The FLOPs accounting treats expert computation as equivalent to dense MLP computation. Soft MoE replaces dense MLP blocks where each token is processed by one MLP. The Soft MoE layer instead processes fewer slots (e.g., 128 slots vs. 196 tokens for S/16), so the expert FLOPs are lower. But the routing adds the $O(m \cdot n \cdot p \cdot d)$ cost of computing $X\Phi$ and the weighted averages. The paper argues this is comparable to self-attention cost, but doesn't provide a breakdown of FLOPs by component (routing vs. expert vs. attention) for representative models. This would help practitioners understand where the compute is spent and whether further optimization of the routing computation is warranted.

The optimal number of MoE layers is likely architecture- and task-dependent. Appendix D's placement ablation uses a 12-layer S/16 backbone and finds 4-8 MoE layers optimal. Whether this generalizes to deeper networks (24, 32, 48 layers), different tasks, or different expert counts is unexplored. The finding that interleaved MoE layers (with gaps) underperform contiguous blocks is noted but not explained — it could be an artifact of the specific architecture or a deeper principle about routing continuity.

The difficulty of matching sequence length to slot count for variable-resolution inputs. Vision Transformers typically use fixed patch sizes but may encounter variable image resolutions at deployment (e.g., finetuning at 384 resolution as the paper does). The paper notes that during finetuning at 384×384 resolution, the number of tokens increases to 576 (for patch size 16) or 752 (for patch size 14), while slots remain fixed at 128 or 256. This means Soft MoE becomes cheaper relative to dense models at higher resolutions (since the token-to-slot ratio increases), which is a benefit. But the routing quality may degrade if the slot count is poorly matched to the token count — the paper doesn't ablate this systematically.

The "one slot per expert is optimal" finding may be specific to the training data scale. All primary experiments use JFT-4B, which provides hundreds of millions of training examples per class on average. With less data, the benefit of many experts with one slot each might diminish because the experts don't have enough data to develop meaningful specializations. The LAION-400M experiments (275M images, much smaller than JFT-4B) still show Soft MoE benefits, but they use fewer experts (128, matching the B/16 configuration) — a sweep of expert count at smaller data scales is not reported.

In summary, the experiments strongly support the paper's core claims within the tested domain (image classification at scale with encoder architectures): Soft MoE is simpler to train than Sparse MoEs, scales better with expert count, achieves superior accuracy/cost Pareto frontiers, and transfers well to contrastive learning. The evidence is unusually comprehensive for a new architectural component, with the Pareto frontier analysis serving as a particularly convincing empirical foundation. However, the claims about generality beyond image classification, applicability to decoder architectures, and practical memory constraints are either unexamined or acknowledged as limitations. The paper makes a compelling case that Soft MoE is the best MoE routing method for vision encoders, but the extension to the broader MoE landscape — particularly language models — remains an open question.

6. Limitations and Trade-offs

6.1 The Method Is Only Demonstrated for Encoder Architectures, Not Auto-Regressive Decoders

The assumption or constraint. Soft MoE constructs each slot as a weighted average of all input tokens in the sequence simultaneously. This requires full bidirectional access to every token, which is available in encoder architectures (Vision Transformers, BERT-style models) but fundamentally incompatible with auto-regressive decoding, where future tokens must be masked during training and generation. The paper acknowledges this explicitly in Section 6:

"One of the key aspects of Soft MoE consists in learning the merging of all tokens in the input. This makes the use of Soft MoEs in auto-regressive decoders difficult, since causality between past and future tokens has to be preserved during training. Although causal masks used in attention layers could be used, one must be careful to not introduce any correlation between token and slot indices, since this may bias which token indices each expert is trained on. The use of Soft MoE in auto-regressive decoders is a promising research avenue that we leave for future work."

The consequence. This scope constraint is significant because auto-regressive decoder architectures (GPT-style models) represent one of the largest and most impactful application domains for MoEs. The Switch Transformer (Fedus et al., 2022) and GShard (Lepikhin et al., 2020) — two of the most prominent MoE deployments — both use decoder-only or encoder-decoder architectures for language modeling. If Soft MoE cannot be extended to these settings, its impact is confined to encoder tasks: image classification, image-text contrastive learning (as demonstrated), and potentially BERT-style masked language modeling. For the many practitioners building and deploying large language models, Soft MoE is currently not a usable option.

The specific failure mode is not just that Soft MoE would be suboptimal under causal masking, but that it could introduce a token-index bias. In an auto-regressive setting where the sequence length varies, if the causal mask creates a correlation between position in the sequence and which experts process which future tokens, the model may learn to route based on position rather than content — defeating the purpose of content-based conditional computation. The paper does not explore this failure mode empirically.

What evidence exists in the paper. None. The paper provides zero experiments on decoder architectures, even with causal masking modifications. The limitation is purely analytic and aspirational. This is not a criticism of the experimental design — the paper's scope is explicitly vision encoders — but it means the claim that Soft MoE "overcomes many of [the] challenges" of Sparse MoEs (abstract) must be understood as conditional on the encoder-only setting.

Mitigation status. The paper frames this as "a promising research avenue" and provides no concrete proposal for adaptation. The note about causal masking suggests a direction but immediately identifies a complication (token-slot index correlation) that would need to be solved. Given the prominence of decoder MoEs, this limitation substantially constrains the paper's claim to have solved general MoE problems.

6.2 Soft MoE Is Not Technically Sparse, and the Memory Cost of Many Experts Is Unbounded

The assumption or constraint. The paper is explicit that Soft MoE achieves parameter scaling without proportional compute scaling, but it does so through a slot bottleneck rather than through sparse activation. Every input token fractionally activates all model parameters — the full dispatch and combine weight matrices are dense. This means that while FLOPs can be controlled (by choosing the total number of slots independently of the number of experts), memory consumption scales with the total number of expert parameters. Section 6 acknowledges this:

"We show in Section 3 that one slot per expert tends to be the optimal choice. In other words, rather than feeding one expert with two slots, it is more effective to use two experts with one slot each. We hypothesize slots that use the same expert tend to align and provide small informational gains, and a expert may lack the flexibility to accommodate very different slot projections. We show this in Appendix I. Consequently, Soft MoE can leverage a large number of experts and—while its cost is still similar to the dense backbone—the memory requirements of the model can grow large."

The consequence. The combination of two findings — (1) one slot per expert is optimal, and (2) more experts improve performance monotonically (Figure 23) — creates an uncomfortable tension for practitioners. The best-performing Soft MoE configuration pushes toward using as many experts as device memory allows, with as few slots per expert as possible. But the memory footprint of a Soft MoE with $n$ experts is roughly $O(n \cdot d^2)$ for the expert MLP parameters, compared to $O(d^2)$ for the dense backbone. A Soft MoE H/14 with 256 experts has 54.1B parameters (Table 1) — this is an enormous model that requires either very large memory capacity per device or model parallelism across many devices. The paper does not report device counts, memory per device, or communication overhead for these large-scale training runs, making it difficult for practitioners to assess hardware feasibility.

This distinguishes Soft MoE from Sparse MoEs in an important way. In a Sparse MoE with capacity factor $C < 1$, only a fraction of expert parameters are needed in fast memory at any given time because any specific input only activates some experts. Sophisticated implementations can load-balance expert parameters across devices and swap them as needed. In Soft MoE, every expert is always fully activated in the sense that all expert parameters are involved in processing every input (since every slot is a mixture of all tokens). You cannot lazily load experts or keep only a subset in memory — the entire model must be available for every forward pass. The paper's "features of both sparse and dense" framing (Section 2.2) is accurate for compute but not for memory: Soft MoE has the memory characteristics of a dense model with $n \times$ the parameter count.

What evidence exists in the paper. The parameter counts in Table 1 are the primary evidence: Soft MoE S/14 with 256 experts is 1.8B parameters (vs. 33M for ViT S/16, a 55× increase), Soft MoE B/16 with 128 experts is 3.7B (vs. 108M, 34× increase), Soft MoE H/14 with 256 experts is 54.1B (vs. 669M, 81× increase). The paper reports training throughput in images/second (Figure 6, bottom) which captures some of the distributed training overhead, but does not report peak memory usage, device count, or the breakdown of time spent in communication vs. computation. Figure 23 shows a 31% throughput reduction when scaling from 32 to 4096 experts (at fixed total slots), attributed to "communication costs due to higher topologies needed for larger models" — but the absolute memory footprint is not quantified.

Mitigation status. The paper acknowledges the memory concern in Section 6 but does not propose mitigations beyond standard model distribution techniques: "we employ the standard techniques to distribute the model across many devices, as in (Lepikhin et al., 2020; Riquelme et al., 2021; Fedus et al., 2022)." It does not investigate whether fewer experts with more slots (which reduces memory at the cost of some performance, per Figure 12) might be a practical compromise. It also does not explore whether expert parameter sharing, distillation from large Soft MoEs into smaller ones, or other memory-reduction strategies could recover the performance of many-expert models at lower memory cost. For practitioners with memory-constrained deployment environments (edge devices, consumer GPUs, single-TPU inference), the headline 54.1B-parameter model with "only 2% increased inference time" is technically true for FLOPs but misleading for practical deployability.

6.3 Only One Task Domain (Image Classification) and One Model Family (ViT) Are Thoroughly Evaluated

The assumption or constraint. Every major quantitative claim in the paper — the Pareto frontier dominance (Section 3.3), the inference-time speed advantages (Section 3.4), the scaling to thousands of experts (Section 3.5), the routing ablations (Section 3.6) — is based on experiments with Vision Transformers on the JFT-4B image classification dataset. The contrastive learning experiments in Section 4 extend evaluation to image-text alignment (frozen vision encoder, trained text tower on WebLI), but still use the same pretrained ViT backbones and evaluate on standard vision benchmarks (ImageNet, CIFAR-100, Pet, COCO). The LAION-400M experiments in Appendix F.1 follow the same paradigm.

The paper does not evaluate Soft MoE on:

• Other vision architectures (e.g., ConvNeXt, Swin, or hybrid CNN-Transformer models).
• Other vision tasks (object detection, semantic segmentation, video understanding, depth estimation).
• Other modalities (text encoders in BERT-style pretraining, speech encoders, multimodal encoders where Soft MoE could be placed in one or both towers).
• Tasks with different sequence-length characteristics (very long sequences where the $O(m^2d)$ routing cost could become dominant, or very short sequences where the slot bottleneck might compress too aggressively).

The consequence. The paper's findings may be specific to the interaction between Soft MoE routing and the ViT architecture, or to the characteristics of image classification at scale. Several aspects of the method could behave differently in other contexts:

• The optimal number of slots relative to sequence length (the compression ratio $m / (n\cdot p)$) was tuned for 196–576 token sequences. For NLP encoders with sequences of 512–2048 tokens, the routing FLOPs ($O(m \cdot (n\cdot p) \cdot d)$) would scale differently relative to expert FLOPs, potentially shifting the efficiency tradeoff.

• The finding that contiguous MoE layers in the second half of the network are optimal (Appendix D) was derived on 12-layer ViTs. Deeper networks (24–48 layers in modern LLM encoders) might benefit from different placement strategies.

• The contrastive learning results (Table 3) show that Soft MoE provides larger improvements on structured classification (ImageNet +7% for S/16) than on retrieval tasks (COCO +2.4% for S/16). This could indicate that the expert routing specializes toward features useful for classification and provides diminishing returns for tasks requiring more open-vocabulary understanding.

• The LAION-400M experiments (Appendix F.1) use only a B/16 backbone with 128 experts — a single configuration. This is not a systematic sweep of whether Soft MoE's advantages scale similarly on smaller public datasets compared to JFT-4B's 4B images.

What evidence exists in the paper. The paper's evaluation is deep within the vision classification domain but narrow outside it. Table 3 and Appendix F.1 are the only non-classification results. The paper presents these as evidence of transfer and generality (Section 4 is titled "Contrastive Learning"), but they test only one additional paradigm (frozen contrastive image encoder) and use the same underlying model architecture.

Mitigation status. The paper makes no claim of domain generality and does not suggest that these results will transfer without modification. The title specifies neither "Vision" nor "Encoder," but the abstract and introduction also avoid domain qualifications — it is described as "a fully-differentiable sparse Transformer" without specifying that only encoder attention patterns have been tested. Section 6 acknowledges the decoder limitation separately but does not discuss the single-task limitation. This is a scope constraint that the paper could have flagged more explicitly: the method is validated for image classification with ViT encoders, and extension to other domains is plausible but unproven.

6.4 The Single-Benchmark, Single-Model Evaluation Limits the Strength of the Pareto Frontier Claim

The assumption or constraint. The paper's central empirical claim — that "Soft MoE greatly outperforms dense Transformers (ViTs) and popular MoEs" (abstract) — is supported by training over 100 model configurations on a single dataset (JFT-4B), evaluated on a single pretraining metric (JFT precision-at-1), a single few-shot transfer task (ImageNet 10-shot), and a single finetuning task (ImageNet-1k at 384 resolution). The claim of Pareto frontier dominance is conditional on these specific evaluation axes. If a deployment setting prioritizes a different capability — say, robustness to distribution shift, calibration quality, or performance on long-tailed classes — the Pareto frontier might look different.

The consequence. Several aspects of the evaluation are worth scrutinizing for how they might favor Soft MoE:

• JFT-4B is a 29k-class dataset with many fine-grained categories. This provides a strong training signal that may encourage the kind of expert specialization Soft MoE benefits from. On smaller datasets with fewer classes or less structured variation, the many-expert advantage might diminish because experts don't receive enough data to develop meaningful specializations. The LAION-400M experiments in Appendix F.1 partially address this (275M images, much less than JFT-4B's 4B+), but use only one expert configuration and don't sweep expert count to test whether the monotonic scaling holds at smaller data scales.

• ImageNet 10-shot and finetuning evaluations use frozen/partially-frozen representations. Soft MoE's advantage on these metrics could partly reflect that it learns more transferable features (as the contrastive learning results suggest), but could also reflect an interaction with the specific transfer protocol. If the few-shot head training or finetuning procedure interacts poorly with Sparse MoE's routing (e.g., small batch sizes in few-shot training causing routing collapse for Tokens Choice), Soft MoE's advantage might be partly a methodological artifact rather than a representation quality difference.

• The finetuning comparison at 384 resolution changes the ratio of tokens (576 or 752) to slots (128 or 256), making Soft MoE relatively cheaper at higher resolution. The paper reports this as a benefit, but it means the comparison is not FLOPs-matched at the finetuning resolution — Soft MoE gets an implicit efficiency advantage from having fewer slots than tokens.

• All evaluation metrics are top-1 accuracy variants with exact-match grading. The paper doesn't evaluate calibration (expected calibration error), robustness to common corruptions, or performance on tail classes — dimensions where the soft routing's "every token contributes to every expert" property might produce different behavior than Sparse MoEs' hard token dropping.

What evidence exists in the paper. The evaluation protocol is consistent and well-specified, and the breadth of configurations (107 models, Appendix J) is unusually comprehensive. The paper is transparent about its metrics and reports both FLOPs and wall-clock time. But the narrowness of the evaluation axes — essentially, three variants of ImageNet transfer accuracy plus upstream JFT accuracy — limits the generality of the claim. There is no evidence about how Soft MoE affects robustness, calibration, or fairness properties compared to dense or Sparse MoE alternatives.

Mitigation status. The paper does not discuss these metric limitations. It treats the evaluated metrics as sufficient to establish Pareto dominance. For practitioners whose deployment priorities align with top-1 accuracy on standard benchmarks, the evaluation is adequate. For practitioners who care about calibration, robustness, or behavior under distribution shift, the paper provides no guidance. This is a common limitation in the MoE literature (few papers systematically evaluate robustness properties of routing mechanisms), but it is worth noting given the paper's strong claims of superiority.

6.5 The Difficulty Estimation Overhead Is Not Amortized

The assumption or constraint. This limitation does not apply to Soft MoE in the same way it applies to some other papers — Soft MoE does not require explicit difficulty estimation at inference time. However, there is an analogous configuration cost: the paper's optimal configurations for a given backbone size and compute budget were found through extensive hyperparameter sweeps (over 100 models trained for the Pareto frontier analysis in Section 3.3, systematic sweeps of expert count and slots-per-expert in Section 3.5). A practitioner wanting to deploy Soft MoE on a new task, dataset, or backbone architecture cannot simply inherit these configurations — they are tuned to ViT backbones on JFT-4B.

The consequence. The headline efficiency numbers (e.g., Soft MoE B/16 matches ViT H/14 at 5.7× faster inference) are for models whose configurations (128 experts, 1 slot per expert, placement in second half of layers) were selected after extensive sweeps. If a practitioner deploys Soft MoE on a new architecture (say, a ViT variant with different depth or width) or a new task domain (say, medical imaging where patch features have different characteristics), the optimal configuration might differ. The paper does not provide a principled method for selecting expert count, slot count, or layer placement without expensive sweeps. This is not unique to Soft MoE — it is true of virtually all MoE and architecture design work — but it means the practical cost of adopting Soft MoE includes the cost of these sweeps, which are not amortized in the paper's reported training cost comparisons.

What evidence exists in the paper. The paper's extensive ablation experiments (Figures 6, 12, 23; Appendix D; Tables 4-6) provide guidance on trends: more experts with one slot each is generally better; placing MoE layers in the second half of the network in contiguous blocks is generally better; L2 normalization of routing logits is important at large dimensions. These heuristics reduce the sweep space but do not eliminate the need for per-task configuration. The paper does not report how sensitive performance is to deviations from the optimal configuration — if a practitioner uses 64 experts instead of 128, how much performance is lost? If they place MoE layers in layers 4-11 instead of 8-11, what is the degradation? The placement ablation (Appendix D) provides some guidance here, showing that 4-layer contiguous placement (72.8% IN/10shot) outperforms 2-layer (71.7%) and 8-layer with gaps (70.5%), but this is for one specific total expert count (512) and one backbone (S/16).

Mitigation status. The paper provides configuration guidelines based on its empirical findings (Section 2.3: "We typically replace the second half of blocks"; Section 3.5: one slot per expert is optimal) but does not develop a configuration selection methodology. This is a practical limitation rather than a scientific one — the paper's contribution is a new architectural component and its empirical characterization, not a method for automatically configuring that component. But for practitioners, the distinction matters: the paper demonstrates Soft MoE can be excellent, but it does not make it trivial to achieve that excellence without similar-scale sweeps.

7. Implications and Future Directions

How This Work Changes the Landscape

Soft MoE does not introduce a new training paradigm, a new scaling law, or a new theoretical framework. It introduces a single architectural component — a drop-in replacement for the MLP block in a Transformer encoder — and demonstrates through exhaustive empirical characterization that this component sidesteps the primary failure modes that have made Sparse MoEs fragile in practice. The paper's contribution is therefore best understood as a methodological consolidation with disproportionate practical consequences: it takes the MoE concept, which the field has long recognized as theoretically attractive but practically temperamental, and shows that a single design change — replacing discrete token-to-expert assignment with continuous slot-based token mixing — resolves the core tensions well enough that MoE scaling becomes almost as straightforward as scaling a dense model.

The magnitude of this shift should be understood in context. Before Soft MoE, a practitioner wanting to deploy a large MoE model faced a combinatorial optimization problem: choose between Tokens Choice and Experts Choice routing (or more exotic alternatives like optimal transport or reinforcement learning), tune the number of experts, the capacity multiplier, the auxiliary loss weight, the number of selected experts per token ($K$), the batch size (which interacts with expert capacity under per-batch routing), and the layer placement — and then hope that the resulting configuration doesn't collapse under distribution shift, small batch sizes, or novel inputs. The paper's experiments (Appendix B, Figures 8-11) demonstrate concretely that even with careful tuning, Tokens Choice routing with $K=1$ and $C=1$ drops 15% of tokens at moderate expert counts and collapses to 25%+ dropping at scale — meaning a substantial fraction of the input receives no expert processing at all. Experts Choice routing fares even worse, with 40-50% of tokens dropped in some layers at large expert counts. These are not edge cases; they are the default behavior of Sparse MoEs under the configurations that make them computationally efficient.

Soft MoE eliminates this complexity. The design has essentially two hyperparameters — number of experts and slots per expert — and the paper provides clear empirical guidance for both: use as many experts as memory allows with one slot each, place them in the second half of layers, and L2-normalize the routing logits. There is no auxiliary loss to tune, no capacity factor to adjust, no batch-size sensitivity to worry about, and no token dropping to monitor. This is not an incremental improvement; it is a reduction in the degrees of freedom of MoE design from a high-dimensional, interdependent space to essentially one knob (total expert count), bounded only by memory rather than by optimization stability.

The paper also resolves a latent contradiction in the MoE literature that has been visible but rarely articulated directly: the tension between the theoretical promise of MoEs (arbitrary parameter scaling at near-constant FLOPs) and the empirical reality that existing routers hit a performance ceiling at relatively modest expert counts. Section 3.5 (Figure 6) shows this ceiling explicitly: Tokens Choice accuracy peaks at 64-128 experts and then degrades; Experts Choice peaks somewhat higher but still shows diminishing returns. Soft MoE's accuracy improves monotonically to at least 4096 experts (Figure 23). This demonstrates that the scaling ceiling was a property of the routers, not of the MoE concept itself. The implication is clear: research effort spent on improving discrete routing heuristics — better auxiliary losses, more sophisticated assignment algorithms, clever capacity allocation — is effort spent optimizing a fundamentally constrained approach. The more productive direction, which Soft MoE exemplifies, is to ask whether discrete routing is necessary at all.

This reframing has implications for how the field thinks about conditional computation more broadly. The dominant mental model has been that efficiency requires sparsity — that to get the benefits of large parameter counts without paying proportional compute, you must activate only a subset of parameters for each input. Soft MoE demonstrates a different mechanism: dimensionality reduction in a learned space can achieve the same efficiency without any hard zeroing-out of computation. The slots act as a computational bottleneck — they compress the sequence into fewer representations before expensive expert processing, then expand back — and this bottleneck, not sparsity, is what decouples parameter count from FLOPs. This is a genuinely different design principle that may generalize beyond the specific dispatch/combine mechanism of Soft MoE. Any learned pooling operation that reduces the number of items processed by expensive downstream modules could, in principle, achieve similar efficiency gains.

The paper also changes the conversation around MoE deployment feasibility. Prior Sparse MoE work (Fedus et al., 2022; Lepikhin et al., 2020) demonstrated impressive results but required careful engineering of expert sharding, load balancing, and capacity management — infrastructure that few organizations outside of large industrial labs can build and maintain. Soft MoE's implementation is a few lines of JAX (Algorithm 1) involving only standard operations (matrix multiply, softmax, einsum). The barrier to entry for experimenting with and deploying MoE models drops substantially. This democratization effect — making large-scale conditional computation accessible to groups without specialized distributed systems expertise — may prove more impactful in aggregate than any single accuracy number in the paper.

A more subtle landscape shift concerns the relationship between training data scale and expert specialization. The paper's LAION-400M experiments (Appendix F.1) show that Soft MoE benefits persist on a dataset two orders of magnitude smaller than JFT-4B, and the Uniform baseline (Table 2) shows that simply having multiple expert views of uniformly averaged tokens already outperforms a dense model. This suggests that expert specialization is not purely a function of massive data — even with modest data, having multiple processing pathways for the same (pooled) information provides a representational benefit. The paper does not deeply analyze this, but it hints that the benefit of MoEs may be partly architectural (more parameters providing more representational capacity even under crude routing) and partly data-driven (learned routing providing genuine specialization). Disentangling these two effects would be a valuable contribution that Soft MoE's simplicity makes newly tractable.

Finally, the paper clarifies which research directions become less attractive. The extensive failure characterization of Sparse MoE routers — token dropping at scale (Appendix B), throughput collapse from routing overhead (Figure 6 bottom), performance ceilings at moderate expert counts (Figure 6 top) — suggests that incremental improvements to discrete routing are unlikely to close the gap with fully-differentiable approaches. The paper does not state this directly, but the empirical case is strong: even with Batch Priority Routing (which substantially improves Tokens Choice at $K=1$, per Table 7), carefully tuned capacity multipliers, and systematic sweeps of expert count, the best Sparse MoE configurations consistently underperform Soft MoE at matched FLOPs and matched wall-clock time (Figures 3, 19-22). Research on discrete routing may still be valuable for decoder architectures where Soft MoE's bidirectional mixing is not directly applicable, but for encoder settings, the paper makes a compelling case that the discrete routing paradigm has been superseded.

Follow-Up Research This Work Enables

Adapting Soft MoE to auto-regressive decoders via causal slot mixing. The most impactful open question the paper identifies (Section 6) is whether Soft MoE can be extended to decoder architectures where full bidirectional token mixing violates causality. A direct adaptation would compute dispatch and combine weights using only past tokens: for token position $t$, slot $j$ would be a weighted average of tokens $1 \ldots t$ rather than $1 \ldots m$. This preserves causality but introduces a position-dependent computational pattern — early tokens contribute to fewer slots, and the slot representations evolve as the sequence grows — which may create the token-index correlation the paper warns about. A concrete first experiment: implement causally-masked Soft MoE in a decoder-only language model at modest scale (e.g., 100M-300M parameters, trained on C4 or the Pile), compare against a Tokens Choice Sparse MoE baseline with matched expert count and FLOPs, and measure both perplexity and the distribution of expert usage across sequence positions. The key diagnostic is whether later tokens systematically route to different experts than early tokens purely due to position rather than content. If causal Soft MoE underperforms Sparse MoEs on language modeling, that would establish a fundamental boundary on the approach. If it works, it would open the largest current MoE application domain to fully-differentiable routing.

Scaling laws for Soft MoE: expert count vs. training tokens vs. backbone size. The paper demonstrates that more experts improve performance monotonically (Figure 23) but does not characterize the shape of this improvement — is it logarithmic (diminishing returns) or power-law (continued gains)? More importantly, the paper varies expert count while holding training data constant (JFT-4B, 4M steps). The interaction between expert count and training data scale is unexplored: at what point do additional experts provide negligible benefit because there isn't enough data to learn meaningful specializations? A systematic study could train Soft MoE models with expert counts ranging from 2 to 4096 at multiple data scales (e.g., 10M, 100M, 1B, 4B images), each trained to convergence, and fit scaling laws relating parameters, data, and loss analogous to Hoffmann et al. (2022). The key question is whether Soft MoE follows the same compute-optimal scaling pattern as dense models (parameters and data should scale together) or whether the decoupling of parameters from FLOPs changes the optimal allocation — perhaps you should always use as many experts as memory allows, and the limiting factor is data, not compute. The paper's memory limitation (Section 6) would also be directly addressed by such a study: if diminishing returns set in at expert counts well below memory limits for realistic data scales, the memory concern is less pressing.

Combining Soft MoE with other efficiency mechanisms to push the accuracy-per-FLOP frontier. The paper evaluates Soft MoE as a standalone replacement for MLP blocks, but modern efficient Transformers often combine multiple mechanisms — sparse attention, token pruning, conditional depth, parameter sharing. Soft MoE is architecturally compatible with all of these: the dispatch/combine mechanism is independent of the attention pattern, and the slot bottleneck is orthogonal to token pruning (you could prune tokens before they enter the Soft MoE layer, reducing both routing and expert cost). A concrete experiment: combine Soft MoE with token merging (e.g., ToMe or TokenLearner) applied before the Soft MoE layers, so the routing operates on a reduced token set, and compare against both standalone Soft MoE and standalone token merging on ImageNet-21k or JFT at matched FLOPs. The hypothesis is that token merging reduces the $O(m^2d)$ routing cost while Soft MoE provides the capacity scaling, and the two mechanisms compound rather than interfere. A negative result — that token merging removes the token-level diversity that Soft MoE routing exploits — would be equally informative, clarifying the conditions under which the slot bottleneck mechanism works.

Stress-testing Soft MoE under distribution shift and adversarial inputs. MoEs introduce a routing step that, in principle, could learn spurious correlations between input features and expert assignment, making the model brittle under distribution shift. In Sparse MoEs, this manifests as routing collapse under small batch sizes or novel inputs (noted in Section 1). Soft MoE's soft routing might be more robust — because every token contributes to every slot, there is no hard assignment to fail — or it might be vulnerable in different ways: an adversary could craft inputs that push the dispatch weights toward degenerate distributions (e.g., all tokens routed to a single slot, effectively bypassing the multi-expert architecture). Concrete experiments: evaluate Soft MoE models on ImageNet-C (common corruptions), ImageNet-R (renditions), and ObjectNet (object pose/viewpoint variation) and compare degradation relative to dense ViTs and Sparse MoEs. Additionally, test whether an adversary with access to the routing logits can craft perturbations that increase routing concentration (measured by the entropy of dispatch weights) and whether this degrades accuracy. This would characterize Soft MoE's robustness properties — currently completely unknown — and potentially motivate routing regularization techniques analogous to the auxiliary losses used in Sparse MoEs but for different reasons (encouraging diverse routing rather than balanced routing).

Soft MoE for modalities beyond vision: BERT-style text encoders and multimodal fusion. The paper is careful to limit its claims to the evaluated domain (vision encoders on image classification), but the mechanism is modality-agnostic: any sequence of token representations can be mixed via learned dispatch and combine weights. A natural extension is to replace the feedforward blocks in a BERT-style text encoder with Soft MoE layers and evaluate on GLUE/SuperGLUE benchmarks. This tests whether the benefits transfer to a domain with different sequence length characteristics (128-512 tokens vs. 196-752), different token semantics (word pieces vs. image patches), and different task structures (single-sentence classification, pairwise similarity, question answering). A more ambitious experiment: place Soft MoE in one or both towers of a multimodal contrastive model (replacing the frozen-vision setup in Section 4 with fully trainable Soft MoE in both encoders) and measure whether the expert specialization leads to better cross-modal alignment, perhaps via probing whether certain experts specialize in attributes relevant to specific text concepts. The LAION-400M experiments in Appendix F.1 provide a starting point — training both towers from scratch with Soft MoE — but only evaluate on standard vision benchmarks, not on text or retrieval tasks that would reveal cross-modal interaction effects.

Distillation and pruning of Soft MoE models to recover the many-expert benefit at lower memory cost. The paper's memory limitation — many experts consume memory even though they don't consume FLOPs — is the primary practical barrier to deploying large Soft MoE models outside of well-resourced environments. A straightforward mitigation is to train a large Soft MoE (many experts, high accuracy) and then distill it into a smaller Soft MoE (fewer experts) or into a dense model, using the large model's predictions as soft targets. The question is whether the many-expert model learns specializations that can be compressed, or whether the benefit is inherently tied to having many separate parameter sets. Concrete experiment: train a Soft MoE H/14 with 256 experts (54.1B parameters, 62.1% JFT P@1 per Table 8), then distill into a Soft MoE H/14 with 32 or 64 experts, and compare against training the smaller model from scratch with the same total compute. If the distilled model substantially outperforms the from-scratch baseline, it suggests that the many-expert model captures useful specializations that can be transferred to a more memory-efficient architecture. A complementary approach is to investigate whether expert parameters can be shared or factorized — e.g., using a low-rank decomposition of the expert MLP weights with a shared basis and per-expert coefficients — to reduce memory while preserving the multi-expert inductive bias.

Practical Applications and Downstream Use Cases

Efficient high-accuracy image classification in resource-constrained inference settings. The paper's inference-time cost comparison (Figure 5, Table 1) provides a direct recipe for practitioners: replace the MLP blocks in the second half of a ViT with Soft MoE layers using 128-256 experts and one slot per expert, and train for extended durations with a long learning rate cooldown. A Soft MoE S/16 achieves 79.8% ImageNet 10-shot at 0.7 ms/img and 8.6 GFLOP/img — faster than a ViT B/16 (1.3 ms/img, 35.1 GFLOP/img) while substantially outperforming it (79.8% vs. 76.8%). This is directly actionable for applications like on-device photo classification, real-time video understanding, or large-scale image tagging where inference latency and energy consumption are primary constraints. The specific configuration numbers (128 experts, second-half placement, 500k-step cooldown) are provided in the paper and can be adopted without architecture search. The main caveat is memory: a 933M-parameter Soft MoE S/16 requires more device memory than a 33M-parameter ViT S/16, but at inference time this is a one-time model loading cost, not a per-image compute cost.

Frozen feature extraction for downstream transfer tasks. Section 4 (Table 3) demonstrates that Soft MoE vision encoders, when frozen and paired with a trained text tower, substantially outperform equivalently-sized dense ViTs on zero-shot transfer: +7.0% on ImageNet, +10.6% on CIFAR-100, +1.8% on Pet for S/16. For practitioners building multimodal systems — image search, content moderation, visual question answering — who want to use a frozen pretrained vision backbone, Soft MoE offers a drop-in replacement that yields better features at lower or equal inference cost. The specific improvement is largest at smaller backbone sizes (S/16, B/16), which is where inference cost is most constrained. A deployment scenario: a content moderation pipeline that screens millions of images per day using a frozen vision encoder paired with lightweight task-specific heads. Replacing ViT-S/16 with Soft MoE-S/16 (128 experts) would improve ImageNet zero-shot from 74.2% to 81.2% — a 7-point accuracy gain — while inference time increases only from 0.5 to 0.7 ms/img. If the task requires higher accuracy, Soft MoE-B/16 provides 82.5% at 1.5 ms/img, matching ViT-L/16's accuracy (82.7%) at 3.3x lower inference time.

Scaling model capacity under fixed inference latency budgets for large-scale batch processing. The paper's finding that expert count can scale to thousands without meaningful throughput reduction (Figure 6 bottom, Figure 23) enables a deployment strategy that is unavailable with Sparse MoEs: train a single model with very many experts, achieving high accuracy, and deploy it at the same inference speed as a much smaller dense model. For batch processing scenarios — generating embeddings for a billion-image dataset, nightly evaluation runs, or offline feature extraction for retrieval indices — the inference cost per image is the dominant economic factor, and model loading time or memory is amortized across millions of images. A Soft MoE H/14 with 256 experts achieves 62.1% JFT P@1 and 84.3% ImageNet 10-shot (Table 8) at 10.9 ms/img — comparable inference time to a dense ViT H/14 (8.6 ms/img) but with substantially higher accuracy. For a batch of 100M images, the total inference time increases by about 64 hours on a single TPUv3 (from 239 to 303 hours), while the accuracy gain may justify the cost depending on the downstream task. The key enabler is that Soft MoE's routing cost is constant with expert count, unlike Sparse MoEs where routing overhead eventually dominates (Figure 6 bottom: Tokens Choice throughput drops 8.1x from 8 to 4096 experts).

Data-efficient fine-tuning with many-expert models. The paper's long-training results (Table 8, Appendix F) show that Soft MoE models trained for extended durations continue to improve, particularly with long cooldowns. This suggests a deployment pattern where a large Soft MoE is pretrained once (amortizing the high training cost) and then fine-tuned on downstream tasks with relatively small datasets, benefiting from the pretrained expert specializations. The finetuning results at 384 resolution (Table 8: 88.5% for Soft MoE B/16 vs. 86.6% for ViT B/16) demonstrate that this transfer works for standard full-dataset fine-tuning. An open question for practitioners is whether the many-expert model also improves few-shot or parameter-efficient fine-tuning (e.g., linear probing, adapter tuning), where the expert specializations learned during pretraining might provide richer features for the downstream head to select from. The paper's 10-shot evaluation protocol addresses this partially, but a systematic comparison of Soft MoE vs. dense ViT under varying few-shot data regimes (1, 5, 10, 25, 100 shots) would provide actionable guidance for the common scenario where downstream labeled data is scarce.

When to Prefer This Method

The paper is explicit in positioning Soft MoE against both dense ViTs and Sparse MoEs (Tokens Choice and Experts Choice), and the experimental results provide clear guidance on the conditions under which Soft MoE is the preferred choice:

• Prefer Soft MoE over dense ViTs when you have the memory budget to hold additional expert parameters and want to maximize accuracy for a given inference FLOPs or latency budget. The crossover point where Soft MoE becomes worth the memory cost depends on the backbone size: at S/16, Soft MoE with 128 experts (933M params) outperforms ViT B/16 (108M params) at lower inference cost (0.7 vs. 1.3 ms/img), making it an unambiguous win. At H/14, Soft MoE with 256 experts (54.1B params) outperforms ViT H/14 (669M params) at similar inference cost (10.9 vs. 8.6 ms/img) but requires 81x more memory, which may be prohibitive without model parallelism.

• Prefer Soft MoE over Sparse MoEs (Tokens Choice and Experts Choice) in essentially all encoder settings. The Pareto frontier analysis (Figures 3, 19-22) shows Soft MoE dominates at every training budget and backbone size, with the gap largest at medium compute budgets and narrowing but persisting at the high end. The only potential exception is if memory constraints are extremely tight and the Sparse MoE can achieve reasonable accuracy with fewer total parameters — but even then, Soft MoE with fewer experts (and thus fewer parameters) still outperforms Sparse MoEs at matched parameter count in the paper's sweeps (Figure 6: Soft MoE with 8 experts achieves 66.6% IN/10shot at 1 slot-per-token, vs. 65.4% for Experts Choice and 64.7% for Tokens Choice).

• Prefer Soft MoE when deployment requires deterministic per-sequence computation. Sparse MoEs operating under capacity constraints typically batch multiple sequences together for routing, making the expert assignment for a given input depend on other inputs in the batch (Section 2.2: "the model is no longer deterministic at the sequence-level, but only at the batch-level"). Soft MoE's dispatch and combine weights are computed independently per sequence. For applications requiring reproducible outputs regardless of batch composition — serving individual user requests, debugging model behavior, or ensuring fairness across inputs — this determinism is a hard requirement that Sparse MoEs violate.

• Prefer Soft MoE when the expert count must scale to very large values (hundreds to thousands of experts). Beyond approximately 128-256 experts, both Tokens Choice and Experts Choice routing degrade in accuracy and throughput (Figure 6). Soft MoE shows monotonic accuracy improvement to at least 4096 experts with near-constant throughput. If the goal is to maximize total parameters under a fixed FLOPs budget, Soft MoE is essentially the only viable option among the methods evaluated.