Residual Connections and RMSNorm

A transformer layer transforms hidden states — but if each layer replaced the input entirely, information from earlier layers could be lost. Worse, gradients would have to pass through every transformation to reach early layers during training, making learning unstable.

A residual connection solves this: the layer's output is added to its input. The layer only needs to learn the change, not the entire new representation. Information passes through even if the layer does nothing useful.

But addition can cause scale to drift — after many layers, the vectors may grow without bound. RMSNorm (Root Mean Square Normalization) fixes this by rescaling vectors to a consistent magnitude before each sub-layer. In modern transformers, the norm is applied before the transform (pre-norm), not after.

The pattern inside a block is: norm → transform → add residual, repeated for both attention and FFN.

RMSNorm first normalizes a vector so its root-mean-square is 1, then applies a learned per-dimension scale. Given a vector x of dimension d:

The division by RMS(x) normalizes the vector magnitude. The learned scale parameter γ (one per dimension) lets the model adjust the normalized output — it can re-amplify dimensions that matter and suppress ones that do not.

Why RMSNorm instead of the more common LayerNorm? LayerNorm also subtracts the mean before normalizing, which requires computing both mean and variance. RMSNorm skips the mean subtraction. In practice, this is slightly faster (one fewer pass over the data) and works equally well for transformer hidden states, which is why most modern LLMs use it.

The original transformer paper applied normalization after the sub-layer: transform → add residual → norm. Most modern LLMs use pre-norm: norm → transform → add residual.

The reason is training stability, specifically gradient flow. In post-norm, the output of each block is norm(x + transform(x)). During backpropagation, gradients from higher layers must pass through the norm layer to reach lower layers. The norm's Jacobian can scale gradients up or down unpredictably, causing training instability in deep networks.

In pre-norm, the output is x + transform(norm(x)). The residual connection provides a direct, unmodified path for gradients: they flow through the addition without passing through any norm or transform. Each sub-layer only needs to learn a useful delta to add — if a layer has nothing useful to contribute, it can learn weights close to zero and the gradient still flows cleanly through the residual. This makes training much more stable for 32+ layer models.

The residual connection is a simple addition:

RMSNorm rescales each vector by its root mean square: