Every parameter, every token
In the block from part 1, attention mixes positions and the feed-forward rewrites each position. Count the matrices. Attention is a handful of projections, each the width of the residual stream. The SwiGLU feed-forward is three projections with an inner width several times larger. In a dense transformer that feed-forward is most of the parameter count, and every token walks through all of it.
Scaling that design has meant wider matrices and a deeper stack. Quality has followed the size. So has the bill. At inference every parameter is read, including the ones this token had no use for.
A mixture of experts grows the feed-forward pool and touches a slice of it.
A pool of experts
Shazeer et al. (2017) replace one feed-forward layer with a sparsely gated pool of them. Attention stays shared by every token. Each expert is an ordinary feed-forward layer with its own weights. A token that needs code and a token that needs ordinary prose can be sent to different experts. The choice can also change from layer to layer. A lower layer may use a broad expert. A higher layer may use one that has absorbed a narrower habit of the data.
Jiang et al. (2024) give the picture in Mixtral 8×7B. Each layer holds 8 feed-forward experts. The router picks 2, and the chosen pair can change from token to token and from layer to layer. The paper reports 47 billion parameters stored and 13 billion active during inference. The headline count is the pool. The active count is the forward pass. The same paper separates the bill: inference compute follows the 13 billion that ran, and serving memory follows the 47 billion that have to be reachable.
The gating network
The gate is a linear layer. It maps the token's vector to one score per expert. Jiang et al. implement it as a softmax over the top-k logits of that layer, the form used in GShard (Lepikhin et al., 2020). Left as a softmax over every expert, the mixture would be soft, and every expert would still run.
Sparse routing keeps the top k scores, sets the rest to zero, and renormalizes the survivors so they sum to one again. For Mixtral, k is 2.
y = g₁ · Expert₁(x) + g₂ · Expert₂(x)
g₁ and g₂ are the renormalized gate weights. Each expert returns a vector the width of the residual stream. The weighted sum is added back the way a single feed-forward output would be. From the rest of the block, the slot looks the same. The selection happened inside it.
The forward choice is discrete. An expert is on or off for this token. The backward pass still needs a signal for the scores that made the choice. Training treats the selected set as fixed on the way forward, and sends gradient through the gate weights of the experts that ran. The router can learn to raise a useful expert and to lower a poor one. Experts that were off get no gradient from that token.
Keeping the gate honest
A gate with no further pressure collapses. Early in training one or two experts pull ahead. The router notices. Soon every token goes there. The idle experts stop learning. You have paid for a large model, and you are training a small one.
Switch Transformers answer that collapse with a penalty and a budget. A later design adds a third piece, an expert that always runs.
A balance penalty
Fedus, Zoph, and Shazeer (2021) write the penalty used in Switch Transformers. Let fᵢ be the fraction of tokens in the batch actually sent to expert i. Let Pᵢ be that expert's average gate probability. With N experts, the added loss is α × N × Σ fᵢPᵢ. It stays small when tokens and probability are spread. It grows when a few experts are both likely and busy. The router is trained on the task and on this penalty together. The task wants the right expert. The penalty wants a used pool. In their runs, α = 0.01 was enough to balance the load without drowning the training loss.
A budget
The same paper gives each expert a capacity, a limited number of tokens in a batch. A capacity factor above 1 leaves slack for an uneven batch. Tokens past the budget are dropped: their expert update is skipped and the representation continues through the residual connection. A balanced router rarely hits the cap. An unbalanced one drops work, and the expert that should have learned from those tokens never sees them.
A shared expert
Dai et al. (2024) isolate the computation every token needs in shared experts, and route the remaining experts as specialists. Common features have a home that is always on. The routed experts are free to be narrow, and the gate no longer has to spend a slot on "everyone needs this."
What you still pay for
Sparse routing cuts arithmetic. Jiang et al. separate that from memory: Mixtral's compute follows the 13 billion active parameters, and its serving memory follows the 47 billion that must be reachable. On many devices the experts are partitioned. The Mixtral paper describes expert parallelism, following GShard: a token travels to the device that holds the expert it was assigned, and the expert's output travels back. That exchange is the new cost.
The savings show up when the exchange costs less than running every expert on every token. That is the regime these models are built for. Part 5 takes up the related serving fact for dense models: once the parameters are chosen, the cost of the next token is dominated by reading them.
What to take with you
The gate is a third selection. Attention selects tokens. SwiGLU selects features inside one feed-forward. The router selects which feed-forwards exist for this token. Mixture of experts scales the parameters a model can store, and the specialists it can train. Sparse routing scales the parameters that fire.
When you read a parameter count, ask how many ran. When a sparse model disappoints next to a dense model of the active size, inspect the router before you blame the data. A gate that ignores half its experts has thrown the extra parameters away. Balance, capacity, and a shared expert are there so the extra ones stay in the training run.
Sources
- Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv:1701.06538, 2017.
- Dmitry Lepikhin and colleagues. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. arXiv:2006.16668, 2020.
- William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961, 2021.
- Albert Q. Jiang and colleagues. Mixtral of Experts. arXiv:2401.04088, 2024.
- Damai Dai and colleagues. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. arXiv:2401.06066, 2024.