Build causal self-attention and inspect intermediate outputs

Problem

When you need this. The story’s final color depends on its first sentence, but a bigram reads only the previous token. Add a layer that can combine earlier information without seeing the future token it must predict.

The bigram baseline in the previous build predicted the story’s final color correctly in 0 of 60 test stories, because it reads one token back. The color is in the first sentence. A model needs a layer that lets each position collect information from every earlier position. That layer is causal self-attention, and this build implements it.

This build writes a whole attention layer as a program: the forward pass for a full sequence, the causal mask, the backward pass that training will need, and tests that check the mask.

The contract: attention_forward(params, x, mask) takes the projection weights and an input of shape (T, d_model), one row per position, and returns the output of the same shape together with every intermediate tensor. attention_backward(params, x, cache, d_out) returns the gradient for each weight matrix and for the input. The recorded run uses the sentence mia has a red ball., tokenized by the series tokenizer, with untrained weights.

Success has three checks. Every row of attention weights adds up to 1 and puts exactly 0 on later positions. Every hand-written gradient matches a finite-difference estimate. Changing a token changes the outputs at that position and later ones, never at earlier ones.

Concepts

A bigram can only look one token back. Self-attention lets each position combine information from earlier positions in the same sequence. It creates three vectors from each input: a query used to score matches, a key that other queries match against, and a value containing the information to combine. Query–key dot products become softmax weights; those weights determine the mixture of values.

A causal mask blocks later positions before softmax. Without it, training could use the very next token it is supposed to predict. The following choices implement that rule and make the gradients inspectable:

Projection: A weight matrix that turns each input row into a query, key or value. This is the original transformer’s multi-head layer with one head. Here W_q, W_k and W_v have shape (4, 2), and an output projection W_o of shape (2, 4) maps the mixed values back to the input width.
Head size: The width of each query, key and value, here 2. Scores are divided by its square root, as in the original transformer.
Masking with −∞: Before softmax, the score of every later position is replaced with −∞. Its exponential is 0, so softmax gives it weight 0 and shares all the weight among the allowed positions.
Softmax backward: Each softmax weight depends on every score in its row. For one row with weights a and incoming gradient da, the gradient of the scores is a × (da − Σ a·da). Masked positions have a = 0, so no gradient flows to them.

Architecture

attention.py · one causal self-attention layer on “mia has a red ball.”

The layer runs top to bottom; the matrix is its recorded weights. A row uses its own column and every earlier column, the cells on and below the diagonal; the cells above it are later tokens, removed by the mask before softmax. Values come from attention-trace.json.

The figure runs the layer on the 6 tokens of the sentence and shows the weights it recorded. Each row is the token whose output is being computed, and each column is a token it may attend to. A row uses its own column and the columns of every earlier token, the allowed cells on and below the diagonal. The masked cells above the diagonal are later tokens. The Causal mask step removes them before Softmax each row.

attention.py holds the layer: attention_forward(), attention_backward(), and small list helpers matmul() and transpose(). It has no batch dimension; a batch is a loop over sequences. trace_attention.py builds the input from an embedding table drawn from a fixed seed, runs the layer, compares three mask modes, checks the gradients and writes attention-trace.json. The mask mode "none" turns the mask off, and "zero" reproduces a common bug; both exist only to be tested.

Tensor	Shape	Computed as
`x`	(6, 4)	One embedding row per token
`q`, `k`, `v`	(6, 2)	`x` times W_q, W_k, W_v
`scores`	(6, 6)	`q` times `k` transposed, divided by √2, later positions set to −∞
`weights`	(6, 6)	Softmax of each row of `scores`
`mixed`	(6, 2)	`weights` times `v`
`out`	(6, 4)	`mixed` times W_o

Working code

Runtime: Python 3.10 or newer. Dependencies: Python standard library only.

Download and extract the complete lab using Download the complete lab above. Open build-causal-self-attention-and-inspect-intermediate-outputs/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.

attention.pySource · pythonDownload

"""Single-head causal self-attention with a hand-written backward pass.

Input x has shape (T, d_model): one row per position. The layer computes
  Q = x Wq, K = x Wk, V = x Wv                  each (T, d_head)
  S = Q K^T / sqrt(d_head)                       (T, T) scores
  S[i][j] = -inf for j > i                        the causal mask
  A = softmax of each row of S                    (T, T) attention weights
  out = A V Wo                                    (T, d_model)
Python 3.10+, standard library only.
"""
import math
import random

MASK_MODES = ("causal", "none", "zero")


def shape(t):
    dims = []
    while isinstance(t, list):
        dims.append(len(t))
        t = t[0] if t else None
    return tuple(dims)


def matmul(a, b):
    (n, k), (k2, m) = shape(a), shape(b)
    if k != k2:
        raise ValueError(f"cannot multiply {shape(a)} by {shape(b)}: inner sizes {k} and {k2} differ")
    bt = transpose(b)
    return [[math.fsum(x * y for x, y in zip(row, col)) for col in bt] for row in a]


def transpose(a):
    return [list(r) for r in zip(*a)]


def init_attention(d_model, d_head, seed):
    rng = random.Random(seed)
    w = lambda r, c: [[round(rng.uniform(-0.5, 0.5), 2) for _ in range(c)] for _ in range(r)]
    return {"Wq": w(d_model, d_head), "Wk": w(d_model, d_head),
            "Wv": w(d_model, d_head), "Wo": w(d_head, d_model)}


def softmax_row(row):
    finite = [v for v in row if v != -math.inf]
    top = max(finite)
    exps = [0.0 if v == -math.inf else math.exp(v - top) for v in row]
    total = math.fsum(exps)
    return [e / total for e in exps]


def attention_forward(params, x, mask="causal"):
    """Return the output and every intermediate tensor."""
    if mask not in MASK_MODES:
        raise ValueError(f"mask must be one of {MASK_MODES}, got {mask!r}")
    if not x:
        raise ValueError("empty sequence: attention needs at least one position")
    d_model = len(params["Wq"])
    if any(len(row) != d_model for row in x):
        raise ValueError(f"each input row must have {d_model} numbers, got {sorted({len(r) for r in x})}")
    T, d_head = len(x), len(params["Wq"][0])
    q, k, v = matmul(x, params["Wq"]), matmul(x, params["Wk"]), matmul(x, params["Wv"])
    scale = 1 / math.sqrt(d_head)
    raw = [[s * scale for s in row] for row in matmul(q, transpose(k))]
    if mask == "causal":
        scores = [[raw[i][j] if j <= i else -math.inf for j in range(T)] for i in range(T)]
    elif mask == "zero":  # a common bug: zeroing future scores instead of excluding them
        scores = [[raw[i][j] if j <= i else 0.0 for j in range(T)] for i in range(T)]
    else:
        scores = raw
    weights = [softmax_row(row) for row in scores]
    mixed = matmul(weights, v)
    out = matmul(mixed, params["Wo"])
    return {"q": q, "k": k, "v": v, "scores": scores, "weights": weights, "mixed": mixed, "out": out}


def attention_backward(params, x, cache, d_out):
    """Gradients of a scalar loss, given d_out = dLoss/d(out), shape (T, d_model)."""
    scale = 1 / math.sqrt(len(params["Wq"][0]))
    d_wo = matmul(transpose(cache["mixed"]), d_out)
    d_mixed = matmul(d_out, transpose(params["Wo"]))
    d_weights = matmul(d_mixed, transpose(cache["v"]))
    d_v = matmul(transpose(cache["weights"]), d_mixed)
    # Softmax backward, row by row: dS = A * (dA - sum(dA * A)). Masked entries have A = 0, so dS = 0.
    d_scores = []
    for a_row, da_row in zip(cache["weights"], d_weights):
        dot = math.fsum(a * da for a, da in zip(a_row, da_row))
        d_scores.append([a * (da - dot) * scale for a, da in zip(a_row, da_row)])
    d_q = matmul(d_scores, cache["k"])
    d_k = matmul(transpose(d_scores), cache["q"])
    grads = {
        "Wq": matmul(transpose(x), d_q),
        "Wk": matmul(transpose(x), d_k),
        "Wv": matmul(transpose(x), d_v),
        "Wo": d_wo,
    }
    d_x = [[math.fsum(t) for t in zip(*rows)] for rows in zip(
        matmul(d_q, transpose(params["Wq"])), matmul(d_k, transpose(params["Wk"])),
        matmul(d_v, transpose(params["Wv"])))]
    return grads, d_x

softmax_row() skips −∞ entries when it finds the largest score and gives them an exponential of exactly 0. attention_backward() walks the forward pass in reverse: output projection, value mixing, softmax, scaling, then the three projections. The input gradient adds the three paths through W_q, W_k and W_v, because the same input row feeds all three.

Read trace_attention.py

trace_attention.pySource · pythonDownload

"""Run causal self-attention on one tokenized sentence and record every intermediate value.

Run: python3 trace_attention.py          (print the trace)
     python3 trace_attention.py --write  (rewrite attention-trace.json)
     python3 trace_attention.py --check  (exit 1 if attention-trace.json is out of date)
"""
import json
import math
import random
import sys
from pathlib import Path

import attention as at
from bpe import Tokenizer

HERE = Path(__file__).parent
D_MODEL, D_HEAD, SEED = 4, 2, 11
SENTENCE, CHANGED = "mia has a red ball.", "mia has a red hat."
EPS = 1e-6


def embeddings(vocab_size, seed=SEED):
    rng = random.Random(seed + 1)
    return [[round(rng.uniform(-1, 1), 2) for _ in range(D_MODEL)] for _ in range(vocab_size)]


def readout(T, seed=SEED):
    rng = random.Random(seed + 2)
    return [[round(rng.uniform(-1, 1), 2) for _ in range(D_MODEL)] for _ in range(T)]


def loss(params, x, r, mask="causal"):
    out = at.attention_forward(params, x, mask)["out"]
    return math.fsum(o * w for orow, rrow in zip(out, r) for o, w in zip(orow, rrow))


def numeric_grad(params, x, r):
    grads = {}
    for name in params:
        g = []
        for i, row in enumerate(params[name]):
            g.append([])
            for j in range(len(row)):
                old = row[j]
                row[j] = old + EPS
                up = loss(params, x, r)
                row[j] = old - EPS
                down = loss(params, x, r)
                row[j] = old
                g[-1].append((up - down) / (2 * EPS))
        grads[name] = g
    return grads


def max_err(a, b):
    return max(abs(p - q) for name in a for ra, rb in zip(a[name], b[name]) for p, q in zip(ra, rb))


def r4(t):
    if isinstance(t, dict):
        return {k: r4(v) for k, v in t.items()}
    if isinstance(t, list):
        return [r4(v) for v in t]
    return None if t == -math.inf else round(t, 4)


def changed_rows(a, b):
    return [i for i, (ra, rb) in enumerate(zip(a, b)) if ra != rb]  # exact: any changed bit counts


def trace():
    tok = Tokenizer.load(HERE / "tokenizer.json")
    table = embeddings(len(tok.vocab))
    params = at.init_attention(D_MODEL, D_HEAD, SEED)
    ids, ids2 = tok.encode(SENTENCE), tok.encode(CHANGED)
    x, x2 = [table[i] for i in ids], [table[i] for i in ids2]
    cache = at.attention_forward(params, x)
    r = readout(len(x))
    grads, _ = at.attention_backward(params, x, cache, r)
    leak = {}
    for mode in at.MASK_MODES:
        a = at.attention_forward(params, x, mode)
        b = at.attention_forward(params, x2, mode)
        leak[mode] = {
            "rows_changed_by_the_edit": changed_rows(a["out"], b["out"]),
            "weight_on_future_row_0": round(math.fsum(a["weights"][0][1:]), 4),
        }
    return {
        "sentence": SENTENCE,
        "tokens": tok.tokens(SENTENCE),
        "ids": ids,
        "changed_sentence": CHANGED,
        "changed_tokens": tok.tokens(CHANGED),
        "d_model": D_MODEL,
        "d_head": D_HEAD,
        "seed": SEED,
        "shapes": {name: list(at.shape(cache[name])) for name in ("q", "k", "v", "scores", "weights", "mixed", "out")},
        "x": r4(x),
        "params": r4(params),
        "q": r4(cache["q"]), "k": r4(cache["k"]), "v": r4(cache["v"]),
        "scores": r4(cache["scores"]),
        "weights": r4(cache["weights"]),
        "out": r4(cache["out"]),
        "row_sums": [round(math.fsum(row), 6) for row in cache["weights"]],
        "masks": leak,
        "gradient_check": {
            "epsilon": EPS,
            "max_error_rounded": round(max_err(grads, numeric_grad(params, x, r)), 7),
        },
    }


