Warmup Guide — Mixture of Experts From Scratch

How to read this. No prior knowledge of MoE, routing, or sparse models is assumed. Everything is built from a dense feed-forward layer upward: what it is, why MoE exists, how the mechanism works underneath, what it costs in production, and the misconception people carry. If you already know MoE, start at Chapter 5 and do not skip Chapter 7 or Chapter 9.

Table of Contents

• Chapter 1: The Problem MoE Solves
• Chapter 2: The Dense Feed-Forward Layer, First
• Chapter 3: The MoE Layer
• Chapter 4: The Router, In Full
• Chapter 5: Router Collapse — the Default Failure
• Chapter 6: The Router Z-Loss
• Chapter 7: Capacity, Dropping, and the Silent Quality Loss
• Chapter 8: Total vs Active Parameters
• Chapter 9: The Communication Wall
• Chapter 10: MoE Scaling Laws and "Running Out of Internet"
• Chapter 11: The Modern Refinements
• Lab Walkthrough
• Success Criteria
• Interview Q&A
• Tips & Takeaways
• References

Chapter 1: The Problem MoE Solves

The welded constraint

In a dense transformer, every parameter participates in every token. Want the model to know more? Add parameters. But now every token costs more to process — during training and forever afterward during serving.

From Phase 00: C = 6ND for training, 2N FLOPs per token for inference. Both scale linearly with N. Capacity and cost are welded together.

That is a genuine problem, because the two things you want are in tension:

• Knowledge and capability want many parameters.
• Latency and serving cost want few parameters.

The observation MoE exploits

Here is the insight: not every token needs every parameter. The weights that help predict the next token in a Python function are largely not the weights that help with 18th-century French poetry. In a dense model, both sets fire for both inputs, and most of that computation is wasted on any individual token.

Conditional computation breaks the weld: store many parameters, but activate only a small, input-dependent subset per token.

DENSE                                MoE
──────                               ───
params = 8 units                     params = 64 units   (8x the knowledge)
FLOPs  = 8 units/token               FLOPs  =  8 units/token   (same cost!)
                                     memory = 64 units   <- the bill arrives here

Why it exists historically

The idea is old (Jacobs et al., 1991). Shazeer et al. (2017) made it work at scale for language, Switch Transformers (2021) simplified it to k = 1 and made it stable, GShard (2020) built the distributed machinery, and by 2024 essentially every frontier model is an MoE — including, as Feinberg notes, the Gemini 2.0 series.

The misconception

"MoE gives you a bigger model for free."

It gives you a bigger model for the same FLOPs. It costs you memory, communication, stability, dropped tokens, and data hunger. This entire warmup is an accounting of where that bill lands.

Chapter 2: The Dense Feed-Forward Layer, First

You cannot understand what MoE replaces without knowing what it replaces.

What it is

Each transformer block has two sub-layers: attention (which routes information between positions) and a feed-forward network / MLP (which transforms each position independently).

def dense_ffn(token, w_in, w_out):
    """Expand to d_ff, apply a nonlinearity, project back to d_model."""
    hidden = matvec(w_in, token)          # d_model -> d_ff   (usually d_ff ~ 4*d_model)
    activated = [max(0.0, h) for h in hidden]
    return matvec(w_out, activated)       # d_ff -> d_model

Why it is the target

From Phase 00's parameter count, for a typical block:

attention projections :  41,943,040   (24%)
MLP                   : 135,266,304   (76%)   <- MoE replaces THIS
norms                 :       8,192   ( 0%)

Two reasons the MLP is the right target:

1. It is where the parameters are — about three-quarters of the block.
2. It is position-independent. Each token goes through the MLP alone, with no interaction between positions. So you can route each token to a different expert without breaking anything. Attention cannot be split this way, because it is precisely the part that mixes positions.

This is the single best reason MoE targets the MLP and not attention, and it is worth being able to say out loud in an interview.

The misconception

"The MLP is just a nonlinearity; attention is where the intelligence is."

Empirically, most factual knowledge lives in MLP weights — the model-editing literature (ROME, MEMIT) locates and edits facts specifically there. Attention decides what to look at; the MLP decides what to think about it.

