Skip to content
Residual Stream
Go back

Mixture of Experts: More Parameters, Same FLOPs

Updated:
Tan Le
ENVI

The Residual Stream found two thirds of a Transformer block’s weights in its MLP, and How LLM Architecture Evolved, 2017–2026 listed the mixture of experts as the substitution that let the largest models grow past what their compute alone would allow. This post opens it up. The idea fits in a sentence: replace the one MLP in each block with EE smaller MLPs, the experts, and let a small router send each token to only kk of them. A model with 256 experts and k=8k = 8 holds 32 times as many routed-expert parameters as it uses for any one token. Everything difficult about MoE is about making that work: the router has to spread tokens well, the experts have to divide the work usefully, and the hardware has to move tokens to experts that live on other GPUs.

One MLP becomes E experts; each token runs only k of them dense block attention MLP every token output MoE block, E = 8, k = 2 attention router: top-2 of 8 e1 e2 e3 e4 e5 e6 e7 e8 Σ weighted by the router output stores 1 MLP computes 1 MLP stores 8 experts computes 2 experts per token DeepSeek-V3: 256 routed experts, top-8 + 1 shared; 671B stored, 37B active

The whole idea of a mixture of experts: the block stores EE experts’ parameters but spends only kk experts’ compute on each token, and a router decides which kk. Schematic.

At a glance
  • A mixture of experts replaces the MLP with EE experts and a router that sends each token to kk of them: EE experts’ parameters, kk experts’ compute. DeepSeek-V3’s routed experts are 97% of its 671B parameters, and a token uses 37B.
  • Without balancing, the router collapses: one expert takes 86 to 100% of tokens in five of the eight layer-seed pairs. A balancing loss or DeepSeek-V3’s bias adjustment keeps the busiest expert at 14 to 17%, against a fair share of 12.5%.
  • Parameters need data. On this 600 KB blog the MoE overfit; on WikiText-2 the same top-1 model beat dense at equal compute by about 0.08 bits per byte.
  • At serving time every expert must sit in memory, tokens cross GPUs twice per layer (up to 168 KiB per token per layer for DeepSeek-V3), and a batch of 64 still gives each touched expert only 2.3 tokens.

Table of contents

Open Table of contents

The router

The router is a single linear layer: for a token’s vector xx it computes one score per expert, xWrx W_r, and a softmax turns the scores into probabilities p1,…,pEp_1, \dots, p_E. The token goes to the kk experts with the highest probabilities, and its MLP output is their outputs, weighted by those probabilities renormalized over the chosen kk:

MoE(x)=∑e∈top-kpe∑e′∈top-kpe′ Experte(x)\text{MoE}(x) = \sum_{e \in \text{top-}k} \frac{p_e}{\sum_{e' \in \text{top-}k} p_{e'}}\, \text{Expert}_e(x)

Worked, with 4 experts and k=2k = 2: a token’s router probabilities are (0.50,0.30,0.15,0.05)(0.50, 0.30, 0.15, 0.05). It goes to experts 1 and 2, with weights 0.50/0.80=0.6250.50 / 0.80 = 0.625 and 0.30/0.80=0.3750.30 / 0.80 = 0.375. Experts 3 and 4 never see it, and cost nothing for it. The gradient reaches the router through those weights: if expert 1’s output helped, the loss pushes p1p_1 up for tokens like this one. (With k=1k = 1 this renormalization would make every weight 1 and cut that path; Switch Transformers, and the code below, use the raw probability instead.)

token x router softmax(x Wr​) 0.50 expert 1 SwiGLU MLP × 0.625 0.30 expert 2 SwiGLU MLP × 0.375 0.15 expert 3 skipped 0.05 expert 4 skipped Σ output top-2 of 4: experts 3 and 4 cost nothing for this token; weights renormalized over the chosen two

The worked example. Top-2 of 4: the chosen experts’ weights are their probabilities renormalized over the two, and the skipped experts cost nothing for this token.