if __name__ == "__main__":
    path = HERE / "attention-trace.json"
    text = json.dumps(trace(), indent=2, ensure_ascii=False) + "\n"
    if "--write" in sys.argv:
        path.write_text(text)
    elif "--check" in sys.argv:
        sys.exit(0 if path.read_text() == text else "attention-trace.json is out of date")
    else:
        print(text, end="")

Run the trace:

Terminal commandsTerminal · bash

python3 trace_attention.py

python3 trace_attention.py --check exits with an error if your output differs from the recorded attention-trace.json.

Inspect the scores and weights

The tokenizer splits the sentence into 6 tokens: "mia", " has", " a", " red", " ball", ".". The recorded scores for the last token, ".", are [-0.0431, 0.1031, -0.0155, 0.0036, -0.0151, 0.0892], and for the first token, [-0.0306, −∞, −∞, −∞, −∞, −∞]: the first token may see only itself.

Token	`"mia"`	`" has"`	`" a"`	`" red"`	`" ball"`	`"."`
`"mia"`	1.0000	0.0000	0.0000	0.0000	0.0000	0.0000
`" has"`	0.5164	0.4836	0.0000	0.0000	0.0000	0.0000
`" a"`	0.3515	0.3064	0.3421	0.0000	0.0000	0.0000
`" red"`	0.2348	0.2697	0.2417	0.2539	0.0000	0.0000
`" ball"`	0.1852	0.2196	0.1919	0.2040	0.1993	0.0000
`"."`	0.1562	0.1808	0.1606	0.1636	0.1606	0.1783

Every row adds up to 1, and every entry right of the diagonal is exactly 0. The weights inside each row are close to equal, because the projection matrices are untrained and every score is close to 0. Attention patterns only become meaningful after training; do not read anything into these.

Check the gradients

The trace defines a loss as a fixed weighted sum of the outputs, so that the gradient arriving at the output is known. It compares every gradient for the four weight matrices with a central finite difference using ε = 0.000001. The largest difference, rounded to seven decimal places, is 0.0000000. The tests also check the gradient for every input value.

Change a later token

The trace runs the layer again on mia has a red hat., which differs only in its fifth token. With the causal mask, the outputs at token positions 5 and 6 change, and the 4 earlier rows are identical to the bit. Without a mask, all 6 rows change: every position read the changed token.

Check your understanding

A common bug masks later positions by setting their scores to 0 instead of −∞. The first token’s own score is -0.0306. Before you answer the check and open the comparison, estimate how much weight the first token puts on later tokens when they are scored 0.

Read the explanation

Masking must remove a position from softmax, which means a score of −∞ and an exponential of exactly 0. A score of 0 leaves the future visible.

Compare your estimate

The recorded weight on later tokens is 0.8375. A score of 0 is an ordinary score: its exponential is 1, almost the same as the first token’s own e^-0.0306. So the first token splits its weight about evenly over all 6 positions, and 5 of them are in the future. Only −∞ gives an exponential of 0. With the zero-score bug, changing the fifth token changes all 6 rows, the same leak as having no mask.

Inspect the recorded trace

attention-trace.jsonRecorded output · jsonDownload

