Attention and the transformer block

What it is and when you reach for it

Attention answers one question: given a sequence of vectors, how should each position gather information from the others? Before attention, the answers were recurrence (carry a state left to right, as in LSTMs) and convolution (mix a fixed local window). Both impose a path length between distant positions that grows with distance, and recurrence resists parallel training because step t depends on step t−1. Attention replaces both with a direct, content-based route: any position can read from any other in a single step, the routing weights are computed from the data rather than fixed by architecture, and the whole thing is a few matrix multiplies, which is exactly the shape of compute GPUs are best at. You reach for it whenever the relationships that matter are not local and not known in advance: which earlier word a pronoun refers to, which image patch contains the object a query cares about, which frame of a trajectory explains the current one. The transformer block is attention plus the machinery that makes stacking it a hundred layers deep trainable: residual connections, normalization, and a position-wise MLP. Almost everything after 2017 that people call "a transformer" is this one block repeated.

The math

Attention as a soft dictionary lookup

Start from an ordinary dictionary. A lookup takes a query, finds the key that exactly matches it, and returns that key's value. Attention relaxes "exactly matches" into "matches by dot product" and returns a weighted blend instead of a single value. Each input vector x_i is linearly projected three ways: a query q_i = x_iW_Q (what position i is looking for), a key k_i = x_iW_K (what position i offers to be found by), and a value v_i = x_iW_V (what position i hands over when selected). The output at position i is

out_i = Σ_j softmax_j(q_i·k_j / √d_k) v_j,

or in matrix form Attention(Q, K, V) = softmax(QK^T/√d_k) V. The softmax turns raw similarity scores into a convex combination, so the output always lies inside the convex hull of the values: a hard lookup in the limit where one score dominates, a uniform average in the limit where all scores tie, and everything in between is differentiable, which is what lets the projections be learned by gradient descent. The separation into three roles is the load-bearing design choice: what a token seeks, what it can be found by, and what it contributes are three different learned functions of the same vector, and conflating them (as older "dot-product over raw embeddings" models did) forces one representation to serve all three jobs.

A worked example with two positions and d_k = 2. Let the query be q = [1, 1], the keys k₁ = [1, 1] and k₂ = [1, −1], the values v₁ = [1, 0] and v₂ = [0, 1]. Raw scores are q·k₁ = 2 and q·k₂ = 0; dividing by √2 gives [1.414, 0]. Softmax of that is [e^1.414, e⁰] / (e^1.414 + e⁰) = [4.113, 1] / 5.113 ≈ [0.804, 0.196]. The output is 0.804·v₁ + 0.196·v₂ = [0.804, 0.196]: mostly the value of the well-matched key, with a differentiable remainder from the other.

Why the √d scaling exists

The √d_k divisor looks cosmetic and is not. Suppose the components of q and k are independent with mean 0 and variance 1, which is roughly what normalization layers and sensible initialization give you. Then the dot product q·k = Σ_m q_mk_m is a sum of d_k terms, each with mean 0 and variance 1, so it has variance d_k and standard deviation √d_k. At d_k = 64 (a typical head size) the raw scores have standard deviation 8, and a softmax over logits that spread out is nearly one-hot. That is bad for learning, not for prediction: the softmax Jacobian is diag(y) − yy^T, and every entry of it goes to zero as y saturates toward a one-hot vector, so gradients through the attention weights vanish and the projections stop training. Dividing by √d_k restores the scores to unit variance regardless of head size, keeping the softmax in the regime where it still passes gradient. This is the same saturation analysis that motivates temperature in any softmax; here the "temperature" is derived from the geometry rather than tuned.

Causal masking

A language model trained to predict token t+1 must not let position t read positions after t, or the training task collapses into copying. The fix acts on the score matrix before the softmax: set every score s_ij with j > i to −∞. After exponentiation those entries become exactly 0, the softmax renormalizes over the surviving prefix, and each row remains a valid distribution over positions 1..i. Two details matter. The mask must be applied before the softmax, not by zeroing weights after it, because zeroing afterwards leaves rows that no longer sum to one and quietly leaks probability mass to the future during normalization. And masking with a very negative finite number instead of −∞ works in float32 but can round to a nonzero weight in float16, so use the dtype's proper −inf and let exp(−inf) = 0 do the work. The triangular structure also means a causal model can cache keys and values for a growing prefix and attend from only the newest query at inference time, which is the entire basis of fast autoregressive decoding.

