Build Attention Before You Trust a Transformer Diagram

¶How do queries, keys, and values select information from a sequence?

A token embedding alone cannot decide which surrounding words matter to its current meaning. This chapter builds the Transformer from the operation that lets token representations exchange information, so you can trace the numbers instead of trusting boxes labelled “attention.”

Consider the study assistant reading “The exam covers graphs; it excludes trees.” The representation at “it” needs information from elsewhere in the sequence. Attention supplies a weighted mixture of available information; it does not perform a hard lookup or guarantee that the pronoun is resolved correctly. I prefer to start with that limitation: “the model focuses on the right word” hides both the arithmetic and the possibility of failure.

¶Match with keys, carry information with values

Each participating position supplies three vectors. A query describes what the receiving position matches against; a key supplies matching features; a value supplies the information contributed to the result. These are learned numerical roles, not English questions and answers.

In self-attention, all three come from the same input sequence, but through separate learned projections: Q = X @ W_Q, K = X @ W_K, and V = X @ W_V. Sharing the input does not make the projections identical. If X has shape (3, 4) and each projection matrix has shape (4, 2), each result has shape (3, 2). For clarity, these shapes omit batch and head dimensions.

A query scores every available key using a dot product. Larger scores receive larger weights after softmax, unless a later restriction excludes that position. The output is the weighted sum of the values—not the keys. Separating matching features from carried information lets a position advertise one feature while contributing another.

Scaled dot-product attention computes softmax(Q @ K.T / sqrt(d_k)) @ V, with softmax applied across keys. Here d_k is the query/key width. Under the simplifying assumption of independent, zero-mean, unit-variance components, a dot product has variance d_k; dividing by its square root controls that growth. Without scaling, wide vectors can produce sharply concentrated weights and small softmax gradients. This is a motivation, not a promise about every learned vector distribution.

¶Inspect one receiving position

I use hand-picked vectors first because random projections obscure the mechanism. The following experiment uses 1 query, 3 keys, and 3 values, each of width 2. It illustrates the calculation, not a trained language model.

import numpy as np

query = np.array([[1.0, 0.0]])
keys = np.array([[2.0, 0.0],
                 [0.0, 2.0],
                 [1.0, 1.0]])
values = np.array([[10.0, 0.0],
                   [0.0, 10.0],
                   [4.0, 6.0]])

if query.shape[1] != keys.shape[1]:
    raise ValueError("query and key widths must match")
if keys.shape[0] != values.shape[0]:
    raise ValueError("each key must have a corresponding value")

scores = query @ keys.T / np.sqrt(keys.shape[1])
shifted = scores - scores.max(axis=-1, keepdims=True)
weights = np.exp(shifted)
weights /= weights.sum(axis=-1, keepdims=True)
output = weights @ values

assert weights.shape == (1, 3)
assert output.shape == (1, 2)
assert np.allclose(weights.sum(axis=-1), 1.0)

print("weights:", weights.round(3))
print("output:", output.round(3))

The weights are approximately [0.576, 0.140, 0.284]; the output is approximately [6.896, 3.104]. The first value contributes most, but the other two still contribute. Replacing a value changes the output without changing these weights, provided the query and keys remain fixed.

Two bugs deserve explicit checks. Mismatched query/key widths cause a matrix-multiplication ValueError. Normalizing over the wrong axis can silently produce incorrect weights without any exception; each receiving position’s weights must sum to 1 across available keys. Subtracting the row maximum also avoids exponential overflow for large finite scores.

For an interview, I would explain selection as “query–key scores determine how values are mixed,” rather than calling attention an explanation of the model’s decision. High weight alone does not establish causal importance. Change the third value to [100.0, 100.0]: why does the output move dramatically even though the selected weights stay identical?

¶Can we implement scaled dot-product attention and verify every shape?

An attention block can return the expected shape and still mix information along the wrong axis. Applying torch.softmax across queries rather than keys is particularly unpleasant: nothing necessarily crashes, but each output no longer represents the intended weighted average. I prefer a tiny, inspectable implementation before using an optimized attention operation. A passing shape check is necessary, not sufficient.

For the study assistant, imagine 2 query positions attending to 3 positions in a short course-note sequence. These are positions inside attention, not a document-search system. Use query and key vectors of width 2, but value vectors of width 1 so that you can inspect each output numerically. Query and key widths must match; value width need not.

