Mixture of Experts (MoE): How Transformers Scale Without Activating Everything.

-->

Mixture of Experts (MoE) is one of the main techniques used to scale modern language models without making every token pay the full computational cost of the model.

The basic idea is surprisingly simple: instead of sending every token through one enormous feed-forward network, we split it into many smaller expert networks and only activate a few experts for each token.

These notes walk through the intuition behind MoE, starting from the role of the feed-forward network (FFN) inside a Transformer, then moving through routing, top-$k$ selection, fine-grained experts, shared experts, expert capacity, dropless MoE, load balancing, and router stability.

The focus is not just on what MoE does, but why each piece exists and what problem it is trying to solve.

I originally learned this material from Jia-Bin Huang’s visual explanation of MoE 1. I have rewritten the ideas here as notes for myself, with the equations and implementation details that I found useful when trying to understand how modern sparse MoE models actually work.

Table of Contents

Feed-Forward Networks in Transformers

RMSNorm

The FFN as a Knowledge Store

Why Mixture of Experts?

Sparse Mixture of Experts

The Router

Top-$k$ Routing

Fine-Grained Experts

Shared Experts

Expert Capacity & Token Overflow

Dropless MoE

The Load Balancing Problem

Noisy Top-$k$ Gating

Importance vs Load

Load Balancing Loss

Device-Level Load Balancing

Auxiliary-Loss-Free Load Balancing

Router Stability & the Router Z-Loss

Putting It All Together

Appendix

References

Citation

Feed-Forward Networks in Transformers

A Transformer layer typically alternates between an attention mechanism and a feed-forward network (FFN) .

Attention allows each token to incorporate information from other tokens in the sequence. The FFN then processes each token independently.

For a single token embedding $x$, a simplified FFN can be written as:

\[z = xW_{\text{up}} + b_{\text{up}}\]

\[a = \sigma(z)\]

\[y = aW_{\text{down}} + b_{\text{down}}\]

where:

$W_{\text{up}}$ projects the token into a larger hidden dimension.

$\sigma$ is the nonlinear activation function.

$W_{\text{down}}$ projects the representation back to the original model dimension.

Typically,

\[d_h \approx 4d\]

where $d$ is the input dimension and $d_h$ is the hidden dimension of the FFN.

RMSNorm

Before the FFN, modern Transformers commonly use a normalization layer such as RMSNorm .

For an input vector $x$:

\[\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2+\epsilon}\]

and

\[\text{RMSNorm}(x) = \gamma \odot \frac{x}{\text{RMS}(x)}\]

where $\gamma$ is a learnable scaling vector.

The important idea is that normalization keeps the magnitude of activations under control while allowing the model to learn a different scale for each dimension.

The FFN as a Knowledge Store

The FFN is more than just a generic nonlinear transformation.

One useful interpretation is that the first projection asks a collection of learned questions about the token representation.

Each row of $W_{\text{up}}$ can be thought of as a learned direction in representation space.

\[z_i = x \cdot W_{\text{up},i} + b_i\]

If $z_i$ is large, the input strongly matches the feature represented by that row.

The activation function then suppresses irrelevant features.

For a simple ReLU example:

\[a_i = \max(0,z_i)\]

The second projection maps the activated features back into the model dimension.

This gives us an intuitive picture:

The FFN can be thought of as a large collection of learned feature detectors that activate different pieces of stored information depending on the input.

This interpretation is useful for understanding why increasing the FFN hidden dimension can improve model capacity.

But there is a problem.

Why Mixture of Experts?

If we simply increase $d_h$, the FFN becomes larger.

That gives the model more capacity, but it also increases:

training computation,

inference computation,

parameter memory,

communication requirements.

And there is another observation:

A token does not need every feature in the FFN.

For example, a token about chemistry probably does not need every feature that might be useful for programming, mathematics, or another language.

So instead of making one enormous FFN that processes every token, we can divide the FFN into multiple smaller networks.

Each one becomes an expert .

The model can then choose which experts should process each token.

This is the central idea behind a sparse Mixture of Experts (MoE) .

Sparse Mixture of Experts

Suppose we have $N$ experts:

\[E_1, E_2, \dots, E_N\]

Each expert is itself an FFN.

Instead of evaluating all $N$ experts for every token, we select only the top $k$ experts.

If $k \ll N$, then the model can contain many more parameters while only activating a small fraction of them for each token.

This creates an important distinction:

Total...