{
  "sentence": "mia has a red ball.",
  "tokens": [
    "mia",
    " has",
    " a",
    " red",
    " ball",
    "."
  ],
  "ids": [
    58,
    36,
    33,
    89,
    85,
    4
  ],
  "changed_sentence": "mia has a red hat.",
  "changed_tokens": [
    "mia",
    " has",
    " a",
    " red",
    " hat",
    "."
  ],
  "d_model": 4,
  "d_head": 2,
  "seed": 11,
  "shapes": {
    "q": [
      6,
      2
    ],
    "k": [
      6,
      2
    ],
    "v": [
      6,
      2
    ],
    "scores": [
      6,
      6
    ],
    "weights": [
      6,
      6
    ],
    "mixed": [
      6,
      2
    ],
    "out": [
      6,
      4
    ]
  },
  "x": [
    [
      -0.75,
      0.28,
      0.69,
      0.22
    ],
    [
      0.47,
      0.13,
      -0.97,
      0.42
    ],
    [
      0.12,
      0.25,
      0.84,
      0.92
    ],
    [
      0.9,
      0.49,
      -0.5,
      -0.16
    ],
    [
      0.65,
      0.51,
      -0.32,
      -0.24
    ],
    [
      -0.88,
      0.64,
      -0.85,
      0.37
    ]
  ],
  "params": {
    "Wq": [
      [
        -0.05,
        0.06
      ],
      [
        0.42,
        -0.03
      ],
      [
        0.01,
        0.09
      ],
      [
        -0.32,
        0.01
      ]
    ],
    "Wk": [
      [
        0.13,
        0.29
      ],
      [
        -0.41,
        -0.2
      ],
      [
        -0.41,
        0.31
      ],
      [
        0.19,
        -0.46
      ]
    ],
    "Wv": [
      [
        0.48,
        0.46
      ],
      [
        0.15,
        0.12
      ],
      [
        -0.34,
        -0.48
      ],
      [
        0.03,
        -0.44
      ]
    ],
    "Wo": [
      [
        -0.31,
        -0.26,
        -0.47,
        -0.04
      ],
      [
        -0.06,
        0.34,
        0.02,
        0.14
      ]
    ]
  },
  "q": [
    [
      0.0916,
      0.0109
    ],
    [
      -0.113,
      -0.0588
    ],
    [
      -0.187,
      0.0845
    ],
    [
      0.207,
      -0.0073
    ],
    [
      0.2553,
      -0.0075
    ],
    [
      0.1859,
      -0.1448
    ]
  ],
  "k": [
    [
      -0.4534,
      -0.1608
    ],
    [
      0.4853,
      -0.3836
    ],
    [
      -0.2565,
      -0.178
    ],
    [
      0.0907,
      0.0816
    ],
    [
      -0.039,
      0.0977
    ],
    [
      0.042,
      -0.8169
    ]
  ],
  "v": [
    [
      -0.546,
      -0.7394
    ],
    [
      0.5875,
      0.5126
    ],
    [
      -0.1629,
      -0.7228
    ],
    [
      0.6707,
      0.7832
    ],
    [
      0.4901,
      0.6194
    ],
    [
      -0.0263,
      -0.0828
    ]
  ],
  "scores": [
    [
      -0.0306,
      null,
      null,
      null,
      null,
      null
    ],
    [
      0.0429,
      -0.0228,
      null,
      null,
      null,
      null
    ],
    [
      0.0503,
      -0.0871,
      0.0233,
      null,
      null,
      null
    ],
    [
      -0.0655,
      0.073,
      -0.0366,
      0.0129,
      null,
      null
    ],
    [
      -0.081,
      0.0896,
      -0.0454,
      0.0159,
      -0.0076,
      null
    ],
    [
      -0.0431,
      0.1031,
      -0.0155,
      0.0036,
      -0.0151,
      0.0892
    ]
  ],
  "weights": [
    [
      1.0,
      0.0,
      0.0,
      0.0,
      0.0,
      0.0
    ],
    [
      0.5164,
      0.4836,
      0.0,
      0.0,
      0.0,
      0.0
    ],
    [
      0.3515,
      0.3064,
      0.3421,
      0.0,
      0.0,
      0.0
    ],
    [
      0.2348,
      0.2697,
      0.2417,
      0.2539,
      0.0,
      0.0
    ],
    [
      0.1852,
      0.2196,
      0.1919,
      0.204,
      0.1993,
      0.0
    ],
    [
      0.1562,
      0.1808,
      0.1606,
      0.1636,
      0.1606,
      0.1783
    ]
  ],
  "out": [
    [
      0.2136,
      -0.1094,
      0.2418,
      -0.0817
    ],
    [
      0.0074,
      -0.0461,
      -0.0037,
      -0.0188
    ],
    [
      0.042,
      -0.1015,
      0.0248,
      -0.0463
    ],
    [
      -0.0493,
      -0.0457,
      -0.076,
      -0.008
    ],
    [
      -0.0789,
      -0.0192,
      -0.1062,
      0.0076
    ],
    [
      -0.0598,
      -0.0213,
      -0.0824,
      0.0032
    ]
  ],
  "row_sums": [
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0
  ],
  "masks": {
    "causal": {
      "rows_changed_by_the_edit": [
        4,
        5
      ],
      "weight_on_future_row_0": 0.0
    },
    "none": {
      "rows_changed_by_the_edit": [
        0,
        1,
        2,
        3,
        4,
        5
      ],
      "weight_on_future_row_0": 0.8379
    },
    "zero": {
      "rows_changed_by_the_edit": [
        0,
        1,
        2,
        3,
        4,
        5
      ],
      "weight_on_future_row_0": 0.8375
    }
  },
  "gradient_check": {
    "epsilon": 1e-06,
    "max_error_rounded": 0.0
  }
}