The operation is softmax(QKᵀ / √d_k)V. Here d_k is the query/key width. With a batch size of 1, the shape ledger is:

Why divide by the square root? Under the simplifying assumption that query and key components are independent, zero-mean variables with variance 1, their dot product has variance d_k. Division by √d_k brings that variance back to 1. Without scaling, wider vectors can produce more extreme scores, pushing softmax toward near-one-hot weights and small gradients. Learned vectors need not satisfy those assumptions exactly; this is the motivation, not a guarantee about every trained layer.

import math
import torch


def scaled_attention(q, k, v):
    if any(t.ndim != 3 for t in (q, k, v)):
        raise ValueError("expected rank-3 q, k, v")
    if not (q.shape[0] == k.shape[0] == v.shape[0]):
        raise ValueError("batch sizes must match")
    if q.shape[-1] != k.shape[-1] or q.shape[-1] == 0:
        raise ValueError("query/key widths must match and be positive")
    if k.shape[1] != v.shape[1] or k.shape[1] == 0:
        raise ValueError("key/value lengths must match and be positive")

    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
    weights = torch.softmax(scores, dim=-1)
    output = weights @ v
    return output, weights


q = torch.tensor([[[math.sqrt(2), 0.0],
                   [0.0, math.sqrt(2)]]])
k = torch.tensor([[[1.0, 0.0],
                   [0.0, 1.0],
                   [0.0, 0.0]]])
v = torch.tensor([[[10.0], [20.0], [30.0]]])

output, weights = scaled_attention(q, k, v)

assert weights.shape == (1, 2, 3)
assert output.shape == (1, 2, 1)
torch.testing.assert_close(weights.sum(dim=-1), torch.ones(1, 2))

expected_weights = torch.softmax(
    torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]),
    dim=-1,
)
torch.testing.assert_close(weights, expected_weights)
torch.testing.assert_close(output, expected_weights @ v)

uniform_output, _ = scaled_attention(torch.zeros_like(q), k, v)
torch.testing.assert_close(uniform_output, torch.full((1, 2, 1), 20.0))

shifted_output, shifted_weights = scaled_attention(q, k, v + 5.0)
torch.testing.assert_close(shifted_weights, weights)
torch.testing.assert_close(shifted_output, output + 5.0)

print("weights:", weights)
print("output:", output)

The first score row is [1, 0, 0], giving weights approximately [0.5761, 0.2119, 0.2119]. Its output is approximately 16.36, not 10: attention mixes values rather than selecting only the strongest match. The second output is 20. Equal zero scores must yield the arithmetic mean, which is why the zero-query test matters.

I would keep the value-shift test even in a small exercise. Adding 5 to every value must add 5 to every output, as long as each weight row sums to 1. It must not change the weights. This catches mistakes that checking dimensions alone misses.

The function deliberately rejects unequal batch sizes with ValueError("batch sizes must match"); otherwise, broadcasting could silently reuse one example across several others. A query/key width mismatch raises ValueError("query/key widths must match and be positive") before matrix multiplication fails. Inputs must also share a floating-point dtype and device, unless you explicitly convert them first. Non-finite inputs can propagate nan; scaling does not repair them.

Interview check: Why does dim=-1 matter? Because the last score dimension indexes keys, and each query needs its own distribution over those keys. As an exercise, change it to dim=-2 and identify which assertions fail. This implementation has no position restrictions yet: what must change when a query is forbidden from reading a particular key?

¶Why do causal masks and positional encodings change what a model can know?

Your study assistant can appear excellent at predicting the next token while secretly reading it. If every query can attend to every input position, the representation used to predict an answer can already contain that answer. A causal mask closes this information leak: position 2 may read positions 0, 1, and 2, but not position 3. The diagonal stays visible because the representation at the current position predicts the following token.

I test this restriction before trusting any attention implementation. A convincing loss curve cannot prove that information flowed legally.

¶Block access before normalization

For a sequence of 4 tokens, attention has a 4-by-4 score matrix. Rows are query positions; columns are key positions. A causal mask blocks the 6 entries strictly above the diagonal. Set those scores to negative infinity before calling torch.softmax, so their probabilities become zero and each permitted row still sums to 1.

Setting blocked scores to zero is a bug: zero remains a valid score and can receive substantial probability. Zeroing probabilities afterward is also wrong unless you renormalize them. Worse, blocking every entry in a row makes this explicit softmax calculation return nan: there is no permitted distribution. Padding masks, which exclude padding tokens rather than future tokens, must not accidentally create such rows.

