SwiGLU feed-forward and Mixture-of-Experts

The usual FFN

A classical Transformer block's feed-forward network (FFN) is:

FFN(x) = Linear_out(GELU(Linear_in(x)))

with Linear_in: d_model → d_ff and Linear_out: d_ff → d_model, where d_ff is typically about 4× d_model. This is where most of the parameters of a Transformer live.

SwiGLU — a gated activation

SwiGLU replaces the two-layer FFN with a gated variant that uses three projections instead of two:

SwiGLU(x) = down_proj( silu(gate_proj(x)) * up_proj(x) )

• silu(z) = z · sigmoid(z) (also called Swish).
• * is elementwise multiplication.
• gate_proj(x) * up_proj(x) is the "gate": one projection decides how much of the other to let through, per feature.

Why bother with three projections? Shazeer (GLU Variants Improve Transformer, 2020) showed gated activations consistently beat GELU on perplexity at matched parameter counts. Every frontier decoder- only model since Llama 2 uses SwiGLU.

Paper: https://arxiv.org/abs/2002.05202.

Mixture-of-Experts, intuitively

Making the FFN bigger makes the model better — but also linearly more expensive per token. Mixture-of-Experts lets us scale parameters without linearly scaling per-token compute:

• Have N copies of the FFN (called "experts"), each with its own weights.
• For each token, route it to only k of those experts, where k ≤ N.
• Sum the k outputs, weighted by a learned gate.

If N = 8 and k = 2, the layer has 8× the FFN parameters but only 2× the per-token FFN compute. Crucially, different experts can specialize in different kinds of tokens (code vs. prose, for example), and the router learns that assignment end-to-end.

The router

A simple linear layer scores each token against each expert:

scores = softmax(W_route @ x)           # shape: (B, T, N)
top_k_idx  = scores.topk(k).indices     # (B, T, k)
top_k_vals = scores.topk(k).values      # (B, T, k)
gates      = top_k_vals / top_k_vals.sum(-1, keepdim=True)   # renormalize to sum=1

For each token we run only the k chosen experts on it, then sum their outputs weighted by gates.

Load balancing

Without intervention the router tends to collapse — it routes every token to the single strongest expert, leaving the rest untrained. The fix is a Switch-Transformer-style auxiliary loss that penalizes imbalanced utilization:

aux_loss = N · Σ_i ( fraction_routed_to[i] · fraction_importance_of[i] )

where fraction_routed_to[i] is the fraction of tokens that pick expert i, and fraction_importance_of[i] is the average gate weight expert i receives across tokens. Minimizing this pushes the router toward uniform utilization.

v0 adds aux_loss to the main cross-entropy loss with weight 0.01.

Shipped configs

Config	num_experts	top_k	d_ff (per expert)
Dev	4	2	small
70B reference	8	2	13824

Papers

• Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, ICLR 2017 — the original top-k routing recipe. https://arxiv.org/abs/1701.06538.
• Fedus et al., Switch Transformers, 2021 — the load-balancing auxiliary loss used here. https://arxiv.org/abs/2101.03961.
• Jiang et al., Mixtral of Experts, 2024 — modern top-2 MoE at production scale. https://arxiv.org/abs/2401.04088.

Reading the code

The SwiGLU expert (gate_proj, up_proj, down_proj, F.silu) and the MoE router both live in the model cell of full_llm_pipeline.ipynb. The auxiliary-loss accumulation happens inside the forward pass; the training loop pulls it out and adds it to the main loss.

Next stage: Training loop — how the optimizer and loss come together.