Build a transformer block with positional information

Problem

When you need this. Attention can combine earlier information, but it is only one part of a language model. Add positions and the remaining transformations to produce next-token scores that the training loop can improve.

The attention layer in the previous build lets each position read earlier positions, but on its own it is not a language model. It has no way to turn token IDs into vectors, no notion of where a token sits, nothing that processes each position after mixing, and no output scores. This build adds those parts and produces the complete model that the next build, Build a small language-model training pipeline, trains.

The model has one transformer block. forward(params, ids) takes a list of token IDs and returns a score for every vocabulary token at every position. loss_and_dlogits() turns those scores into the mean next-token cross-entropy and its gradient, and backward() returns a gradient for every one of the model’s 7,440 parameters at the series size: d = 16, a feed-forward width of 64, and 64 positions.

Success has three checks. Every gradient of a smaller copy of the model matches a finite-difference estimate. Changing a later token never changes the scores at earlier positions. And a recorded test shows what the position embeddings are for: without them, two sentences with the same words in a different order get identical scores at the last position.

Concepts

Attention combines information across positions, but a language model also needs to distinguish word order and transform each position’s features. A transformer block combines attention with a feed-forward network, normalization and residual additions. This build makes four choices:

Learned position embeddings

A second table with one row of 16 numbers per position, 64 rows in all. Row t is added to the token embedding at position t. Training adjusts these rows like any other weights. The original transformer added fixed sine and cosine patterns instead, and reported that learned embeddings gave nearly identical results. The model cannot read a position beyond the last row.

Pre-norm

Layer normalization runs at the input of each sublayer, before attention and before the feed-forward network, and once more before the output scores. GPT-2 uses this order; the original transformer normalized after each addition instead.

Skip path

Each sublayer’s output is added to its input, a residual connection. It preserves a direct path for both the input and its gradient. In the backward pass, the gradient arriving at the sum flows to both branches unchanged, so the input of a sublayer receives its gradient twice: once through the sublayer and once directly.

Feed-forward network (MLP)

Two weight matrices with a ReLU between them, applied to each position separately. ReLU returns max(0, x), introducing a nonlinear transformation; without it, the two matrices would reduce to one linear transformation. Its hidden width, 64, is four times d, the same ratio as the original transformer’s 2,048 and 512.

Layer normalization subtracts a row’s mean and divides by √(variance + ε), then applies a learned gain and bias. This controls the scale of the values entering a sublayer. Here ε = 10⁻⁵ keeps nearly constant rows from causing division by zero; because of ε, the normalized variance is not exactly 1.

Architecture

Start at Token + position, which adds the two embedding rows for each token. The first Layer norm feeds Attention, and the first Add sums the attention output with the block’s input, which arrives on the skip path that runs from Token + position. The second Layer norm feeds the MLP, and the second Add sums its output with the first sum, which arrives on the skip path that runs from the first Add. Norm, then scores normalizes once more and multiplies by the output matrix to give one score per vocabulary token.

model.py holds everything in one file: list helpers, layer normalization forward and backward, the attention layer from the previous build widened to d, and forward() and backward() for the whole model. The model processes one sequence at a time; the training build loops over a batch. trace_block.py runs the model on the first story of the corpus and on two short sentences, checks gradients on a smaller copy and writes block-trace.json.

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-a-transformer-block-with-positional-information/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.

"""A one-block transformer language model in plain Python, with a hand-written backward pass.

For a sequence of token IDs (length T) the forward pass computes
  x0 = token embedding + position embedding            (T, d)
  x1 = x0 + Attention(LayerNorm1(x0))                  residual around causal self-attention
  x2 = x1 + MLP(LayerNorm2(x1))                         residual around a two-layer ReLU network
  logits = LayerNormF(x2) W_out + b_out                 (T, V)
and the loss is the mean cross-entropy of each next token. Normalization comes
before each sublayer (pre-LN). Every float sum uses math.fsum, so results are
identical on every Python version. Python 3.10+, standard library only.
"""
import math
import random
from operator import mul

EPS = 1e-5  # added to the variance inside layer normalization


# ---- list helpers ---------------------------------------------------------------

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


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


def matmul(a, b):
    """(n, k) @ (k, m) -> (n, m)."""
    if len(a[0]) != len(b):
        raise ValueError(f"cannot multiply {shape(a)} by {shape(b)}: inner sizes {len(a[0])} and {len(b)} differ")
    bt = transpose(b)
    return [[math.fsum(map(mul, row, col)) for col in bt] for row in a]


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


def add_bias(a, bias):
    return [[x + y for x, y in zip(row, bias)] for row in a]


def col_sums(a):
    return [math.fsum(col) for col in zip(*a)]


def zeros(r, c):
    return [[0.0] * c for _ in range(r)]


# ---- parameters -----------------------------------------------------------------

def init_params(vocab, d, hidden, max_len, seed):
    """Small random weights from a fixed seed; layer-norm gains start at 1 and biases at 0."""
    rng = random.Random(seed)
    def mat(r, c, scale):
        return [[rng.gauss(0.0, scale) for _ in range(c)] for _ in range(r)]
    s = 0.02
    return {
        "tok_emb": mat(vocab, d, s), "pos_emb": mat(max_len, d, s),
        "ln1_g": [1.0] * d, "ln1_b": [0.0] * d,
        "Wq": mat(d, d, s), "Wk": mat(d, d, s), "Wv": mat(d, d, s), "Wo": mat(d, d, s),
        "ln2_g": [1.0] * d, "ln2_b": [0.0] * d,
        "W1": mat(d, hidden, s), "b1": [0.0] * hidden, "W2": mat(hidden, d, s), "b2": [0.0] * d,
        "lnf_g": [1.0] * d, "lnf_b": [0.0] * d,
        "W_out": mat(d, vocab, s), "b_out": [0.0] * vocab,
    }


def count_params(params):
    return sum(len(v) * (len(v[0]) if isinstance(v[0], list) else 1) for v in params.values())


# ---- layer normalization ----------------------------------------------------------

def layernorm_forward(x, g, b):
    out, cache = [], []
    for row in x:
        n = len(row)
        mu = math.fsum(row) / n
        var = math.fsum((v - mu) ** 2 for v in row) / n
        inv = 1.0 / math.sqrt(var + EPS)
        xhat = [(v - mu) * inv for v in row]
        out.append([gi * xi + bi for gi, xi, bi in zip(g, xhat, b)])
        cache.append((xhat, inv))
    return out, cache


def layernorm_backward(dy, g, cache):
    dx, rows_dg, rows_db = [], [], []
    for dyr, (xhat, inv) in zip(dy, cache):
        n = len(dyr)
        dxhat = [d * gi for d, gi in zip(dyr, g)]
        s1, s2 = math.fsum(dxhat), math.fsum(map(mul, dxhat, xhat))
        dx.append([inv / n * (n * d - s1 - xh * s2) for d, xh in zip(dxhat, xhat)])
        rows_dg.append(list(map(mul, dyr, xhat)))
        rows_db.append(dyr)
    dg, db = col_sums(rows_dg), col_sums(rows_db)
    return dx, dg, db


# ---- causal self-attention (single head, width d) ------------------------------------

def softmax_row(row):
    top = max(v for v in row if v != -math.inf)
    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(p, x):
    T, d = len(x), len(p["Wq"][0])
    q, k, v = matmul(x, p["Wq"]), matmul(x, p["Wk"]), matmul(x, p["Wv"])
    scale = 1 / math.sqrt(d)
    weights = []
    for i in range(T):
        scores = [math.fsum(map(mul, q[i], k[j])) * scale if j <= i else -math.inf for j in range(T)]
        weights.append(softmax_row(scores))
    mixed = matmul(weights, v)
    out = matmul(mixed, p["Wo"])
    return out, {"x": x, "q": q, "k": k, "v": v, "weights": weights, "mixed": mixed}