¶Make position distinguishable

Access control does not fully describe order. Without positional information or an order-dependent mask, rearranging the input rows simply rearranges self-attention’s output rows. Content-only attention cannot distinguish “student teaches model” from the corresponding rearrangement by sequence order alone.

The original Transformer adds a fixed positional vector to each token embedding. Its even coordinates use sine; odd coordinates use cosine, with different frequencies across coordinate pairs. For coordinate pair i, the angle is position / 10000 ** (2 * i / d_model). Both vectors have width d_model, so addition preserves the tensor shape.

A causal mask itself introduces ordering asymmetry, so “no positional encoding means absolutely no order information” is too strong for a causal model. But it does not provide the same explicit position signal. I use sinusoidal positions here because you can inspect every value, not because they are universally best. Other architectures learn position vectors or rotate queries and keys; those mechanisms are not interchangeable additions.

import math
import torch

torch.manual_seed(7)
length, width = 4, 4
tokens = torch.randn(length, width)

position = torch.arange(length, dtype=torch.float32)[:, None]
frequency = torch.exp(
    torch.arange(0, width, 2, dtype=torch.float32)
    * (-math.log(10000.0) / width)
)
encoding = torch.zeros(length, width)
encoding[:, 0::2] = torch.sin(position * frequency)
encoding[:, 1::2] = torch.cos(position * frequency)

blocked = torch.triu(
    torch.ones(length, length, dtype=torch.bool), diagonal=1
)

def attend(inputs):
    # Identity projections isolate masking from learned parameters.
    scores = inputs @ inputs.T / math.sqrt(width)
    weights = torch.softmax(
        scores.masked_fill(blocked, float("-inf")), dim=-1
    )
    return weights @ inputs, weights

inputs = tokens + encoding
output, weights = attend(inputs)

changed = inputs.clone()
changed[2:] += 100
changed_output, _ = attend(changed)

assert torch.all(weights[blocked] == 0)
assert torch.allclose(weights.sum(dim=-1), torch.ones(length))
assert torch.allclose(output[:2], changed_output[:2])
print(weights)
print("Future-token isolation passed.")

The final assertion changes the last 2 input vectors while requiring the first 2 outputs to remain unchanged. Unless another operation mixes sequence positions outside this attention block, future content has no path into those outputs. This checks behavior, not merely whether the mask looks triangular. Remove masked_fill and the isolation assertion fails with AssertionError.

Interview check: “Can positional encoding replace a causal mask?” No. Position vectors label locations; they do not prohibit reading them. Knowing that a token belongs to tomorrow does not stop attention from using it today.

Exercise: Reverse the triangle in blocked and identify which positions become visible. Then ask the architectural question: should a representation of an already-complete course note have the same visibility restrictions as a representation generating an answer?

¶What do multiple heads, feed-forward layers, residual paths, and normalization add?

A single attention operation mixes token information, but it does not give you a complete Transformer block. Without the surrounding layers, you lose separate information-routing patterns, nonlinear processing after mixing, and a direct path for preserving the input. I find it more useful to read the block as “gather, transform, and add” than to memorize boxes in a diagram.

¶Multiple heads gather different mixtures

For the study assistant, one attention head might connect a term to its definition while another connects it to a nearby qualification. That is a possible learned behavior, not an assigned job. I avoid labeling heads “grammar” or “reasoning”: separate parameters allow specialization, but do not guarantee interpretable roles.

With model width 64 and 4 heads, each head uses query, key, and value width 16 in the conventional arrangement. Each receives its own learned projections of the full input; you do not simply assign each head a different quarter of the original features. For input shape (2, 5, 64), the projected queries can be rearranged to (2, 4, 5, 16). Each head computes its own attention weights. Concatenation restores width 64, and an output projection mixes the heads’ results.

This arrangement requires the model width to divide evenly by the head count. In PyTorch, an incompatible configuration raises an AssertionError containing embed_dim must be divisible by num_heads. More heads are not automatically better: at fixed model width, each head gets fewer features.

¶The feed-forward layer transforms each position

Attention moves information between positions; the feed-forward network transforms the resulting features at each position independently, using shared weights. A conventional example expands width 64 to 256, applies ReLU, then projects back to 64. Without a nonlinear activation, those two affine layers collapse into one affine transformation. Expansion alone is not the point.