Each expert here is an ordinary MLP, SwiGLU in current models. The rest of the block, attention, the norms, the residual stream, is unchanged: MoE is a substitution for one sublayer. Per token, the MLP’s cost is kk experts’ worth; its parameters are EE experts’ worth. (Mixtures of experts date from Jacobs et al., 1991; Shazeer et al., 2017, introduced sparse top-kk gating at scale; Switch Transformers, Fedus et al., 2022, showed that even k=1k = 1 works. DeepSeek-V3 scores experts with a sigmoid rather than a softmax, normalized over the chosen ones.)

Keeping the load even

Left alone, a router collapses. Early in training a few experts happen to be slightly better; the router sends them more tokens; they get more gradient and become better still; the others starve. The model then has the parameters of EE experts and uses a few of them. Three remedies:

  • A balancing loss. Switch Transformers adds, for each MoE layer,
Lbal=E∑e=1Efe Pe\mathcal{L}_{\text{bal}} = E \sum_{e=1}^{E} f_e\, P_e

where fef_e is the fraction of the batch’s tokens routed to expert ee and PeP_e is the router’s average probability for expert ee. Both sum to 1 over experts. The loss is E⋅E⋅(1/E)2=1E \cdot E \cdot (1/E)^2 = 1 when both are uniform and EE when one expert takes everything. Only PeP_e has a gradient (fef_e comes from a top-kk, which doesn’t), and that gradient, EfeE f_e, is largest for the busiest experts, so the loss pushes probability away from them until the load evens out. It’s added to the language-model loss with a small weight, 0.01 here as in the Switch paper.

  • Bias balancing (Wang et al., 2024; used in DeepSeek-V3). Instead of a loss, keep a bias beb_e per expert, added to the scores only when choosing the top kk, never in the weights that combine the outputs. After every step, raise beb_e by a fixed small amount if expert ee got fewer than its share of tokens and lower it if it got more. Nothing is added to the loss, so nothing pulls the router away from what the language-model loss wants, and the experts’ outputs are weighted exactly as the router scored them. (DeepSeek-V3 does keep a tiny sequence-wise balancing loss, weight 0.0001, as a guard against extreme imbalance within one sequence.)
  • A capacity limit. Each expert accepts at most a fixed number of tokens per batch, the capacity, set as a capacity factor times its fair share kN/EkN/E; tokens beyond it are dropped from that expert and pass through the block on the residual stream alone (Lepikhin et al., 2021). This bounds the memory and compute of the busiest expert, which is what decides a step’s speed when experts run in parallel, at the price of some tokens skipping the MLP.

The widget below routes 256 tokens to 8 experts, with a router skewed toward the low-numbered experts, the way an untrained or collapsing router often is. Raise the skew to see the busiest expert take a large share; press the bias balancing step a few times to see the loads even out; set a capacity factor to see how many assignments would be dropped:

Routing 256 tokens to 8 experts

experts per token (k)

balancing steps: 0 · busiest expert: 29% (fair share 13%)

Bars: tokens sent to each expert. The bias changes only which experts are chosen, not how their outputs are weighted.

01234567

Bias balancing brings the busiest expert to within a few points of the fair share in 5 to 15 steps, depending on the skew, without changing how any chosen expert’s output is weighted; with a fixed step it then hovers around even rather than landing on it. A capacity factor of 1.0 drops whatever exceeds the fair share; systems that use one mostly train at 1.0 to 1.25 and evaluate at up to 2, and many newer ones, DeepSeek-V3 among them, drop no tokens in training.

Shared and fine-grained experts

DeepSeekMoE (Dai et al., 2024) changed two things about the experts themselves, and most large MoE models since follow it:

  • Fine-grained experts. Split each expert into smaller ones and route to more of them: instead of 8 experts with top-2, 32 experts a quarter the size with top-8, the same parameters and the same compute per token. The number of possible combinations of experts grows enormously, from 28 pairs to over 10 million sets of 8, which lets experts specialize more narrowly. (How much they specialize is debated: Mixtral’s authors found routing followed syntax and position more than topic.)
  • Shared experts. Add one or two experts that every token uses, for knowledge every token needs, so the routed experts don’t each have to relearn it. DeepSeek-V3 has 256 routed experts of width 2,048, top-8, plus one shared expert, in each of its 58 MoE layers.