def attention_backward(p, c, d_out):
    scale = 1 / math.sqrt(len(p["Wq"][0]))
    grads = {"Wo": matmul(transpose(c["mixed"]), d_out)}
    d_mixed = matmul(d_out, transpose(p["Wo"]))
    d_weights = matmul(d_mixed, transpose(c["v"]))
    d_v = matmul(transpose(c["weights"]), d_mixed)
    d_scores = []
    for a_row, da_row in zip(c["weights"], d_weights):
        dot = math.fsum(map(mul, 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, c["k"])
    d_k = matmul(transpose(d_scores), c["q"])
    xt = transpose(c["x"])
    grads["Wq"], grads["Wk"], grads["Wv"] = matmul(xt, d_q), matmul(xt, d_k), matmul(xt, d_v)
    d_x = add(add(matmul(d_q, transpose(p["Wq"])), matmul(d_k, transpose(p["Wk"]))),
              matmul(d_v, transpose(p["Wv"])))
    return grads, d_x


# ---- the full model ---------------------------------------------------------------------

def check_ids(p, ids):
    vocab, max_len = len(p["tok_emb"]), len(p["pos_emb"])
    if not ids:
        raise ValueError("empty sequence")
    if len(ids) > max_len:
        raise ValueError(f"sequence of {len(ids)} tokens is longer than the {max_len} positions the model has")
    bad = [i for i in ids if not 0 <= i < vocab]
    if bad:
        raise ValueError(f"token IDs outside the vocabulary of size {vocab}: {bad[:5]}")


def forward(p, ids, use_positions=True):
    """Return the logits (T, V) and everything backward() needs."""
    check_ids(p, ids)
    x0 = [list(p["tok_emb"][i]) for i in ids]
    if use_positions:
        x0 = add(x0, p["pos_emb"][:len(ids)])
    h1, ln1 = layernorm_forward(x0, p["ln1_g"], p["ln1_b"])
    a, att = attention_forward(p, h1)
    x1 = add(x0, a)
    h2, ln2 = layernorm_forward(x1, p["ln2_g"], p["ln2_b"])
    pre = add_bias(matmul(h2, p["W1"]), p["b1"])
    act = [[v if v > 0 else 0.0 for v in row] for row in pre]
    m = add_bias(matmul(act, p["W2"]), p["b2"])
    x2 = add(x1, m)
    hf, lnf = layernorm_forward(x2, p["lnf_g"], p["lnf_b"])
    logits = add_bias(matmul(hf, p["W_out"]), p["b_out"])
    cache = {"ids": ids, "use_positions": use_positions, "ln1": ln1, "att": att, "h2": h2, "ln2": ln2,
             "pre": pre, "act": act, "hf": hf, "lnf": lnf, "x0": x0, "x1": x1, "x2": x2}
    return logits, cache


def loss_and_dlogits(logits, targets):
    """Mean cross-entropy over positions and its gradient with respect to the logits."""
    if len(targets) != len(logits):
        raise ValueError(f"{len(logits)} positions but {len(targets)} targets")
    losses, d = [], []
    n = len(targets)
    for row, t in zip(logits, targets):
        probs = softmax_row(row)
        losses.append(-math.log(probs[t]))
        d.append([(pv - (1.0 if j == t else 0.0)) / n for j, pv in enumerate(probs)])
    return math.fsum(losses) / n, d


def backward(p, cache, d_logits):
    g = {}
    g["W_out"] = matmul(transpose(cache["hf"]), d_logits)
    g["b_out"] = col_sums(d_logits)
    d_hf = matmul(d_logits, transpose(p["W_out"]))
    d_x2, g["lnf_g"], g["lnf_b"] = layernorm_backward(d_hf, p["lnf_g"], cache["lnf"])
    # MLP branch; the residual passes d_x2 straight through to x1.
    g["W2"] = matmul(transpose(cache["act"]), d_x2)
    g["b2"] = col_sums(d_x2)
    d_act = matmul(d_x2, transpose(p["W2"]))
    d_pre = [[d if v > 0 else 0.0 for d, v in zip(dr, vr)] for dr, vr in zip(d_act, cache["pre"])]
    g["W1"] = matmul(transpose(cache["h2"]), d_pre)
    g["b1"] = col_sums(d_pre)
    d_h2 = matmul(d_pre, transpose(p["W1"]))
    d_x1_ln, g["ln2_g"], g["ln2_b"] = layernorm_backward(d_h2, p["ln2_g"], cache["ln2"])
    d_x1 = add(d_x2, d_x1_ln)
    # Attention branch; the residual passes d_x1 straight through to x0.
    att_g, d_h1 = attention_backward(p, cache["att"], d_x1)
    g.update(att_g)
    d_x0_ln, g["ln1_g"], g["ln1_b"] = layernorm_backward(d_h1, p["ln1_g"], cache["ln1"])
    d_x0 = add(d_x1, d_x0_ln)
    g["tok_emb"] = zeros(len(p["tok_emb"]), len(p["tok_emb"][0]))
    for i, row in zip(cache["ids"], d_x0):
        g["tok_emb"][i] = [a + b for a, b in zip(g["tok_emb"][i], row)]
    g["pos_emb"] = zeros(len(p["pos_emb"]), len(p["pos_emb"][0]))
    if cache["use_positions"]:
        for t, row in enumerate(d_x0):
            g["pos_emb"][t] = list(row)
    return g


def sequence_loss(p, ids, targets, use_positions=True):
    logits, _ = forward(p, ids, use_positions)
    return loss_and_dlogits(logits, targets)[0]

The implementation uses math.fsum for accumulation to reduce rounding differences between sums. The tests compare the recorded trace and check the position experiment exactly on the supported runtime; this does not guarantee identical floating-point results on every platform or Python version. The original source comment makes that broader claim, but the recorded checks establish results only for the tested runtime.

layernorm_backward() first multiplies the incoming output gradient by the learned gain to get dxhat, called u here. For normalized values x̂ and σ = √(variance + ε), the input gradient is (1/σ)(u − mean(u) − x̂ · mean(u · x̂)). Omitting the gain gives the wrong gradient once it differs from 1. The test with non-unit gains checks input, gain and bias gradients against finite differences. backward() adds the two gradients at each skip path explicitly, with add(d_x2, d_x1_ln) and add(d_x1, d_x0_ln).

"""Run the one-block model on story text and record shapes, normalization and the position test.

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

import model as m
from bpe import Tokenizer

HERE = Path(__file__).parent
SERIES = {"d": 16, "hidden": 64, "max_len": 64, "seed": 3}  # the size later builds train
TINY = {"d": 4, "hidden": 8, "max_len": 8, "seed": 3}        # small enough to check every gradient
ORDER_A, ORDER_B = "now mia gives the hat to leo.", "now leo gives the hat to mia."
EPS = 1e-6


def build(tok, cfg):
    return m.init_params(len(tok.vocab), cfg["d"], cfg["hidden"], cfg["max_len"], cfg["seed"])


def gradient_error(p, x, y):
    logits, cache = m.forward(p, x)
    _, d = m.loss_and_dlogits(logits, y)
    g = m.backward(p, cache, d)
    worst = 0.0
    for name, value in p.items():
        rows, grow = (value, g[name]) if isinstance(value[0], list) else ([value], [g[name]])
        for row, gr in zip(rows, grow):
            for j in range(len(row)):
                old = row[j]
                row[j] = old + EPS
                up = m.sequence_loss(p, x, y)
                row[j] = old - EPS
                down = m.sequence_loss(p, x, y)
                row[j] = old
                worst = max(worst, abs((up - down) / (2 * EPS) - gr[j]))
    return worst


def position_gaps(p, tok, use_positions):
    """Largest score difference at each position between the two sentences."""
    a, _ = m.forward(p, tok.encode(ORDER_A), use_positions)
    b, _ = m.forward(p, tok.encode(ORDER_B), use_positions)
    return [round(max(abs(x - y) for x, y in zip(ra, rb)), 4) for ra, rb in zip(a, b)]


def last_logits_gap(p, tok, use_positions):
    a, _ = m.forward(p, tok.encode(ORDER_A), use_positions)
    b, _ = m.forward(p, tok.encode(ORDER_B), use_positions)
    return max(abs(x - y) for x, y in zip(a[-1], b[-1]))


def trace():
    tok = Tokenizer.load(HERE / "tokenizer.json")
    story = (HERE / "corpus.txt").read_text().split("\n")[0] + "\n"
    ids = tok.encode(story)
    p = build(tok, SERIES)
    logits, cache = m.forward(p, ids[:-1])
    loss, _ = m.loss_and_dlogits(logits, ids[1:])
    h1 = m.layernorm_forward(cache["x0"], p["ln1_g"], p["ln1_b"])[0]
    tiny = build(tok, TINY)
    short = tok.encode("mia has a red ball.")
    return {
        "story": story,
        "story_tokens": len(ids),
        "series_config": SERIES,
        "vocab_size": len(tok.vocab),
        "parameter_count": m.count_params(p),
        "parameters_by_tensor": {k: list(m.shape(v)) for k, v in p.items()},
        "shapes": {"x0": list(m.shape(cache["x0"])), "attention_weights": list(m.shape(cache["att"]["weights"])),
                   "mlp_hidden": list(m.shape(cache["act"])), "logits": list(m.shape(logits))},
        "initial_loss": round(loss, 4),
        "uniform_loss": round(math.log(len(tok.vocab)), 4),
        "layernorm_row0": {
            "mean": round(math.fsum(h1[0]) / len(h1[0]), 6) + 0.0,
            "variance": round(math.fsum(v * v for v in h1[0]) / len(h1[0]), 4),
        },
        "order_test": {
            "a": ORDER_A, "b": ORDER_B,
            "tokens_a": tok.tokens(ORDER_A), "tokens_b": tok.tokens(ORDER_B),
            "last_logit_gap_without_positions": last_logits_gap(p, tok, False),
            "last_logit_gap_with_positions": round(last_logits_gap(p, tok, True), 6),
            "gaps_by_position_without_positions": position_gaps(p, tok, False),
        },
        "gradient_check": {
            "config": TINY,
            "parameters_checked": m.count_params(tiny),
            "max_error_rounded": round(gradient_error(tiny, short[:-1], short[1:]), 7),
        },
    }


if __name__ == "__main__":
    path = HERE / "block-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 "block-trace.json is out of date")
    else:
        print(text, end="")

Run the trace:

python3 trace_block.py

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

Count the parameters

The total is 7,440. The two embedding tables and the output matrix hold most of it, because they have a row or column for every vocabulary token or position.

Check the starting loss

With small random weights, the model gives every token roughly the same score. Its loss on the first story is 4.5805, close to ln 96 = 4.5643, the uniform model’s loss from the baseline build. A starting loss far from that value usually means the initial weights are too large.

The first layer normalization turns the first position’s embedding into a row with mean 0.0 and variance 0.9852, not exactly 1. The embeddings start tiny, so ε is about 1.5% of the variance plus ε, and dividing by √(variance + ε) pulls the result below 1. The larger a row’s variance, the less ε matters.

Check the gradients

gradient_error() nudges each of the 1,060 parameters of a copy with d = 4 in both directions and compares the loss difference with backward(). The largest difference, rounded to seven decimal places, is 0.0000000.

Check your understanding

The trace runs two sentences through the series-size model: now mia gives the hat to leo. and now leo gives the hat to mia.. They contain the same tokens, "n", "ow", " mia", " gives", " the", " hat", " to", " leo", ".", with the two names swapped. Before you open the comparison, predict how different the scores at the last position are, first without position embeddings and then with them.

{
  "story": "mia has a red hat. mia goes to the garden. at the garden, mia gives the hat to ava. now ava has the red hat.\n",
  "story_tokens": 37,
  "series_config": {
    "d": 16,
    "hidden": 64,
    "max_len": 64,
    "seed": 3
  },
  "vocab_size": 96,
  "parameter_count": 7440,
  "parameters_by_tensor": {
    "tok_emb": [
      96,
      16
    ],
    "pos_emb": [
      64,
      16
    ],
    "ln1_g": [
      16
    ],
    "ln1_b": [
      16
    ],
    "Wq": [
      16,
      16
    ],
    "Wk": [
      16,
      16
    ],
    "Wv": [
      16,
      16
    ],
    "Wo": [
      16,
      16
    ],
    "ln2_g": [
      16
    ],
    "ln2_b": [
      16
    ],
    "W1": [
      16,
      64
    ],
    "b1": [
      64
    ],
    "W2": [
      64,
      16
    ],
    "b2": [
      16
    ],
    "lnf_g": [
      16
    ],
    "lnf_b": [
      16
    ],
    "W_out": [
      16,
      96
    ],
    "b_out": [
      96
    ]
  },
  "shapes": {
    "x0": [
      36,
      16
    ],
    "attention_weights": [
      36,
      36
    ],
    "mlp_hidden": [
      36,
      64
    ],
    "logits": [
      36,
      96
    ]
  },
  "initial_loss": 4.5805,
  "uniform_loss": 4.5643,
  "layernorm_row0": {
    "mean": 0.0,
    "variance": 0.9852
  },
  "order_test": {
    "a": "now mia gives the hat to leo.",
    "b": "now leo gives the hat to mia.",
    "tokens_a": [
      "n",
      "ow",
      " mia",
      " gives",
      " the",
      " hat",
      " to",
      " leo",
      "."
    ],
    "tokens_b": [
      "n",
      "ow",
      " leo",
      " gives",
      " the",
      " hat",
      " to",
      " mia",
      "."
    ],
    "last_logit_gap_without_positions": 0.0,
    "last_logit_gap_with_positions": 0.001282,
    "gaps_by_position_without_positions": [
      0.0,
      0.0,
      0.3037,
      0.0149,
      0.0137,
      0.0125,
      0.0232,
      0.3097,
      0.0
    ]
  },
  "gradient_check": {
    "config": {
      "d": 4,
      "hidden": 8,
      "max_len": 8,
      "seed": 3
    },
    "parameters_checked": 1060,
    "max_error_rounded": 0.0
  }
}

Failure cases

Each case has a named test in test_model.py.

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

python3 -m unittest -v test_model.py

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

import model as m
import trace_block as tb
from bpe import Tokenizer

HERE = Path(__file__).parent
TOK = Tokenizer.load(HERE / "tokenizer.json")


def series():
    return tb.build(TOK, tb.SERIES)


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


class Mechanism(unittest.TestCase):
    def test_gradients_match_finite_differences(self):
        ids = TOK.encode("mia has a red ball.")
        self.assertLess(tb.gradient_error(tb.build(TOK, tb.TINY), ids[:-1], ids[1:]), 1e-7)

    def test_layernorm_centers_and_scales_each_row(self):
        x = [[1.0, 2.0, 3.0, 6.0], [0.5, 0.5, 0.5, 2.5]]
        out, _ = m.layernorm_forward(x, [2.0] * 4, [1.0] * 4)
        for row in out:
            self.assertAlmostEqual(math.fsum(row) / 4, 1.0, places=12)  # the bias sets the mean
            var = math.fsum((v - 1.0) ** 2 for v in row) / 4
            self.assertAlmostEqual(var, 4.0, places=3)  # the gain squared sets the variance

    def test_layernorm_backward_with_nonunit_gains(self):
        x = [[0.2, -0.4, 1.1], [2.0, 0.3, -0.8]]
        gain, bias = [2.0, 0.5, -1.0], [0.1, -0.2, 0.3]
        dy = [[0.4, -0.7, 0.2], [-0.3, 0.8, 0.5]]
        _, cache = m.layernorm_forward(x, gain, bias)
        dx, dg, db = m.layernorm_backward(dy, gain, cache)

        def loss():
            y, _ = m.layernorm_forward(x, gain, bias)
            return math.fsum(a * b for row, grad in zip(y, dy) for a, b in zip(row, grad))

        def check(values, gradients):
            for i, expected in enumerate(gradients):
                old, epsilon = values[i], 1e-6
                values[i] = old + epsilon
                higher = loss()
                values[i] = old - epsilon
                lower = loss()
                values[i] = old
                self.assertAlmostEqual(expected, (higher - lower) / (2 * epsilon), places=7)

        for row, grad in zip(x, dx):
            check(row, grad)
        check(gain, dg)
        check(bias, db)

    def test_later_token_does_not_change_earlier_logits(self):
        p = series()
        a, _ = m.forward(p, TOK.encode("mia has a red ball."))
        b, _ = m.forward(p, TOK.encode("mia has a red hat."))
        self.assertEqual(a[:4], b[:4])
        self.assertNotEqual(a[4], b[4])

    def test_unused_positions_get_zero_gradient(self):
        p = series()
        ids = TOK.encode("mia has a red ball.")
        logits, cache = m.forward(p, ids[:-1])
        g = m.backward(p, cache, m.loss_and_dlogits(logits, ids[1:])[1])
        self.assertTrue(any(v != 0.0 for v in g["pos_emb"][0]))
        self.assertTrue(all(v == 0.0 for row in g["pos_emb"][len(ids) - 1:] for v in row))


class FailureCases(unittest.TestCase):
    def test_without_positions_the_last_position_cannot_see_word_order(self):
        self.assertEqual(tb.last_logits_gap(series(), TOK, use_positions=False), 0.0)
        self.assertGreater(tb.last_logits_gap(series(), TOK, use_positions=True), 0.0)

    def test_earlier_positions_still_differ_without_positions(self):
        gaps = tb.position_gaps(series(), TOK, use_positions=False)
        self.assertEqual(gaps[-1], 0.0)
        self.assertTrue(all(g > 0 for g in gaps[2:-1]))

    def test_sequence_longer_than_positions(self):
        with self.assertRaisesRegex(ValueError, "longer than the 64 positions"):
            m.forward(series(), [5] * 65)

    def test_token_id_outside_vocabulary(self):
        for ids in ([5, 96], [5, -1]):
            with self.subTest(ids=ids), self.assertRaisesRegex(ValueError, "outside the vocabulary"):
                m.forward(series(), ids)

    def test_empty_sequence(self):
        with self.assertRaisesRegex(ValueError, "empty sequence"):
            m.forward(series(), [])

    def test_targets_must_match_positions(self):
        logits, _ = m.forward(series(), [5, 6, 7])
        with self.assertRaisesRegex(ValueError, "3 positions but 2 targets"):
            m.loss_and_dlogits(logits, [6, 7])


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

test_gradients_match_finite_differences fails if the layer normalization backward drops its last term, if ReLU passes gradient through negative inputs, or if a skip-path gradient is missing. test_later_token_does_not_change_earlier_logits fails if the mask lets attention see one token ahead.

Deployment

There is nothing to deploy yet. The next build, Build a small language-model training pipeline, trains this model and saves its parameters.

Scaling

This section is unmeasured here; the next build records training time on a stated machine. The work per sequence grows with the square of its length in attention and with d times the vocabulary size in the output layer. Every multiplication runs in a Python loop, which is enough for 7,440 parameters and far too slow for a model of millions.

Key takeaways

• A transformer language model wraps attention and a feed-forward network in layer normalization and skip paths, then turns the last vectors into one score per vocabulary token.
• At each skip path, the backward pass adds the gradient that comes through the sublayer to the gradient that skips it.
• In this one-block model, the last position cannot distinguish two orders of the same earlier tokens without position embeddings. Learned position embeddings distinguish those inputs; the experiment does not claim the same result for every multilayer architecture.

Best practices

• Check the starting loss against ln(vocabulary size) before training.
• Run a full-model gradient check on a small copy with the same code, not only on each layer separately.
• Keep a word-order test such as the one here, so that a change that drops position information fails loudly.

References

Sources checked October 20, 2026.

• Vaswani et al., Attention Is All You Need (NIPS 2017), arXiv 1706.03762v7: residual connections with layer normalization after each sublayer (section 3.1), the two-layer ReLU feed-forward network with d_ff = 2048 and d_model = 512 (3.3), positional encodings added to the embeddings and learned position embeddings giving nearly identical results (3.5).
• Radford et al., Language Models are Unsupervised Multitask Learners (GPT-2, 2019), section 2.3: layer normalization moved to the input of each sub-block, with one more after the final block.
• Ba, Kiros and Hinton, Layer Normalization, arXiv 1607.06450: statistics over the units of one layer for one training case, with a learned gain and bias.
• PyTorch 2.14: torch.nn.LayerNorm: y = (x − E[x]) / √(Var[x] + ε) · γ + β with ε = 10⁻⁵ by default, the biased variance, γ starting at 1 and β at 0.