Transformers: attention explained and implemented

Read a sample · Chapter 01 of 06

Each row has a role. Write it down before multiplying anything.

Create a folder containing attention.py, block.py, demo.py and test_attention.py. The complete files appear in this workbook. Run them with Python 3.12; the exercised version was 3.12.10. There are no dependencies to install. Keep these files in one folder so the imports resolve. Work on the tiny fixture first: increasing its size will make the arithmetic less visible without making it a useful language model.

A query is a vector used to compare against key vectors. Each key has a corresponding value vector. The comparison produces one score per key; the normalised scores determine how much of each value reaches the output. Q contains query rows, K key rows and V value rows. If Q is L by dk, K is S by dk and V is S by dv, the weights have shape L by S and the result has shape L by dv.

Our first fixture has two query rows, two key rows and two value features. Q and K are both [[1, 0], [0, 1]], while V is [[10, 0], [0, 20]]. These numbers are chosen for hand calculation, not obtained from text or training. The first query matches the first key more strongly, but the second key still receives positive weight when both are visible. An attention row is a weighted mixture rather than a mandatory single-item selection.

Self-attention uses representations from one sequence to form queries, keys and values. This does not mean their numbers must be identical: trainable projections usually give them different roles. Cross-attention can take queries from one sequence and keys and values from another. Our attention function permits different query and key lengths in noncausal mode. Its causal mode deliberately requires equal lengths, avoiding an ambiguous cache-offset convention in this small lab.

Start attention.py with the imports, input checker and matrix multiplication below, placing the second part after the first with blank lines between functions. The checker accepts only rectangular, bounded finite numeric matrices. The 64-row, 64-column and magnitude limits are teaching constraints, not transformer architecture limits. Booleans are rejected even though Python treats bool as a subclass of int. The checks keep a malformed shape from silently passing through zip.

"""Small forward-only attention lab; no training or tensor dependency."""
import math


def matrix(value, name):
    if type(value) not in (list, tuple) or not 1 <= len(value) <= 64:
        raise ValueError(f"{name}: expected 1 to 64 rows")
    if any(type(row) not in (list, tuple) for row in value):
        raise ValueError(f"{name}: expected rows")
    width = len(value[0])
    if not 1 <= width <= 64 or any(len(row) != width for row in value):
        raise ValueError(f"{name}: expected a non-empty rectangle")
    if any(type(x) not in (int, float) or abs(x) > 1e6
           or not math.isfinite(x) for row in value for x in row):
        raise ValueError(f"{name}: expected bounded finite numbers")
    return [[float(x) for x in row] for row in value]

def matmul(left, right):
    a, b = matrix(left, "left"), matrix(right, "right")
    if len(a[0]) != len(b):
        raise ValueError("Incompatible multiplication shapes")
    columns = list(zip(*b))
    return [[math.fsum(x*y for x, y in zip(row, col))
             for col in columns] for row in a]

For matmul, the left width must equal the right row count. Multiplying [[1, 2]] by [[3, 4], [5, 6]] gives [[13, 16]]: the first result is 1*3 + 2*5, and the second is 1*4 + 2*6. Our explicit loops prioritise traceability. They do not offer the execution speed or automatic differentiation of a tensor library. Matrix validation also does not establish that a chosen model architecture is useful.

Read a sample · Chapter 02 of 06

Scale the comparisons, stabilise the exponential and preserve the row sum.

The scaled dot-product rule uses Q times the transpose of K, divides scores by the square root of dk, applies a row-wise softmax, and multiplies the resulting weights by V. The original transformer paper describes this operation and multi-head composition. Here we implement the arithmetic directly. A softmax row assigns nonnegative weights summing to one before any dropout; this lab implements no dropout.

In the first fixture, dk is two. Query [1, 0] has dot products 1 and 0 with the two keys. Dividing by sqrt(2) gives approximately 0.707107 and 0. Exponentiating and normalising gives approximately 0.669762 and 0.330238. The output is therefore [10*0.669762, 20*0.330238], approximately [6.697615, 6.604769]. Keep full precision in the program; rounding intermediate values changes the last digits.