8 experts, top-2 28 possible pairs 32 quarter-size experts, top-8, + 1 shared shared expert: every token split experts: same cost, 10M+ possible sets DeepSeekMoE’s two changes (illustrative sizes; DeepSeek-V3: 256 routed of width 2,048, top-8, 1 shared)

DeepSeekMoE’s two changes, at illustrative sizes: split experts finer and route to more of them, and add an expert that every token shares.

An MoE tiny GPT

Build a Tiny GPT, End to End, and Train It on This Blog’s tiny GPT, with its MLP replaced by 8 SwiGLU experts, each as wide as the dense MLP was: top-1 routing then costs the same compute per token as the dense model while holding 8 times its MLP parameters. The layer, with all three balancing methods:

import torch
import torch.nn.functional as F

# A mixture-of-experts MLP: E small SwiGLU experts and a router. Each token
# goes to its top-k experts, weighted by the router's softmax probabilities
# (renormalized over the chosen k). Balancing, one of:
#   "aux":  the Switch Transformer loss, E * sum_e f_e * P_e, where f_e is the
#           share of tokens routed to expert e and P_e its mean router
#           probability; minimized when both are uniform (Fedus et al., 2022);
#   "bias": DeepSeek-V3's auxiliary-loss-free method: a per-expert bias added
#           to the scores for *choosing* experts only, nudged after each step
#           toward under-used experts (Wang et al., 2024);
#   "none": nothing.
class MoE(torch.nn.Module):
    def __init__(self, D, F_, E, k, balance, bias_rate=1e-3, renorm_top1=False):
        super().__init__()
        self.E, self.k, self.balance, self.bias_rate = E, k, balance, bias_rate
        self.renorm_top1 = renorm_top1
        self.router = torch.nn.Linear(D, E, bias=False)
        self.gate = torch.nn.Parameter(torch.randn(E, D, F_) * 0.02)
        self.up = torch.nn.Parameter(torch.randn(E, D, F_) * 0.02)
        self.down = torch.nn.Parameter(torch.randn(E, F_, D) * 0.02)
        self.register_buffer("bias", torch.zeros(E))
        self.aux_loss, self.load = torch.tensor(0.0), torch.zeros(E)

    def forward(self, x):
        B, T, D = x.shape
        x = x.reshape(-1, D)
        probs = self.router(x).softmax(-1)                                # (tokens, E)
        choose_by = probs + self.bias if self.balance == "bias" else probs
        top = choose_by.topk(self.k, -1).indices                          # (tokens, k)
        weight = probs.gather(1, top)
        if self.k > 1 or self.renorm_top1:
            weight = weight / weight.sum(-1, keepdim=True)        # renormalize over the chosen k
        # With k = 1, renormalizing makes every weight exactly 1, so the
        # router would get no gradient from the language-model loss at all;
        # Switch Transformers uses the raw probability instead.
        out = torch.zeros_like(x)
        counts = torch.zeros(self.E, device=x.device)
        for e in range(self.E):                                           # run each expert on its tokens
            rows, slot = (top == e).nonzero(as_tuple=True)
            counts[e] = len(rows)
            if len(rows):
                h = x[rows]
                y = (F.silu(h @ self.gate[e]) * (h @ self.up[e])) @ self.down[e]
                out.index_add_(0, rows, y * weight[rows, slot, None])
        f = counts / counts.sum()
        self.load = f.detach()
        self.aux_loss = self.E * (f * probs.mean(0)).sum() if self.balance == "aux" else torch.tensor(0.0, device=x.device)
        if self.balance == "bias" and self.training:                      # raise the bias of under-loaded experts
            with torch.no_grad():
                self.bias += self.bias_rate * torch.sign(1 / self.E - f)
        return out.reshape(B, T, D)

Full script: moe.py

Two details in it are easy to get wrong. With k=1k = 1, renormalizing the chosen probability makes every weight exactly 1, and the router then gets no gradient from the language-model loss at all; Switch Transformers weights the expert’s output by the raw probability instead. And the experts’ weights are stacked into 3-D tensors, which matters below.