Chapter 3: The MoE Layer

The structure

Replace the one MLP with E parallel MLPs ("experts") plus a small router.

                        token
                          │
                 ┌────────▼────────┐
                 │     ROUTER      │   E scores, one per expert
                 │  (d_model x E)  │
                 └────────┬────────┘
                          │ softmax, take top-k, renormalize
         ┌────────┬───────┼───────┬────────┐
         ▼        ▼       ▼       ▼        ▼
        E0       E1      E2      ...      E7
         ·        ✓       ·       ·        ✓        only 2 of 8 run
                  └───────┬────────┘
                          ▼   weighted sum by gate
                        output

The forward pass, exactly

def moe_layer(token, router_weights, experts, top_k=2):
    # 1. score every expert
    logits = matvec(router_weights, token)          # E scores
    probs = softmax(logits)

    # 2. pick the best k, renormalize their gates to sum to 1
    chosen = sorted(range(len(probs)), key=lambda i: -probs[i])[:top_k]
    total = sum(probs[i] for i in chosen)
    gates = {i: probs[i] / total for i in chosen}

    # 3. run ONLY those k, and combine by gate weight
    out = [0.0] * len(token)
    for i, g in gates.items():
        contribution = dense_ffn(token, experts[i]["w_in"], experts[i]["w_out"])
        out = [o + g * c for o, c in zip(out, contribution)]
    return out

The k choice

Why renormalize the gates

This is subtle and worth being precise about. Suppose the softmax gives expert 6 a probability of 0.575 and expert 4 a probability of 0.151, and you take the top 2.

• Renormalized: gates become 0.79 and 0.21, summing to 1. The layer output has a consistent scale regardless of how confident the router was.
• Raw: gates stay 0.575 and 0.151, summing to 0.726. The layer output is scaled down by 27% — and by a different amount for every token, depending on how much probability mass leaked to the experts you did not pick.

That varying scale interacts badly with the residual stream and layer norms. Almost all production implementations renormalize. The lab implements both so you can see the difference, and a test asserts it.

Chapter 4: The Router, In Full

What it is

The smallest and most important component. Just d_model × E parameters — for a model with d_model = 7168 and E = 256, that is 1.8M parameters, versus ~44M for a single expert. It is 0.04% of the layer and it decides everything.

Numerical stability is not optional here

def softmax(xs):
    """Max-subtraction is what stops exp() overflowing to inf."""
    m = max(xs)
    exps = [math.exp(x - m) for x in xs]
    s = sum(exps)
    return [e / s for e in exps]

Why this matters more for a router than elsewhere: router logits are unconstrained and there is a positive feedback loop pushing them up (Chapter 5). In bf16, exp(90) already overflows. A naive softmax gives NaN, the NaN propagates through the residual stream, and your 40-day training run is dead. The lab tests this with logits of 1000.

Subtracting the max changes nothing mathematically — softmax(x) = softmax(x - c) for any constant c — and changes everything numerically. A test asserts this shift-invariance.

Ties and determinism

Two experts can score identically (especially early in training, when weights are near zero). If your tie-break is dict ordering or an unstable sort, the same input routes differently on different runs, and you lose reproducibility — which you need for debugging a run that costs $30M. Break ties by expert index, always.

Chapter 5: Router Collapse — the Default Failure

This is the chapter that matters most. A router left alone will collapse, and it is not a rare pathology — it is the natural dynamic.

The mechanism

  expert i is randomly slightly better at step 0
            │
            ▼
  it wins slightly more tokens
            │
            ▼
  it receives more gradient, so it trains faster
            │
            ▼
  it becomes genuinely better
            │
            ▼
  it wins even more tokens ──────────┐
            ▲                        │
            └────────────────────────┘
                  RICH GET RICHER

Within a few thousand steps you have one expert doing everything and E − 1 dead ones. You paid for E× the memory, E× the checkpoint size, and E× the sharding complexity, and you are running a dense model.

The fix: an auxiliary load-balancing loss

def load_balance_loss(assignments, all_probs, n_experts):
    """Switch-style:  L_aux = E * sum_i f_i * P_i

    f_i = fraction of token-SLOTS routed to expert i   (discrete -> NO gradient)
    P_i = mean router PROBABILITY on expert i          (continuous -> differentiable)
    """
    total_slots = sum(len(a) for a in assignments)
    f = [0.0] * n_experts
    for row in assignments:
        for e in row:
            f[e] += 1.0 / total_slots

    P = [0.0] * n_experts
    for probs in all_probs:
        for i, p in enumerate(probs):
            P[i] += p / len(all_probs)

    return n_experts * sum(fi * Pi for fi, Pi in zip(f, P))

Why the product f_i · P_i, and not something simpler?

This design is genuinely clever and the reasoning is a great interview answer.

You want to penalize imbalance in f — the actual token counts. But f comes from a top-k operation, which is discrete: it has zero gradient almost everywhere, so you cannot backpropagate through it.

P — the mean softmax probability — is differentiable, but on its own it is a weak signal: the router could keep P uniform while its top-k choices remain lopsided.

Multiplying them gives you both: the loss is scaled by the real imbalance (f) while the gradient flows through P. When expert i is over-subscribed, f_i is large, so the gradient pushing P_i down is large.

The calibration point — memorize this

Under perfect balance, f_i = P_i = 1/E for every i:

$$ L_{\text{aux}} = E \sum_{i=1}^{E} \frac{1}{E}\cdot\frac{1}{E} = E \cdot E \cdot \frac{1}{E^2} = 1.0 $$

1.0 means perfectly balanced. Larger is worse. That single fact makes an MoE training dashboard readable at a glance — and the lab has a test asserting exactly 1.0.

Measured in the lab:

balanced : L_aux=1.0000  max/mean=1.00
collapsed: L_aux=7.4400  max/mean=8.00  dead=7

Watching it happen

The lab simulates the feedback loop directly:

aux_weight=0.000: max/mean after 600 steps =    4.10  COLLAPSED
aux_weight=0.002: max/mean after 600 steps =    1.43  healthy
aux_weight=0.005: max/mean after 600 steps =    1.06  healthy
aux_weight=0.010: max/mean after 600 steps =    1.00  healthy

The coefficient is a real hyperparameter

L_total = L_task + alpha * L_balance + gamma * L_z

alpha ≈ 0.01 is the common value, and it is finicky:

• Too small → collapse. You lose the entire benefit of MoE.
• Too large → you force tokens to experts that are wrong for them, purely to satisfy the balance constraint. Quality drops. The router is now optimizing bookkeeping instead of prediction.

What to put on the dashboard: max_over_mean utilization and L_balance. If max_over_mean climbs past ~3, intervene now — not at the next checkpoint. This is exactly the sort of thing the five-person Flash 2.0 rotation was watching for.

Chapter 6: The Router Z-Loss

What it is

$$ L_z = \frac{1}{T}\sum_{t} \left(\log \sum_i e^{z_{t,i}}\right)^2 $$

A penalty on the magnitude of the router logits, introduced in ST-MoE (Zoph et al., 2022).

def router_z_loss(all_logits):
    return sum(logsumexp(row) ** 2 for row in all_logits) / len(all_logits)

Why it exists — two distinct failures

1. Numerical overflow. Router logits are unconstrained and the collapse dynamic pushes them up. In bf16, exp(90) overflows. NaN in the router means NaN in the residual stream means a dead run.

2. A saturated softmax has no gradient. This one is subtler and more damaging. If logits are [100, 2, 1], the softmax is [1.0, 0.0, 0.0] to floating-point precision. The derivative of softmax is p_i(δ_ij − p_j) — when p_i is 1 or 0, that is zero. The router stops learning. Routing freezes into whatever pattern it happened to have, permanently.

The lab shows the magnitude growing sharply:

logit scale x 1.0 -> z_loss=    9.999
logit scale x 5.0 -> z_loss=  160.348
logit scale x20.0 -> z_loss= 2527.342

gamma ≈ 1e-3 is typical — small, because you are only trying to keep logits in a sane range, not to shape the distribution.

The general lesson, which transfers well beyond MoE: whenever a small number of unconstrained logits control a discrete decision, add a magnitude penalty. The same reasoning shows up in attention (QK-norm) for the same reason.