Softmax is unchanged mathematically when the same constant is subtracted from every finite score in a row: the common exponential factor cancels from numerator and denominator. Subtracting the largest score makes the largest exponential one. Thus [10000, 10000] can yield [0.5, 0.5] without evaluating exp(10000). Small terms may still underflow towards zero. This routine is a controlled floating-point teaching implementation, not an arbitrary-precision calculator.

Append softmax and attention to attention.py below. Negative infinity represents a disallowed key; its exponential contribution is zero. An entirely masked row is rejected because there is no visible value to mix. NaN and positive infinity are invalid scores. The public attention function validates bounded finite matrices before constructing scores; softmax is a small internal helper with a narrower input contract.

def softmax(scores):
    if not scores or any(math.isnan(x) or x == math.inf for x in scores):
        raise ValueError("Invalid scores")
    peak = max(scores)
    if peak == -math.inf:
        raise ValueError("No visible key in this row")
    values = [math.exp(x - peak) for x in scores]
    total = math.fsum(values)
    return [x / total for x in values]

def attention(query, key, value, *, causal=False, allow=None):
    q, k, v = (matrix(x, name) for x, name in
               ((query, "query"), (key, "key"), (value, "value")))
    nq, nk, dk, dv = len(q), len(k), len(q[0]), len(v[0])
    if len(k[0]) != dk or len(v) != nk:
        raise ValueError("Q/K widths and K/V row counts must agree")
    if type(causal) is not bool or (causal and nq != nk):
        raise ValueError("Causal mode requires equal sequence lengths")
    if allow is None:
        allow = [[True] * nk for _ in range(nq)]
    if (type(allow) not in (list, tuple) or len(allow) != nq
            or any(type(row) not in (list, tuple) or len(row) != nk
                   or any(type(x) is not bool for x in row)
                   for row in allow)):
        raise ValueError("Expected a boolean allow matrix")
    weights, output = [], []
    for i, row in enumerate(q):
        scores = [math.fsum(a*b for a, b in zip(row, col)) / math.sqrt(dk)
                  if allow[i][j] and (not causal or j <= i) else -math.inf
                  for j, col in enumerate(k)]
        probs = softmax(scores)
        weights.append(probs)
        output.append([math.fsum(probs[j] * v[j][d] for j in range(nk))
                       for d in range(dv)])
    return output, weights

The division uses the query/key width dk, not the number of tokens, the value width or the number of heads. Changing the sequence length therefore does not change this scaling factor when dk stays fixed. For our first row, dropping the scale would produce about 0.731059 instead of 0.669762 for the first key. That is a different distribution even though it still sums to one. A row-sum test alone would miss the defect.

Each output feature is a weighted average of that feature across visible value rows. It must lie between their minimum and maximum, up to floating-point tolerance. This is a useful local invariant for the attention output. It need not hold after an output projection, residual addition or feed-forward transformation, where other arithmetic changes the values. A correct invariant must name the stage to which it applies.

Read a sample · Chapter 03 of 06

Mask the scores before normalisation and distinguish causal order from padding.

Hide the future before mixing values

The query rows are [1, 0] and [0, 1], in that order. Keys and values are paired by position. These are original fixed teaching numbers, not learned word embeddings or measured model performance.

Exact input keys and paired value vectors

Position	Key	Value
0	[1, 0]	[10, 0]
1	[0, 1]	[0, 20]

For query zero, dot products are [1, 0]. Divide by sqrt(2), the square root of the query/key feature width, giving scores [0.707107, 0] rounded to six decimals. Noncausal softmax gives [0.669762, 0.330238]. In causal mode, only key zero is visible to this query; key one's score becomes negative infinity before softmax, giving weights [1, 0]. Zeroing the score would not hide a key because exp(0)=1.

Every attention weight: six-decimal display, unrounded calculation

Mode / query	Key 0 weight	Key 1 weight
Noncausal, query 0	0.669762	0.330238
Noncausal, query 1	0.330238	0.669762
Causal, query 0	1.000000	0.000000
Causal, query 1	0.330238	0.669762