The network does not directly mix sequence positions. Nevertheless, its input already contains information gathered by attention, so “position-wise” does not mean “unaware of context.”

¶Residual paths preserve; normalization controls scale

A residual connection adds a sublayer’s output to its input. This preserves an identity path for information and gradients; it does not guarantee that gradients cannot vanish or explode. Addition requires compatible shapes. If the feed-forward network returns width 256 instead of 64, the residual addition throws a RuntimeError rather than repairing the mismatch.

LayerNorm normalizes across features within each token, then applies learned scale and bias. It does not normalize across sequence positions when configured with the feature width, so it does not expose future tokens. Its small eps term prevents division by zero; a hand-written version without it can produce NaN values on constant inputs.

The original Transformer applies normalization after residual addition. I reach for pre-normalization in a small experimental block: normalize before each sublayer, leaving the residual stream’s identity path unobstructed by normalization. This is a useful training choice, not a universal replacement for every architecture’s specified ordering. The following runnable PyTorch example uses pre-normalization and omits dropout deliberately.

import torch
from torch import nn

torch.manual_seed(7)

class StudyBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm1 = nn.LayerNorm(64)
        self.attention = nn.MultiheadAttention(
            embed_dim=64, num_heads=4, dropout=0.0, batch_first=True
        )
        self.norm2 = nn.LayerNorm(64)
        self.feed_forward = nn.Sequential(
            nn.Linear(64, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
        )

    def forward(self, x):
        length = x.size(1)
        blocked = torch.triu(
            torch.ones(length, length, dtype=torch.bool, device=x.device),
            diagonal=1,
        )
        normalized = self.norm1(x)
        attended, _ = self.attention(
            normalized, normalized, normalized,
            attn_mask=blocked, need_weights=False,
        )
        x = x + attended
        return x + self.feed_forward(self.norm2(x))

block = StudyBlock().eval()
x = torch.randn(2, 5, 64)

with torch.no_grad():
    output = block(x)
    changed = x.clone()
    changed[:, 4, :] += torch.randn_like(changed[:, 4, :]) * 10
    changed_output = block(changed)

assert output.shape == x.shape
assert torch.isfinite(output).all().item()
torch.testing.assert_close(output[:, :4], changed_output[:, :4])
print(tuple(output.shape))

The last assertion checks a boundary, not model intelligence: changing the fifth token must not change the first 4 outputs under this causal mask. As an exercise, remove the mask and rerun it; the assertion should fail. For an interview, explain why: neither residual addition nor position-wise processing restores causality after attention has mixed in future information. What changes when another block is allowed to attend to a separate sequence?

¶How do the original Transformer, encoder-only models, and decoder-only models differ?

A model can produce perfectly shaped tensors while receiving information it should never see. Remove a decoder’s causal mask, for example, and a position can read the answer token it is supposed to predict. The code may run without an exception. The experiment is still broken. I compare Transformer families by drawing their permitted information paths first; counting layers comes second.

¶The original Transformer: two stacks, three attention paths

The original 2017 Transformer paired an encoder with a decoder. Its base configuration used 6 layers in each stack and a model width of 512. Those numbers describe that configuration, not requirements for calling something a Transformer.

The encoder reads the source sequence using bidirectional self-attention: each source position can attend to every source position, excluding padding when present. The decoder has causal self-attention over the target sequence, followed by cross-attention to the encoder’s output. Cross-attention connects two sequences: decoder states supply Q, while encoder outputs supply K and V. The feed-forward layers and residual paths remain as described earlier.

Suppose a translation input has 12 source tokens and the current target sequence has 3 tokens. For one attention head, encoder self-attention scores have shape (12, 12), decoder self-attention scores have shape (3, 3), and cross-attention scores have shape (3, 12), omitting the batch dimension. Each target position asks a separate question of the same source representation.

The source is fully available before translation begins, so reading its final token is allowed. Reading a future target token is not. Applying a triangular mask indiscriminately to both stacks silently restricts the encoder; leaving the target unmasked leaks future answers. Neither mistake needs to cause a shape error.

¶Encoder-only: represent the available input

BERT is the familiar encoder-only example. Its ordinary attention pattern lets input positions exchange information in both directions. There is no separate autoregressive decoder stack. An output head can turn its representations into classifications or token-level predictions, provided that head has been trained for the task.

For the study assistant, an encoder could eventually help represent a question or classify a passage; the retrieval machinery comes later. It does not follow that every encoder produces useful sentence vectors without additional choices. I would not select an encoder merely because its output has the vector shape an application expects. Shape compatibility is not semantic suitability.

¶Decoder-only: generate from a shared prefix

GPT-style decoder-only models use causal self-attention over one sequence, without the original architecture’s separate encoder and encoder-decoder cross-attention. Notes, a question, and an answer prefix can occupy that same sequence. An answer position can attend to earlier notes because they precede it, not because a hidden encoder supplied them.

“Decoder-only” therefore does not mean an untouched original decoder with its encoder disconnected. It names a family of causal Transformer stacks. Unless an architecture explicitly adds another information path, do not invent one from the label. A decoder-only model can process supplied text, but its earlier positions cannot read later question tokens within that forward pass.

¶A useful interview check

Question: Does answering questions about an input document require an encoder-decoder model?

Answer: No. An encoder-decoder model can encode the document separately; a decoder-only model can receive it before the question and generate an answer afterward. The distinguishing mechanism is where source information enters attention, not whether the application accepts input and produces output.

As a tracing exercise, place a question after a passage and mark which positions can read it in each family. Once those arrows are correct, the next question is practical: how much computation does each attention matrix require as the passage grows?

¶Where does attention spend memory, and how would you defend the architecture at a whiteboard?

Doubling the course notes fed into your study assistant can quadruple an attention allocation, even though the model’s weights have not changed. That is the first calculation I would put on a whiteboard—not a diagram of stacked rectangles. An architecture explanation earns trust when it predicts what breaks.

¶Count the large intermediate, not just the weights

In the explicit attention implementation from earlier, the score tensor has shape (B, H, N, N): batch size, head count, and two sequence dimensions. Its storage is B * H * N * N * element_bytes. Queries, keys, and values each have shape (B, H, N, d_h); their storage grows linearly with sequence length when the other dimensions stay fixed.

For a batch of 1, 8 heads, and 2,048 tokens, one score tensor contains 33,554,432 elements. At 2 bytes per element, that is 64 MiB. At 4,096 tokens, it becomes 256 MiB. These are calculated tensor sizes, not measured peak memory: separate probabilities, gradients, temporary buffers, and other layers can add allocations, depending on the implementation.

I prefer an allocation-free estimate before a large experiment. This Python script uses score_bytes and element_bytes to make the assumptions visible; it does not require a GPU.

batch_size = 1
heads = 8
element_bytes = 2

for sequence_length in (2048, 4096):
    score_bytes = (
        batch_size * heads * sequence_length**2 * element_bytes
    )
    print(f"{sequence_length} tokens: {score_bytes / 1024**2:.0f} MiB")

A causal mask does not automatically halve this allocation. If your code creates the full score matrix and then masks its upper triangle, those entries still occupy memory. Padding can waste it too: when every sequence is padded to the longest example, the dense allocation follows that padded length, not the average length.

The resulting failure may appear as torch.OutOfMemoryError. Reducing the batch size can help, but shortening the sequence attacks the squared term. A different failure, NaN attention probabilities, warrants checking for a fully masked row: softmax over only negative infinities is undefined. More GPU memory will not repair that bug.

¶Defend the mechanism, including its limits

Interview question: “Is attention necessarily quadratic in memory?” My answer: “The explicit score matrix is quadratic in sequence length. Dense attention arithmetic also scales quadratically in sequence length, with head width held fixed. But an implementation need not store the entire matrix simultaneously.”

FlashAttention-style kernels process tiles and maintain the quantities needed for softmax without materializing the full score matrix in device memory. They do not make dense attention’s arithmetic linear. I avoid promising a speedup without a measurement on the target hardware; kernel availability, dimensions, and numerical precision affect the result.

Follow-up: “Would adding heads always increase memory?” Only under stated assumptions. At fixed model width, adding heads reduces each head’s width, leaving the combined query, key, and value sizes unchanged. The explicit score tensor still grows with head count. Saying “more heads means proportionally more memory everywhere” misses that distinction.

For an exercise, change the script to 16 heads and predict the output before running it. Then explain why a dense causal mask does not change those printed sizes. That explanation is a better architecture defense than reciting component names.

The next chapter needs this separation between information flow, stored tensors, and computation to explain how training changes a Transformer’s weights. Keep the mask in view: it determines which information a prediction can use while those weights are learned.