The Transformer Architecture End-to-End

Autocomplete has read function returns the and needs one more token. result feels plausible, but picking it requires each token to carry its own information while the consults its history without peeking at the answer. What changes between that text and the final choice?

Follow one small decoder-only Transformer forward pass. Three token IDs become positioned vectors, attention exchanges information along legal paths, a block updates every row without changing its width, and the final row becomes vocabulary scores. Small NumPy experiments expose individual operations; three linked PyTorch cells then assemble the complete model. Hand-set weights make the arithmetic readable; a plausible continuation alone says nothing about language understanding.

The preceding autoencoder lesson squeezed an input through a narrow latent code. Here we keep one vector per token and let those vectors communicate. You'll reuse matrix multiplication, nonlinear layers, and the softmax calculation from Softmax & Cross-Entropy. The same differentiation rules support this forward pass.

Before reading the pipeline, predict two shapes: with three tokens and width four, the hidden state should stay [3, 4]; with six vocabulary entries, the selected row should produce six logits. Keep those predictions in view as the data moves.

View larger

That path is what we mean by a Transformer here: a stack of attention and position-wise feed-forward updates that carries a same-width representation forward.

The architecture introduced in 2017 paired an encoder with a decoder for translation. A text-generation decoder-only model keeps the decoder's masked self-attention path and predicts each token from preceding context.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762[2]Reference 2Improving Language Understanding by Generative Pre-Training.https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf}

Tokens become positioned vectors

Start at the input boundary with a tiny autocomplete vocabulary. These toy IDs keep the forward pass readable; production tokenizers define their own IDs:

Our decoder maps each token ID to a learned vector, then adds the vector for its position. Moving the changes which position vector it receives.

Unmasked self-attention without position information is permutation equivariant: shuffle input rows, and its output rows follow the same shuffle. A fixed causal mask changes that argument because shuffling content changes which tokens precede which queries. Explicit positional mechanisms still give the decoder a direct way to distinguish positions and distances.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Examples of positional mechanisms include:

