Decoder-Only Transformers and Deep-Learning Foundations

Build causal self-attention, multi-head projections, LayerNorm, and a pre-norm residual block in NumPy before checking the reusable PyTorch decoder.

A fixed-context MLP combines a predetermined neighborhood of tokens. A decoder-only Transformer learns how to mix every earlier position while preserving the autoregressive boundary. The NumPy cells keep the attention equations and axis conventions visible; the cumulative PyTorch implementation then packages the same contracts into the ProofLM decoder used by training and evaluation.

Shape checks, a future-token intervention, a local gradient check, and a depth experiment make the architectural contracts observable before the reusable model is trusted.

Self-attention

Let X\in\mathbb{R}^{T\times d} contain one hidden vector per position. Learned projections produce

Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V,

and scaled dot products produce a score from query position i to key position j:

S_{i,j}=\frac{Q_iK_j^\top}{\sqrt{d_k}}.

The causal condition is a lower-triangular visibility matrix. Set the scores for j>i to -\infty before the row-wise softmax. The resulting weights sum to one over visible keys, and the weighted values are the information passed to each query position.

Start with one head and one sequence. That version isolates the mask from the bookkeeping needed by multiple heads.

import numpy as np


SEED = 41


def row_softmax(values):
    values = np.asarray(values, dtype=float)
    shifted = values - np.max(values, axis=-1, keepdims=True)
    return np.exp(shifted) / np.exp(shifted).sum(axis=-1, keepdims=True)


def causal_mask(time):
    if time < 1:
        raise ValueError("time must be positive")
    return np.tril(np.ones((time, time), dtype=bool))


def single_head_attention(hidden, wq, wk, wv, wo):
    hidden = np.asarray(hidden, dtype=float)
    query = hidden @ wq
    key = hidden @ wk
    value = hidden @ wv
    scores = query @ key.T / np.sqrt(query.shape[-1])
    visible = causal_mask(len(hidden))
    masked_scores = np.where(visible, scores, -np.inf)
    weights = row_softmax(masked_scores)
    attended = weights @ value
    return attended @ wo, weights, attended


rng = np.random.default_rng(SEED)
time, width = 5, 6
hidden = rng.normal(size=(time, width))
projections = [rng.normal(scale=0.25, size=(width, width)) for _ in range(4)]
attention_output, attention_weights, attended = single_head_attention(hidden, *projections)
print("output shape:", attention_output.shape)
print("attention weights at the last position:", np.round(attention_weights[-1], 3))
assert attention_output.shape == (time, width)
assert attended.shape == (time, width)
assert np.allclose(attention_weights.sum(axis=1), 1.0)
assert np.allclose(attention_weights[np.triu_indices(time, k=1)], 0.0)

future_changed = hidden.copy()
future_changed[-1] += 100.0
changed_output = single_head_attention(future_changed, *projections)[0]
np.testing.assert_allclose(changed_output[:-1], attention_output[:-1], atol=1e-10)
print("future-token mask test: earlier positions unchanged")

output shape: (5, 6)
attention weights at the last position: [0.276 0.127 0.206 0.256 0.135]
future-token mask test: earlier positions unchanged

The future-token intervention changes the query, key, and value at the last position, but earlier rows remain identical because their visible key set ends at their own position. The zero upper triangle is the direct causal invariant. A row-wise softmax is essential: normalizing across all scores would allow a future position to receive nonzero weight even after masking.

Multi-head attention

With h heads and width d, choose d_h=d/h. The same projected sequence is reshaped from (T,d) to (h,T,d_h), and each head computes its own score matrix. Concatenate the head outputs back to (T,d) and apply an output projection. The split does not create more information by itself; it gives different subspaces separate attention weights.

Keep the head axis first in the next implementation. This makes the score contraction explicit and gives shape assertions a single convention.

def split_heads(values, heads):
    time, width = values.shape
    if width % heads:
        raise ValueError("width must be divisible by heads")
    head_width = width // heads
    return values.reshape(time, heads, head_width).transpose(1, 0, 2)


def merge_heads(values):
    heads, time, head_width = values.shape
    return values.transpose(1, 0, 2).reshape(time, heads * head_width)


