LM10-T4 — SwiGLU, RMSNorm and Mixture of Experts

Learning objectives

After this lesson, you should be able to:

• Write RMSNorm and compare it with LayerNorm on a worked vector.
• Write the SwiGLU layer, count its three matrices and explain the 8d/3 hidden width.
• Explain top-k expert routing and distinguish total from active parameters.
• Implement a small mixture-of-experts layer and test it against a dense computation.

The idea: simpler norms, gated flow, sparse capacity

Normalisation. LayerNorm (Module 4) makes each token’s vector have mean 0 and variance 1 and then learns a scale and a shift. It does two things: it re-centres, and it re-scales. Experiments on large models found that the re-scaling carries nearly all the benefit. RMSNorm keeps only that part: divide the vector by its root mean square and multiply by a learned per-feature scale. There is no mean subtraction and no learned shift. It is fewer operations and fewer parameters (d instead of 2d per norm) with no measurable loss in quality for large language models.

Gating. The classic feed-forward layer applies one fixed nonlinearity to every hidden unit: up-project, activate, down-project. A gated linear unit uses two up-projections. One, the gate, is passed through a nonlinearity; the other, the up value, is not. Their element-wise product goes to the down-projection. The gate decides, per unit and per token, how much of the up value gets through: a learned, input-dependent switch. In the “SwiGLU” variant the gate’s nonlinearity is SiLU, x·sigmoid(x), a smooth function that is close to zero for negative inputs and close to the identity for large positive ones. In comparisons by Shazeer (2020) at equal compute, gated variants gave lower loss than the plain GELU layer; that is why they became the default.

Sparse experts. A larger feed-forward layer stores more knowledge but costs more arithmetic for every token. A mixture of experts (MoE) breaks the feed-forward layer into E separate smaller layers, the experts, and adds a tiny linear router that scores them for each token. Only the top k experts (typically 1 or 2) are run, and their outputs are combined with the router’s weights. The model has E times the feed-forward parameters, but the arithmetic per token is that of k experts. Attention layers, embeddings and norms are shared by all tokens as before.

What the experts learn

Experts are not hand-assigned to topics. The router and experts are trained together, and what each expert specialises in is an outcome, often surprisingly syntactic or token-level rather than “this is the chemistry expert”. Do not describe a trained MoE as a committee of subject specialists.

Total versus active

The mathematics: three layers with their counts

RMSNorm

LayerNorm in comparison: γ_i · (x_i − μ)/√(σ² + ε) + β_i. Worked example with x = (3, 4, 0, −5) and γ = 1: the mean is 0.5, so LayerNorm gives (0.7143, 1.0, −0.1429, −1.5714) and RMSNorm, with root mean square √(50/4) = 3.5355, gives (0.8485, 1.1314, 0.0, −1.4142). For a vector whose mean is zero the two coincide. RMSNorm keeps the vector’s direction exactly and only rescales its length.

SwiGLU feed-forward

W_gate, W_up

both d → h; the gate passes through SiLU, the up value is linear.

⊙

element-wise product: the gate scales each hidden unit.

h

hidden width; set so that 3·d·h ≈ 8d², that is h ≈ 8d/3.

The classic layer has 2 · d · 4d = 8d² parameters (ignoring biases). The gated layer has three matrices, 3dh. To keep the budget, choose h = 8d/3, then round to a multiple of 64 or 256 for hardware efficiency. For d = 64 this rounding rule gives h = 192 (the rule in Lab LM10-L2); for d = 4,096 the exact 8d/3 is 10,922.67 and the 7B model uses 11,008 (the Llama 2 family chooses a multiple of 256). Parameters for d = 64, h = 192: 3 · 64 · 192 = 36,864, against 32,768 for the classic layer.

Values of SiLU: SiLU(−4) = −0.0719, SiLU(−1) = −0.2689, SiLU(0) = 0, SiLU(1) = 0.7311, SiLU(4) = 3.9281. It is not exactly zero for negative inputs, it has a small dip, and it is smooth.