Trained four ways on the blog’s characters, two seeds each, 4,000 steps, scored on a held-out slice of the training posts (lower is better). train.py and gpt.py are the training script and model from Build a Tiny GPT, End to End, and Train It on This Blog, with experts, top_k and balance options that swap the MoE layer above in for the MLP. balancing.py runs train.py once per configuration:

# 8 experts on the blog's characters, top-1 with each balancing method and
# top-2 with the balancing loss, two seeds each, 4,000 steps on the dev slice.
BASE = dict(kind="char", steps=4000, dropout=0.2, checkpoints=[1000, 2000, 3000, 4000])
for seed in (0, 1):
    train(f"8 experts, top-1, balancing loss   s{seed}", **BASE, seed=seed, experts=8, top_k=1, balance="aux")
    train(f"8 experts, top-1, no balancing     s{seed}", **BASE, seed=seed, experts=8, top_k=1, balance="none")
    train(f"8 experts, top-1, bias balancing   s{seed}", **BASE, seed=seed, experts=8, top_k=1, balance="bias")
    train(f"8 experts, top-2, balancing loss   s{seed}", **BASE, seed=seed, experts=8, top_k=2, balance="aux")

Full script: balancing.py

8 experts, top-1, balancing loss   s0: 1000: 2.091 / 2000: 2.100 / 3000: 2.279 / 4000: 2.527 / busiest expert: 17% 15% 14% 15%  (61 passes)
8 experts, top-1, no balancing     s0: 1000: 2.146 / 2000: 2.059 / 3000: 2.059 / 4000: 2.102 / busiest expert: 61% 98% 99% 75%  (61 passes)
8 experts, top-1, bias balancing   s0: 1000: 2.110 / 2000: 2.120 / 3000: 2.311 / 4000: 2.561 / busiest expert: 15% 14% 15% 14%  (61 passes)
8 experts, top-2, balancing loss   s0: 1000: 2.050 / 2000: 2.116 / 3000: 2.315 / 4000: 2.567 / busiest expert: 15% 14% 15% 14%  (61 passes)
8 experts, top-1, balancing loss   s1: 1000: 2.106 / 2000: 2.109 / 3000: 2.303 / 4000: 2.549 / busiest expert: 14% 16% 16% 15%  (61 passes)
8 experts, top-1, no balancing     s1: 1000: 2.143 / 2000: 2.055 / 3000: 2.095 / 4000: 2.132 / busiest expert: 100% 86% 100% 41%  (61 passes)
8 experts, top-1, bias balancing   s1: 1000: 2.129 / 2000: 2.137 / 3000: 2.339 / 4000: 2.541 / busiest expert: 15% 15% 15% 14%  (61 passes)
8 experts, top-2, balancing loss   s1: 1000: 2.077 / 2000: 2.124 / 3000: 2.317 / 4000: 2.548 / busiest expert: 15% 15% 16% 15%  (61 passes)
busiest expert’s share of tokens, four layers (seed 0); fair share 12.5% balancing loss 17% 15% 14% 15% bias balancing 15% 14% 15% 14% no balancing 61% 98% 99% 75% layer 1 layer 2 layer 3 layer 4

Balancing, measured. Both methods keep every expert near its fair share; without either, the router sends most tokens to one expert, nearly all in some layers.

Data: the output above

For reference, the dense model with the same settings scored 2.14 and 2.15 at step 1,000 and 2.03 and 2.01 at step 4,000 (How LLM Architecture Evolved, 2017–2026).

  • Balancing works. With the balancing loss, the busiest expert takes 14 to 17% of routed assignments in every layer (measured on the final training batch), against a fair share of 12.5%. Bias balancing does the same with nothing added to the loss.
  • Without balancing, the router collapses. One expert takes 86 to 100% of tokens in five of the eight layers across the two seeds, and 41 to 75% in the rest: the model has 8 experts and mostly uses one.
  • On this data, more parameters overfit. The balanced models learn fastest, 2.05 to 2.13 at step 1,000 against the dense model’s 2.14 and 2.15, and then get worse, to 2.53 to 2.57 at step 4,000, while the dense model keeps improving to 2.01 to 2.03. The training text is 600 KB, and 4,000 steps read it 61 times, and each expert sees only about an eighth of the tokens. The collapsed models, which in effect use one expert, overfit least (2.10 to 2.13). Here collapse looks like a regularizer, which is the wrong lesson: it only means the extra parameters weren’t being used.

Overfitting is a claim about two numbers, so it is worth measuring both: the same runs again, on the CPU, scoring 60,000 characters of the text the models train on at each checkpoint as well (overfit.py):

# The dense model and the top-1 MoE again, seed 0, on the CPU, also scoring
# 60,000 characters of the training text at each checkpoint (track_fit).
BASE = dict(kind="char", steps=4000, dropout=0.2, checkpoints=[1000, 2000, 3000, 4000], seed=0, device="cpu",
            track_fit=True)
train("dense                              s0", **BASE)
train("8 experts, top-1, balancing loss   s0", **BASE, experts=8, top_k=1, balance="aux")

Full script: overfit.py

dense                              s0: 1000: 2.145 / 1000 on training text: 1.807 / 2000: 2.047 / 2000 on training text: 1.576 / 3000: 2.051 / 3000 on training text: 1.455 / 4000: 2.029 / 4000 on training text: 1.365  (61 passes)
8 experts, top-1, balancing loss   s0: 1000: 2.095 / 1000 on training text: 1.658 / 2000: 2.104 / 2000 on training text: 1.228 / 3000: 2.292 / 3000 on training text: 0.915 / 4000: 2.553 / 4000 on training text: 0.642 / busiest expert: 17% 15% 14% 15%  (61 passes)

The MoE model drives its training-text loss down to 0.64 bits per byte while its held-out loss climbs from 2.10 to 2.55: it is memorizing. The dense model memorizes too, more slowly (1.37 on training text), and its held-out loss is still falling at 4,000 steps. Switch Transformers met the same effect fine-tuning on small datasets, and countered it with heavier dropout inside the experts.

The counterfactual is data that doesn’t run out. The same dense and MoE models on WikiText-2’s 10.9 million characters, where 4,000 steps read the text about 3 times. wikitext2.py downloads the text, and wikitext2_runs.py trains on it:

Full script: wikitext2.py

Full script: wikitext2_runs.py

dense (WikiText-2)                 s0: 1000: 2.202 / 2000: 2.063 / 3000: 2.006 / 4000: 1.974  (3 passes)
8 experts, top-1, balancing loss   s0: 1000: 2.146 / 2000: 1.990 / 3000: 1.922 / 4000: 1.885 / busiest expert: 20% 18% 16% 14%  (3 passes)
8 experts, top-2, balancing loss   s0: 1000: 2.083 / 2000: 1.943 / 3000: 1.905 / 4000: 1.860 / busiest expert: 15% 14% 13% 15%  (3 passes)
dense (WikiText-2)                 s1: 1000: 2.217 / 2000: 2.088 / 3000: 2.009 / 4000: 1.979  (3 passes)
8 experts, top-1, balancing loss   s1: 1000: 2.154 / 2000: 2.015 / 3000: 1.937 / 4000: 1.898 / busiest expert: 19% 15% 15% 14%  (3 passes)
8 experts, top-2, balancing loss   s1: 1000: 2.095 / 2000: 1.966 / 3000: 1.900 / 4000: 1.862 / busiest expert: 14% 13% 13% 14%  (3 passes)
0.5 1 1.5 2 2.5 1k 2k 3k 4k steps bits per byte the blog, 600 KB (seed 0) MoE held-out MoE training 1.8 1.9 2 2.1 2.2 1k 2k 3k 4k steps WikiText-2, 10.9 MB (mean of 2) dense 8 experts, top-1 (same compute) top-2 dashed (left): training text

The same models on two corpora. On the small blog the MoE memorizes; with enough data, its extra parameters pay off at the same compute per token.

Data: the two runs above