def multi_head_attention(hidden, wq, wk, wv, wo, heads):
    query = split_heads(hidden @ wq, heads)
    key = split_heads(hidden @ wk, heads)
    value = split_heads(hidden @ wv, heads)
    scores = np.einsum("htd,hsd->hts", query, key) / np.sqrt(query.shape[-1])
    visible = causal_mask(hidden.shape[0])
    weights = row_softmax(np.where(visible[None, :, :], scores, -np.inf))
    attended = np.einsum("hts,hsd->htd", weights, value)
    return merge_heads(attended) @ wo, weights, merge_heads(attended)


head_rng = np.random.default_rng(SEED + 1)
heads = 2
head_weights = [head_rng.normal(scale=0.2, size=(width, width)) for _ in range(4)]
mult_head_output, mult_head_weights, merged_context = multi_head_attention(
    hidden, *head_weights, heads
)
print("head weights shape:", mult_head_weights.shape)
print("merged context shape:", merged_context.shape)
assert mult_head_output.shape == (time, width)
assert mult_head_weights.shape == (heads, time, time)
assert merged_context.shape == (time, width)
assert np.allclose(mult_head_weights.sum(axis=-1), 1.0)
assert np.allclose(mult_head_weights[:, np.triu_indices(time, k=1)[0], np.triu_indices(time, k=1)[1]], 0.0)

head weights shape: (2, 5, 5)
merged context shape: (5, 6)

The output projection restores the model width after the heads are concatenated. The shape assertions catch three common mistakes: treating the head axis as a batch axis, contracting keys along the wrong time dimension, and returning one output per head instead of one output per position. Each head still has a zero upper triangle, so splitting channels does not weaken the causal boundary.

Pre-norm residual blocks

For a hidden vector x\in\mathbb{R}^d, LayerNorm computes

\operatorname{LN}(x)=\gamma\odot\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta,

where \mu and \sigma^2 are the mean and variance across the channel axis. It normalizes each position independently, so sequence length and other examples do not change the statistics of a position. The learnable \gamma and \beta restore a trainable scale and offset.

A pre-norm residual block applies normalization inside each branch:

H_1=X+\operatorname{MHA}(\operatorname{LN}(X)),\qquad H_2=H_1+\operatorname{FFN}(\operatorname{LN}(H_1)).

The identity terms provide a direct path for both activations and derivatives. The feed-forward branch expands the channel width, applies a pointwise nonlinearity, and projects back to the model width.

def layer_norm(values, gamma, beta, epsilon=1e-5):
    values = np.asarray(values, dtype=float)
    mean = values.mean(axis=-1, keepdims=True)
    variance = ((values - mean) ** 2).mean(axis=-1, keepdims=True)
    normalized = (values - mean) / np.sqrt(variance + epsilon)
    return gamma * normalized + beta


def gelu(values):
    coefficient = np.sqrt(2.0 / np.pi)
    return 0.5 * values * (1.0 + np.tanh(coefficient * (values + 0.044715 * values ** 3)))


def feed_forward(values, w1, b1, w2, b2):
    return gelu(values @ w1 + b1) @ w2 + b2


def decoder_block(values, attention_weights, feed_forward_weights, heads):
    gamma1, beta1, wq, wk, wv, wo = attention_weights
    gamma2, beta2, w1, b1, w2, b2 = feed_forward_weights
    normalized = layer_norm(values, gamma1, beta1)
    attended, _, _ = multi_head_attention(normalized, wq, wk, wv, wo, heads)
    after_attention = values + attended
    normalized_again = layer_norm(after_attention, gamma2, beta2)
    return after_attention + feed_forward(
        normalized_again, w1, b1, w2, b2
    )


def tiny_transformer_forward(token_ids, token_embedding, positional_embedding, blocks, output_projection, heads):
    sequence_length = len(token_ids)
    hidden = token_embedding[token_ids] + positional_embedding[:sequence_length]
    for attention_weights, feed_forward_weights in blocks:
        hidden = decoder_block(hidden, attention_weights, feed_forward_weights, heads)
    return hidden @ output_projection


model_rng = np.random.default_rng(SEED + 2)
vocab_size = 9
heads = 2
feed_forward_width = 4 * width
attention_parameters = (
    np.ones(width),
    np.zeros(width),
    model_rng.normal(scale=0.2, size=(width, width)),
    model_rng.normal(scale=0.2, size=(width, width)),
    model_rng.normal(scale=0.2, size=(width, width)),
    model_rng.normal(scale=0.2, size=(width, width)),
)
feed_forward_parameters = (
    np.ones(width),
    np.zeros(width),
    model_rng.normal(scale=0.15, size=(width, feed_forward_width)),
    np.zeros(feed_forward_width),
    model_rng.normal(scale=0.15, size=(feed_forward_width, width)),
    np.zeros(width),
)
token_embedding = model_rng.normal(scale=0.2, size=(vocab_size, width))
positional_embedding = model_rng.normal(scale=0.2, size=(time, width))
output_projection = model_rng.normal(scale=0.2, size=(width, vocab_size))
block_output = tiny_transformer_forward(
    np.arange(time),
    token_embedding,
    positional_embedding,
    [(attention_parameters, feed_forward_parameters)],
    output_projection,
    heads,
)
print("decoder logits shape:", block_output.shape)
print("LayerNorm means:", np.round(layer_norm(hidden, np.ones(width), np.zeros(width)).mean(axis=-1), 6))
assert block_output.shape == (time, vocab_size)
assert np.allclose(layer_norm(hidden, np.ones(width), np.zeros(width)).mean(axis=-1), 0.0, atol=1e-5)
assert np.allclose(layer_norm(hidden, np.ones(width), np.zeros(width)).var(axis=-1), 1.0, atol=1e-4)

decoder logits shape: (5, 9)
LayerNorm means: [-0.  0. -0.  0.  0.]

The block preserves the (T,d) shape through both residual additions, while the feed-forward branch temporarily uses 4d channels. LayerNorm’s mean and variance checks are per position, not over the whole sequence. The output projection turns the final hidden states into one vocabulary logit vector per position; a later causal loss can compare position t with token t+1.

Attention backpropagation

A full Transformer backward pass combines matrix products, softmax, normalization, and the nonlinearity. A useful intermediate check isolates a projection that appears inside the complete forward path. If C is the merged attention context, O=CW_O, and U=\partial L/\partial O, then

\frac{\partial L}{\partial W_O}=C^\top U.

Compare this hand-derived slice with finite differences. Once this check passes, the same matrix-calculus rule can be composed with the gradient of the attention context.

def output_projection_loss_and_grad(context, output_projection, upstream):
    output = context @ output_projection
    loss = float(np.sum(output * upstream))
    gradient = context.T @ upstream
    return loss, gradient


grad_rng = np.random.default_rng(SEED + 3)
context = grad_rng.normal(size=(4, width))
projection = grad_rng.normal(scale=0.2, size=(width, width))
upstream = grad_rng.normal(size=(4, width))
loss, analytic_gradient = output_projection_loss_and_grad(context, projection, upstream)
epsilon = 1e-5
numeric_gradient = np.zeros_like(projection)
for index in np.ndindex(projection.shape):
    original = projection[index]
    projection[index] = original + epsilon
    plus = output_projection_loss_and_grad(context, projection, upstream)[0]
    projection[index] = original - epsilon
    minus = output_projection_loss_and_grad(context, projection, upstream)[0]
    projection[index] = original
    numeric_gradient[index] = (plus - minus) / (2.0 * epsilon)
maximum_error = np.max(np.abs(analytic_gradient - numeric_gradient))
print("maximum output-projection gradient error:", maximum_error)
np.testing.assert_allclose(analytic_gradient, numeric_gradient, atol=1e-8, rtol=1e-6)

maximum output-projection gradient error: 2.1735502286901465e-11

The local gradient agrees at finite-difference precision. This test does not certify every derivative in the block; it certifies one interface that is easy to get wrong when head outputs are merged. The remaining attention and LayerNorm derivatives can be tested by applying the same finite-difference pattern to a small selected parameter set.

Depth and gradient flow

