What a mixture of experts really is

Setup#

pip install torch

Written against PyTorch 2.5.1, Python 3.10. CPU, instant.

A mixture-of-experts layer, complete#

import torch, torch.nn as nn, torch.nn.functional as F

torch.manual_seed(0)
d_model, d_ff, n_exp, top_k, n_tok = 16, 32, 8, 2, 6

class MoE(nn.Module):
    def __init__(self):
        super().__init__()
        self.router = nn.Linear(d_model, n_exp, bias=False)
        self.experts = nn.ModuleList(
            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
            for _ in range(n_exp))

    def forward(self, x):
        logits = self.router(x)                             # one score per expert per token
        probs = F.softmax(logits, dim=-1)
        weight, idx = probs.topk(top_k, dim=-1)             # keep the best few only
        weight = weight / weight.sum(-1, keepdim=True)      # renormalise the survivors
        out = torch.zeros_like(x)
        for slot in range(top_k):
            for e in range(n_exp):
                hit = idx[:, slot] == e
                if hit.any():                               # this expert runs only if chosen
                    out[hit] += weight[hit, slot, None] * self.experts[e](x[hit])
        return out, idx, weight

moe = MoE()
x = torch.randn(n_tok, d_model)
out, idx, weight = moe(x)

print("token -> experts chosen (and their weights)")
for t in range(n_tok):
    picks = ", ".join(f"E{idx[t,s].item()} {weight[t,s]:.2f}" for s in range(top_k))
    print(f"  token {t}: {picks}")

one_expert = sum(p.numel() for p in moe.experts[0].parameters())
total = sum(p.numel() for p in moe.parameters())
active = one_expert * top_k + moe.router.weight.numel()
print(f"\nparameters in one expert     : {one_expert:,}")
print(f"total parameters in the layer: {total:,}")
print(f"used for any single token    : {active:,}   ({active/total:.1%})")

print("\nsame layer, more experts, same work per token")
print(f"{'experts':>8}{'total params':>15}{'active params':>15}{'fraction':>10}")
for n in (8, 64, 256, 512):
    tot = n * one_expert + d_model * n
    act = top_k * one_expert + d_model * n
    print(f"{n:>8}{tot:>15,}{act:>15,}{act/tot:>10.2%}")

print("\nreal models (from each config.json on the Hub)")
print(f"{'model':<24}{'experts':>8}{'top-k':>7}{'shared':>8}")
for nm, e, k, s in [("Mixtral-8x7B-v0.1", 8, 2, 0), ("DeepSeek-V3", 256, 8, 1),
                    ("Qwen3-30B-A3B", 128, 8, 0), ("Qwen3-Next-80B-A3B", 512, 10, 1),
                    ("gpt-oss-20b", 32, 4, 0)]:
    print(f"{nm:<24}{e:>8}{k:>7}{s:>8}")

token -> experts chosen (and their weights)
  token 0: E5 0.53, E7 0.47
  token 1: E7 0.60, E0 0.40
  token 2: E3 0.66, E5 0.34
  token 3: E7 0.58, E0 0.42
  token 4: E6 0.57, E7 0.43
  token 5: E6 0.60, E4 0.40

parameters in one expert     : 1,072
total parameters in the layer: 8,704
used for any single token    : 2,272   (26.1%)

same layer, more experts, same work per token
 experts   total params  active params  fraction
       8          8,704          2,272    26.10%
      64         69,632          3,168     4.55%
     256        278,528          6,240     2.24%
     512        557,056         10,336     1.86%

real models (from each config.json on the Hub)
model                    experts  top-k  shared
Mixtral-8x7B-v0.1              8      2       0
DeepSeek-V3                  256      8       1
Qwen3-30B-A3B                128      8       0
Qwen3-Next-80B-A3B           512     10       1
gpt-oss-20b                   32      4       0

Reading the output#

Different tokens went to different experts. Token 2 chose E3 and E5; token 5 chose E6 and E4. That divergence, on an untrained router with random weights, is only routing noise. In a trained model it is the mechanism the whole design rests on.

Expert 7 appears in four of six tokens. Experts 1 and 2 appear in none. With a random router this is chance. In training, this exact pattern becomes self-reinforcing and is the failure mode of the entire technique — see load balancing and expert collapse.

The scaling table is the whole argument. From 8 experts to 512, total parameters grow 64x while active parameters grow 4.5x — and that growth is entirely the router, which must produce one score per expert. The expert compute per token is fixed at top_k experts, forever.

Real models sit far along that curve. Qwen3-Next-80B-A3B activates 10 of 512 experts. Naming a model "80B-A3B" is the industry's shorthand: 80 billion total, 3 billion active.

The loop is teaching code, not production code#

for slot ... for e ... runs top_k * n_exp iterations with a boolean mask each time. It is readable and it is slow. Production kernels sort tokens by expert assignment once and then run one grouped matrix multiply per expert.

transformers uses a MixtralTopKRouter that computes softmax then topk then renormalises, exactly as above, and hands indices to a fused MixtralExperts module. The routing logic is what you see here; the dispatch is not.

What is and is not replaced#

An MoE transformer layer is:

x -> attention (unchanged, still GQA or MLA) -> MoE feed-forward -> out

Attention stays dense. Everything in Fast Attention and Long Context applies unchanged to an MoE model. Some models keep the first few layers dense: DeepSeek-V3's first_k_dense_replace: 3 means layers 0 to 2 use an ordinary feed-forward block, because routing is unstable in the earliest layers.

Common mistakes#

Expecting a memory saving. Every expert is resident. A 671B-parameter model needs 671B parameters' worth of memory whatever its active count. MoE trades memory for compute, in that direction only.

Expecting interpretable specialisation. Published analyses find experts specialising by token identity, syntax and position far more than by human topic. Do not promise a "legal expert".

Forgetting to renormalise the top-k weights. Without weight / weight.sum(), output magnitude varies with how confident the router was, and training becomes unstable.

Comparing an MoE and a dense model by parameter count. Mixtral-8x7B has around 47B parameters and about 13B active. It behaves like neither a 47B nor a 13B dense model. Compare on measured quality and measured cost, never on the headline number.

Assuming any training framework handles it. MoE needs balancing losses, capacity handling and expert-parallel communication. Naive data-parallel training of the code above will run, and will collapse onto a few experts.

Try it yourself#

Set top_k = 8 with n_exp = 8. Every expert now runs for every token, which is a dense layer with extra steps. Compare the active-parameter fraction. Then set top_k = 1 and see how much the output changes when one router decision flips.

What to learn next#

• The router — how the choice is made, and why it is hard to train.
• Active parameters vs total parameters — reading model names honestly.
• Transformers — the architecture an MoE layer plugs into.