L9.1 · LESSON 1 OF 6

Mixture of experts

Explain expert routing, long-context costs, state-space models and multimodal token streams under explicit assumptions.

THE IDEA

Many experts, few used

In a mixture-of-experts (MoE) layer the single MLP is replaced by several parallel MLPs, the experts, plus a small router. For each token the router scores the experts, the top k are chosen, and the layer output is the weighted sum of just those experts' outputs.

So the model can hold many more parameters than it uses for any one token. With 2 billion shared parameters and eight 1-billion experts the total is 10 billion, but with two experts per token only about 4 billion are on that token's path. The unused experts still have to be stored or loaded.

shared parameters (billions)2experts8parameters per expert (billions)1experts used per token (k)2

total = 2 + 8 × 1 = 10 billion (all must be stored or loaded) active per token = 2 + 2 × 1 = 4 billion (what that token's computation uses) active share = 40%

The router scores each expert for a token, keeps the top k and mixes only those experts' outputs with softmax weights.

router score, expert 12router score, expert 21router score, expert 30router score, expert 4-1

top-2 experts: 1, 2 weights 0.731, 0.269 the other experts do no work for this token, but their weights still have to exist

Change the expert counts and k. Then look at the router.

Predict first. Does activating two of eight experts mean storing only two?

Chapter context

New architectures alter where computation and memory are spent. Compare them by tracing one request through their actual operations rather than assuming every improvement reduces every cost.

TRY IT YOURSELF

Part 1. Do it hands-on

Count total and active parameters and compute router weights by hand, then write a routing layer.

shared parameters (billions)2experts8parameters per expert (billions)1experts used per token (k)2

total = 2 + 8 × 1 = 10 billion (all must be stored or loaded) active per token = 2 + 2 × 1 = 4 billion (what that token's computation uses) active share = 40%

The router scores each expert for a token, keeps the top k and mixes only those experts' outputs with softmax weights.

router score, expert 12router score, expert 21router score, expert 30router score, expert 4-1

top-2 experts: 1, 2 weights 0.731, 0.269 the other experts do no work for this token, but their weights still have to exist

total (2 shared + 8 × 1) active with 2 experts

weights for scores [2,1,0,−1]:

To do: 2 billion shared parameters plus eight experts of 1 billion each. Total parameters (in billions)?
To do: Each token is routed to 2 of the 8 experts. Parameters on a token's path (billions)?
To do: The router gives four experts the scores [2, 1, 0, −1]. It keeps the top 2 and takes a softmax over just those two scores. Type their two weights (3 decimals)

Part 2. Do it in code

moe_params(shared, expert, n_experts): shared + n_experts · expert.
active_params(shared, expert, k): shared + k · expert. k must be between 1 and n_experts when n_experts is given as a fourth argument (ValueError otherwise).
route(logits, k): the indices of the k largest router scores (ties: lower index first) and the softmax weights over just those k scores. Return (indices list, weights array). k must be at least 1 (ValueError).
moe_layer(x, experts, router_W, k): x is a vector, experts is a list of matrices (each maps x to a vector of the same size), router_W has one row per expert. Scores = router_W @ x. Output = Σ weight_i · (experts[i] @ x) over the chosen experts.

import numpy as np

def moe_params(shared, expert, n_experts):
    pass

def active_params(shared, expert, k, n_experts=None):
    pass

def route(logits, k):
    pass

def moe_layer(x, experts, router_W, k):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Does activating two of eight experts mean storing only two?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a toy MoE with shared weights S, E equal experts of P weights and k experts active per token:

N_{total}=S+EP,\quad N_{active}\approx S+kP.

This is a parameter-accounting model, not a full FLOP or communication model. Naively materialized attention scores use $T^2$ entries per head.

Worked example

With 2 billion shared parameters and eight experts of 1 billion each, total parameters are 10 billion. If two experts are selected per token, approximately 4 billion parameters participate along that token's path. The other experts still require storage or loading. A 4096×4096 naive score matrix contains 16,777,216 entries per head.

Reference code

The following example uses the Python standard library and does not call a remote service.

shared, experts, per_expert, active=2,8,1,2
assert shared+experts*per_expert == 10
assert shared+active*per_expert == 4
assert 4096**2 == 16777216

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Active parameters and resident parameters differ. Linear sequence complexity does not make every model property linear. Context-window capacity and demonstrated context use are separate claims.