Transformer Init, Causal Mask & Forward Pass

All you need is attention. And layer normalization. And positional encodings. And residual connections. And feed-forward networks. But mostly attention.

Norm, Perpetually Normalizing AI Agent

This continuation of Section 3.1 picks up after the structural elements of a single transformer block. It assembles those parts into a working decoder: how weights are initialized so signal flows cleanly, how the causal mask turns self-attention into autoregressive prediction, what the full forward pass looks like in code, and how to count parameters so you can predict memory before you allocate it.

3.2.1 Weight Initialization

Proper initialization is critical for training deep Transformers. The standard approach uses Xavier (Glorot) initialization for most weights: values are drawn from a uniform or normal distribution with variance $2 / (fan_{in} + fan_{out})$. This ensures that the variance of activations stays roughly constant as they propagate forward through layers.

A subtle but important refinement, used in GPT-2 and later models, is to scale the initialization of the output projection in the residual path by $1 / \sqrt{2N}$, where $N$ is the number of layers. The factor of 2 comes from having two residual sub-layers per block (attention and FFN). This keeps the residual stream variance from growing as $O(N)$ through the network.

from torch import nn
# GPT-2 style weight initialization: N(0, 0.02) for most layers,
# scaled by 1/sqrt(2*n_layers) for residual projections to stabilize depth.
def init_weights(module, n_layers):
    """GPT-2 style initialization."""
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            def scale_residual_init(module, n_layers):
                """Scale output projections in residual blocks."""
                for name, param in module.named_parameters():
                    if name.endswith('W_o.weight') or name.endswith('net.3.weight'):
                        # net.3 is the second linear layer in the FFN
                        nn.init.normal_(param, mean=0.0, std=0.02 / (2 * n_layers) ** 0.5)

3.2.2 The Causal Mask (Decoder Self-Attention)

In auto-regressive language models (and in the decoder of the original Transformer), each position can only attend to itself and to earlier positions. This is enforced by a causal mask: an upper-triangular matrix of -inf values added to the attention scores before the softmax. The softmax converts -inf to zero, effectively blocking information flow from future tokens.

import torch
# Build a causal (lower-triangular) boolean mask that blocks each token
# from attending to future positions, enforcing autoregressive generation.
def causal_mask(seq_len, device):
 """Returns a boolean mask: True = allowed, False = blocked."""
return torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool))
# Usage in attention:
# scores.masked_fill_(~mask, float('-inf'))

Figure 3.2.5a: The residual stream as a branch-and-merge highway. Each sub-layer (Attention, FFN) reads the current stream, computes its contribution, and adds it back. The stream is never overwritten, sub-layers can only add. This is why information from early layers reaches late layers directly, and why deleting a layer often degrades the model less than expected.

The mask ensures that the prediction for position $t$ depends only on tokens at positions $0, 1, ..., t$. This is what makes the model auto-regressive: during generation, each new token can be produced by conditioning only on the tokens generated so far.

3.2.2.1 The KV cache: Avoiding Redundant Computation

Consider what happens when a decoder-only model generates text one token at a time. To produce token $t+1$, the model runs a forward pass over the entire sequence $[x_0, x_1, ..., x_t]$. To produce token $t+2$, it must run over $[x_0, x_1, ..., x_t, x_{t+1}]$. Notice that the Key and Value projections for positions $0$ through $t$ are identical in both passes. Recomputing them from scratch at every step is extremely wasteful, turning generation into an O(n²) operation in sequence length when it could be O(n).

The solution is the KV cache: after computing the Key and Value vectors for each position, the model stores them in a cache. When generating the next token, only the new token's Query, Key, and Value need to be computed. The new Key and Value are appended to the cache, and the attention for the new Query is computed against all cached Keys and Values. This reduces each generation step from processing the full sequence to processing a single token (plus the cached context), providing a dramatic speedup in practice.