Every weighted output feature, before any projection or residual path

Mode / query	Feature 0	Feature 1
Noncausal, query 0	6.697615	6.604769
Noncausal, query 1	3.302385	13.395231
Causal, query 0	10.000000	0.000000
Causal, query 1	3.302385	13.395231

Each output is w0 * [10, 0] + w1 * [0, 20]. Use unrounded weights when calculating; rounding the displayed weights first changes the last digits. Query one's results are unchanged because both positions remain visible. Our causal rule is j <= i: the current position is included.

The explicit allow mask is a separate visibility condition, with True meaning visible. If its intersection with the causal rule removes every key in a row, the lab raises ValueError; it does not fabricate a probability distribution. Attention weights describe this numeric mixture, not the truth of an answer. The diagram covers the two-position attention primitive; the course's later three-position, two-head block is a separate example.

In our causal mode, query position i may see key positions j less than or equal to i. The diagonal is visible. This convention fits a next-token training arrangement where an input position predicts the following token: seeing the current input token does not expose the next target token. An incorrectly aligned target can still leak the answer even with a correct triangular mask. The mask and the input/target shift must be designed together.

For the two-position fixture, the first query sees only value [10, 0], so its weights become [1, 0] and its output is exactly [10, 0]. The second query sees both values, so its result matches the second row of the noncausal example. Setting a prohibited score to zero would be wrong: exp(0) is one, so the supposedly hidden position could still contribute. Mask before softmax, then normalise over the positions that remain.

allow is an explicit boolean matrix with True meaning visible. It can exclude a padding key or an application-defined position. Effective visibility combines allow with the causal condition. If their intersection removes every key in a query row, the function raises ValueError. Padding keys and padding queries are different issues: our function masks keys but does not silently discard query outputs or training-loss positions. Those policies belong in the surrounding pipeline.

Save demo.py below. It runs both visibility modes, then uses the fixed block introduced in the next chapter. Create block.py before running python demo.py. The output below is captured from those exact files. It records rounded values for inspection, while the tests compare unrounded calculations with numerical tolerance. Rows and columns always retain their original positional order.

from attention import attention
from block import X, transformer_block

Q = [[1, 0], [0, 1]]
K = [[1, 0], [0, 1]]
V = [[10, 0], [0, 20]]


def show(label, rows):
    print(label)
    for row in rows:
        print(" ".join(f"{x: .6f}" for x in row))


if __name__ == "__main__":
    for causal in (False, True):
        output, weights = attention(Q, K, V, causal=causal)
        show(f"causal={causal} weights", weights)
        show("output", output)
    output, maps = transformer_block(X)
    show("fixed block output", output)
    for index, weights in enumerate(maps, 1):
        show(f"head {index}", weights)

causal=False weights
 0.669762  0.330238
 0.330238  0.669762
output
 6.697615  6.604769
 3.302385  13.395231
causal=True weights
 1.000000  0.000000
 0.330238  0.669762
output
 10.000000  0.000000
 3.302385  13.395231
fixed block output
 0.999998 -0.999998  0.999998 -0.999998
-0.999998  0.999998 -0.999998  0.999998
 1.347147  0.577349 -0.962248 -0.962248
head 1
 1.000000  0.000000  0.000000
 0.330238  0.669762  0.000000
 0.248255  0.248255  0.503490
head 2
 1.000000  0.000000  0.000000
 0.330238  0.669762  0.000000
 0.333333  0.333333  0.333333

Framework mask conventions must be checked at the actual API boundary. The reviewed PyTorch scaled_dot_product_attention documentation uses True for participating elements in its boolean attention mask. MultiheadAttention uses True for excluded positions in its boolean masks. Those meanings are opposite. The functional API also applies the dropout probability supplied to it, so evaluation code must supply zero when dropout should be disabled. No PyTorch call is executed in this workbook.

A future-change experiment gives stronger evidence than a triangular-looking display. Change only a later input, run the causal block again, and compare the earlier outputs. They should remain unchanged for this deterministic fixture. Run the same idea without the causal constraint and the later input may affect earlier outputs. This test checks the direction of information flow; it does not prove that a trained model follows every application instruction.