Backpropagation multiplies a Jacobian from every block. To isolate that multiplication, use a controlled linearized branch: each layer receives an orthogonal branch Jacobian with a prescribed gain. A residual block has Jacobian I+A; a plain stacked branch has Jacobian A. The normalized flag reduces the branch gain, representing the scale control that normalization provides around a typical operating point. This is not a replacement for the exact LayerNorm Jacobian; it is a small experiment for the path-length mechanism.

Measure the norm of a fixed upstream gradient after traversing increasing depth. Use the same random branch directions for all four configurations.

def gradient_norm_by_depth(depth, residual, normalized, width=8, seed=SEED):
    random_generator = np.random.default_rng(seed)
    total_jacobian = np.eye(width)
    for _ in range(depth):
        random_matrix = random_generator.normal(size=(width, width))
        orthogonal, _ = np.linalg.qr(random_matrix)
        branch_gain = 0.15 if normalized else 0.70
        branch_jacobian = branch_gain * orthogonal
        block_jacobian = branch_jacobian + (np.eye(width) if residual else 0.0)
        total_jacobian = block_jacobian @ total_jacobian
    upstream = np.ones(width) / np.sqrt(width)
    return float(np.linalg.norm(total_jacobian.T @ upstream))


configurations = {
    "plain": (False, False),
    "normalized branch": (False, True),
    "residual": (True, False),
    "pre-norm residual": (True, True),
}
depths = (1, 2, 4, 8, 12)
experiment = {}
for name, (residual, normalized) in configurations.items():
    values = [gradient_norm_by_depth(depth, residual, normalized) for depth in depths]
    experiment[name] = values
    print(f"{name:>18}: " + " ".join(f"{value:6.3f}" for value in values))

assert experiment["pre-norm residual"][-1] > experiment["plain"][-1]
assert experiment["plain"][-1] < experiment["residual"][-1]

             plain:  0.700  0.490  0.240  0.058  0.014
 normalized branch:  0.150  0.022  0.001  0.000  0.000
          residual:  1.125  1.341  0.621  2.672 10.149
 pre-norm residual:  0.987  0.982  0.855  0.777  0.927

Tensor shapes

A training call usually carries a batch axis in addition to time and channel. The single-sequence function above is easy to inspect, so wrap it over a batch without changing its per-sequence mask. Assert every axis after stacking; a hidden broadcast over the batch axis would otherwise look like a valid matrix multiplication.

def batched_multi_head_attention(hidden_batch, wq, wk, wv, wo, heads):
    hidden_batch = np.asarray(hidden_batch, dtype=float)
    if hidden_batch.ndim != 3:
        raise ValueError("hidden_batch must have shape (batch, time, width)")
    outputs = []
    weights = []
    for sequence in hidden_batch:
        output, sequence_weights, _ = multi_head_attention(
            sequence, wq, wk, wv, wo, heads
        )
        outputs.append(output)
        weights.append(sequence_weights)
    return np.stack(outputs), np.stack(weights)


batch_hidden = np.stack([hidden, hidden + 0.1])
batch_output, batch_attention_weights = batched_multi_head_attention(
    batch_hidden, *head_weights, heads
)
print("batched output shape:", batch_output.shape)
print("batched attention shape:", batch_attention_weights.shape)
assert batch_output.shape == (2, time, width)
assert batch_attention_weights.shape == (2, heads, time, time)
assert np.allclose(batch_attention_weights.sum(axis=-1), 1.0)

batched output shape: (2, 5, 6)
batched attention shape: (2, 2, 5, 5)

The batched wrapper preserves the order (batch, time, channel) for hidden states and (batch, head, time, time) for attention weights. Each batch item receives its own causal mask; no sequence can read another sequence because the loop keeps their score matrices separate. The assertion complements the earlier single-sequence checks rather than replacing them.

Reusable PyTorch decoder

The NumPy attention path exposes the mathematics and the causal intervention. The model artifact used by later chapters is ProofLM: it adds RoPE, tied embeddings, masked loss, generation, and state-dict serialization behind a stable forward interface. The checks below compare the same invariants at the module boundary.

import io

import torch

from proof_lm.model import DecoderConfig, ProofLM

