Mixture of Experts (MoE) Intuition

Capacity, active arithmetic and resident memory

A sparse mixture-of-experts layer routes a token to a subset of feed-forward networks. The hospital analogy has a useful distinction: every specialist occupies an office, but only selected specialists work on this patient. It fails if it implies one clean semantic specialty per expert, equal quality to a much larger dense model, or identical wall-clock speed to a smaller dense model. Those are empirical questions.

A common decoder keeps attention and replaces selected feed-forward sublayers with a router and experts. Some layers remain dense; shared experts may run for every token. Total parameters, activated parameters, FLOPs, resident bytes and network traffic are distinct quantities. In particular, top_k/num_experts applies to the routed expert bank, not the whole model.

Prerequisites are feed-forward blocks, softmax, gradients, parameter counting and basic distributed execution. We implement sparse routing on CPU and a complete bounded training ablation; distributed dispatch and large-model serving remain unexecuted procedures.

Router and expert equations

For a token x of width d, a bias-free router produces E logits g=x W_gate, with W_gate of shape (d,E). Select top-k indices, normalize the selected logits and combine their expert outputs. A SwiGLU expert computes W_down(silu(W_gate_expert x) * W_up x), using three d-by-f matrices and 3df parameters without biases.

With logits [.2,-1.3,2.5,.1,1.8,-.5,0,.7], experts 2 and 4 win. Their weights are softmax(2.5,1.8)≈(.668188,.331812). If their outputs are [1,3] and [−1,1], the mixture is [.336376,2.336376]. Full softmax followed by truncation without renormalization is a different operator whose selected weights sum below one.

Hard top-k membership is discontinuous at a tie. Gradients flow through the selected logits/weights and selected expert computations, not through a smooth choice over expert identities. In particular, top-1 selected-softmax is identically one and gives no main-loss router gradient through that weight. Switch-style full-softmax selected gates have different gradient semantics. Name the operator before comparing training claims.

Complete sparse CPU implementation

import torch
from torch import nn
from torch.nn import functional as F
 
torch.manual_seed(77)
torch.set_num_threads(1)
class Expert(nn.Module):
    def __init__(self,d,f):
        super().__init__()
        self.gate=nn.Linear(d,f,bias=False)
        self.up=nn.Linear(d,f,bias=False)
        self.down=nn.Linear(f,d,bias=False)
    def forward(self,x):
        return self.down(F.silu(self.gate(x))*self.up(x))
 
class SparseMoE(nn.Module):
    def __init__(self,d=4,f=8,e=4,k=2):
        super().__init__()
        if any(type(v) is not int or v <= 0 for v in (d,f,e,k)) or k > e:
            raise ValueError('positive integer dimensions and 1 <= k <= e')
        self.router=nn.Linear(d,e,bias=False)
        self.experts=nn.ModuleList([Expert(d,f) for _ in range(e)])
        self.k=k
    def forward(self,x):
        if (x.ndim < 2 or x.shape[-1] != self.router.in_features or x.numel() == 0
                or not torch.isfinite(x).all()):
            raise ValueError('nonempty finite token tensor with configured width required')
        flat=x.reshape(-1,x.shape[-1])
        logits=self.router(flat)
        values,indices=logits.topk(self.k,dim=-1)
        weights=values.softmax(-1)
        out=torch.zeros_like(flat)
        for expert_id,expert in enumerate(self.experts):
            token,slot=torch.where(indices==expert_id)
            if len(token):
                contribution=expert(flat[token])*weights[token,slot,None]
                out=out.index_add(0,token,contribution)
        # Assignment fraction sums to one, not k.
        f=F.one_hot(indices,len(self.experts)).float().mean((0,1)).detach()
        p=logits.softmax(-1).mean(0)
        aux=len(self.experts)*(f*p).sum()
        return out.reshape_as(x),aux,indices,weights
 
m=SparseMoE().double()
x=torch.randn(2,3,4,dtype=torch.float64,requires_grad=True)
y,aux,indices,weights=m(x)
assert y.shape==x.shape and torch.allclose(weights.sum(-1),torch.ones_like(weights[...,0]))
# Independent dense-all-expert selection oracle; same operator, different dispatch.
flat=x.reshape(-1,4)
all_outputs=torch.stack([e(flat) for e in m.experts],dim=1)
selected=all_outputs.gather(1,indices[...,None].expand(-1,-1,4))
expected=(selected*weights[...,None]).sum(1).reshape_as(x)
torch.testing.assert_close(y,expected)
y.square().mean().backward()
assert torch.isfinite(x.grad).all() and m.router.weight.grad.norm()>0
assert sum(p.numel() for p in m.parameters())==4*4+4*3*4*8
print('sparse/dense routing parity, gradients and parameter count pass')