Mixture of experts

With a router matrix W_r of shape E × d, for one token x:

For d = 64, h = 192, E = 4, k = 2: P_total = 4 · 36,864 + 256 = 147,712 and P_active = 2 · 36,864 + 256 = 73,984. For the 7B feed-forward width (d = 4,096, h = 11,008) with E = 8 and k = 2, the feed-forward parameters of one layer grow from 135,266,304 to 1,082,130,432 (plus a router of 32,768), but the active ones are 270,532,608 (plus the router).

Load balancing

Without a safeguard the router can send most tokens to a few experts and leave the others untrained (“collapse”), losing the capacity gain and making some devices idle. A common remedy adds an auxiliary loss. Let f_e be the fraction of tokens routed to expert e and P_e the mean router probability of e over a batch:

It is added to the language-modelling loss with a small coefficient. Systems also give each expert a capacity (a maximum number of tokens per batch) and drop or reroute the overflow.

A worked example: RMSNorm, SwiGLU and a routed layer

The code reproduces the numbers above, builds a four-expert layer with top-2 routing, and tests it two ways: its parameter counts equal the formulas, and its grouped computation (each expert processes only its tokens) equals a slow token-by-token dense computation of the same mixture.

import torch
import torch.nn.functional as F

x = torch.tensor([[3.0, 4.0, 0.0, -5.0]])
mean = x.mean(-1, keepdim=True)
ln = (x - mean) / torch.sqrt(((x - mean) ** 2).mean(-1, keepdim=True) + 1e-5)
rms = x / torch.sqrt((x ** 2).mean(-1, keepdim=True) + 1e-5)
print([round(v, 4) for v in ln[0].tolist()])     # [0.7143, 1.0, -0.1429, -1.5714]
print([round(v, 4) for v in rms[0].tolist()])    # [0.8485, 1.1314, 0.0, -1.4142]