Multi-head attention: parallel subspace attention

One attention pattern per layer is a bottleneck: the softmax produces a single distribution per query, so a position that needs to simultaneously look at "the previous verb" and "the matching open bracket" has to average the two needs into one blurry lookup. Multi-head attention runs h independent attentions in parallel, each in a d_model/h dimensional subspace: the projections W_Q, W_K, W_V are split into h smaller ones, each head computes softmax(Q_hK_h^T/√(d/h))V_h with its own weights and its own attention pattern, the h outputs are concatenated back to d_model, and a final output projection W_O mixes them. Because the subspaces are slices of the same projections, the parameter count and FLOPs are essentially identical to one full-width head; what changes is expressiveness, h separate distributions instead of one. In trained models the heads do specialize measurably: some track positional offsets, some track syntax, some are induction heads that implement in-context copying. In code, heads are nothing more than a reshape of (B, T, C) into (B, h, T, C/h) so the batched matmul treats the head axis as more batch.

The transformer block

Attention mixes information across positions but is itself almost linear: values are averaged, and averaging cannot compute per-position nonlinear functions of what was gathered. The block therefore pairs it with a position-wise MLP (two linear layers with a GELU between, hidden width conventionally 4×d_model) that processes each position independently. A useful reading: attention moves information between positions, the MLP transforms information within a position. Both sublayers sit on a residual stream:

x ──────────────────────────────┐
│                               │
├─► LayerNorm ─► Attention ─►(+)┤   x = x + Attn(LN(x))
│                               │
├─► LayerNorm ─► MLP ────────►(+)   x = x + MLP(LN(x))
│                               │
▼                               ▼
        (repeated N times)

The residual connections make the identity the default behavior of every layer: a block that has learned nothing yet passes its input through unchanged, so gradients flow to the bottom of a deep stack from step one. Normalization placement is the one real fork in the road. The original transformer used post-LN, x = LN(x + sublayer(x)), which normalizes the residual stream itself; it can reach slightly better final loss but is famously touchy at depth, needing learning-rate warmup to avoid early divergence because gradient magnitudes vary sharply across layers. GPT-2 popularized pre-LN, x = x + sublayer(LN(x)), where the residual path is never touched and each sublayer merely reads a normalized view of the stream and writes a delta back. Pre-LN trains stably at large depth without warmup gymnastics and is the default in essentially all modern LLMs (usually with RMSNorm in place of LayerNorm); the one wrinkle is that the stream itself is never normalized, so a final norm after the last block is required before the output head. The implementations below are pre-LN.

Implementation, twice

First the core function and the multi-head causal self-attention module. The PyTorch version mirrors the CausalSelfAttention module in nanoGPT (my annotated walkthrough of that codebase lives at /oss/nanogpt): one fused linear produces q, k, and v for all heads at once, and heads are a reshape, not a loop. The JAX version says the same thing with einsum, with every shape written out.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def scaled_dot_product_attention(q, k, v, causal=False):
    """q, k: (..., T, d); v: (..., T, dv). Returns (..., T, dv)."""
    d = q.size(-1)
    # (..., Tq, d) @ (..., d, Tk) -> (..., Tq, Tk)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d)
    if causal:
        Tq, Tk = scores.shape[-2:]
        future = torch.triu(torch.ones(Tq, Tk, dtype=torch.bool,
                                       device=scores.device), diagonal=1)
        # -inf BEFORE softmax: exp(-inf) = 0, rows renormalize
        # over the visible prefix and still sum to 1.
        scores = scores.masked_fill(future, float('-inf'))
    w = F.softmax(scores, dim=-1)
    return w @ v                          # convex combination of values