This CPU lesson uses nonempty finite floating-point tensors with the model's dtype/device and configured last-axis width; dimensions and k are positive integers, not booleans. Keep logits/activations in a moderate finite range: finite input does not prove intermediate arithmetic cannot overflow. At equal logits, topk tie order is not a portable expert-identity guarantee; test mixture/operator invariants or specify a tie policy, rather than asserting one arbitrary expert wins.

This loops over experts and uses index_add; it is an inspectable sparse implementation, not a GPU kernel-performance claim. A production implementation groups tokens by expert, performs grouped/block-sparse matrix operations and restores original token order. A scatter permutation error can produce plausible shapes with wrong outputs, hence the independent dense oracle.

Load balancing, capacity and a controlled ablation

If popular experts get more training examples, their improvement can reinforce router preference. This is a plausible positive feedback mechanism, not a theorem that every unregularized MoE collapses within 100 steps. Unselected experts can later receive traffic if routing changes; a temporarily unused expert is not necessarily an absorbing state.

Switch's top-1 auxiliary term uses E times the dot product of hard assignment fractions and average full-softmax probabilities. For top-k, declare normalization: here f_e=count_e/(N*k) sums to one, while P_e=mean softmax(logits)_e also sums to one. Uniform f and P give auxiliary value 1. If instead f counts tokens routed to each expert divided by N, it sums to k and the uniform reference is k. Hard f is treated as constant in the gradient. This objective encourages balance through a surrogate; uniformity is not a guaranteed unique global optimizer of the coupled routing system.

Capacity C per expert limits dispatched assignments. With eight tokens, top-2, four experts and capacity factor 1, C=ceil(16/4)=4. If all tokens select the same pair, each of those experts receives eight assignments and four overflow. Dropping them changes outputs and training; rerouting changes the routing operator; padding wastes compute. Dropless block-sparse implementations such as MegaBlocks address this tradeoff differently. Report pre-capacity demand, post-policy execution and dropped fraction separately.

The implementation above is dropless: the following separate first-come capacity fixture counts assignments, not distinct tokens, and does not silently modify SparseMoE. It admits the first C assignments for each expert in token order. On the skewed example it drops eight of sixteen assignments; the last four tokens lose both contributions. Partial survivors would retain their original weights in this drop-only policy, not be renormalized. A residual path, if present, is separate from the zero MoE branch.

def capacity_mask(routes, experts, capacity):
    if (type(experts) is not int or experts <= 0
            or type(capacity) is not int or capacity < 0):
        raise ValueError('positive experts and nonnegative integer capacity')
    width=len(routes[0]) if routes else 0
    counts=[0]*experts
    kept=[]
    for row in routes:
        if (not row or len(row)!=width or any(type(e) is not int or not 0<=e<experts for e in row)
                or len(set(row))!=len(row)):
            raise ValueError('fixed-width distinct valid expert IDs')
        flags=[]
        for e in row:
            flags.append(counts[e] < capacity)
            counts[e]+=1
        kept.append(flags)
    return kept,counts
 
keep,demand=capacity_mask([[0,1]]*8,4,4)
assert demand==[8,8,0,0]
assert keep==[[True,True]]*4+[[False,False]]*4
assert sum(map(sum,keep))==8  # 8 executed / 16 requested assignments

Auxiliary fractions above are pre-capacity demand. Recomputing them after dropping would hide overload: an expert could appear balanced only because its excess traffic vanished. With fixed detached f, the derivative with respect to token logit z[t,j] is (E/N) * p[t,j] * (f[j] - sum_i(f[i]*p[t,i])). This is a local derivative away from hard membership changes, not a gradient through argmax/top-k or a proof of a unique balanced optimum. The scalar uniform reference is 1 before multiplying by the coefficient, and changes to k under count/N normalization.

import copy
base=SparseMoE()
inputs=torch.randn(24,4)
target=inputs @ torch.tensor([[1.,0,0,0],[0,.5,0,0],[0,0,-.5,0],[0,0,0,1.]])
records=[]
for coefficient in (0.,.02):
    net=copy.deepcopy(base)
    opt=torch.optim.Adam(net.parameters(),lr=.01)
    history=[]
    for step in range(60):
        opt.zero_grad()
        prediction,balance,chosen,w=net(inputs)
        task=F.mse_loss(prediction,target)
        loss=task+coefficient*balance
        loss.backward(); opt.step()
        if step%20==0 or step==59:
            counts=torch.bincount(chosen.flatten(),minlength=4)
            shares=counts/counts.sum()
            entropy=-(shares[shares>0]*shares[shares>0].log()).sum().item()
            history.append((step,task.item(),balance.item(),counts.tolist(),entropy))
    assert all(torch.isfinite(p).all() for p in net.parameters())
    assert history[-1][1] < history[0][1]
    records.append((coefficient,history))