class SwiGLU(torch.nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.gate = torch.nn.Linear(d, h, bias=False)
        self.up = torch.nn.Linear(d, h, bias=False)
        self.down = torch.nn.Linear(h, d, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

torch.manual_seed(0)
d, h = 64, 192
ff = SwiGLU(d, h)
print(sum(p.numel() for p in ff.parameters()), 2 * d * 4 * d)         # 36864 32768
print([round(v, 4) for v in F.silu(torch.tensor([-4.0, -1.0, 0.0, 1.0, 4.0])).tolist()])
# [-0.0719, -0.2689, 0.0, 0.7311, 3.9281]

class MoE(torch.nn.Module):
    def __init__(self, d, h, E, k):
        super().__init__()
        self.router = torch.nn.Linear(d, E, bias=False)
        self.experts = torch.nn.ModuleList(SwiGLU(d, h) for _ in range(E))
        self.k = k
    def forward(self, x):                                  # x: (N, d), one row per token
        probs = torch.softmax(self.router(x), dim=-1)      # (N, E)
        top_p, top_i = probs.topk(self.k, dim=-1)          # (N, k)
        top_p = top_p / top_p.sum(-1, keepdim=True)        # renormalise over the chosen experts
        out = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            rows, slot = (top_i == e).nonzero(as_tuple=True)
            if rows.numel():
                out[rows] += top_p[rows, slot, None] * expert(x[rows])
        return out, top_i

moe = MoE(64, 192, E=4, k=2)
tokens = torch.randn(32, 64)
y, chosen = moe(tokens)
print(y.shape, sum(p.numel() for p in moe.parameters()))              # torch.Size([32, 64]) 147712
print(sum(p.numel() for p in moe.router.parameters()) + 2 * sum(p.numel() for p in moe.experts[0].parameters()))   # 73984
print(torch.bincount(chosen.flatten(), minlength=4).tolist())         # [17, 13, 17, 17]  (64 routings over 4 experts)

probs = torch.softmax(moe.router(tokens), dim=-1)
tp, ti = probs.topk(2, dim=-1)
tp = tp / tp.sum(-1, keepdim=True)
ref = sum(tp[:, j, None] * torch.stack([moe.experts[int(e)](tokens[i:i + 1])[0] for i, e in enumerate(ti[:, j])]) for j in range(2))
print(torch.allclose(y, ref, atol=1e-5))                              # True

frac = torch.bincount(chosen.flatten(), minlength=4).float() / chosen.numel()
print(round(4 * (frac * probs.mean(0)).sum().item(), 3))              # 0.995  (balance loss; 1 is perfectly even)

The routing counts are close to even because the router is untrained and random; in training the balance loss keeps them that way. The “grouped” loop, with each expert processing its tokens in one batch, is how real implementations avoid running every expert on every token; production systems add capacity limits and spread the experts across devices.

Which change, and why

Engineering implications and limitations

Check the activation's name. “SwiGLU” means silu(gate) * up followed by a down-projection. Mixing up which of the two projections goes through SiLU gives a model that runs and is wrong; compare with reference logits.

Count parameters to detect mismatches. A SwiGLU layer with the wrong number of matrices will not reproduce the published parameter count (Lesson LM10-T1).

Compute norms in 32 bits. Squares of large activations overflow or lose precision in 16-bit formats. Upcast inside the norm and cast back.

Do not judge an MoE by its active parameters alone. Memory follows the total. An 8-expert model with 2 active needs the memory of the whole set; if experts are spread over devices, the communication of tokens to experts can dominate.

Monitor expert usage. Log the fraction of tokens per expert and the balance loss in training; a collapsing router is visible long before the final loss shows it.

Fine-tune carefully. LoRA on an MoE can target the attention projections (cheap and stable) or the experts (many adapters). Start with attention only (Module 8).

For robots. Memory-limited on-board hardware rarely benefits from MoE: you pay for all experts in memory. A dense small model with gated layers is usually the better fit; MoE is a tool for servers with plenty of memory and a compute budget to save.

Common misconceptions

Engineering takeaways

1. RMSNorm = γ · x / √(mean(x²) + ε): it rescales without re-centring, with d instead of 2d parameters.
2. SwiGLU = W_down (SiLU(W_gate x) ⊙ W_up x): three matrices; h ≈ 8d/3 keeps the budget near the classic 8d².
3. A mixture of experts with E experts and top-k routing stores E·3dh + dE parameters and uses k·3dh + dE per token.
4. Balance routing with an auxiliary loss E·Σ f_e P_e (equal to 1 for perfectly even routing) and capacity limits.
5. Memory follows total parameters, arithmetic follows active parameters.

References and further reading

The numbers in the examples were computed by the code in this lesson. The cited papers report the quality comparisons; they are not reproduced here. The balance loss shown follows the common formulation of later mixture-of-experts work.

• Zhang, B. and Sennrich, R. (2019). Root mean square layer normalization. arXiv:1910.07467.
• Shazeer, N. (2020). GLU variants improve transformer. arXiv:2002.05202.
• Shazeer, N. et al. (2017). Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv:1701.06538.
• Touvron, H. et al. (2023). Llama 2: Open foundation and fine-tuned chat models. arXiv:2307.09288.
• Shazeer (2020), GLU Variants (research paper)
• Shazeer et al. (2017), Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (research paper)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: KV caching, prefill and decode: what the cache stores, why it is exact, and what it costs.

Related content

Nothing below is required to finish this lesson.

• Before
  LM10-T3: Multi-Query and Grouped-Query Attention

  Share key and value heads to cut KV-cache size, and count the saving.

• Lab
  LM10-L2: Compare Attention Variants

  Count parameters, cache size and compute for multi-head, multi-query and grouped-query attention.