class CausalSelfAttention(nn.Module):
    """Multi-head causal self-attention, shaped like nanoGPT's.

    One fused linear emits q, k, v for every head at once;
    splitting into heads is a view + transpose, so multi-head
    costs the same parameters and FLOPs as one full-width head.
    """

    def __init__(self, n_embd, n_head):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)   # q, k, v together
        self.c_proj = nn.Linear(n_embd, n_embd)       # mixes heads (W_O)

    def forward(self, x):                             # x: (B, T, C)
        B, T, C = x.shape
        q, k, v = self.c_attn(x).split(C, dim=2)      # each (B, T, C)
        # (B, T, C) -> (B, n_head, T, C // n_head): head axis is batch
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        y = scaled_dot_product_attention(q, k, v, causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # concat heads
        return self.c_proj(y)

import jax
import jax.numpy as jnp

def scaled_dot_product_attention(q, k, v, causal=False):
    """q, k: (..., T, d); v: (..., T, dv). Returns (..., T, dv)."""
    d = q.shape[-1]
    # (..., Tq, d) x (..., Tk, d) -> (..., Tq, Tk)
    scores = jnp.einsum('...qd,...kd->...qk', q, k) / jnp.sqrt(d)
    if causal:
        Tq, Tk = scores.shape[-2:]
        visible = jnp.tril(jnp.ones((Tq, Tk), dtype=bool))
        # -inf BEFORE softmax so masked weights are exactly zero
        # and each row renormalizes over its visible prefix.
        scores = jnp.where(visible, scores, -jnp.inf)
    w = jax.nn.softmax(scores, axis=-1)   # rows sum to 1
    # (..., Tq, Tk) x (..., Tk, dv) -> (..., Tq, dv)
    return jnp.einsum('...qk,...kv->...qv', w, v)

def init_attention(key, n_embd):
    k1, k2 = jax.random.split(key)
    s = 1.0 / jnp.sqrt(n_embd)
    return {
        'w_qkv': jax.random.normal(k1, (n_embd, 3 * n_embd)) * s,
        'b_qkv': jnp.zeros(3 * n_embd),
        'w_out': jax.random.normal(k2, (n_embd, n_embd)) * s,
        'b_out': jnp.zeros(n_embd),
    }

def causal_self_attention(params, x, n_head):
    B, T, C = x.shape                             # (batch, time, model dim)
    qkv = x @ params['w_qkv'] + params['b_qkv']   # (B, T, 3C)
    q, k, v = jnp.split(qkv, 3, axis=-1)          # each (B, T, C)

    def heads(t):                                 # (B, T, C) -> (B, H, T, C//H)
        return t.reshape(B, T, n_head, C // n_head).swapaxes(1, 2)

    y = scaled_dot_product_attention(heads(q), heads(k), heads(v),
                                     causal=True)  # (B, H, T, C//H)
    y = y.swapaxes(1, 2).reshape(B, T, C)          # concat heads
    return y @ params['w_out'] + params['b_out']

Then the full pre-LN block: two sublayers, two norms, two residual additions, and a 4× MLP. Note how short the forward is once the sublayers exist; the block is plumbing, and the plumbing is the point.

class MLP(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.c_fc = nn.Linear(n_embd, 4 * n_embd)   # conventional 4x widen
        self.c_proj = nn.Linear(4 * n_embd, n_embd)

    def forward(self, x):
        return self.c_proj(F.gelu(self.c_fc(x)))

class Block(nn.Module):
    """Pre-LN transformer block. The residual stream is never
    normalized in place: each sublayer reads a normalized view
    and writes a delta back, so the identity path stays clean."""

    def __init__(self, n_embd, n_head):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head)
        self.ln_2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # mix across positions
        x = x + self.mlp(self.ln_2(x))    # transform within positions
        return x

def layer_norm(x, gain, bias, eps=1e-5):
    mu = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    return gain * (x - mu) / jnp.sqrt(var + eps) + bias

def init_block(key, n_embd):
    ka, k1, k2 = jax.random.split(key, 3)
    s = 1.0 / jnp.sqrt(n_embd)
    return {
        'attn': init_attention(ka, n_embd),
        'ln1': (jnp.ones(n_embd), jnp.zeros(n_embd)),
        'ln2': (jnp.ones(n_embd), jnp.zeros(n_embd)),
        'w_fc': jax.random.normal(k1, (n_embd, 4 * n_embd)) * s,
        'b_fc': jnp.zeros(4 * n_embd),
        'w_proj': jax.random.normal(k2, (4 * n_embd, n_embd)) * s,
        'b_proj': jnp.zeros(n_embd),
    }

def block(params, x, n_head):
    """Pre-LN: x = x + Attn(LN(x)); x = x + MLP(LN(x))."""
    x = x + causal_self_attention(
        params['attn'], layer_norm(x, *params['ln1']), n_head)
    h = layer_norm(x, *params['ln2'])
    h = jax.nn.gelu(h @ params['w_fc'] + params['b_fc'])
    return x + h @ params['w_proj'] + params['b_proj']

Using it on a real shape of problem

GPT-2 small dimensions are the standard smoke test: batch 8, context 1024, model width 768, 12 heads of size 64. The block is shape-preserving, (B, T, C) in and (B, T, C) out, which is what makes stacking trivial.

torch.manual_seed(0)
blk = Block(n_embd=768, n_head=12)
x = torch.randn(8, 1024, 768)
y = blk(x)
print(y.shape)                    # torch.Size([8, 1024, 768])
print(sum(p.numel() for p in blk.parameters()))  # ~7.1M per block

key = jax.random.PRNGKey(0)
params = init_block(key, n_embd=768)
x = jax.random.normal(key, (8, 1024, 768))
apply = jax.jit(lambda p, x: block(p, x, n_head=12))
y = apply(params, x)              # first call compiles, later calls are fast
print(y.shape)                    # (8, 1024, 768)

The thing to internalize from running this is the memory shape, not the output. The materialized score matrix is (B, h, T, T): at these dimensions that is 8 × 12 × 1024 × 1024 float32 values, about 400 MB for a single block's forward, growing quadratically with context. That number, more than any FLOP count, is why the fused kernels in the next-to-last section exist. Exact timings are machine-dependent, but on any GPU you will see the naive version's memory climb with T² while a fused kernel's stays flat.

Applications

The honest summary is "everything", but it pays to say it precisely. Decoder-only language models (GPT-4, Llama, Claude, Gemini) are exactly the causal block above repeated 30 to 100+ times with a token embedding at the bottom and a softmax over the vocabulary at the top; the differences between them live in normalization flavor, positional encoding, and scale, not in the mechanism. Encoder-style bidirectional attention (no causal mask) powers BERT-family retrieval and embedding models. Vision Transformers cut an image into 16×16 patches, embed each patch as a token, and run the identical block; ViT and its descendants replaced convnets at the top of most image benchmarks and are the visual encoder inside CLIP and most multimodal LLMs. Cross-attention, where queries come from one sequence and keys and values from another, is the standard conditioning mechanism: in Stable Diffusion, U-Net layers attend from image feature positions to CLIP text embeddings, which is literally how the prompt steers the image. The same cross-attention pattern conditions Whisper's speech decoder on audio features and lets AlphaFold's structure module attend over residue pairs. When one mechanism serves language, vision, audio, protein structure, and robot control, the interesting question stops being "where does it apply" and becomes "what are its costs", which is the next section.

Against the real libraries

The from-scratch version above materializes the (T, T) score matrix in global memory. Production kernels do not, and that is the entire gap.

torch.nn.functional.scaled_dot_product_attention is the API to reach for first in PyTorch. It is one call with the exact semantics of the reference function above, and behind it a dispatcher picks among fused backends: a FlashAttention-2 kernel, a memory-efficient kernel derived from the xFormers line of work, a cuDNN backend, and a plain math fallback that behaves like the naive version. You can pin a backend for benchmarking with the torch.nn.attention.sdpa_kernel context manager. nanoGPT itself switched its inner loop to this call; my notes on the surrounding codebase are at /oss/nanogpt.

flash-attention (Dao et al.) is the kernel most of those backends trace back to. The algorithmic core is online softmax: because a softmax can be computed in one streaming pass carrying only a running (max, sum) pair, the score matrix can be processed tile by tile in on-chip SRAM and never written to global memory, turning attention from memory-bound to compute-bound and making 100k-token contexts affordable. I derive that rescaling identity on the softmax page, and my own step-by-step CUDA reconstruction of the forward pass is written up at /oss/flash-attention. The production repo adds what a reference implementation never needs: dropout inside the kernel, paged KV-cache support for serving, head dimensions beyond 128, and a backward pass that recomputes tiles instead of storing them.

On the JAX side, flax's MultiHeadDotProductAttention module is the production analogue of the functional version above, adding the things real training runs need: separate q/k/v/output kernels with configurable initializers, attention dropout, an autoregressive decoding cache, and pluggable attention functions so a fused TPU or Pallas kernel can be swapped in. Recent JAX also ships jax.nn.dot_product_attention with a cuDNN flash path. And huggingface/transformers (notes at /oss/transformers) is where you see the block at industrial scale: every model file in that repo is some dialect of the block on this page, with per-model choices of norm, positional scheme, and attention backend selected by one attn_implementation flag.

The from-scratch version is genuinely enough for research at small scale: short contexts, model surgery, interpretability work where you want the attention weights as an inspectable tensor (the fused kernels never form them). The correctness check is one comparison, run in float64 so tolerance failures mean logic bugs rather than accumulation order:

torch.manual_seed(0)
q, k, v = torch.randn(3, 4, 8, 128, 64, dtype=torch.float64).unbind(0)

ref = F.scaled_dot_product_attention(q, k, v, is_causal=True)
ours = scaled_dot_product_attention(q, k, v, causal=True)

print((ref - ours).abs().max())   # ~1e-16 in float64
assert torch.allclose(ref, ours, atol=1e-12)
# In float32 expect ~1e-6: same math, different summation order.

key = jax.random.PRNGKey(0)
kq, kk, kv = jax.random.split(key, 3)
# jax.nn.dot_product_attention uses (B, T, H, d) layout
q = jax.random.normal(kq, (4, 128, 8, 64))
k = jax.random.normal(kk, (4, 128, 8, 64))
v = jax.random.normal(kv, (4, 128, 8, 64))

ref = jax.nn.dot_product_attention(q, k, v, is_causal=True)
# our function wants (B, H, T, d): move the head axis
ours = scaled_dot_product_attention(
    q.swapaxes(1, 2), k.swapaxes(1, 2), v.swapaxes(1, 2),
    causal=True).swapaxes(1, 2)

print(jnp.abs(ref - ours).max())  # ~1e-6 in float32
assert jnp.allclose(ref, ours, atol=1e-4)

Traps and misconceptions

Masking after the softmax. Zeroing forbidden weights after normalization leaves rows that do not sum to one and lets future positions influence the normalizer. The mask must be −∞ on the scores, before the softmax, so the surviving prefix renormalizes cleanly. Relatedly, a row where everything is masked (padding-only rows in an encoder) makes softmax produce NaN, since it normalizes over an empty set; production code guards those rows explicitly.

"More heads means more compute." Heads split the model width; they do not multiply it. Twelve heads of size 64 cost the same parameters and matmul FLOPs as one head of size 768. What you buy is twelve independent attention distributions per layer; what you pay is a smaller dot-product subspace per head, which is why head size rarely drops below about 64.

Forgetting the scaling, or scaling by d. Dropping the √d_k often still trains at small width and then quietly degrades at larger d_k as the softmax saturates and attention gradients vanish. The divisor is √d, not d: the standard deviation of the dot product grows as √d, and overcorrecting by d pushes the softmax toward uniform, which blurs every lookup.

"Attention is all you need, so the MLP is optional." In parameter terms the block is mostly MLP: at 4× widening the MLP holds about two thirds of the block's weights, and ablating it cripples the model. Attention routes; the MLP computes. Interpretability work locates most factual-recall behavior in the MLPs, not the attention maps.

Reading attention weights as explanations. The weights say where a head looked, in one layer, in one subspace, before the output projection and the other heads mix everything back together. They are a useful debugging signal and a weak explanation of model behavior, and several papers have shown you can perturb them substantially without changing predictions. Treat weight visualizations as hypotheses, not evidence.