The tradeoff is memory: the KV cache stores two tensors (K and V) per layer, per attention head, for every token in the sequence. For a model with $L$ layers, $H$ heads, and head dimension $d_h$, the cache for a sequence of length $n$ requires $2 \times L \times H \times d_h \times n$ elements. For large models serving long contexts, this cache can consume tens of gigabytes of GPU memory. Managing KV cache memory efficiently is one of the central challenges of LLM inference, and techniques like grouped-query attention (GQA), paged attention (vLLM), and cache quantization have been developed to address it. We will explore these optimizations in detail in Chapter 9: Inference Optimization.

3.2.3 The Complete Forward Pass

Let us trace a single forward pass through a decoder-only Transformer, the architecture used by GPT and most modern LLMs:

1. Tokenize the input text into a sequence of integer token IDs.
2. Embed the tokens: look up each ID in the embedding table and scale by $\sqrt{d}$.
3. Add positional encoding (sinusoidal, learned, or RoPE).
4. For each of the $N$ Transformer blocks:
   1. Apply Layer Normalization (Pre-LN).
   2. Compute Masked Multi-Head Self-Attention (with the causal mask).
   3. Add the residual (skip connection).
   4. Apply Layer Normalization (Pre-LN).
   5. Apply the Feed-Forward Network.
   6. Add the residual.
5. Apply a final Layer Normalization.
6. Project to vocabulary size with a linear layer (often weight-tied with the embedding matrix).
7. Apply softmax to obtain next-token probabilities.

Algorithm: Decoder-Only Transformer (next-token logits)
Input:  token IDs t in {0,..,V-1}^{B x T}
Output: logits L in R^{B x T x V}

  // 1. Embed + add positional information
  x := Embed[t] * sqrt(d_model)                 // (B, T, d_model)
  x := x + PositionalEncoding(T)                // sinusoidal, learned, or RoPE-applied later

  // 2. Build causal mask
  M[i, j] := 0 if j <= i else -inf              // (T, T)

  // 3. Stack of N Pre-LN blocks
  For l = 1..N:
      h := LayerNorm(x)                         // pre-attention norm
      a := MultiHeadAttention(h, mask = M)      // Algorithm 3.2.1, applied per head
      x := x + a                                // residual

      h := LayerNorm(x)                         // pre-FFN norm
      f := FFN(h)                               // two-layer MLP, often SwiGLU
      x := x + f                                // residual

  // 4. Final norm + unembedding
  x := LayerNorm(x)
  L := x @ E^T                                  // weight-tied with embedding E, shape (B, T, V)
  Return L

3.2.4 Information Flow Through the Residual Stream

A powerful mental model for understanding Transformers (popularized by Elhage et al. at Anthropic) is the residual stream perspective. Instead of viewing the Transformer as a sequence of layers, imagine a single stream of vectors (one per position) flowing from the input embedding to the output. Each attention layer and each FFN layer reads from and writes to this stream additively:

Each sub-layer sees the accumulated state of the residual stream up to that point, performs some computation, and adds its contribution back. This means that earlier layers can communicate with later layers directly through the residual stream, without the information needing to "pass through" every intermediate layer. It also means that deleting a layer from a trained Transformer may be less catastrophic than you might expect; the residual stream carries forward most of the information regardless.

Figure 3.2.6: The residual stream perspective (Elhage et al., 2021). The Transformer's residual path acts as a shared communication channel. Each attention and FFN sub-layer reads from the stream and adds its output back, rather than sequentially transforming a single representation.

This perspective is central to the field of mechanistic interpretability, where researchers decompose the behavior of trained Transformers into the contributions of individual heads and MLP layers. We will return to this in a later module.

3.2.5 Putting It All Together: Parameter Counts

For a Transformer with $N$ layers, model dimension $d$, feed-forward dimension $d_{ff} = 4d$, $h$ attention heads, and vocabulary size $V$, the parameter count is approximately:

For GPT-3 (N=96, d=12288, V=50257), this gives roughly 175 billion parameters. The FFN contributes about twice as many parameters as the attention layers, a ratio that remains consistent across model sizes.

Exercises