print('synthetic task ablations',records)

This is a fixed synthetic regression task, not language pretraining. It deliberately does not assert that coefficient zero collapses or that .02 wins. For a language experiment, insert the layer into a causal decoder, freeze identical data/split/initialization budgets, train multiple seeds, and plot task validation loss, expert demand, entropy, capacity overflow, router gradient norms and step time. Too strong a balance incentive can hurt task loss; balanced aggregate counts do not preclude specialization because different inputs can route to different experts while totals stay equal.

Selection-bias balancing uses adjusted scores to choose experts, then original affinity scores for mixture weights. Increase biases for underused experts and decrease for overused ones using a documented batch/window statistic and update rule. DeepSeek-V3's report retains a complementary sequence-wise balance loss: “auxiliary-loss-free” does not mean no auxiliary balance term anywhere. Do not substitute an invented EMA rule for its published method without labelling the variant.

Exact parameter and memory arithmetic

Let L sparse layers have routed expert size A, E experts, k selected, shared expert size S per layer, dense per-layer parameters D and outside-layer parameters B. Then total is B+L*(D+E*A+S+router) and an activated-parameter accounting convention is B+L*(D+k*A+S+router). If A/D already denote whole-model totals, do not multiply them by L again. Embedding lookup touches selected rows while output logits may use the full vocabulary matrix; “active parameters” is a reporting convention, not exact per-token FLOPs.

For a Mixtral-8×7B-shaped configuration, d=4096, f=14336, 32 query heads, 8 KV heads of width128, E=8, k=2, L=32 and vocabulary32000:

• Q and O each contain d² parameters; K and V each contain d*(8*128). Attention totals 41,943,040 per layer, not 4d².
• One SwiGLU expert has 176,160,768 parameters; the router adds 32,768.
• Two RMSNorm scales add 8192 per layer. With separate input/output vocabulary matrices add 2*32000*4096, then one final norm of4096.
• Total is 46,702,792,704; replacing E by k in the expert term gives 12,879,925,248 under this active-count convention. These explain the rounded 47B/13B labels, not equivalent quality to any dense model.

d,f,L,E,k,V=4096,14336,32,8,2,32000
attn=2*d*d+2*d*8*128
expert=3*d*f
outside=2*V*d+d
common=attn+d*E+2*d
total=outside+L*(common+E*expert)
active=outside+L*(common+k*expert)
assert total==46702792704 and active==12879925248
assert 3*64*128==24576
print('total/active',total,active,'ideal BF16 GB',total*2/1e9,
      'ideal 4bit GB',total/2/1e9)

Ideal BF16 weights are about93.406GB (86.99GiB); ideal four-bit payload is23.351GB before scales, metadata, unquantized tensors and runtime memory. The old “45GB INT4” confused bit widths. All experts must exist somewhere, but not necessarily all in GPU memory: sharding or offloading trades resident memory for communication and latency. Sizing a device from active count alone is still wrong.

For the small d64/f128/e8/k2 example, stored expert parameters are196608; selected experts contribute49152 and router512. Add one always-active expert: total experts221184 and active experts73728. Keep attention/embeddings separate. The earlier Kimi example's 104B active count cannot be derived by multiplying its whole total by16/896; selected routed experts, shared experts and dense components must all be counted under the model's declared convention.

Serving, parallelism and quantization procedure

Expert parallelism places expert weights on different devices. Dispatch token activations to destinations, compute expert outputs, return weighted contributions, restore token order. Traffic depends on token count, width, selected experts, dtype, layout and overlap. A rough all-remote forward dispatch+return budget is 2*N*k*d*bytes per MoE layer before protocol/replication details. Backward adds its own transfers. NVLink versus PCIe alone does not determine whether communication dominates: topology, collective implementation, grouping, skew and batch size matter.

For an authorized large-model experiment, pin model/runtime/kernel revisions and expert placement. Compare tensor-parallel, expert-parallel and hybrid layouts under identical prompt/output distributions and precision. Record router load, per-expert token counts, all-to-all bytes/time, grouped-GEMM shapes, memory high-water mark, TTFT, decode latency and goodput. Test a deliberately skewed workload and a balanced fixture. Replication and tensor parallelism are different: replicating the full model on each GPU is data parallelism, not “use tensor parallelism instead.”

Quantization may affect routing and experts differently. Preserve a calibration/evaluation set, compare gate decisions and task loss before/after, inspect per-expert outliers/scales and test rare domains. LoRA on attention can change hidden states reaching a frozen router; freezing router weights does not freeze assignments. Monitor utilization and quality after fine-tuning. No universal one-point/dense versus five-point/MoE perplexity penalty or guaranteed expert-specific cure is asserted.