torch.manual_seed(SEED)
decoder_config = DecoderConfig(
    vocab_size=32,
    context_length=16,
    n_layers=2,
    d_model=64,
    n_heads=4,
    d_ff=128,
)
decoder = ProofLM(decoder_config)
decoder.eval()
token_ids = torch.tensor([[1, 4, 7, 9, 3]])
decoder_logits, decoder_loss = decoder(token_ids)
changed_ids = token_ids.clone()
changed_ids[0, -1] = 12
changed_logits = decoder(changed_ids)[0]
labels = torch.tensor([[4, 7, 9, 3, 2]])
loss_mask = torch.ones_like(labels, dtype=torch.bool)
_, masked_loss = decoder(token_ids, labels=labels, loss_mask=loss_mask)
masked_loss.backward()
buffer = io.BytesIO()
torch.save(decoder.state_dict(), buffer)
reloaded_decoder = ProofLM(decoder_config)
reloaded_decoder.load_state_dict(
    torch.load(io.BytesIO(buffer.getvalue()), weights_only=False)
)
reloaded_decoder.eval()
reloaded_logits = reloaded_decoder(token_ids)[0]
finite_gradients = all(
    parameter.grad is None or torch.isfinite(parameter.grad).all()
    for parameter in decoder.parameters()
)

print("PyTorch decoder logits:", tuple(decoder_logits.shape))
print("parameter count:", decoder.parameter_count)
print("tied embeddings:", decoder.embeddings_are_tied)
print("masked loss:", round(float(masked_loss.detach()), 6))
print("CUDA available:", torch.cuda.is_available())
assert decoder_loss is None
assert decoder_logits.shape == (1, 5, decoder_config.vocab_size)
assert decoder.embeddings_are_tied
assert masked_loss is not None and torch.isfinite(masked_loss)
torch.testing.assert_close(decoder_logits[:, :-1], changed_logits[:, :-1])
torch.testing.assert_close(decoder_logits, reloaded_logits)
assert finite_gradients

PyTorch decoder logits: (1, 5, 32)
parameter count: 69120
tied embeddings: True
masked loss: 3.640021
CUDA available: False

/var/folders/jq/9vsvd9252_349lsng_5gc_jw0000gn/T/ipykernel_56589/3606104176.py:43: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
Consider using tensor.detach() first. (Triggered internally at /Users/runner/work/pytorch/pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:823.)
  print("masked loss:", round(float(masked_loss), 6))

The plain stack repeatedly multiplies a branch with gain below one, so its upstream signal shrinks rapidly. Residual addition retains an identity component, and reducing the branch gain keeps the residual Jacobian closer to an identity map across depth. The exact values depend on width and initialization; the robust mechanism is the additive identity path, while normalization controls the scale of the learned branch that is added to it.

Summary

• Causal self-attention masks future keys before softmax, and a future-token intervention verifies the mask behavior directly.
• Multi-head attention splits the channel axis, computes independent score matrices, and restores the model width after merging.
• LayerNorm normalizes channels per position; the pre-norm residual block preserves shape and creates a direct gradient path.
• A hand-derived output-projection gradient agrees with finite differences, providing a local backward check for the block.
• A controlled depth experiment shows why residual paths and branch-scale control affect gradient propagation.

Chapter 05 places a small causal model in an end-to-end training loop with validation, checkpoints, and reproducible resumption.

Exercises

Use the exercises to test the chapter’s invariants and connect the derivations to the reusable implementation. Solutions are hidden in the notebook source and are available through the course tooling when needed.

[P4.1] Causal intervention

Causal intervention. Describe a test that changes only the final hidden state and verifies that attention outputs at earlier positions do not change. State why the same test would fail without a causal mask.

[P4.2] Residual-depth gradient

Residual-depth analysis. Explain the difference between the Jacobian of a plain branch and a residual branch. Use the depth experiment to identify which configuration keeps the upstream gradient closer to its initial norm and why.

[P4.3]

Decoder invariants. A serialized decoder must preserve its output contract. List three checks for a causal decoder module covering (1) future-token visibility, (2) parameter sharing, and (3) save/load identity. Then state why the first check is stronger than inspecting only the triangular attention weights.

def decoder_checks(model, original_logits, changed_logits, reloaded_logits):
    # Return three booleans for the requested module contracts.
    pass