1. Fixed sinusoidal encodings: The original 2017 Transformer generated sine and cosine waves of varying frequencies, allowing relative shifts to be expressed as linear transformations without adding learned parameters.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}
2. Learned absolute positional embeddings: The original GPT and BERT use an embedding table $W_{\text{pos}} \in \mathbb{R}^{T_{\text{max}} \times d_{\text{model}}}$ and add row $t$ to the token vector.^{[2]Reference 2Improving Language Understanding by Generative Pre-Training.https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf[3]Reference 3BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.https://arxiv.org/abs/1810.04805} An unmodified table has no row beyond that range. Extending it requires a model change, and usable performance at longer lengths needs evaluation.
3. Rotary Position Embedding (RoPE): LLaMA's 2023 architecture rotates pairs of query and key coordinates by position-dependent angles.^{[4]Reference 4RoFormer: Enhanced Transformer with Rotary Position Embedding.https://arxiv.org/abs/2104.09864[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971} Their dot product depends on the content vectors and a rotation for relative distance $(m-n)$. Evaluate model quality when using positions beyond the trained context length.

Other positional designs include T5's learned relative-position biases on attention scores.^{[6]Reference 6Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.https://arxiv.org/abs/1910.10683}

Our arithmetic uses small additive vectors to keep every coordinate inspectable by hand:

import numpy as np

tokens = ["function", "returns", "the"]
ids = np.array([0, 1, 2])
embedding_table = np.array([
    [1.0, 0.0, 0.4, 0.0],  # function
    [0.2, 1.0, 0.5, 0.0],  # returns
    [0.0, 0.3, 0.8, 1.0],  # the
    [0.1, 0.5, 1.0, 0.9],  # result
    [0.0, 0.1, 0.0, 0.1],  # None
    [0.2, 0.0, 0.4, 0.5],  # value
])
position = np.array([
    [0.00, 0.00, 0.00, 0.00],
    [0.05, 0.00, 0.00, 0.05],
    [0.10, 0.00, 0.00, 0.10],
])

x = embedding_table[ids] + position
print("tokens:", tokens)
print("ids shape:", ids.shape)
print("hidden shape:", x.shape)
print("last positioned vector:", np.round(x[-1], 2))

tokens: ['function', 'returns', 'the']
ids shape: (3,)
hidden shape: (3, 4)
last positioned vector: [0.1 0.3 0.8 1.1]

A row belongs to one token position; a column is one feature of that position. Real models use far wider vectors and batches, but the axes mean the same thing:

The full vocabulary is function, returns, the, result, None, value, in that ID order. The output head will score those same six tokens. Input and output weights needn't be tied, but a generated ID must be meaningful when fed back into the embedding table.

Attention lets a token read earlier tokens

At the, the model needs a useful read of its history. An RNN transports information through one recurrent state. Self-attention takes a different route: every position computes a read from the visible positions, then blends their value vectors.

For a decoder, "visible" means the current and earlier tokens only. That rule lets the model process a whole training sequence without letting a later answer flow backward.

Each token produces three learned projections:

• A query asks what information this position needs.
• A key advertises what a position contains.
• A value supplies information if the query matches its key.

Why separate queries, keys, and values instead of using $X X^T$? Direct self-matching gives symmetric raw scores, but row-wise softmax can already give asymmetric weights because rows have different denominators. A causal mask adds another asymmetry. Separate learned query and key maps let the model choose a richer matching rule; the value map independently chooses what information to mix.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Check the distinction with two rows. Before masking, the off-diagonal scores are both one. Will their normalized weights also match?

import numpy as np

x = np.array([[1.0, 0.0], [1.0, 1.0]])
scores = x @ x.T
exp = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = exp / exp.sum(axis=-1, keepdims=True)
print("symmetric scores:", bool(np.allclose(scores, scores.T)))
print("symmetric weights:", bool(np.allclose(weights, weights.T)))
print("opposite weights:", np.round([weights[0, 1], weights[1, 0]], 3))

symmetric scores: True
symmetric weights: False
opposite weights: [0.5   0.269]

For one head, $W_q$ and $W_k$ map d_model to key width $d_k$. $W_v$ maps to value width $d_v$, which needn't equal $d_k$. Our examples choose equal head widths for all three projections.

Before checking the numbers, predict the score shape: three query rows compared with three key rows should produce a [3, 3] matrix. The the row can score all three keys at this stage; the causal mask will remove any illegal entries next.

To see the multiplication, take one two-feature attention head. These handpicked arrays isolate the dot product; they aren't the projections used by the complete model later:

import numpy as np

q = np.array([
    [1.0, 0.0],  # function position
    [0.5, 1.0],  # returns position
    [0.2, 1.2],  # the position
])
k = np.array([
    [1.0, 0.0],
    [0.4, 1.0],
    [0.0, 0.8],
])
v = np.array([
    [1.0, 0.0],
    [0.3, 1.0],
    [0.0, 0.6],
])

raw_scores = q @ k.T
print("Q shape:", q.shape, "K shape:", k.shape, "V shape:", v.shape)
print("Q @ K.T shape:", raw_scores.shape)
print("raw score row for 'the':", np.round(raw_scores[2], 2))

Q shape: (3, 2) K shape: (3, 2) V shape: (3, 2)
Q @ K.T shape: (3, 3)
raw score row for 'the': [0.2  1.28 0.96]

There is one score for every query-key pair. That square T x T matrix also explains the quadratic cost of standard attention. Doubling T creates four times as many pair scores before the value mix.

Scores only decide where to read. After scaling and masking, softmax turns each row into weights, and multiplying those weights by V produces one mixed value row for each position. The the row becomes a two-feature update that the heads will merge back into model width.

Scale before softmax

Scaled dot-product attention divides query-key scores by the square root of the head width before applying softmax:

Why divide by $\sqrt{d_k}$? Consider the statistical distribution of raw scores. If the components of query $q \in \mathbb{R}^{d_k}$ and key $k \in \mathbb{R}^{d_k}$ are independent random variables with mean zero and variance one, their dot product is:

Each individual product $q_i k_i$ has mean $\mathbb{E}[q_i k_i] = \mathbb{E}[q_i]\mathbb{E}[k_i] = 0$ and variance:

Because the components are independent, their variances sum directly:

Under those assumptions, the raw dot product's standard deviation is $\sqrt{d_k}$: eight at width 64, or about 11.31 at width 128. Score ranges depend on $d_k$ and the learned projections. Large score gaps make softmax nearly one-hot, where its Jacobian entries $\frac{\partial s_i}{\partial z_j} = s_i(\delta_{ij}-s_j)$ become small and can approach zero, allowing gradients to shrink.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Dividing by $\sqrt{d_k}$ gives variance one in this statistical model. Learned, correlated query and key components may violate those assumptions. Scaling controls a width-related variance effect; gradients still depend on the learned projections and surrounding layers.

Their base model used $d_{\text{model}}=512$ with $h=8$ heads, so each head had $d_k=64$. The calculation below compares the unscaled and scaled versions of concrete scores [12.4, 9.1, 3.8]:

import numpy as np

def softmax(values):
    shifted = values - values.max()
    exp = np.exp(shifted)
    return exp / exp.sum()

scores = np.array([12.4, 9.1, 3.8])
scaled = scores / np.sqrt(64)

print("unscaled weights:", np.round(softmax(scores), 3))
print("scaled scores:", np.round(scaled, 3))
print("scaled weights:", np.round(softmax(scaled), 3))

unscaled weights: [0.964 0.036 0.   ]
scaled scores: [1.55  1.138 0.475]
scaled weights: [0.499 0.33  0.17 ]

The unscaled row puts nearly all attention mass on one key, while the scaled row spreads it across visible choices. Scaling reduces the width-related sharpening in the preceding statistical model. Learned projections can still produce sharp scores and confident choices.

Use the earlier the row to finish one attention read. Its raw scores [0.2, 1.28, 0.96] become [0.141, 0.905, 0.679] after division by $\sqrt{2}$. Softmax gives weights about [0.21, 0.44, 0.35]; applying them to the three value rows produces the contextual vector [0.34, 0.65]. Vocabulary scores come later.

View larger

The causal mask prevents answer leakage

Training can process the positions in function returns the result in parallel within each layer, but position the is still responsible for predicting result. If its query can read the result representation, it sees the answer. A causal mask sets future scores to negative infinity before softmax; those entries receive probability zero. Setting them to zero wouldn't work: $e^0 = 1$, so forbidden keys would still receive active attention weight. Using $-\infty$ guarantees $e^{-\infty} = 0$.

Predict the allowed cells before using the table: the the row should include function, returns, and the, but not result. The final row can see all four inputs because no later input is present in this example.

Here is a separate four-position score matrix that includes the future the → result cell. The earlier dot-product example uses three positions.

import numpy as np

scores = np.array([
    [2.0, 1.5, 0.5, 1.0],
    [1.0, 2.5, 1.5, 0.5],
    [0.5, 1.0, 3.0, 2.0],
    [1.5, 0.5, 1.0, 2.5],
])
future = np.triu(np.ones_like(scores, dtype=bool), k=1)
masked_scores = np.where(future, -np.inf, scores)

shifted = masked_scores - masked_scores.max(axis=-1, keepdims=True)
weights = np.exp(shifted)
weights /= weights.sum(axis=-1, keepdims=True)
assert np.all(weights[future] == 0.0)
assert np.allclose(weights.sum(axis=-1), 1.0)

print("masked weights:")
print(np.round(weights, 3))
print("future attention mass:", float(weights[future].sum()))

masked weights:
[[1.    0.    0.    0.   ]
 [0.182 0.818 0.    0.   ]
 [0.067 0.111 0.821 0.   ]
 [0.213 0.078 0.129 0.579]]
future attention mass: 0.0

The output shows future attention mass: 0.0. Read row the (row index 2) as [0.067, 0.111, 0.821, 0.000]: its unrounded weights sum to one, and the result column is zero. If this next-token decoder assigns weight to a future input, inspect its mask and position alignment. The model could be reading an answer that won't exist at generation time. The mask stopped that leak. It still left the with only one weighted average of the visible values. What if that position needs a second, separate read of the same history?

Heads run in parallel; blocks run in sequence

One head is that single average. A token often needs several different reads at once: one head might track syntactic dependencies (matching a verb to its subject function), while another tracks typing or semantic scoping, and a third tracks the immediate prior token.

Multi-head attention gives a block several such views: it projects model features into heads, applies attention independently in each head, concatenates their outputs, and projects back to d_model.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Each projection can use every input feature. Head 1 can draw on all original embedding coordinates before its projection selects a subspace.

Suppose our tiny model width is four and it has two heads. Each head owns two features. Predict the shape trace before running the cell: [1, 3, 4] should become [1, 2, 3, 2] after the split, then return to [1, 3, 4] after the merge. Splitting and merging must preserve every token position:

import numpy as np

batch, tokens, d_model, heads = 1, 3, 4, 2
assert d_model % heads == 0
d_head = d_model // heads
projected = np.arange(batch * tokens * d_model).reshape(batch, tokens, d_model)

split = projected.reshape(batch, tokens, heads, d_head).transpose(0, 2, 1, 3)
merged = split.transpose(0, 2, 1, 3).reshape(batch, tokens, d_model)
assert np.array_equal(projected, merged)

print("all-heads shape:", split.shape)
print("merged shape:", merged.shape)
print("merge restores values:", bool(np.array_equal(projected, merged)))

all-heads shape: (1, 2, 3, 2)
merged shape: (1, 3, 4)
merge restores values: True

The head axis describes independent work inside one block. The block axis is sequential depth: output from block 1 enters block 2. Stacking repeats the structure, not the weights; each block has its own projections. Mixing heads with blocks loses that dependency.

A block alternates communication and local transformation

Attention lets token positions exchange information across sequence steps. The next sublayer changes each row on its own, using the same weights at every position. That feed-forward network (FFN) is where a token's mixed context gets a per-position nonlinear transformation. Geva et al. interpreted FFNs in their studied language models as key-value memories: keys correlate with textual patterns, while values influence the output vocabulary distribution. In this interpretation, a neuron can respond to patterns and shape many output scores; it needn't store one fact.^{[7]Reference 7Transformer Feed-Forward Layers Are Key-Value Memories.https://aclanthology.org/2021.emnlp-main.446/}

The original Transformer used two linear layers with a ReLU nonlinearity between them, widening from $d_{\text{model}}=512$ to $d_{\text{ff}}=2048$ (a $4\times$ expansion) and back:^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

The first projection widens the feature dimension; the second returns it to d_model, which keeps blocks stackable. This isolated FFN demo feeds in our positioned vectors to show [3, 4] -> [3, 6] -> [3, 4]. In a complete pre-LN block, the FFN instead receives the normalized stream after attention's residual addition. Bias terms are omitted here so the matrix sizes stay visible.

import numpy as np

x = np.array([
    [1.00, 0.00, 0.40, 0.00],
    [0.25, 1.00, 0.50, 0.05],
    [0.10, 0.30, 0.80, 1.10],
])
w1 = np.array([
    [1.0, 0.0, 0.5, 0.0, 0.0, 0.2],
    [0.0, 1.0, 0.0, 0.5, 0.2, 0.0],
    [0.4, 0.2, 1.0, 0.0, 0.0, 0.5],
    [0.0, 0.2, 0.0, 1.0, 0.4, 0.0],
])
w2 = np.array([
    [0.3, 0.0, 0.0, 0.1],
    [0.0, 0.3, 0.1, 0.0],
    [0.2, 0.0, 0.3, 0.0],
    [0.0, 0.2, 0.0, 0.3],
    [0.1, 0.0, 0.0, 0.2],
    [0.0, 0.1, 0.2, 0.0],
])

expanded = np.maximum(0.0, x @ w1)
ffn_output = expanded @ w2
print("input shape:", x.shape)
print("expanded shape:", expanded.shape)
print("returned shape:", ffn_output.shape)
print("last-token update:", np.round(ffn_output[-1], 3))

input shape: (3, 4)
expanded shape: (3, 6)
returned shape: (3, 4)
last-token update: [0.346 0.496 0.407 0.517]

Each FFN multiplication maps token rows independently, producing token-by-feature matrices without row-to-row interaction. Attention's Q @ K.T multiplication instead produces token-to-token scores.

Gated FFNs add another way to transform features. The 2023 LLaMA architecture uses SwiGLU (Swish Gated Linear Unit):^{[8]Reference 8GLU Variants Improve Transformerhttps://arxiv.org/abs/2002.05202[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971}

Here, $\operatorname{Swish}(z)=z\sigma(z)$ is Swish with its parameter set to one, also called SiLU. Two upward projections run in parallel; the activated branch multiplies the other branch elementwise before $W_{\text{down}}$ returns to model width. The gate is an unconstrained scalar that can reverse a branch's sign or scale it above one.

Ignoring biases, three matrices at intermediate width $d_{\text{ff}}$ have $3d_{\text{model}}d_{\text{ff}}$ parameters. Matching a two-matrix FFN at width $4d_{\text{model}}$ gives $d_{\text{ff}}=\frac{8}{3}d_{\text{model}}$. LLaMA uses this sizing idea; other models choose other widths.^{[8]Reference 8GLU Variants Improve Transformerhttps://arxiv.org/abs/2002.05202[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971} The original Gemma report specifies GeGLU, with a GELU gate.^{[9]Reference 9Gemma: Open Models Based on Gemini Research and Technologyhttps://arxiv.org/abs/2403.08295}

Residual paths and normalization

A residual connection adds a sublayer result back to its input. Because attention and the FFN return to d_model, addition keeps the running stream at the same shape. It creates an identity path through deep stacks, following the residual-network concept from He et al.^{[10]Reference 10Deep Residual Learning for Image Recognition.https://arxiv.org/abs/1512.03385}

Placement of normalization changes gradient paths and can affect optimization:

• Post-LayerNorm (2017 original): Normalization sits directly in the residual stream after each addition:^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} $$x_{l+1} = \operatorname{LayerNorm}(x_l + \operatorname{SubLayer}(x_l))$$ For one update, its Jacobian has the form $J_{\text{LN}}(I+J_{\text{SubLayer}})$. Xiong et al. found large expected parameter gradients near the output at initialization in their Post-LN analysis, making large initial learning rates unstable in their experiments. Warmup helped. Their findings describe the evaluated setups; depth, optimization, and architecture determine behavior in other models.^{[11]Reference 11On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745}
• Pre-LayerNorm (used by our decoder and GPT-2): Normalization precedes each sublayer, leaving the main residual addition untouched:^{[12]Reference 12GPT-2 Source Implementation.https://github.com/openai/gpt-2/blob/master/src/model.py[11]Reference 11On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745} $$x_{l+1} = x_l + \operatorname{SubLayer}(\operatorname{LayerNorm}(x_l))$$ For one update, its Jacobian is $I+J_{\text{SubLayer}}J_{\text{LN}}$. This supplies a direct identity contribution, but branch derivatives can reinforce or cancel it. Products across updates still determine the full gradient. Xiong et al. trained Pre-LN variants without warmup in their evaluated settings. LLaMA uses pre-normalization and still reports 2,000 warmup steps, so warmup depends on the training setup.^{[11]Reference 11On Layer Normalization in the Transformer Architecture.https://arxiv.org/abs/2002.04745[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971}

Here $l$ indexes a residual update; an attention-plus-FFN block contains two updates. Our GELU FFN is a two-matrix layer. SwiGLU would add a second upward branch, so the width trace alone wouldn't fully specify its structure.

Architectures also differ in the normalization statistic itself:

• LayerNorm: Centers by subtracting mean $\mu$, then scales by standard deviation $\sqrt{\sigma^2 + \epsilon}$ with learned scale $\gamma$ and shift $\beta$.^{[13]Reference 13Layer Normalization.https://arxiv.org/abs/1607.06450}
• RMSNorm: Divides by a root-mean-square statistic without subtracting the feature mean: $$\operatorname{RMSNorm}(x) = \frac{x}{\operatorname{RMS}(x)} \odot \gamma, \quad \text{where } \operatorname{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}$$ The usual formulation learns scale but has no additive shift. A constant nonzero row becomes zero under LayerNorm before its affine terms, whereas RMSNorm can retain it. Zhang and Sennrich found comparable quality with lower overhead in their experiments. Speed depends on framework, hardware, model, and other computation, so measure the target workload before choosing normalization for performance.^{[14]Reference 14Root Mean Square Layer Normalization.https://arxiv.org/abs/1910.07467} LLaMA and the original Gemma report use RMSNorm.^{[5]Reference 5LLaMA: Open and Efficient Foundation Language Models.https://arxiv.org/abs/2302.13971[9]Reference 9Gemma: Open Models Based on Gemini Research and Technologyhttps://arxiv.org/abs/2403.08295}

Before running the snippet, predict two checks: normalization should put each row near mean 0 and variance 1, while adding the update should leave the residual output at shape [2, 4]. The code simulates that pre-normalization order with a small update.

import numpy as np

def normalize(x, eps=1e-5):
    mean = x.mean(axis=-1, keepdims=True)
    variance = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
    return (x - mean) / np.sqrt(variance + eps)

x = np.array([
    [1.0, 2.0, 0.0, 1.0],
    [0.2, 0.4, 0.8, 0.6],
])
normalized = normalize(x)
sublayer_update = 0.1 * normalized
residual_output = x + sublayer_update

print("normalized means:", np.round(normalized.mean(axis=-1), 6))
print("normalized variances:", np.round(normalized.var(axis=-1), 4))
print("residual output shape:", residual_output.shape)

normalized means: [0. 0.]
normalized variances: [1.     0.9998]
residual output shape: (2, 4)

The small eps keeps division defined even for a constant row. It also explains why the variances are slightly below one. Learned scale and offset can change those statistics afterward; mean zero and variance one describe the normalized values before those learned affine terms. In the block view, follow the attention update into x₁, then the FFN update into x₂. Both additions preserve [3, 4].

View larger

Implement masked attention in PyTorch

The NumPy example exposed the mask entry by entry. Now keep the same contract in one causal multi-head self-attention layer. The projections are learned linear maps; the mask is created from token positions. Calling masked_fill enforces that visibility constraint: it turns every future score into negative infinity before softmax.

Trace the shapes before running it: [1, 4, 8] becomes [1, 2, 4, 4] for each of q, k, and v; scores and weights stay [1, 2, 4, 4]; merging heads returns [1, 4, 8]. The largest future weight should be exactly zero.

import math
import torch
from torch import nn

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model=8, heads=2):
        super().__init__()
        if d_model <= 0 or heads <= 0 or d_model % heads != 0:
            raise ValueError("d_model and heads must be positive; heads must divide d_model")
        self.heads = heads
        self.d_head = d_model // heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        if x.ndim != 3 or x.shape[-1] != self.qkv.in_features:
            raise ValueError("expected [batch, tokens, d_model] input")
        if x.shape[1] == 0:
            raise ValueError("at least one token is required")
        batch, tokens, d_model = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        def split_heads(tensor):
            return tensor.view(batch, tokens, self.heads, self.d_head).transpose(1, 2)
        q, k, v = map(split_heads, (q, k, v))

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)
        future = torch.triu(
            torch.ones(tokens, tokens, dtype=torch.bool, device=x.device),
            diagonal=1,
        )
        weights = torch.softmax(scores.masked_fill(future, float("-inf")), dim=-1)
        mixed = weights @ v
        merged = mixed.transpose(1, 2).contiguous().view(batch, tokens, d_model)
        return self.out(merged), weights

torch.manual_seed(4)
x = torch.randn(1, 4, 8)
attention = CausalSelfAttention()
output, weights = attention(x)
future = torch.triu(torch.ones(4, 4, dtype=torch.bool), diagonal=1)
assert torch.all(weights[0, :, future] == 0)
assert torch.allclose(weights.sum(dim=-1), torch.ones_like(weights.sum(dim=-1)))

print("attention output shape:", tuple(output.shape))
print("weights shape:", tuple(weights.shape))
print("largest future weight:", f"{weights[0, :, future].max().item():.6f}")

attention output shape: (1, 4, 8)
weights shape: (1, 2, 4, 4)
largest future weight: 0.000000

The output shape matches the input shape, and the [batch, heads, tokens, tokens] weight tensor exposes each head's routing decisions. The zero future weight is the causal invariant in PyTorch form.

The PyTorch cell uses random inputs to test shapes and visibility separately from the hand-worked values. It omits padding, dropout, and a KV cache. Every query has a real token on its diagonal; with finite unmasked scores, no row is fully masked. Input validation also rejects zero width, incompatible shapes, and empty token sequences.

torch.nn.functional.scaled_dot_product_attention (SDPA) implements the same attention operation and chooses a backend based on its inputs and device. Eligible fused GPU kernels can avoid storing full score and weight matrices; a math fallback may store them. Backend changes can alter floating-point rounding. This API returns the mixed values; inspectable attention weights require a separate calculation.^{[15]Reference 15torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

In SDPA, a boolean mask's True means allow, the inverse of our future mask passed to masked_fill. Its dropout_p argument must explicitly be zero for evaluation; calling a surrounding module's .eval() doesn't override a nonzero argument. Keep the preceding Python session open for this comparison:

from torch.nn import functional as F

torch.manual_seed(11)
q, k, v = [torch.randn(1, 2, 4, 4) for _ in range(3)]
future = torch.triu(torch.ones(4, 4, dtype=torch.bool), diagonal=1)
scores = q @ k.transpose(-2, -1) / math.sqrt(4)
dense = scores.masked_fill(future, float("-inf")).softmax(-1) @ v
causal = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
allowed = F.scaled_dot_product_attention(q, k, v, attn_mask=~future, dropout_p=0.0)

# One query at absolute position 3 may read all cached keys at positions 0..3.
last_query = F.scaled_dot_product_attention(
    q[:, :, -1:, :], k, v,
    attn_mask=torch.ones(1, 4, dtype=torch.bool), dropout_p=0.0,
)
misaligned = F.scaled_dot_product_attention(
    q[:, :, -1:, :], k, v, is_causal=True, dropout_p=0.0,
)
only_first_key = torch.allclose(misaligned, v[:, :, :1, :], atol=1e-6)
assert only_first_key and not torch.allclose(misaligned, last_query)
checks = [torch.allclose(causal, dense, atol=1e-6),
          torch.allclose(allowed, dense, atol=1e-6),
          torch.allclose(last_query, dense[:, :, -1:, :], atol=1e-6)]
assert all(checks)
print("causal API matches dense:", checks[0])
print("True-means-allow mask matches:", checks[1])
print("cached last query matches:", checks[2])
print("misaligned cached query reads only key zero:", only_first_key)

causal API matches dense: True
True-means-allow mask matches: True
cached last query matches: True
misaligned cached query reads only key zero: True

For a non-square query-key layout, SDPA's is_causal=True uses upper-left alignment. A [1,4] mask then exposes only key zero, which is wrong for our query at absolute position three. Its sole softmax weight is one, so the returned vector is exactly that first key's value, whatever the query scores. The final check catches this silent error, which still has a valid output shape. Use a mask or causal-bias object aligned to the actual query and cached-key positions. Chunked decoding needs the same care.^{[15]Reference 15torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Assemble one decoder block

Attention now returns a same-width update, so it can sit inside a full pre-normalized block. Read the sequence before the code: normalize, attend, add; normalize, widen through the FFN, return to width eight, and add again. For input [2, 4, 8] and d_ff=24, every residual merge should still be [2, 4, 8].

PyTorch's LayerNorm includes the learned scale and offset that the earlier NumPy calculation omitted. This block uses GELU instead of the original Transformer's ReLU, the same activation family as GPT-2's MLP. GPT-2 uses a tanh approximation; nn.GELU() below uses the default exact formulation. Either way, the FFN's structural role stays the same: widen, transform, and return to d_model.^{[16]Reference 16Gaussian Error Linear Units (GELUs).https://arxiv.org/abs/1606.08415[12]Reference 12GPT-2 Source Implementation.https://github.com/openai/gpt-2/blob/master/src/model.py}

Continue in the same Python session so CausalSelfAttention remains defined. The attention module returns both its update and its weights; the block needs only the update for the residual addition.

class DecoderBlock(nn.Module):
    def __init__(self, d_model=8, heads=2, d_ff=24):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attention = CausalSelfAttention(d_model, heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        update, _ = self.attention(self.norm1(x))
        x = x + update
        x = x + self.ffn(self.norm2(x))
        return x

torch.manual_seed(5)
x = torch.randn(2, 4, 8, requires_grad=True)
block = DecoderBlock()
output = block(x)
output.square().mean().backward()
assert output.shape == x.shape
assert x.grad is not None and torch.isfinite(x.grad).all()

print("input -> output:", tuple(x.shape), "->", tuple(output.shape))
print("input gradient is finite:", bool(torch.isfinite(x.grad).all()))

input -> output: (2, 4, 8) -> (2, 4, 8)
input gradient is finite: True

The finite-gradient check catches a broken backward pass for this input. A deep stack's training behavior depends on its depth, initialization, data, and optimization settings. The block produces contextual vectors with the same [B, T, d_model] shape as its input, so another block can consume them.

Project the final position into a vocabulary

The block stack has finished its same-width work. The vocabulary head changes the stream's output width from model features to token scores. GPT-2 applies a final ln_f normalization immediately before computing logits.^{[12]Reference 12GPT-2 Source Implementation.https://github.com/openai/gpt-2/blob/master/src/model.py}

A learned output matrix converts each hidden vector into one raw score per vocabulary item. These scores are logits. Vaswani et al. share the output matrix with token embeddings; GPT-2 reuses wte similarly.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762[12]Reference 12GPT-2 Source Implementation.https://github.com/openai/gpt-2/blob/master/src/model.py} With a common vocabulary and width, weight tying replaces two $V \times d_{\text{model}}$ weight arrays with one shared array. That halves those two arrays' parameters, not the whole model, activation memory, or vocabulary-projection computation. The toy below keeps a separate output_weight so the multiply is visible.

Now connect the actual prompt to the block. Continue in the same session with the DecoderBlock class above. This time use width four, two heads, and an FFN hidden width of six. The embedding table is exactly the one from the opening NumPy example.

The matrix multiply applies the vocabulary head to every position: [1, 3, 4] @ [4, 6] → [1, 3, 6]. Generation selects the last six scores. Training can use all three rows, with a shifted next-token target for each. The output weights below are deliberately chosen for a readable ranking; the random block hasn't learned English.

vocab = ["function", "returns", "the", "result", "None", "value"]
ids = torch.tensor([[0, 1, 2]])
embedding_table = torch.tensor([
    [1.0, 0.0, 0.4, 0.0],
    [0.2, 1.0, 0.5, 0.0],
    [0.0, 0.3, 0.8, 1.0],
    [0.1, 0.5, 1.0, 0.9],
    [0.0, 0.1, 0.0, 0.1],
    [0.2, 0.0, 0.4, 0.5],
])
position = torch.tensor([
    [0.00, 0.00, 0.00, 0.00],
    [0.05, 0.00, 0.00, 0.05],
    [0.10, 0.00, 0.00, 0.10],
])
output_weight = torch.tensor([
    [1.0, 0.2, 0.0, 0.0,  0.0, 0.0],
    [0.0, 0.0, 0.0, 0.0, -1.0, 0.0],
    [0.0, 0.0, 0.0, 0.5,  0.0, 1.0],
    [0.0, 0.0, 0.2, 1.0,  0.0, 0.0],
])

torch.manual_seed(7)
toy_block = DecoderBlock(d_model=4, heads=2, d_ff=6).eval()
final_norm = nn.LayerNorm(4).eval()
with torch.no_grad():
    positioned = embedding_table[ids] + position
    contextual = toy_block(positioned)
    normalized = final_norm(contextual)
    logits = normalized @ output_weight
    probabilities = logits[0, -1].softmax(dim=-1)
    best = probabilities.argmax().item()

assert logits.shape == (1, 3, len(vocab))
assert not torch.allclose(contextual, positioned)
assert torch.allclose(probabilities.sum(), torch.tensor(1.0))

print("last row after block:", [round(v, 3) for v in contextual[0, -1].tolist()])
print("normalized last row:", [round(v, 3) for v in normalized[0, -1].tolist()])
print("all-position logits shape:", tuple(logits.shape))
print("last-position logits:", [round(v, 3) for v in logits[0, -1].tolist()])
print("top token:", vocab[best])
print("top probability:", f"{probabilities[best]:.3f}")

last row after block: [0.037, 0.267, 0.994, 0.87]
normalized last row: [-1.26, -0.686, 1.128, 0.818]
all-position logits shape: (1, 3, 6)
last-position logits: [-1.26, -0.252, 0.164, 1.382, 0.686, 1.128]
top token: result
top probability: 0.353

The vector sent to the output head has passed through attention, both residual additions, the FFN, and final normalization. Its values differ from the original embedding. The overview at the start marks the shape change from [1, 3, 4] to [1, 3, 6]. Because the vocabulary projection acts independently at each position, slicing the last position before or after that projection gives the same six logits.

Choosing the largest logit is greedy decoding. It appends result to make function returns the result.

Alternatively, sampling introduces temperature scaling $T > 0$ to soften or sharpen the logits before softmax:

For fixed finite logits with a unique maximum, $T\to0^+$ concentrates probability on that token. Tied maxima share the limiting mass; greedy argmax uses a tie-breaking rule. At $T=1$, softmax is unscaled. Raising $T$ increases entropy unless the logits are already equal, and $T\to\infty$ approaches uniform probabilities over this finite vocabulary. Randomness changes sampling diversity; continuation quality depends on the model, prompt, and decoding policy. Use max-subtraction as in the softmax lesson, preferably before dividing by a very small temperature.

To run another step, add a fourth row to the positional table as well as appending the new token ID. Our three-position demo isn't an unlimited text generator. Optimized serving reuses earlier attention keys and values through a KV cache instead of recomputing them for every new token.^{[17]Reference 17KV Cache Strategieshttps://huggingface.co/docs/transformers/main/en/kv_cache} We haven't implemented that optimization here. The next lesson turns this single step into a generation loop and a training objective.

Encoder, decoder, and encoder-decoder visibility

Decoder-only generation is one member of the Transformer family. Once the job changes, the visibility rule must change with it: classification can use the whole input, while translation needs a full source plus a causal target. Compare which positions each variant may read:

View larger

For this autocomplete task, a decoder-only design offers a convenient training and generation path. These properties describe its interface; evaluate architectures on the target task:

1. Targets from ordinary text: Shifting a sequence supplies next-token labels without human annotation. Which positions contribute depends on boundaries and loss masks. Encoder-decoder pre-training needn't require paired translations: T5 corrupts spans of ordinary text and predicts the missing text.^{[6]Reference 6Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.https://arxiv.org/abs/1910.10683}
2. Prompt conditioning: Instructions and examples can precede generation in the same stream. This adds no separate cross-attention module, but prompt prefill, cache storage, and attending to a longer prefix still cost computation. Useful instruction following also depends on training.
3. Reusable attention state: Cached decoder-only inference stores keys and values for self-attention layers. Cached encoder-decoder inference can store growing target self-attention state and fixed projected source state for cross-attention. These are logical cache roles; both architectures also run without a cache.^{[17]Reference 17KV Cache Strategieshttps://huggingface.co/docs/transformers/main/en/kv_cache}
4. Measured scaling trends: Kaplan et al. found power-law trends for language-model loss in their studied regimes. Such trends can guide resource planning. Evaluate reasoning skills, individual task results, and architecture choices on the target workload.^{[18]Reference 18Scaling Laws for Neural Language Modelshttps://arxiv.org/abs/2001.08361}

In self-attention, queries, keys, and values come from one sequence. In encoder-decoder cross-attention, decoder queries read keys and values produced by the encoder. A plain decoder-only text generator has no separate source sequence, so its core path is causal self-attention.

Diagnose a broken forward pass

When a forward pass fails, don't inspect every weight first. Ask which contract broke: visibility, shape, device placement, or normalization. The symptom usually narrows the search:

Use a deliberately bad mask to make leakage visible. Predict the two totals before running the cell: an unmasked matrix should assign positive mass to future cells, while the causal version should assign 0.000.

import numpy as np

scores = np.array([
    [2.0, 5.0, 4.0],
    [1.0, 2.0, 6.0],
    [0.5, 1.0, 2.0],
])
future = np.triu(np.ones_like(scores, dtype=bool), k=1)

def softmax_rows(matrix):
    exp = np.exp(matrix - matrix.max(axis=-1, keepdims=True))
    return exp / exp.sum(axis=-1, keepdims=True)

unmasked = softmax_rows(scores)
masked = softmax_rows(np.where(future, -np.inf, scores))
print("unmasked future mass across rows:", f"{unmasked[future].sum():.3f}")
print("masked future mass across rows:", f"{masked[future].sum():.3f}")

unmasked future mass across rows: 1.940
masked future mass across rows: 0.000

The unmasked total exceeds one because it adds probability mass from several separately normalized rows. Each row still sums to one. A zero future total checks the mask. Also change a future input and confirm that earlier outputs stay fixed; this checks for leakage elsewhere in the block.

Test causality through the whole block

Continue the PyTorch session. Replace the final token's features while keeping the first two positions fixed. Which outputs may change? Causality says the first two must be unchanged, even after the FFN and residual additions. The last output should change, so the test also catches a block that ignores its input.

changed = positioned.clone()
changed[:, -1, :] = torch.tensor([2.0, -1.0, 0.0, 0.5])
with torch.no_grad():
    original_output = toy_block(positioned)
    changed_output = toy_block(changed)

prefix_unchanged = torch.allclose(original_output[:, :-1], changed_output[:, :-1])
last_changed = not torch.allclose(original_output[:, -1], changed_output[:, -1])
assert prefix_unchanged and last_changed
print("earlier outputs unchanged:", prefix_unchanged)
print("last output changed:", last_changed)

earlier outputs unchanged: True
last output changed: True

Try removing masked_fill in CausalSelfAttention. The earlier outputs should now change, causing this test to fail. Keep dropout disabled for this comparison; otherwise random dropout masks can change outputs even when causality is correct.