With enough data the picture flips. Top-1 MoE, at the same compute per token as the dense model, ends at 1.885 and 1.898 against the dense model’s 1.974 and 1.979, a gain of about 0.08 to 0.09 bits per byte on both seeds, and it is ahead at every checkpoint. Top-2 gains more, 1.860 and 1.862, but it runs two experts per token, twice the MLP compute, so it isn’t an equal-cost comparison. (The scores are on a 120,000-character held-out slice of WikiText-2’s training text.) Same models, same code: whether the extra parameters help depends on whether there is enough data to fill them, the same trade How LLM Architecture Evolved, 2017–2026’s scaling laws make between parameters and tokens.

A bug worth knowing about

The first version of these runs scored the 8-expert model at 2.73 at step 1,000, far worse than the dense model’s 2.15, and it took a sanity check to see why. A one-expert MoE is mathematically the dense MLP, so it should match the dense model; it also scored 2.73 (and, once fixed, 2.14). Its forward pass and gradients matched the dense model’s exactly. The fault was in the optimizer setup, written for the dense model: weight matrices (2-D) went into the group with weight decay, norm gains (1-D) into the group without, and the experts’ stacked 3-D weights into neither. They were never updated. The fix is one character, p.dim() >= 2, and the lesson is an assertion: every parameter must be in exactly one optimizer group.

Counting real models

# Mixture-of-experts arithmetic from published configurations.
# Each expert is a SwiGLU MLP: 3 * D * F weights (gate, up, down).
def expert(D, F):
    return 3 * D * F

# Mixtral 8x7B: 32 layers, D = 4,096, 8 experts of F = 14,336, top-2.
D, L, F, E, k = 4096, 32, 14336, 8, 2
print(f"Mixtral 8x7B: one expert {expert(D, F) / 1e6:.0f}M weights; all experts, all layers "
      f"{expert(D, F) * E * L / 1e9:.1f}B; used per token {expert(D, F) * k * L / 1e9:.1f}B")

# DeepSeek-V3: 61 layers, the first 3 dense (F = 18,432); the other 58 have
# 256 routed experts of F = 2,048, top-8, plus 1 shared expert every token uses.
D, L, dense, F_dense, E, F, k, shared = 7168, 61, 3, 18432, 256, 2048, 8, 1
routed_all = expert(D, F) * E * (L - dense)
used = (expert(D, F) * (k + shared)) * (L - dense) + expert(D, F_dense) * dense
print(f"DeepSeek-V3: one routed expert {expert(D, F) / 1e6:.0f}M weights; routed experts, all layers "
      f"{routed_all / 1e9:.0f}B of the reported 671B; MLP weights used per token {used / 1e9:.1f}B of the reported 37B active")

# Expert parallelism: experts live on different GPUs, so every token's
# hidden vector travels to its k experts and the results travel back.
# DeepSeek-V3 dispatches in FP8 (1 byte) and combines in BF16 (2 bytes).
per_layer = k * D * (1 + 2)
print(f"all-to-all per token: {per_layer / 1024:.0f} KiB per MoE layer, {per_layer * (L - dense) / 2**20:.1f} MiB "
      f"across the {L - dense} MoE layers (at most; tokens whose experts share a GPU send less)")

# Decoding: a batch of B tokens, each picking k of E experts. How many
# experts does a step touch, and how many tokens does each one get? An
# expert that processes t tokens does 2 * t FLOPs per weight it reads.
for B in (1, 8, 64, 512, 4096):
    touched = E * (1 - (1 - k / E) ** B)
    print(f"batch {B:>5,}: about {touched:5.1f} of {E} experts touched per layer, "
          f"{B * k / touched:6.1f} tokens per touched expert")

Full script: moe_arithmetic.py

Mixtral 8x7B: one expert 176M weights; all experts, all layers 45.1B; used per token 11.3B
DeepSeek-V3: one routed expert 44M weights; routed experts, all layers 654B of the reported 671B; MLP weights used per token 24.2B of the reported 37B active
all-to-all per token: 168 KiB per MoE layer, 9.5 MiB across the 58 MoE layers (at most; tokens whose experts share a GPU send less)
batch     1: about   8.0 of 256 experts touched per layer,    1.0 tokens per touched expert
batch     8: about  57.4 of 256 experts touched per layer,    1.1 tokens per touched expert
batch    64: about 222.4 of 256 experts touched per layer,    2.3 tokens per touched expert
batch   512: about 256.0 of 256 experts touched per layer,   16.0 tokens per touched expert
batch 4,096: about 256.0 of 256 experts touched per layer,  128.0 tokens per touched expert

