Final norm

The last RMSNorm in the network — applied once, immediately before the output projection turns the residual stream into logits.

After the transformer blocks (each with attention and feed-forward), the residual stream holds a dim-dimensional vector (e.g. 64) representing the model's full understanding of the sequence so far. Final norm is a single RMSNorm applied to this vector before anything else reads it. Its job is the same as the per-block norms: measure the root-mean-square of the vector — rms = √(mean(xᵢ²) + ε) — then scale each element by xᵢ / rms · gᵢ, where g is a learned per-dimension gain of shape [64]. The result has controlled magnitude regardless of how many additions the stream has accumulated across all the layers.

Without this step, the output of the last block could have a very different scale than the output projection was trained to expect. Normalising here is cheap (no matrix multiply, just a rescale) and keeps training stable as the model depth grows. It also keeps the logit magnitudes predictable: if the norm is absent, the output projection can produce extremely large or small logits, causing the softmax distribution to collapse or flatten chaotically.

In the diagram, the final-norm node appears on the residual stream just before the output projection at the right end. The feature-rail below shows the 64-element vector at that stage: before norm the values can span a wide range; after, every dimension is rescaled to a consistent neighbourhood around ±1 (modulated by the learned gain g). The individual gains are themselves a small set of learned parameters — one float per dimension, a tiny fraction of the model's total weights.

Related: RMSNorm · Logits · The transformer block