Mixture of Experts - Mathematical Foundations and Scaling

For a long time, making a language model larger came with a blunt bargain: add more parameters, then pay to use all of them for every token. That works, but it gets expensive quickly. A sentence about cooking and a function written in Rust both travel through the same full stack of feed-forward weights, whether every part of that capacity is useful or not.

Mixture of Experts changes that bargain. Instead of one feed-forward network doing every job, a layer keeps a collection of them and learns which few should handle each token. The name “expert” sounds more orderly than the reality—one expert does not necessarily become the poetry department while another learns Python—but the division of work can still be useful.

This is what makes MoE attractive: the model can hold far more parameters without using every expert for every token. It is also what makes MoE difficult. Tokens have to be routed, popular experts become crowded, quiet experts learn slowly, and distributed hardware spends time shuffling activations around. The mathematics is fairly compact; making the system behave well is the interesting part.

In a dense transformer, every token passes through the same feed-forward network in a layer. A sparse MoE layer keeps several feed-forward networks and routes each token to only a small subset. Adding experts increases total parameter capacity, while the expert computation per token can stay roughly fixed when the selected count stays fixed. Capacity grows roughly linearly with the number of experts, not quadratically.

From a Dense Feed-Forward Layer to MoE

A transformer feed-forward block applies a function such as:

\operatorname{FFN}(x)=W_2\,\phi(W_1x+b_1)+b_2

An MoE layer stores $M$ expert functions $E_1,\ldots,E_M$ and a router that produces one score per expert:

r(x)=W_rx

g(x)=\operatorname{softmax}(r(x))

For top- $k$ routing, let $S_k(x)$ be the selected expert indices. The selected weights are commonly renormalized:

\tilde g_i(x)= \frac{g_i(x)}{\sum_{j\in S_k(x)}g_j(x)}, \qquad i\in S_k(x)

The layer output is:

y=\sum_{i\in S_k(x)}\tilde g_i(x)E_i(x)

Architectures differ. Some use top-1 routing, some top-2 or a larger $k$ , and some do not renormalize in exactly this way. The implementation needs to state its convention because it changes both output scale and gradients.

What “Sparse” Means Here

Sparse MoE does not mean most individual expert weights are zero. It means each token activates only $k$ of the $M$ experts.

If every expert has $P_e$ parameters, the expert portion has roughly $MP_e$ total parameters but uses roughly $kP_e$ of them for one token. Shared attention, embeddings, normalization, and routing parameters remain active as well.

This is why active parameter count is not the same as total parameter count or FLOPs. It is also not a complete predictor of quality: the total capacity, training data, routing, shared layers, and optimization all contribute.

Expert Capacity

Tokens in a batch do not arrive evenly at every expert. Distributed implementations therefore reserve a finite number of token slots per expert.

For $N$ routed tokens, top- $k$ selection, $M$ experts, and capacity factor $C$ , a common capacity is:

\operatorname{capacity} =\left\lceil C\frac{kN}{M}\right\rceil

The factor $k$ matters because top-2 routing produces twice as many token-to-expert assignments as top-1 routing.

When an expert is full, an implementation may reroute the token, drop that expert assignment, or let the residual path carry the token without an expert update. Setting the token representation to an all-zero vector is not a universal rule.

A larger capacity factor wastes memory and compute on empty slots. A smaller one risks overflow. The right value depends on batch shape, routing balance, and the collective communication implementation.

Load Balancing

Without an auxiliary objective, a router may send most tokens to a few experts. Those experts overflow while others sit idle.

There are two related but different ideas:

Per-token routing confidence: Is one token’s router distribution sharp or uniform?
Marginal expert load: Across a batch, how much traffic does each expert receive?

The entropy of one token’s router distribution is:

H(g(x))=-\sum_{i=1}^{M}g_i(x)\log g_i(x)

$H\approx\log M$ means the router is nearly uniform and uncertain for that token. It does not mean confident specialization. Good batch-level balance can coexist with confident per-token choices if different tokens choose different experts.

Switch Transformer uses an auxiliary loss based on the fraction of tokens dispatched to each expert and the average router probability assigned to it. Other systems use entropy, variance, or differentiable balancing terms. These losses encourage usable traffic distribution, but too much pressure can prevent meaningful specialization.

A capacity example

Suppose a batch contains 4,096 routed tokens, the layer has 8 experts, routing uses top-2, and the capacity factor is 1.25. There are 8,192 token-to-expert assignments, so the mean load is 1,024 assignments per expert. The reserved capacity is:

\left\lceil1.25\frac{2\cdot4096}{8}\right\rceil=1280.

If one expert receives 1,500 assignments, a fixed-capacity implementation has 220 assignments it cannot place there even though other experts may have empty slots. The global average looked safe; the tail of the load distribution caused the failure.

This is why mean utilization is a weak routing metric. Report the maximum load, percentiles across experts, overflow or dropped-assignment rate, and how those values change by token domain and training step.

Router diagnostics in a few lines

The following function does not implement a distributed MoE layer. It checks the part that should be inspected before dispatch: which experts were selected and whether a fixed capacity would overflow.

import math

import torch


def routing_report(logits: torch.Tensor, top_k: int, capacity_factor: float):
    """logits: [tokens, experts]"""
    probabilities = logits.softmax(dim=-1)
    weights, expert_ids = probabilities.topk(top_k, dim=-1)

    tokens, experts = logits.shape
    capacity = math.ceil(capacity_factor * top_k * tokens / experts)
    counts = torch.bincount(expert_ids.flatten(), minlength=experts)
    overflow = (counts - capacity).clamp_min(0)

    return {
        "capacity": capacity,
        "counts": counts,
        "overflow_assignments": int(overflow.sum()),
        "max_to_mean_load": float(counts.max() / counts.float().mean()),
        "mean_router_probability": probabilities.mean(dim=0),
        "selected_weights": weights,
    }

max_to_mean_load exposes a hot expert that an average hides. The mean router probability and the hard assignment count should both be inspected because a router can have balanced probability mass while top- $k$ choices remain uneven.

Router Gradients

Top- $k$ selection is discontinuous when the membership of the selected set changes. Within a region where the same experts remain selected, gradients flow through their gate values and expert outputs. The discrete choice itself is usually handled by the top- $k$ operation plus auxiliary router losses.

This is not generally the classic straight-through estimator used for binary or quantized variables. Some research routers do use estimators or relaxations, but that is an architectural choice rather than a defining property of MoE.

Expert Parallelism and Communication

MoE layers are often distributed with expert parallelism. Tokens are grouped by destination expert, exchanged across devices with an all-to-all collective, processed, and then returned to their original sequence positions.

There is no rule that one expert must live on one GPU. A device may hold several experts, one expert may be sharded, and expert parallelism may be combined with data, tensor, pipeline, or sequence parallelism.

Communication cost depends on token count, hidden width, selected experts, placement, topology, precision, and implementation. A single neat formula cannot capture network contention, padding, overlap, or hierarchical collectives. In practice, MoE can be compute-efficient and still be slower than expected because tokens spend time moving between devices.

The dispatch path is worth spelling out:

compute router logits locally;
group token representations by destination expert;
exchange them with an all-to-all collective;
run grouped expert matrix multiplications;
exchange expert outputs back;
restore original token order and combine selected outputs.

The payload scales with routed assignments and hidden width, not with the number of expert parameters. Small token counts create tiny expert batches that underuse GPUs; large counts improve matrix efficiency but increase activation traffic and latency. Placement should follow the physical topology: crossing a slow inter-node link is not equivalent to moving tokens over an on-package fabric.

Systems such as MegaBlocks avoid the usual choice between padding every expert to a fixed capacity and dropping excess tokens. They express the dynamic work as block-sparse operations. That is a kernel and layout solution to irregular routing; it does not make skew or communication irrelevant.

Top-1 and Top-2 Routing

Top-1

Each token uses one expert.

Advantages: lower expert computation and less communication.

Costs: greater sensitivity to routing mistakes and expert overflow.

Top-2

Each token combines two experts.

Advantages: a smoother mixture and a backup signal from another expert.