Chapter 7: Capacity, Dropping, and the Silent Quality Loss

The constraint

Accelerators want fixed-shape tensors. You cannot allocate "however many tokens happened to route here" — the shape has to be known ahead of time. So each expert gets a fixed buffer:

$$ \text{capacity} = \text{capacity_factor} \times \frac{\text{tokens} \times k}{E} $$

The fraction is the count each expert would get under perfect balance. capacity_factor (typically 1.0–2.0) is the safety margin for imperfect balance.

The two failure directions

        capacity per expert  ────────────────────────►
        ┌──────────────────────────────────────┐
        │████████████████░░░░░░░░░░░░░░░░░░░░░░│  under-subscribed:
        │  real tokens      PADDING (waste)    │  you compute on zeros
        └──────────────────────────────────────┘

        ┌──────────────────────────────────────┐
        │██████████████████████████████████████│▓▓▓▓  over-subscribed:
        │        buffer full                   │DROP  overflow is discarded
        └──────────────────────────────────────┘

Measured in the lab on a realistically imbalanced batch:

capacity_factor=1.00: cap= 16  dropped= 12 ( 9.4%)  padded= 12  buffer_util=90.6%
capacity_factor=1.25: cap= 20  dropped=  3 ( 2.3%)  padded= 35  buffer_util=78.1%
capacity_factor=2.00: cap= 32  dropped=  0 ( 0.0%)  padded=128  buffer_util=50.0%

There is no setting that avoids both. Going from 1.0 to 2.0 eliminates dropping and halves your effective compute utilization. That trade-off is the hyperparameter.

Why dropping is the dangerous one

Trace what happens to a dropped token in the lab's forward pass:

if not got_any:
    # Every slot dropped: the token skips the FFN entirely and rides the
    # residual. No error, no log line. This is the silent failure.
    acc = list(token)
    fully_dropped += 1

No exception. No warning. No log line. The token passes through the layer unchanged on the residual connection. Training continues. Your loss is very slightly worse than it should be, and nothing tells you why.

This is a genuine, live source of quality loss in production MoE models — and it is why drop_rate belongs on the dashboard next to max_over_mean.

What a shared expert buys

