insideout Get Pro

15 · Mixture of experts

Pro

Transformers, lesson 15

Mixture of experts

A router sends each token to a few experts, so most weights sit idle.

A transformer block runs every token through the same MLP, so a bigger model means more work on every token. A mixture of experts stores many MLPs and runs only a few of them for each token. You need the four parts of a block from "Norm, attention, MLP, residual", and softmax.

By the end, you can:

  1. Send one token through a mixture-of-experts layer by hand.
  2. Count an MoE's total and active parameters, and tell which sets memory and which sets compute.
  3. Work out a load-balancing loss, and say what capacity does to an overloaded expert.
  4. Count the expert mixes that finer experts allow, and say what experts end up specialising in.

Part 1. Many MLPs, a few per token

Send one token through a mixture-of-experts layer by hand.

In GPT-2 small, each block's MLP holds 4.7 million weights, twice what attention holds, and every token runs through all of them. To add weights without adding work per token, keep several MLPs side by side and let each token use a few.

Example 1.1

One token, four experts

A layer has four MLPs, called experts, and keeps the best 2 for each token. A small router scores them 1.0, 3.0, 0.0 and 2.0 for this token. Expert 2 would output (2, 1) and expert 4 would output (−1, 3). What does the layer send out?

  1. Softmax the four scores(0.09, 0.64, 0.03, 0.24)

    e^1 = 2.718, e^3 = 20.09, e^0 = 1 and e^2 = 7.389, out of 31.19.

  2. Keep the top 2E2: 0.64, E4: 0.24

    Experts 1 and 3 never run for this token.

  3. Renormalise0.64 ÷ 0.88 = 0.73

    The kept weights add up to 0.88. Divide both by it: 0.73 and 0.27.

  4. Blend the two outputs(1.46 − 0.27, 0.73 + 0.81)

    Each output times its weight: (1.46, 0.73) and (−0.27, 0.81).

Answer(1.19, 1.54), from 2 of the 4 experts.

Experts 1 and 3 got a weight of 0, so their weights stay in memory, unread. The router is one small grid, a column per expert, so scoring is cheap. Because the blend uses the weights, training adjusts them, and the router learns which experts help which tokens.

token xscoreweight1.003.00.730.002.00.27router keeps the top 2E1E2E3E4(2, 1)(−1, 3)E1 and E3 never run0.73 × (2, 1) + 0.27 × (−1, 3)= (1.19, 1.54)
Figure 1The router scores all four experts, but only the top 2 run. Their outputs are blended by the renormalised weights, 0.73 and 0.27.

Mixture of experts (MoE)

Several MLPs, the experts, plus a router that scores them for each token. Only the top k run, blended by their weights renormalised over the kept ones. It takes the place of the block's MLP.

MoE(x) = sum of gᵢ × Eᵢ(x), over the top k

Only the MLP is copied: attention, the norms and the residual stream stay one per block. Each token, in each block, makes its own choice.

Your turn

Another token scores the four experts 3.0, 0.0, 1.0 and 0.5. Expert 1 would output (1, 0) and expert 3 would output (0, 1). What does the layer send out?

The top 2 are expert 1 (3.0) and expert 3 (1.0). Softmax over those two: e^3 = 20.09 and e^1 = 2.718 give 0.88 and 0.12. (0.79, 0.11) skips the renormalising, so it adds up to 0.90. (0.75, 0.25) divides the scores by their total and skips e.

Part 2. Total and active parameters

Count an MoE's total and active parameters, and tell which sets memory and which sets compute.

Keep going with Pro

You've done the first part. Pro opens the rest of this lesson, and every other one.

Still to come
  1. Total and active parameters
  2. Keeping the experts busy
  3. Smaller experts, and a shared one
See Pro plansI have Pro: sign in
Next lessonLinear attention and hybrid modelsMost layers keep a fixed-size state, and a few keep full attention. Start lesson 16

Transformers overviewSpotted a mistake? Tell us.