Recomputing the published sizes from the configurations is a good check on understanding them. Mixtral’s eight experts per layer hold 45B of its 47B parameters; a token uses two, 11B of MLP weights, which with attention and embeddings make its 13B active parameters. DeepSeek-V3’s routed experts hold 654B of its 671B parameters, 97%; a token uses 9 experts in each MoE layer, 24B of MLP weights, plus attention (multi-head latent attention, How LLM Architecture Evolved, 2017–2026) and embeddings for the reported 37B active. Per token it computes like a 37B model while storing a 671B one.

DeepSeek-V3: 671B parameters routed experts 654B (97%) per token: 37B active, about 5.5% of the whole computes like a 37B model, stores a 671B one

Where DeepSeek-V3’s parameters are, and how few of them one token uses.

Data: the counts above

What experts cost at serving time

The parameters a token doesn’t use still have to be somewhere, and moving tokens to them costs time. Three costs decide whether MoE wins in practice.

  • Memory holds every expert. A 671B-parameter model needs memory for 671B parameters, about 670 GB in 8-bit weights, whatever its 37B active parameters cost in compute. MoE trades compute for memory: it’s a good deal when memory is plentiful relative to compute, as it is in a large multi-GPU deployment, and a poor one on a single small GPU.
  • Experts live on different GPUs. With expert parallelism each GPU holds some of the experts, so every token’s vector must travel to the GPUs holding its kk experts, and the results must travel back: an all-to-all exchange in every MoE layer, both directions. The output above puts it at up to 168 KiB per token per layer for DeepSeek-V3, 9.5 MiB per token over a forward pass, which is why its training limited each token to experts on at most 4 nodes and overlapped this communication with computation.
  • Experts need a batch to be efficient. Decoding one token at a time, each of the 8 experts it touches processes that one token: every weight read from memory is used for 2 FLOPs, deep in the memory-bound region of How a GPU Runs Deep Learning. With a batch of 512 tokens, every expert is touched and each processes about 16 tokens per step; at 4,096, 128. Dense models face the same problem, but MoE spreads a batch across E/kE / k times as many weights, so it needs a batch E/kE / k times larger to use them as well. That makes MoE a better fit for high-throughput serving with large batches than for a single user’s low-latency chat.
1 10 100 1 8 64 512 4,096 tokens in the batch tokens per touched expert 8 of 256 used 57 of 256 used 222 of 256 used 256 of 256 used 256 of 256 used

Why MoE wants big batches. Until a batch is in the hundreds, each expert that’s read from memory does work for only one or two tokens.

Data: the output above

Where this leads

Takeaways

  • A mixture of experts replaces the MLP with EE experts and a router that sends each token to kk of them: EE experts’ parameters, kk experts’ compute per token.
  • Without balancing, a router collapses onto a few experts; the Switch balancing loss or DeepSeek-V3’s bias adjustment keeps loads even, the second almost without touching the loss.
  • Parameters need data. On a 600 KB corpus the 8-expert tiny GPT learned fastest and then overfit, ending worse than dense (2.53 to 2.57 against 2.01 to 2.03); on WikiText-2, read about 3 times, the same top-1 model beat dense at equal compute per token by about 0.08 bits per byte.
  • DeepSeek-V3 holds 671B parameters and uses 37B per token; its experts are 97% of its weights.
  • At serving time every expert must be in memory, tokens cross GPUs twice per MoE layer, and experts need large batches to be used efficiently: MoE suits high-throughput serving more than single-user latency.
  • Check that every parameter is being trained: a 3-D weight tensor slipped past an optimizer filter written for 2-D matrices, and the experts never learned.

References

Start with the Switch Transformers paper for the balancing loss and capacity, then DeepSeekMoE for fine-grained and shared experts. Hugging Face’s Mixture of Experts Explained is a readable tour of the history.


Written by

Tan Le

AI engineer. I write these notes to learn LLM internals and AI engineering in depth, and to keep what I learn where I can find it again.