Read a sample · Chapter 04 of 06

Keep the separate heads visible, then trace the two residual paths.

Save block.py using the four consecutive parts below. X has three positions and four features. Treat these rows as representations already supplied to the block, with any positional information handled upstream. They are not word embeddings learned from a corpus. The first fixed projection selects features zero and one; the second selects features two and three. For transparency, each head reuses its selected vectors as Q, K and V. A trainable implementation would normally provide distinct learned projections.

"""A fixed two-head, post-normalisation teaching block, not a trained model."""
import math
from attention import attention, matmul, matrix

X = [[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 0, 0]]
P1 = [[1, 0], [0, 1], [0, 0], [0, 0]]
P2 = [[0, 0], [0, 0], [1, 0], [0, 1]]
WO = [[1, 0, 0, 0], [0, 1, 0, 0],
      [0, 0, 1, 0], [0, 0, 0, 1]]
W1 = [[1, -1, 0, 0], [0, 1, -1, 0],
      [0, 0, 1, -1], [-1, 0, 0, 1]]
W2 = [[.5, 0, 0, 0], [0, .5, 0, 0],
      [0, 0, .5, 0], [0, 0, 0, .5]]

Each head returns three rows with two features. Concatenating corresponding rows gives three rows with four features. WO is an explicit identity output projection, so it preserves this fixture’s concatenated values; retaining the multiplication shows where a learned output projection would act. Changing WO is a separate experiment from changing attention. It should not be confused with adding another attention head.

The block adds the attention output to its input and normalises each position across its four features. It then applies W1, a ReLU activation and W2 independently to each row, adds that feed-forward result to the hidden representation, and normalises again. This is a post-normalisation teaching layout. The original transformer describes residual connections and layer normalisation; architectures can choose different normalisation placement. The fixed implementation here specifies its own exact sequence.

def normalise(rows):
    result = []
    for row in rows:
        mean = math.fsum(row) / len(row)
        variance = math.fsum((x-mean)**2 for x in row) / len(row)
        result.append([(x-mean) / math.sqrt(variance + 1e-5) for x in row])
    return result

def add(a, b):
    return [[x+y for x, y in zip(ra, rb)] for ra, rb in zip(a, b)]

def transformer_block(rows, *, causal=True):
    x = matrix(rows, "block input")
    if len(x[0]) != 4:
        raise ValueError("This fixed block expects four features")
    heads, maps = [], []
    for projection in (P1, P2):
        projected = matmul(x, projection)
        head, weights = attention(projected, projected, projected,
                                  causal=causal)
        heads.append(head)
        maps.append(weights)
    joined = [a+b for a, b in zip(*heads)]
    mixed = matmul(joined, WO)
    hidden = normalise(add(x, mixed))
    expanded = matmul(hidden, W1)
    activated = [[max(0.0, x) for x in row] for row in expanded]
    feed_forward = matmul(activated, W2)
    return normalise(add(hidden, feed_forward)), maps

normalise uses the population variance across a row and adds 1e-5 inside the square root. Its affine gain is fixed to one and bias to zero, so it is a simplified layer normalisation with no trainable affine parameters. A constant row maps to zeros rather than dividing by zero. Because epsilon is positive, a normalised row’s variance is generally slightly below one. Checking exact unit variance would be an inappropriate test for this implementation.

The feed-forward width remains four in this deliberately compact example. There are no bias vectors, dropout, batches, learned positions or gradients. ReLU retains a positive number and replaces a negative one with zero. W2 then halves each activated feature. Together with the residual path, these operations change the representation without mixing positions. The attention operation is the part of this block that moves information between positions.

The two head maps differ on the third row. Head one has weights about [0.248255, 0.248255, 0.503490]; head two has [1/3, 1/3, 1/3], because its last projected query is [0, 0]. A uniform row is a predictable result here, not a malfunction. The final third output is approximately [1.347147, 0.577349, -0.962248, -0.962248]. These results are from fixed numbers, not evidence of language understanding.

