Soft Mixture-of-Experts in Vision Transformers

Soft Mixture-of-Experts (Soft MoE) is an advanced neural network architecture that resolves the training instabilities of traditional sparse Mixture-of-Experts in Vision Transformers (ViTs). This article provides an overview of how Soft MoE replaces discrete token routing with a soft assignment mechanism. By computing weighted combinations of all input tokens to create fused "slots" for each expert, the model eliminates non-differentiable operations, prevents token dropping, and maintains full end-to-end differentiability during backpropagation.

The Challenge with Traditional Sparse MoE

In standard Mixture-of-Experts models, a routing network assigns individual image patch tokens to specific experts using discrete operations like top-k selection. Because choosing a discrete path is a non-differentiable step, it breaks the continuous flow of gradients required for standard backpropagation. Furthermore, traditional routing often leads to load-balancing issues where some experts are overwhelmed while others remain unutilized, forcing the model to drop tokens or rely on complex auxiliary loss functions to stabilize training.

Computing Weighted Token Combinations

Soft MoE completely removes discrete routing. Instead of sending individual tokens to experts, it provides each expert with a fixed number of "slots," where each slot is a weighted average of all input tokens in the sequence.

The computation follows a strict, continuous mathematical progression:

  1. Calculating Dispatch Weights: Let the input tokens from the Vision Transformer be represented as a matrix \(X\). The model introduces a set of learnable parameters, \(\Phi\), corresponding to the slots available across all experts. The model computes the dot product between the input tokens \(X\) and the slot parameters \(\Phi\) to generate a matrix of unnormalized logits.
  2. Applying Column-Wise Softmax: A softmax function is applied across the token dimension (columns) of the logit matrix. This produces the dispatch weights. Each column now represents a probability distribution over all input tokens for a specific slot.
  3. Generating Expert Slots: Each slot vector is computed by taking a linear combination of all input tokens, weighted by the dispatch weights. Mathematically, this is a simple matrix multiplication between the dispatch weights and the token matrix \(X\). Because every token contributes a fractional amount to the slot based on its weight, no tokens are ever dropped.

Processing and Recombining

Once the weighted token combinations (slots) are formed, they are distributed to their designated experts—typically standard Multi-Layer Perceptrons (MLPs). Each expert processes its assigned slots independently and in parallel.

After the experts generate their outputs, the information must be mapped back to the original token space so the Vision Transformer can continue processing the image representation:

  1. Calculating Combine Weights: The model uses the same initial logit matrix created from the interaction of \(X\) and \(\Phi\), but this time applies a softmax function across the slot dimension (rows). This yields the combine weights, representing a probability distribution over all slots for each original token.
  2. Reconstructing Tokens: The final output tokens are computed by taking a weighted average of all expert-processed slots, using the combine weights. This projects the transformed slot data back into the original sequence length of the input tokens.

Preserving Full Differentiability

Full differentiability is preserved in Soft MoE because the entire routing and recombination process relies exclusively on continuous, smooth mathematical operations: matrix multiplications and softmax activations.

In calculus and deep learning, both matrix multiplication and the softmax function have well-defined derivatives. When the Vision Transformer calculates its loss at the end of a forward pass, the gradients can flow backward through the combine weights, through the expert MLPs, and directly through the dispatch weights back to the input tokens and the learnable slot parameters \(\Phi\).

Because there are no step functions, argmax, or top-k thresholds blocking the gradient path, the routing parameters \(\Phi\) are updated smoothly via standard gradient descent alongside the rest of the network. This allows Vision Transformers to scale to massive capacities with high computational efficiency and zero routing instability.