JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

Inside the block: RMSNorm, SwiGLU, and sparse experts

The feed-forward half of a transformer quietly carries most of its knowledge and most of its parameters. Tour the modern block — pre-norm with RMSNorm, gated SwiGLU — then make it sparse with mixture-of-experts.

The 2024-era block, part by part

A modern transformer block is a small, opinionated recipe: normalize, attend, add the residual, normalize again, run a feed-forward block, add the residual. Two refinements separate a Llama-era block from the 2017 original. First, normalization moved before each sub-layer (pre-norm) and switched from LayerNorm to RMSNorm. Second, the feed-forward block adopted a gated SwiGLU nonlinearity. Both are small on paper and large in practice.

The modern pre-norm block: RMSNorm to attention to residual, then RMSNorm to feed-forward to residual.

Diagram of a transformer block showing normalization, attention, residual connections, and a feed-forward sub-layer.

It helps to remember the budget: in a dense transformer the feed-forward sub-layers hold roughly two-thirds of the parameters. Attention gets the headlines, but the feed-forward block is where most of the model's factual capacity lives — which is exactly why making it sparse is the highest-leverage way to scale.

RMSNorm: drop the mean, keep the scale

Vol I's LayerNorm does two things to an activation vector: subtracts its mean (re-centering) and divides by its standard deviation (re-scaling). RMSNorm keeps only the second: it divides by the root-mean-square of the vector and skips mean subtraction entirely. Empirically the re-centering contributes little to stability, so dropping it costs almost no quality.

\bar{x}_i = \frac{x_i}{\sqrt{\frac{1}{n}\sum_{j=1}^{n} x_j^2 + \epsilon}}\; g_i

RMSNorm rescales by the root-mean-square alone — no mean is subtracted, only a learned gain g remains.

Why bother? Speed and simplicity. RMSNorm has fewer operations and no mean to reduce across the feature dimension, which matters when normalization runs twice per layer across a very deep model. It is the kind of change that looks like a rounding error per call and adds up to real training-time savings at scale — and it has become the default in essentially every recent open LLM.

# LayerNorm:  y = g * (x - mean(x)) / std(x) + b
# RMSNorm:    y = g * x / sqrt(mean(x*x) + eps)
# note: no mean subtraction, no bias term -> fewer ops per call
RMSNorm keeps the rescaling and discards the recentering of LayerNorm.

SwiGLU: a gated feed-forward

A classic feed-forward block is project-up, apply a nonlinearity, project-down. SwiGLU splits the up-projection into two paths: one passes through a smooth SiLU activation and acts as a gate, the other is the value, and they are multiplied element-wise before projecting down. The gate lets the network modulate each value channel multiplicatively rather than just thresholding it the way a bare ReLU would.

SwiGLU's gate runs one projection through the smooth SiLU activation, x·σ(x) — a self-gated relative of the curves shown here.

Activation curves for ReLU, sigmoid, and tanh.

Gating costs a third matrix, so implementations shrink the hidden width (often to about two-thirds) to keep the parameter count fixed. Even at matched parameters SwiGLU consistently improves loss over a plain block — a rare upgrade that is nearly pure benefit, and the reason it is standard in the Llama, Mistral, and PaLM families.

Sparse mixture-of-experts: more parameters, same FLOPs

Here is the big lever. Vol I introduced the mixture-of-experts idea; the modern realization is the sparse MoE block. Replace the single feed-forward block with N independent expert blocks plus a small router network. For each token the router picks just the top-k experts (often k=2), and only those run. The model now holds N-times the feed-forward parameters but spends compute for only k of them per token — capacity decoupled from cost.

y = \sum_{i \in \mathrm{TopK}(x)} g_i(x)\, E_i(x), \qquad g(x) = \mathrm{softmax}(x W_g)

Sparse MoE: a softmax router picks the top-k experts and mixes only their outputs, so FLOPs stay fixed as total parameters grow.

The catch is load balance. If the router falls in love with a few experts, the rest sit idle and effective capacity collapses, so MoE training adds an auxiliary loss that pushes tokens to spread evenly. Switch Transformer simplified this lineage by routing each token to a single expert (top-1), proving that aggressive sparsity trains stably and scales to trillions of parameters while keeping per-token compute flat.

Token-chooses-expert routing must hope for balance and patch it with a loss. Expert-choice routing flips the direction: each expert selects a fixed quota of the tokens it wants most. Balance becomes a property of the construction rather than something coaxed by regularization — no token storms one expert, no expert starves — at the price of letting some tokens be processed by more experts than others.