Read a sample · Chapter 05 of 06

Use known answers and information-flow checks as well as numerical invariants.

Save test_attention.py and run python -m unittest -v test_attention from the folder containing all four files. The twelve methods cover the hand calculation, causal direction, mask renormalisation, empty visibility, stable exponentials, matrix shapes, invalid inputs, a joint key/value permutation, constant values, multiplication, block shape and future independence. The runner should finish with OK. Keep the failing assertion when a test disagrees; do not round the program until a test happens to pass.

import math
import unittest
from attention import attention, matmul, softmax
from block import X, transformer_block
from demo import Q, K, V


class AttentionTests(unittest.TestCase):
    def test_hand_calculated_weights_and_values(self):
        out, weights = attention(Q, K, V)
        a = math.exp(1/math.sqrt(2)) / (math.exp(1/math.sqrt(2)) + 1)
        self.assertAlmostEqual(weights[0][0], a)
        self.assertAlmostEqual(out[0][0], 10*a)
        self.assertAlmostEqual(out[0][1], 20*(1-a))

    def test_causal_direction_and_diagonal(self):
        out, weights = attention(Q, K, V, causal=True)
        self.assertEqual(weights[0], [1, 0])
        self.assertEqual(out[0], [10, 0])
        self.assertGreater(weights[1][0], 0)
        self.assertGreater(weights[1][1], 0)

    def test_mask_renormalises_visible_keys(self):
        out, weights = attention(Q, K, V, allow=[[False, True]]*2)
        self.assertEqual(weights, [[0, 1], [0, 1]])
        self.assertEqual(out, [[0, 20], [0, 20]])

    def test_no_visible_key_is_an_error(self):
        for mask in ([[False, False]]*2, [[False, True], [True, True]]):
            with self.assertRaises(ValueError):
                attention(Q, K, V, causal=True, allow=mask)

    def test_softmax_shift_and_large_scores(self):
        self.assertEqual(softmax([10000, 10000]), [.5, .5])
        for a, b in zip(softmax([1, 2, 3]), softmax([1001, 1002, 1003])):
            self.assertAlmostEqual(a, b)

    def test_shape_contracts(self):
        with self.assertRaises(ValueError):
            attention(Q, [[1]], V)
        with self.assertRaises(ValueError):
            attention(Q, K, [[1]])
        out, _ = attention([[0, 0]], K, [[1, 2, 3], [3, 4, 5]])
        self.assertEqual(out, [[2, 3, 4]])
        with self.assertRaises(ValueError):
            attention([[0, 0]], K, V, causal=True)

    def test_invalid_numbers_and_masks(self):
        for bad in ([], [[True, 0]], [[math.nan, 0]], [[math.inf, 0]],
                    [[1e7, 0]], [[1, 2], [3]]):
            with self.assertRaises(ValueError):
                attention(bad, K, V)
        for mask in ([[1, 1]]*2, [[True]], "all"):
            with self.assertRaises(ValueError):
                attention(Q, K, V, allow=mask)

    def test_joint_key_value_permutation(self):
        out, _ = attention(Q, K, V)
        permuted, _ = attention(Q, K[::-1], V[::-1])
        for row, other in zip(out, permuted):
            for a, b in zip(row, other):
                self.assertAlmostEqual(a, b)

    def test_constant_values_remain_constant(self):
        out, _ = attention(Q, K, [[3, -2], [3, -2]], causal=True)
        for row in out:
            for a, b in zip(row, [3, -2]):
                self.assertAlmostEqual(a, b)

    def test_matrix_multiplication(self):
        self.assertEqual(matmul([[1, 2]], [[3, 4], [5, 6]]), [[13, 16]])
        with self.assertRaises(ValueError):
            matmul([[1, 2]], [[3, 4]])

    def test_block_shape_and_normalisation(self):
        out, maps = transformer_block(X)
        self.assertEqual((len(out), len(out[0]), len(maps)), (3, 4, 2))
        self.assertAlmostEqual(out[0][0], 1, places=5)
        self.assertAlmostEqual(out[0][1], -1, places=5)
        for row in out:
            self.assertAlmostEqual(sum(row), 0)
            self.assertTrue(all(math.isfinite(x) for x in row))

    def test_future_change_cannot_change_block_prefix(self):
        changed = [row[:] for row in X]
        changed[2] = [10, -9, 8, -7]
        original, _ = transformer_block(X)
        altered, _ = transformer_block(changed)
        self.assertEqual(original[:2], altered[:2])
        self.assertNotEqual(original[2], altered[2])


