The Switch layer: one expert per token

A Switch layer replaces one feed-forward block with N copies of it plus a router. The router is the smallest thing in the layer: a single weight matrix W_r: [d, N] followed by a softmax.

h(x) = W_r^T x        logits, shape [N]
p_i(x) = exp(h_i) / Σ_j exp(h_j)
i* = argmax_i p_i(x)
y   = p_i*(x) · E_i*(x)

Two details carry the whole design. First, only E_i* is evaluated — the other N-1 experts do no work for this token, which is where the FLOP saving lives. Second, the expert output is scaled by its own gate probability p_i*(x), a number in (0, 1]. That scalar looks decorative next to a full FFN, but it is the only path by which the router receives any learning signal at all, and it is what makes top-1 viable.

Advertisement

Does k = 1 break the router gradient?

The 2017 sparsely-gated MoE argued that k ≥ 2 was necessary: with one expert, the reasoning went, there is nothing to compare against and so no meaningful gradient to the gate. Switch’s counter is a one-line derivative. Because the output is y = p_i*(x) · E_i*(x), the loss reaches the router through the multiplicative scalar:

∂L/∂W_r = (∂L/∂y · E_i*(x)) · ∂p_i*/∂W_r

And p_i* is a softmax output, so it depends on every logit through the shared denominator — ∂p_i*/∂h_j ≠ 0 for all j. One selected expert still updates all N router rows: if the chosen expert helped, its logit rises and the rest fall. Fedus et al. list three payoffs for the simplification — less router computation, at least halved per-expert buffers, and reduced communication.

Advertisement

Expert capacity: a fixed buffer per expert

Routing is data-dependent, but the hardware is not. Experts live on different devices, and an all-to-all needs tensors of a shape known before the routing decisions exist. So each expert gets a fixed buffer, sized by a capacity factor:

C = CF · (tokens_per_batch / N)     [Switch, k = 1]
CF = 1.0  →  exactly the perfectly balanced share
CF > 1.0  →  slack for skew, at the cost of idle slots

This is the generic MoE capacity CF · T·k/N with k collapsed to 1 — the reason Switch can afford buffers half the size of a top-2 model at the same slack. The capacity factor is a pure memory-for-robustness dial: N · C slots are allocated and padded whether or not tokens fill them, so CF = 2.0 means half the expert compute is spent on padding even when routing is perfect.

Token dropping and what it costs

When more than C tokens choose the same expert, the overflow is dropped: those tokens are simply not processed by any FFN. They are not re-routed and not queued. Their representation passes to the next layer through the residual connection unchanged, as though the layer had been skipped for them.

At k = 2 this is a partial loss — a dropped token usually still has its other expert, and the combine just reweights. At k = 1 there is no second expert, so a dropped token receives zero feed-forward contribution from that layer. This is the sharpest edge of the top-1 design and the reason capacity and balancing matter more here than in any other MoE variant. A small drop rate is survivable and even acts a little like structured dropout; a large one silently removes a whole layer’s worth of computation from a biased subset of the vocabulary.

Worked example: capacity, skew, drops

Take one Switch layer with N = 64 experts serving a batch of T = 8192 tokens across the expert-parallel group (global, not per-device). At k = 1 there are 8192 dispatches.

perfect share μ = 8192 / 64 = 128
CF = 1.25  →  C = 1.25 · 128 = 160

skewed routing:  4 experts ← 300 tokens each
                 8 experts ← 200 tokens each
                52 experts ← ~104 tokens each

dropped = 4·(300-160) + 8·(200-160) = 560 + 320 = 880
drop_rate = 880 / 8192 = 10.7%

Nearly eleven percent of tokens skip the FFN entirely. Raise CF to 2.0 and C = 256: only the four hottest experts overflow, dropped = 4·44 = 176, a 2.1% drop rate — but you now allocate 64 × 256 = 16384 slots for 8192 tokens, wasting half the expert compute on padding. That trade is the whole tuning problem.

The load-balancing auxiliary loss

Capacity handles skew; the auxiliary loss tries to prevent it. Switch adds a product of two per-expert fractions over each batch B:

f_i = fraction of tokens dispatched to expert i      (hard counts)
P_i = mean_{x in B} p_i(x)                        (soft mass)

L_aux = α · N · Σ_i f_i · P_i,   α = 1e-2
uniform: N · Σ (1/N)(1/N) = 1  → the minimum

The product form is what makes it usable. f_i comes from an argmax and a count — piecewise constant, gradient identically zero. P_i is a mean of softmax outputs and is perfectly smooth. Treating f_i as a constant weight leaves ∇L_aux = αN Σ_i f_i · ∇P_i: the observed overload becomes the coefficient, and the gradient pushes probability mass off exactly the experts that were overloaded. The scale α must be large enough to balance and small enough not to fight the task loss.