Dense often makes the first capstone easier to debug. MoE may improve capacity per compute under some budgets, but there is no universal 10B/30B cutoff or rule that a homogeneous corpus cannot benefit. Under a fixed serving memory budget compare dense, MoE and a smaller dense model with more cache headroom using actual task quality and SLOs. Sparse decode may read fewer selected weight bytes; claiming it buys “nothing” whenever bandwidth matters is too strong.

Architecture survey with bounded claims

Mixtral 8×7B/8×22B are concrete sparse-decoder families, not eight independent complete7B/22B models. The8×7B paper documents top-2 of eight per layer and rounded47B/13B accounting. Its benchmark comparisons belong to that report, not all later workloads.

DeepSeek-V2 combines fine-grained routed/shared experts with Multi-head Latent Attention. MLA caches compressed latent state plus the necessary positional part; its published cache reduction has a specific comparator and does not transfer unchanged to every GQA model. V3 reports671B total/37B active,256 routed experts with eight selected and one shared in its sparse layers, plus MTP. Its attention/cache architecture is distinct from the expert bank. The report's2.788M H800 GPU-hours is an attributed training accounting boundary, not total organizational/R&D cost or a controlled comparison to Llama's hardware budget. Some operations use higher precision; “end-to-end everything FP8” is false.

MTP uses sequential additional prediction modules conditioned on prior information/next-token embeddings in training, not just an independent next-next-token head. It can support speculative proposal generation, but exact target verification and acceptance rate determine inference benefit. R1 changes the reasoning/post-training recipe on a V3-family base; it does not imply all R1 distilled models share that MoE architecture.

Qwen MoE, DBRX and Llama4 illustrate that expert granularity, k, shared paths and sparse-layer placement vary. Read the actual config/card before applying Mixtral's formula. Qwen3-235B-A22B's model identity is a total/active convention, not a promise of22B resident memory. DBRX uses16 experts/top4 in its published architecture. Llama4 Scout and Maverick differ in routed-expert counts; do not repeat the original confused “16–16 up to128 Behemoth” table. Native multimodal and interleaved positional choices are separate axes from MoE.

Scaling laws for fine-grained MoE explicitly include granularity, tokens and model size. The source does not establish a universal tokens≈20*sqrt(active*total) rule. Fit/evaluate a budget model within the studied regime; include routing/communication overhead in compute estimates. MegaBlocks shows a dropless block-sparse approach; vLLM/SGLang runtime options are version-specific integration choices, not guaranteed2–4× accelerators for any sparse model. No giant model or distributed run was executed here.

Recall and worked answers

1. What is traded? More stored capacity for sparse per-token expert computation, with routing/communication overhead. Mixtral's rounded total/active ratio is47/13, not8/2 for the whole network.
2. Sixteen experts, top2, expert7B, backbone3B,40layers? First resolve units. If7B and3B are whole-network aggregate expert/backbone sizes, total=3+167=115B and active=3+27=17B. If7B is per expert per layer and3B is outside the sparse layers, total=4483B and active=563B, before unspecified router/dense components. The original question is under-specified, so both interpretations must be named.
3. Why balance, and V3's alternative? Counteract concentrated load/unequal training and device hot spots. Selection biases adjust demand without entering combination weights; V3 also retains a sequence-wise auxiliary term. Collapse is a risk, not guaranteed fate without one named loss.
4. All-to-all bottleneck? Activations move to expert owners and back. Measure bytes, topology and overlap; neither PCIe disaster nor NVLink safety is a theorem.
5. When not MoE? When total state will not fit the permitted memory/transfer budget, when routing complexity prevents a verifiable baseline, or when a measured dense alternative already meets quality/SLOs at lower total cost. No parameter-count cutoff is universal.
6. How does MTP relate to speculation? Extra training objectives can improve representations and provide draft tokens. A verifier must accept/correct proposals according to target semantics; useful speed depends on acceptance and scheduling, not merely having a second head.

Tomorrow, reproduce the8-logit routing example and distinguish sum-one from sum-k load fractions. Next week, compare two synthetic routing histograms without inventing language-quality conclusions.

Sources and progression

Sparsely gated MoE, Switch, Mixtral, DeepSeek-V2, V3, R1, fine-grained scaling, MegaBlocks. Parameter and dispatch arithmetic are local checks, not reproduced paper benchmarks.

Previous: RAG · Next: long context

Optional reading

Explore these supplementary resources after the exercises. Video titles and availability have not been independently verified; use the primary sources above for technical claims.

• Mixtral 8x7B Explained
• Mixture of Experts from Scratch
• Mixtral of Experts (blog + paper)
• Switch Transformer