Home / Fundamentals / Transformer Architecture

Positions, normalization and the residual stream

Sinusoidal, learned, RoPE and ALiBi positions; context extension; LayerNorm vs RMSNorm; pre-norm vs post-norm; SwiGLU; and the residual-stream view of a transformer.

15 min read

Attention by itself is permutation-equivariant: shuffle the input tokens and the outputs shuffle the same way. Position information has to be added. The rest of the block (normalization, residual connections and the MLP) determines whether a deep stack trains stably.

Absolute positions

  • Sinusoidal (original Transformer): add PE(pos, 2i) = sin(pos / 10000^(2i/d)) and PE(pos, 2i+1) = cos(…) to the token embedding. Fixed, no parameters, but models extrapolate poorly beyond training lengths.
  • Learned absolute: a trainable vector per position (BERT, GPT-2). Simple, but undefined beyond the maximum trained position.

Rotary position embeddings (RoPE)

RoPE encodes position by rotating query and key vectors rather than adding to them. Split a d_k-dimensional vector into d_k/2 pairs. For position m, rotate pair i by angle m·θ_i, where θ_i = base^(−2i/d_k) and base is typically 10,000 (larger in long-context models):

[x'_{2i}  ]   [cos mθ_i   −sin mθ_i] [x_{2i}  ]
[x'_{2i+1}] = [sin mθ_i    cos mθ_i] [x_{2i+1}]

Because rotations compose, the score between a query at position m and a key at position n depends only on their offset:

⟨R_m q, R_n k⟩ = ⟨q, R_mᵀ R_n k⟩ = ⟨q, R_{n−m} k⟩

So RoPE gives relative position information while still being applied position-by-position, which also makes it compatible with the KV cache: keys are rotated once when written.

def rope(x, positions, base=10000.0):
    # x: (T, d_k) with d_k even; positions: (T,)
    d_k = x.shape[-1]
    inv_freq = base ** (-np.arange(0, d_k, 2) / d_k)          # (d_k/2,)
    angles = positions[:, None] * inv_freq[None, :]           # (T, d_k/2)
    cos, sin = np.cos(angles), np.sin(angles)
    x1, x2 = x[:, 0::2], x[:, 1::2]
    out = np.empty_like(x)
    out[:, 0::2] = x1 * cos - x2 * sin
    out[:, 1::2] = x1 * sin + x2 * cos
    return out

The pairs rotate at geometrically spaced frequencies: high-frequency pairs distinguish nearby positions, low-frequency pairs vary slowly and carry long-range information.

Extending context with RoPE

A model trained at 8k positions sees rotation angles it has never seen when run at 32k. Common fixes, usually followed by a short fine-tune on long sequences:

  • Position interpolation: scale positions by L_train / L_target so angles stay in the trained range. Simple, but compresses nearby positions together.
  • NTK-aware scaling: increase the base, which stretches low frequencies more than high ones, preserving local resolution.
  • YaRN: interpolate different frequency bands differently and adjust the attention temperature; works with short fine-tunes.
  • Train with a large base from the start, or progressively increase context during pre-training.

ALiBi

ALiBi adds a fixed, head-specific linear penalty to attention scores, −slope_h · |i − j|, and no position embedding at all. It's simple and extrapolates to longer inputs gracefully, though RoPE has become more common in large models.

Normalization

LayerNorm normalizes each token's vector to zero mean and unit variance, then applies a learned scale and shift:

LN(x) = γ ⊙ (x − μ) / √(σ² + ε) + β

RMSNorm drops the mean-centering and the bias:

RMSNorm(x) = γ ⊙ x / √(mean(x²) + ε)

It's cheaper and works as well in practice, so most recent LLMs use it.

Pre-norm vs. post-norm

Post-norm (original):   x = Norm(x + Sublayer(x))
Pre-norm (modern):      x = x + Sublayer(Norm(x))
  • In post-norm, every residual path passes through a normalization, so gradients to early layers are reshaped at every layer. Deep post-norm stacks need careful learning-rate warmup and are prone to divergence.
  • In pre-norm, the residual stream is an identity path from input to output. Gradients flow directly, and deep models train stably with less warmup. A final normalization before the output head is required, because the stream itself is never normalized.
  • The trade-off: with pre-norm, the stream's magnitude grows with depth, so later layers' contributions can become relatively small. Mitigations include scaled initialization of output projections (e.g. by 1/√(2L)), alternative schemes such as DeepNorm, or extra norms.
  • QK-norm (normalizing queries and keys before the dot product) prevents attention logits from growing without bound, a common source of loss spikes at scale.

The MLP block

The position-wise MLP holds about two thirds of a dense model's parameters. Modern models use a gated variant:

SwiGLU(x) = W_2 · ( SiLU(W_1 x) ⊙ (W_3 x) )

Three matrices instead of two, so the hidden size is usually set to about 8d/3 (rounded) to keep parameters near 8d², the same as a 4× GELU MLP. Gated MLPs consistently improve quality per parameter.

The residual-stream view

Scroll sideways to see the whole diagram.

The residual streamComponents: Token embedding (+ positions via RoPE); Residual stream x ∈ ℝ^d; RMSNorm → Attention (mixes across positions); RMSNorm → SwiGLU MLP (per position); Final norm → LM head (logits over vocabulary). Flow: Token embedding to Residual stream x ∈ ℝ^d; Residual stream x ∈ ℝ^d to Final norm → LM head.× L LAYERSToken embedding+ positions via RoPEResidual stream x ∈ ℝ^dRMSNorm → Attentionmixes across positionsRMSNorm → SwiGLU MLPper positionFinal norm → LM headlogits over vocabularyreadaddreadadd
The residual stream. Each position carries a d-dimensional residual stream. Every attention and MLP block reads it through a normalization and adds its output back; nothing overwrites it. A final normalization and the LM head turn the stream into next-token logits.

A useful way to think about a pre-norm transformer: there is one residual stream per position, a d-dimensional vector. Each attention and MLP block reads from it (through a normalization) and adds its output back. Nothing overwrites the stream; blocks only add to it. This is why:

  • information can pass through many layers unchanged,
  • removing a single layer often degrades a model gracefully,
  • interpretability work can attribute the final logits to contributions from individual heads and MLPs.

A modern decoder block

class Block(nn.Module):
    def __init__(self, d, n_heads, n_kv_heads, d_ff):
        super().__init__()
        self.norm1 = RMSNorm(d)
        self.attn  = GroupedQueryAttention(d, n_heads, n_kv_heads, rope=True)
        self.norm2 = RMSNorm(d)
        self.mlp   = SwiGLU(d, d_ff)

    def forward(self, x, kv_cache=None):
        x = x + self.attn(self.norm1(x), kv_cache=kv_cache)   # read, then add
        x = x + self.mlp(self.norm2(x))
        return x

Interview angles

  • Derive RoPE's relative-position property from rotation composition.
  • Compare interpolation, NTK scaling and YaRN for context extension.
  • Explain why pre-norm trains more stably and what it costs.
  • Give the RMSNorm formula and why it replaced LayerNorm.
  • Explain why SwiGLU's hidden size is 8d/3.