A shared expert (DeepSeek-V3's design) runs for every token, unconditionally, in addition to the routed ones:

if shared_expert is not None:
    acc = [a + c for a, c in zip(acc, expert_forward(shared_expert, token))]

Two things fall out:

1. No token can ever be fully dropped. The failure mode above is structurally eliminated.
2. The routed experts can specialize harder, because the shared expert absorbs the common, general-purpose transformation that every token needs. You are no longer forcing every expert to independently learn the basics.

The cost: that expert's FLOPs are paid on every token, always. It is a small, always-on tax that buys a large reduction in variance.

Chapter 8: Total vs Active Parameters

The arithmetic most commonly botched in modern LLM discussion, and it is off by 10–20×.

def moe_parameter_counts(d_model, d_ff, n_experts, top_k, shared_experts=0):
    one_expert = 2 * d_model * d_ff
    router = d_model * n_experts
    total = router + (n_experts + shared_experts) * one_expert    # -> MEMORY
    active = router + (top_k + shared_experts) * one_expert       # -> FLOPs
    return {"total": total, "active": active, "sparsity_ratio": total / active}

The rule, and there are no exceptions:

A 700B-total / 52B-active model trains like a 52B model and stores like a 700B one. Get this backwards and your capacity plan is wrong by an order of magnitude.

From the lab, at toy scale:

top_k=1: total=  2112  active=   320  sparsity=6.60x
top_k=2: total=  2112  active=   576  sparsity=3.67x
top_k=8: total=  2112  active=  2112  sparsity=1.00x

Note the last row: top_k == n_experts gives sparsity exactly 1.0. A dense model wearing an MoE costume. The lab tests this boundary explicitly, because it is the sanity check that proves your accounting is right.

Chapter 9: The Communication Wall

Now the cost that shaped Feinberg's Flash 2.0 story.

Why sharding is forced

All E experts must be resident in HBM. A 700B-parameter MoE in bf16 is 1.4 TB. A TPU v5e has 16 GB; an H100 has 80 GB. The model cannot fit on one chip, so experts are distributed across chips — expert parallelism.

What that costs, per layer

Layer ℓ:  token lives on chip 0
          router says "expert 37", which lives on chip 4
          ──> send the activation to chip 4      [NETWORK]
          chip 4 computes
          ──> send the result back to chip 0     [NETWORK]

Layer ℓ+1: router says "expert 12" (chip 1) ...  [NETWORK] [NETWORK]

... for every one of ~60 layers.

Feinberg's description: "that token might live on the first TPU, but it needs to go to the last TPU." The collective is an all-to-all — every chip sends a different slice to every other chip — which is the most expensive collective there is, and its cost, he notes, "increases dramatically with N."

The arithmetic

def expert_parallel_comm_bytes(n_tokens, d_model, n_layers, top_k, bytes_per_elem=2):
    per_hop = n_tokens * d_model * bytes_per_elem * top_k
    return 2 * n_layers * per_hop        # dispatch + combine, every layer

From the lab, for a realistic prefill:

8192 tokens, d_model=8192, 60 layers, top_k=2
bytes across the interconnect :     32.2 GB
time at 100 Gbps              :    2.578 s

Two and a half seconds of pure network time before a single useful FLOP. For an interactive product with a sub-second budget (Phase 00's napkin math), that is fatal.

The fix, in one sentence

Stop sharding experts across chips. Shard layers instead, and stream chunks of the prompt through the resulting pipeline so the transfers overlap with computation. That is pipelined prefill, credited in the interview to Geng Yan, and it is what made an MoE Gemini Flash servable.

The full treatment — including why it works for prefill and not for decode — is in the transcript dissection, Claim 12.

The transferable lesson, and it is the real one: the fix was not a better kernel or a better model. It was changing which axis you shard along, chosen with knowledge of which phase of inference you are in. That is what "inference co-design" means in practice.

Chapter 10: MoE Scaling Laws and "Running Out of Internet"

Feinberg's slide states both halves:

"MoE scaling laws are better, but have implications for token hunger. We're running out of internet! ... Notice relative data hunger compared to dense! At same active param count and fixed 100B token training, MoE 64E improves on dense."

Decoded

The good half. For a fixed compute budget, an MoE reaches a lower loss than a dense model. The scaling law is strictly better. At equal active parameters and equal tokens, 64 experts beats dense. This is why everyone switched.

The bill. The compute-optimal D for an MoE is larger than for a dense model. MoE has more capacity to fill, and filling capacity takes tokens. So MoE converts "we have compute" into "we need more unique tokens" — and unique, high-quality tokens are the resource that is actually scarce.

Hence: MoE trades a compute problem for a data problem. Which is a good trade right up until you run out of data, and then it is not.

That is the bridge to the rest of Feinberg's agenda — multimodal data, synthetic data, and the data-constrained scaling law L(N, U, R) where U is unique tokens and R is repeats. He notes that data work is where he spends "probably half my focus this year so far." The pre-training lead's time goes to data, not architecture.

The reference for the routed-model scaling law is Clark et al., Unified Scaling Laws for Routed Language Models (2022). The data-constrained law is Muennighoff et al. (2023). Both are on his slide's resource list.

Chapter 11: The Modern Refinements

What production MoE looks like in 2024–2025, beyond the Switch baseline.

Notice the pattern across that table. Four of the six refinements exist to fix problems created by the two mechanisms in Chapters 5 and 7 — the auxiliary loss and the capacity buffer. Understand those two deeply and the rest of the literature reads as commentary.

Lab Walkthrough

Lab 01 — Router, Load Balancing, Capacity & the MoE Layer

Implement in this order:

1. softmax, logsumexp, matvec, relu. Get test_softmax_survives_huge_logits passing first — max-subtraction is the lesson, not a detail.
2. router_logits, route_token, route_batch. Watch the renormalize flag; both behaviours are tested. Break ties by expert index.
3. load_balance_loss. Aim for the calibration point: test_balanced_routing_gives_aux_loss_of_exactly_one. If you do not get exactly 1.0, the formula is wrong.
4. router_z_loss, expert_utilization.
5. expert_capacity, apply_capacity. Note that capacity_factor < 1.0 must raise — it drops tokens even under perfect balance, which is never intended.
6. moe_forward. The residual fallback for fully-dropped tokens is the important line.
7. moe_parameter_counts, expert_parallel_comm_bytes, comm_seconds.
8. simulate_collapse. The money test.

The traps:

• A naive softmax gives NaN on large logits and the test will catch you.
• The load-balance loss divides f by total slots (tokens × k), not by token count.
• apply_capacity is first-come-first-served in token order — deterministic, and the tests depend on it.
• top_k == n_experts must give sparsity exactly 1.0; if it does not, your router or parameter count is wrong.
• Everything random goes through a seeded random.Random(seed).

Success Criteria

• LAB_MODULE=solution pytest test_lab.py -v → 53 passed.
• Your lab.py reaches 53 passed.
• python solution.py runs and you can explain all eight sections.
• Balanced routing gives a load-balance loss of exactly 1.0 in your code.
• You can explain why the aux loss multiplies f by P rather than penalizing f alone.
• You have watched a router collapse at aux_weight = 0 and recover at 0.01.
• You can state the capacity formula and the drop/pad trade-off from memory.
• You can compute total vs active parameters and say which goes into 6ND.
• You can compute all-to-all bytes for a realistic MoE and convert to seconds.

Interview Q&A

Q: What is a Mixture of Experts and what does it buy you? Replace each transformer block's MLP with E parallel MLPs plus a small router that picks the top-k per token. Parameters scale with E; FLOPs scale with k. So you get far more capacity at the same compute per token. It targets the MLP rather than attention for two reasons: the MLP is ~76% of the block's parameters, and it is position-independent, so routing each token separately does not break anything — attention is precisely the part that mixes positions and cannot be split this way.

Q: What breaks first when you train an MoE? The router collapses. It is a rich-get-richer loop: one expert is randomly slightly better, so it wins more tokens, so it trains more, so it becomes genuinely better, so it wins more. Within a few thousand steps you have one live expert and E−1 dead ones — you paid for E× the memory to run a dense model. The fix is an auxiliary load-balancing loss, L = E·Σ f_i·P_i, with a coefficient around 0.01.

Q: Why does that loss multiply a token fraction by a probability? Because f, the fraction of tokens routed to each expert, comes from a top-k, which is discrete and has no gradient. P, the mean softmax probability, is differentiable but a weak signal on its own — the router could keep P uniform while its actual choices stay lopsided. Multiplying gives you a loss scaled by the real imbalance with gradient flowing through the differentiable part. And it is calibrated: perfect balance gives exactly 1.0.

Q: What is capacity factor and what happens if you get it wrong? Each expert gets a fixed buffer of cf × tokens × k / E slots, because hardware needs fixed-shape tensors. Too low and tokens overflow and are dropped — silently: the token skips the FFN, rides the residual, no error is raised, and your quality is quietly worse. Too high and buffers are padded with zeros and you waste compute. There is no setting that avoids both, so you tune it against measured drop rate. A shared expert that runs for every token eliminates the fully-dropped case structurally.

Q: A 700B-parameter MoE with 8 of 256 experts active. How much compute to train it? 6ND with N = active parameters, not total. If active is ~52B, it trains like a 52B dense model. But it stores like a 700B model — 1.4 TB in bf16 — which is what dictates your sharding, your checkpoint size, and your HBM budget. Active for FLOPs, total for memory.

Q: Why is MoE hard to serve? All experts must be HBM-resident, so a large MoE cannot fit on one chip and you shard experts across chips. Then every layer, every token must be sent to whichever chip holds its expert and the result sent back — two all-to-all collectives per layer. For an 8k-token prefill on a 60-layer model that is tens of gigabytes and seconds of pure network time, which is fatal for an interactive product. The fix Google used for Flash 2.0 was to shard layers instead of experts and stream prompt chunks through the resulting pipeline, so the transfers hide behind computation. Notably it works for prefill — which is compute-bound and has thousands of tokens to chunk — and not for decode, which produces one token at a time and has nothing to overlap with.

Q: If MoE is strictly better per FLOP, why isn't everything an MoE? Three reasons. The memory and sharding complexity is real. The training instability is real — you have added a discrete decision to a system that otherwise has none. And the compute-optimal token count is larger for an MoE, so it converts a compute problem into a data problem — and unique high-quality data is the resource actually running out.

Q: What is the router z-loss for? It penalizes the magnitude of router logits. Two failures it prevents: numerical overflow — router logits are unconstrained, the collapse dynamic pushes them up, and exp() overflows in bf16, putting NaN through the residual stream; and softmax saturation — once the softmax is effectively one-hot, its derivative p(1−p) is zero, the router stops learning, and routing freezes permanently.

Tips & Takeaways

Tips

• Put max_over_mean and drop_rate on the dashboard. They are the two numbers that tell you an MoE is failing, and both fail silently otherwise.
• Memorize L_aux = 1.0 means balanced. It makes the metric readable instantly.
• Always ask "total or active?" the moment someone quotes an MoE parameter count.
• Break router ties deterministically. Reproducibility on a $30M run is not negotiable.
• Never write a softmax without max-subtraction. Especially in a router.
• When you see a new MoE paper, ask which of the two core problems it is fixing — balance or capacity. Most of the literature is one or the other.
• Compute the all-to-all bytes before you commit to an architecture, not after.

Takeaways

1. Parameters scale with E; FLOPs scale with k. That is MoE, entire.
2. MoE targets the MLP because it holds ~76% of the parameters and is position-independent.
3. Renormalizing the top-k gates keeps the output scale consistent; not doing so silently attenuates the layer by a per-token amount.
4. Routers collapse by default. The auxiliary loss is what prevents it, and 1.0 is balanced.
5. The f·P product exists to get a gradient through a discrete decision.
6. The z-loss prevents overflow and softmax saturation — two different failures.
7. Capacity trades dropping against wasted compute. Dropping is silent, which makes it worse.
8. A shared expert eliminates fully-dropped tokens and lets the routed experts specialize.
9. Active for FLOPs, total for memory. Off by 10–20× if you swap them.
10. Naive expert parallelism costs seconds of network time. The fix is a different sharding axis, not a better kernel.
11. MoE trades a compute problem for a data problem — and data is what is running out.

References

Foundational

• Jacobs et al., Adaptive Mixtures of Local Experts, 1991 — the original idea
• Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer, 2017 — https://arxiv.org/abs/1701.06538
• Lepikhin et al., GShard: Scaling Giant Models with Conditional Computation, 2020 — https://arxiv.org/abs/2006.16668 (expert parallelism, capacity, the all-to-all)
• Fedus, Zoph & Shazeer, Switch Transformers, 2021 — https://arxiv.org/abs/2101.03961 (k=1, the load-balance loss, capacity factor)

Stability and refinements

• Zoph et al., ST-MoE: Designing Stable and Transferable Sparse Expert Models, 2022 — https://arxiv.org/abs/2202.08906 (the router z-loss)
• Zhou et al., Mixture-of-Experts with Expert Choice Routing, 2022 — https://arxiv.org/abs/2202.09368
• Gale et al., MegaBlocks: Efficient Sparse Training with Mixture-of-Experts, 2022 — https://arxiv.org/abs/2211.15841 (dropless MoE via block-sparse kernels)
• Komatsuzaki et al., Sparse Upcycling, 2022 — https://arxiv.org/abs/2212.05055

Production models

• Jiang et al., Mixtral of Experts, 2024 — https://arxiv.org/abs/2401.04088
• DeepSeek-AI et al., DeepSeek-V3 Technical Report, 2024 — https://arxiv.org/abs/2412.19437 (fine-grained + shared experts, auxiliary-loss-free balancing)

Scaling

• Clark et al., Unified Scaling Laws for Routed Language Models, 2022 — https://arxiv.org/abs/2202.01169
• Muennighoff et al., Scaling Data-Constrained Language Models, 2023 — https://arxiv.org/abs/2305.16264

Primary source for this phase's framing

• Feinberg, Gemini Pretraining, Princeton, Apr 2025 — slides
• Developing Dev interview · video