The Shape of a Transformer

A transformer is a residual stream with repeated processing stages.

Tokens enter as embeddings. Each block reads from the stream, computes an update, and adds that update back. The final representation is the accumulated sum of all block contributions. The shape follows from one decision: never replace the stream, only refine it.

The simplest shape: an embedding table and a sum#

The naive model maps each token ID to a vector and averages the vectors.

import numpy as np
 
vocab_size, d_model = 1000, 128
embed = np.random.randn(vocab_size, d_model).astype("float32")
 
def naive_forward(token_ids: list[int]) -> np.ndarray:
    vectors = embed[token_ids]          # shape: (seq_len, d_model)
    return vectors.mean(axis=0)         # one vector per sequence
 
ids = [12, 45, 45, 789]
print(naive_forward(ids).shape)  # (128,)

The bottleneck is order. Averaging discards position: "dog bites man" and "man bites dog" produce the same representation. There is also no memory between tokens, so earlier tokens cannot influence how later tokens are interpreted. Two fixes follow from these bottlenecks: add position information explicitly, and give tokens a mechanism to exchange information.

Position is a signal, not a byproduct#

Order must enter the representation before any mixing happens. A common choice adds a sinusoidal vector per position.

import numpy as np
 
def positional_encoding(seq_len: int, d_model: int) -> np.ndarray:
    position = np.arange(seq_len)[:, None]
    div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
    pe = np.zeros((seq_len, d_model), dtype="float32")
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)
    return pe
 
pe = positional_encoding(seq_len=16, d_model=128)
print(pe.shape)  # (16, 128)

Position vectors are added to token embeddings. Two sequences with the same tokens in different orders now produce different vectors at each position. The bottleneck shifts from order to context: each position still represents only itself. The next step adds communication between positions.

Attention is the communication channel#

Self-attention computes, for each position, a weighted average of all positions. Three projections produce queries, keys, and values.

import numpy as np
 
def scaled_dot_product_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray) -> np.ndarray:
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)              # (seq_len, seq_len)
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights /= weights.sum(axis=-1, keepdims=True)  # softmax rows
    return weights @ V                            # (seq_len, d_v)
 
seq_len, d_model = 8, 64
X = np.random.randn(seq_len, d_model).astype("float32")
W_q = W_k = W_v = np.random.randn(d_model, d_model).astype("float32") * 0.1
 
out = scaled_dot_product_attention(X @ W_q, X @ W_k, X @ W_v)
print(out.shape)  # (8, 64)

The softmax rows answer a routing question: for this position, which other positions matter. A single head learns one routing pattern. The bottleneck is that one pattern cannot express every dependency at once. Syntax, coreference, and long-range retrieval compete inside one matrix of weights. The fix is to run several heads in parallel.

Multi-head attention splits the routing problem#

Each head uses a smaller dimension and learns a different pattern. Outputs concatenate and project back to the stream width.

import numpy as np
 
def multi_head_attention(X: np.ndarray, num_heads: int) -> np.ndarray:
    seq_len, d_model = X.shape
    d_k = d_model // num_heads
    outs = []
    for h in range(num_heads):
        W_qh = np.random.randn(d_model, d_k).astype("float32") * 0.1
        W_kh = np.random.randn(d_model, d_k).astype("float32") * 0.1
        W_vh = np.random.randn(d_model, d_k).astype("float32") * 0.1
        Q, K, V = X @ W_qh, X @ W_kh, X @ W_vh
        scores = Q @ K.T / np.sqrt(d_k)
        w = np.exp(scores - scores.max(axis=-1, keepdims=True))
        w /= w.sum(axis=-1, keepdims=True)
        outs.append(w @ V)
    concat = np.concatenate(outs, axis=-1)  # (seq_len, d_model)
    W_o = np.random.randn(d_model, d_model).astype("float32") * 0.1
    return concat @ W_o
 
X = np.random.randn(16, 128).astype("float32")
print(multi_head_attention(X, num_heads=8).shape)  # (16, 128)

Attention alone only mixes information across positions. It does not transform each position's content. The bottleneck is a missing nonlinearity and a missing memory. A linear mixing plus linear projection is still linear. The fix is a per-position feed-forward network.

The MLP transforms each position independently#

The feed-forward network expands the dimension, applies a nonlinearity, and contracts back.

import torch
import torch.nn as nn
 
class MLP(nn.Module):
    def __init__(self, d_model: int, d_hidden: int | None = None):
        super().__init__()
        d_hidden = d_hidden or 4 * d_model
        self.up = nn.Linear(d_model, d_hidden)
        self.act = nn.GELU()
        self.down = nn.Linear(d_hidden, d_model)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down(self.act(self.up(x)))
 
mlp = MLP(d_model=128)
x = torch.randn(16, 128)
print(mlp(x).shape)  # torch.Size([16, 128])

The expansion creates room for feature detection. Keys act as pattern detectors, values act as stored associations, and a key-value pair with a strong match writes its value into the stream. The expansion factor of four means a 128-wide stream passes through a 512-wide hidden layer. Per-position processing gives the MLP a role distinct from attention: attention moves information between tokens, the MLP transforms information within a token. The bottleneck now is composition: one block performs one round of mixing and one round of transformation. Depth solves that.

The block is attention, residual, MLP, residual#

A transformer block combines the two sublayers with residual connections and normalizations.

import torch
import torch.nn as nn
 
class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = MLP(d_model)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.ln1(x)
        attn_out, _ = self.attn(h, h, h)
        x = x + attn_out              # residual 1: add the routing update
        x = x + self.mlp(self.ln2(x)) # residual 2: add the transformation update
        return x
 
block = TransformerBlock(d_model=128, num_heads=8)
x = torch.randn(2, 16, 128)
print(block(x).shape)  # torch.Size([2, 16, 128])

The residual stream is the structure doing the most work. Every block adds a contribution; the stream carries the sum forward. The additive form creates direct paths for gradients, so each block only needs to improve the representation rather than preserve it. The stream is a highway. Attention and the MLP are detours that return their results to the highway.

Stacking blocks: depth composes, width stores#

A full model stacks blocks and finishes with a final normalization and a projection to vocabulary size.

import torch
import torch.nn as nn
 
class Transformer(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, num_heads: int, num_layers: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.ModuleList(
            [TransformerBlock(d_model, num_heads) for _ in range(num_layers)]
        )
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
 
    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        x = self.embed(token_ids)
        for block in self.blocks:
            x = block(x)
        return self.head(self.ln_f(x))
 
model = Transformer(vocab_size=1000, d_model=128, num_heads=8, num_layers=4)
ids = torch.randint(0, 1000, (2, 16))
print(model(ids).shape)  # torch.Size([2, 16, 1000])

Depth composes transformations. More layers allow more steps of routing, retrieval, and refinement, and the composition supports more complex behavior. Width expands capacity inside each step. Wider streams and wider MLPs store more features, and redundant capacity improves robustness under approximation.

The trade is computational. Attention cost grows with sequence length and depth, while width cost grows with representational size. Hardware favors the trade that matches parallel structure and memory access. The parameter count makes the division concrete.

def block_params(d_model: int, d_hidden: int | None = None) -> int:
    d = d_model
    h = d_hidden or 4 * d
    attn = 4 * d * d  # Q, K, V, O projections
    mlp = 2 * d * h   # up and down
    return attn + mlp
 
for d in [128, 256, 512, 1024]:
    print(f"d={d}: per-block params ≈ {block_params(d):,}")
 
# Attention share at d=512, h=2048
d = 512
attn, mlp = 4 * d * d, 2 * d * 4 * d
print(f"attention share: {attn / (attn + mlp):.0%}, MLP share: {mlp / (attn + mlp):.0%}")

At the standard expansion ratio, the MLP holds about two-thirds of each block's parameters. Early layers tend to handle local and syntactic patterns; later layers handle abstraction and task-specific features; the division emerges from training rather than from a specification.

The compute bottleneck: quadratic attention#

Attention computes an $L \times L$ score matrix for sequence length $L$. Cost per layer per head is $O(L^2 \cdot d_k)$ in arithmetic and $O(L^2)$ in memory if the matrix is materialized. Doubling the sequence length quadruples the memory.

def attention_flops_per_head(seq_len: int, d_k: int) -> int:
    # QK^T and attention-weighted V: two matmuls of shape (L, d_k) x (d_k, L)
    return 2 * (seq_len * seq_len * d_k * 2)
 
for L in [512, 2048, 8192, 32768]:
    print(f"L={L}: scores matrix = {L * L / 1e6:.0f}M entries, "
          f"flops/head ≈ {attention_flops_per_head(L, 64) / 1e9:.1f} GFLOP")

Two fixes address the quadratic term without removing it. Flash attention tiles the computation so the full $L \times L$ matrix never sits in high-bandwidth memory; it changes memory traffic, not asymptotic arithmetic. Sliding-window and sparse patterns restrict which positions each query attends to, changing the arithmetic itself at the cost of direct long-range links. Neither fix eliminates the cost; both move it.

What to measure#

Four measurements separate the components.

import time
import torch
 
# Time one block; attention dominates at long sequences, MLP at short ones
block = TransformerBlock(d_model=512, num_heads=8)
for L in [128, 1024, 4096]:
    x = torch.randn(1, L, 512)
    start = time.perf_counter()
    for _ in range(5):
        block(x)
    print(f"L={L}: {(time.perf_counter() - start) / 5 * 1000:.1f} ms/block")

• Parameter fraction per component: attention versus MLP, measured from the module tree.
• Time split at the target sequence length: attention time and MLP time per block.
• Memory scale with $L$: verify whether the $L \times L$ term is materialized or tiled.
• Gradient norm per block: residual depth can stall if early blocks receive no signal.

A checklist for reading a transformer checkpoint:

• Identify $d_{\text{model}}$, head count, layer count, and expansion ratio.
• Compute the parameter split; expect the MLP to dominate.
• Check the positional scheme: fixed, learned, or relative.
• Note the attention variant: full, chunked, sliding-window, or hashed.
• Verify the residual order: pre-norm versus post-norm.
• Measure time and memory at the intended sequence length, not at length one.

The shape explains the behavior: a highway, a router, and a transformer bank, repeated at scale.