Failure cases

Each case has a named test in test_attention.py.

Case	Result with this implementation	What it tells you
Attention without a mask	Changing the fifth token changes all 6 output rows	During training, every position could read the token it is asked to predict
Future scores set to 0 instead of −∞	The first token puts `0.8375` of its weight on later tokens	Zero is a score, not an exclusion
An unknown mask mode	`attention_forward()` raises `ValueError`	A typo in the mask setting must not silently turn masking off
An empty sequence	`attention_forward()` raises `ValueError`	Softmax over zero positions is undefined
Input rows of the wrong width	`attention_forward()` raises `ValueError` naming the expected width	Shape mistakes fail at the boundary, not deep inside a matrix product

The complete lab includes these tests and their fixtures. Run the 12 standard-library tests:

Terminal commandsTerminal · bash

python3 -m unittest -v test_attention.py

Read the tests

test_attention.pySource · pythonDownload

"""Tests for attention.py. Run: python3 -m unittest -v test_attention.py"""
import math
import subprocess
import sys
import unittest
from pathlib import Path

import attention as at
import trace_attention as tr
from bpe import Tokenizer

HERE = Path(__file__).parent
TOK = Tokenizer.load(HERE / "tokenizer.json")
TABLE = tr.embeddings(len(TOK.vocab))


def setup(sentence=tr.SENTENCE):
    params = at.init_attention(tr.D_MODEL, tr.D_HEAD, tr.SEED)
    return params, [list(TABLE[i]) for i in TOK.encode(sentence)]


class Recorded(unittest.TestCase):
    def test_recorded_trace(self):
        result = subprocess.run([sys.executable, "trace_attention.py", "--check"], cwd=HERE)
        self.assertEqual(result.returncode, 0)