if __name__ == "__main__":
    unittest.main()

Reordering key rows and their corresponding value rows together should preserve a noncausal output, apart from small floating-point effects, because the same pairings contribute to the same query. Reordering keys alone changes those pairings and is not a valid invariance. A causal mask refers to positional order, so the same permutation argument cannot be transferred unchanged to causal attention. State the assumptions behind each test.

The fixed-value test checks another useful property. If every visible value row is [3, -2], their weighted mixture is [3, -2] regardless of the query scores. The future-independence test exercises the complete block, including both residual paths and normalisation. Per-row normalisation preserves that information-flow boundary; normalising across sequence positions could break it. A passing attention primitive alone would not detect every error in a larger block.

The release verification also compared 120 small cases against a separate 50-digit Decimal calculation, with 3,722 additional assertions. It varied lengths one to five, query/key widths one to four, value widths one to three, visibility patterns and causal mode. The reference used unshifted high-precision exponentials rather than the float implementation’s max-subtraction path. This checks forward arithmetic over a bounded fixture family; it is not a test of gradients, GPU kernels or model quality.

Three fault experiments ran in separate copies: remove the square-root scale, reverse the causal inequality, and ignore the allow mask. Each produced actual assertion failures. An import error or syntax error would not establish that the numerical rule was tested. Restore the baseline after an experiment and rerun the complete suite. The delivered source contains the passing implementation, with the failed mutation logs retained only as review evidence.

Read a sample · Chapter 06 of 06

List the missing interfaces before adding data, gradients or a larger machine.

This block returns one four-feature representation per input position. It does not return token probabilities. To make a small language model, a surrounding system needs token IDs, an embedding lookup, a positional scheme, trainable parameters, an output projection to vocabulary logits, an appropriate next-token loss and a gradient-based update. A generation loop must feed the produced token back into the next step. Those components are planned follow-on work, not hidden capabilities of these four files.

Write the input/target alignment explicitly. For a toy sequence [a, b, c, d], one training pair could use inputs [a, b, c] and targets [b, c, d]. Position zero can see a while predicting b; it must not receive b through an input feature, leaked target or incorrect mask. Special start/end tokens, padding and loss masking need their own definitions. An attention mask does not automatically remove padding from the loss.

A visibility mask provides an order-dependent boundary, but it is not a substitute for a deliberate representation of position. Decide how position enters the model and keep that decision consistent between training and generation. This course’s X is an explicit representation fixture; there is no positional encoding algorithm to infer from its numbers. A future trainable implementation should test repeated tokens at different positions and document the selected scheme.

Our function constructs an L by S weight table. With equal sequence lengths n, this table has n*n entries per head: four times as many when n doubles. Two heads and a batch dimension multiply the count further. This arithmetic describes the explicit table in the teaching code, not every optimised implementation’s peak memory. Python object overhead is also very different from a packed tensor. Do not turn this small script into a GPU purchasing estimate.

Attention weights reveal how this forward operation mixes value vectors. They are not a probability that a source statement is true and do not establish a complete explanation of a model’s answer. In our fixture the values are numbers with no textual claim attached at all. A language model adds many transformations before and after attention. Keep numerical inspection, predictive evaluation and factual evaluation as separate questions.

For a framework port, record tensor dimension order, parameter initialisation, mask polarity, dropout behaviour, normalisation axes, epsilon and dtype. Compare a tiny fixed forward case first, then add gradient and optimisation tests. Use a held-out data split for model evaluation and record data provenance. This release ran no framework port, gradient check, training benchmark or text generation. Its completed result is a reproducible attention calculation and a precise specification for the next step.