Costs: approximately twice as many expert assignments and more communication.

Whether the second expert is worth the cost is an empirical deployment decision.

Shared and Routed Experts

Some designs keep one or more shared experts active for every token and use routed experts for additional capacity. Shared experts can learn broadly useful transformations while routed experts handle more conditional patterns.

This does not prove that an expert corresponds to a human-readable topic. Experts may specialize by syntax, token position, language, frequency, or a distributed feature that is difficult to name. Router visualizations are evidence about traffic, not a full explanation of the model’s reasoning.

Balancing Without a Large Auxiliary Loss

An auxiliary loss is not the only way to manage traffic. DeepSeek-V3 describes an auxiliary-loss-free strategy that adjusts per-expert bias terms used during routing while leaving the model’s main routing scores to carry the learned affinity. The aim is to correct persistent load imbalance without forcing a large balancing term into the training objective.

That should not be reduced to “auxiliary losses are obsolete.” Bias updates introduce another control loop, and the result depends on its update rule, granularity, and interaction with the router. A useful comparison keeps model, data, capacity, and hardware fixed, then reports quality, overflow, load variance, and end-to-end step time for both strategies.

Fine-Grained Experts

DeepSeekMoE divides the feed-forward capacity into smaller experts and activates several of them per token. The published 16.4-billion-parameter model activates about 2.8 billion parameters for each token. It is not a 1.3-trillion-parameter model, and its paper is arXiv:2401.06066.

Finer experts give the router more combinations but also increase routing and systems complexity. Parameter counts should always state whether they are total, active, or limited to the expert layers.

Soft MoE Is a Different Routing Pattern

Soft MoE does not choose experts with a sigmoid threshold. It learns soft assignments between input tokens and a fixed set of expert slots. Tokens are combined into slot inputs, experts process those slots, and another set of soft weights combines the outputs back to token positions.

This keeps the dispatch operation differentiable and avoids hard token dropping, but expert slots still have fixed computational cost. It is a distinct design, not simply top- $k$ routing with the corners rounded off.

MoE Beyond Feed-Forward Experts

The phrase “Mixture of Experts” is sometimes applied to adapters, convolutional branches, multimodal modules, or attention-related components. Sparse attention is not automatically MoE. Routing Transformer, for example, sparsifies attention by clustering queries and keys; it should not be described as ordinary expert routing inside attention heads.

Stability Problems

MoE training adds several failure modes:

Expert collapse: most traffic goes to a small subset;
overflow: popular experts exceed capacity;
router instability: small logit changes repeatedly switch the selected experts;
expert under-training: rarely selected experts receive too little data;
communication stalls: all-to-all exchange dominates layer time;
numerical sensitivity: router logits can be fragile in low precision.

Common responses include load-balancing losses, router-logit regularization, adequate batch diversity, careful capacity factors, routing in higher precision, expert dropout, and monitoring the actual token distribution.

What to Measure

A useful MoE evaluation reports more than validation loss:

total and active parameters;
expert FLOPs per token;
end-to-end latency and throughput;
all-to-all communication time;
capacity utilization and overflow rate;
marginal load per expert;
routing stability across batches;
quality on rare domains and languages;
memory per device.

An MoE can have attractive theoretical FLOPs and poor wall-clock performance. Conversely, a carefully placed and fused implementation can use much larger capacity without a proportional increase in per-token computation.

Scaling Without the Mythology

MoE provides conditional computation. It does not show that routing has become program synthesis, that experts form a set of independent minds, or that specialization is automatically modular cognition. Those are research hypotheses, not consequences of the routing equation.

The defensible claim is narrower and still useful: sparse routing can increase total learned capacity faster than it increases expert computation per token. Whether that capacity improves a model depends on training, data, architecture, and the system that keeps the experts fed.

References

Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017).
Fedus, Zoph, and Shazeer, Switch Transformers (2021).
Du et al., GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (2021).
Puigcerver et al., From Sparse to Soft Mixtures of Experts (2023).
Dai et al., DeepSeekMoE (2024).
DeepSeek-AI, DeepSeek-V3 Technical Report (2024).
Gale et al., MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (2022).