class Mechanism(unittest.TestCase):
    def test_shapes(self):
        params, x = setup()
        c = at.attention_forward(params, x)
        T = len(x)
        self.assertEqual(at.shape(c["weights"]), (T, T))
        self.assertEqual(at.shape(c["out"]), (T, tr.D_MODEL))

    def test_rows_sum_to_one_and_future_weights_are_zero(self):
        params, x = setup()
        w = at.attention_forward(params, x)["weights"]
        for i, row in enumerate(w):
            self.assertAlmostEqual(math.fsum(row), 1.0, places=12)
            self.assertEqual(row[i + 1:], [0.0] * (len(row) - i - 1))

    def test_gradients_match_finite_differences(self):
        params, x = setup()
        r = tr.readout(len(x))
        grads, d_x = at.attention_backward(params, x, at.attention_forward(params, x), r)
        self.assertLess(tr.max_err(grads, tr.numeric_grad(params, x, r)), 1e-7)
        for i in range(len(x)):
            for j in range(tr.D_MODEL):
                old = x[i][j]
                x[i][j] = old + tr.EPS
                up = tr.loss(params, x, r)
                x[i][j] = old - tr.EPS
                down = tr.loss(params, x, r)
                x[i][j] = old
                self.assertAlmostEqual(d_x[i][j], (up - down) / (2 * tr.EPS), places=7)

    def test_each_position_attends_to_itself(self):
        params, x = setup()
        w = at.attention_forward(params, x)["weights"]
        self.assertEqual(w[0][0], 1.0)  # the first position can only see itself
        self.assertTrue(all(row[i] > 0 for i, row in enumerate(w)))

    def test_scores_are_scaled_by_the_square_root_of_the_head_size(self):
        params, x = setup()
        c = at.attention_forward(params, x)
        raw = at.matmul(c["q"], at.transpose(c["k"]))
        for i, row in enumerate(c["scores"]):
            for j in range(i + 1):
                self.assertAlmostEqual(row[j], raw[i][j] / math.sqrt(tr.D_HEAD), places=15)

    def test_future_token_does_not_change_earlier_rows(self):
        params, x = setup()
        _, x2 = setup(tr.CHANGED)
        changed = tr.changed_rows(at.attention_forward(params, x)["out"], at.attention_forward(params, x2)["out"])
        first = next(i for i, (a, b) in enumerate(zip(x, x2)) if a != b)
        self.assertEqual(changed, list(range(first, len(x))))


class FailureCases(unittest.TestCase):
    def test_no_mask_leaks_the_future(self):
        params, x = setup()
        _, x2 = setup(tr.CHANGED)
        changed = tr.changed_rows(at.attention_forward(params, x, "none")["out"],
                                  at.attention_forward(params, x2, "none")["out"])
        self.assertEqual(changed, list(range(len(x))))

    def test_zeroed_scores_still_leak(self):
        params, x = setup()
        w = at.attention_forward(params, x, "zero")["weights"]
        self.assertGreater(math.fsum(w[0][1:]), 0.5)

    def test_unknown_mask_mode(self):
        params, x = setup()
        with self.assertRaisesRegex(ValueError, "mask must be one of"):
            at.attention_forward(params, x, "upper")

    def test_empty_sequence(self):
        params, _ = setup()
        with self.assertRaisesRegex(ValueError, "empty sequence"):
            at.attention_forward(params, [])

    def test_wrong_input_width(self):
        params, x = setup()
        with self.assertRaisesRegex(ValueError, "must have 4 numbers"):
            at.attention_forward(params, [row[:3] for row in x])


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

test_gradients_match_finite_differences fails if the softmax backward drops its Σ a·da term, if the scale is left out of the backward pass, or if a projection receives another projection’s gradient. test_future_token_does_not_change_earlier_rows fails if the mask lets a position see even one token ahead.

Deployment

There is nothing to deploy yet. This layer becomes part of the transformer block in the next build, Build a transformer block with positional information, and Build a small language-model training pipeline calls it after that.

Scaling

This section is unmeasured. The scores and weights have one entry for every pair of positions, so their size and the work to compute them grow with the square of the sequence length. The loops here also recompute the full layer for every sequence; Build a KV cache for LLM generation revisits that cost.

Key takeaways

Causal self-attention projects each input row to a query, key and value, scores queries against keys, masks later positions with −∞, applies softmax and mixes the values.
The backward pass through softmax needs the whole row: a × (da − Σ a·da). A finite-difference check catches a missing term.
Test causality by changing a later token. Earlier outputs must not change at all.

Best practices

Mask with −∞ before softmax and confirm the masked weights are exactly 0, not small.
Keep a perturbation test for causality in the suite of every model that uses attention.
Check gradients for the input as well as the weights; the next layer down depends on them.

References

Sources checked October 20, 2026.

Vaswani et al., Attention Is All You Need (NIPS 2017), arXiv 1706.03762v7: scaled dot-product attention, softmax(QKᵀ/√d_k)V, in section 3.2.1; the learned query, key, value and output projections in section 3.2.2; decoder self-attention masked by setting illegal softmax inputs to −∞ in section 3.2.3.
PyTorch 2.14: torch.nn.functional.scaled_dot_product_attention: is_causal applies a lower-triangular mask that keeps the diagonal; the reference code adds −∞ to masked scores before softmax and scales by 1/√E by default.