One head and four heads

This notebook extends the Part II attention model. It changes the number of heads, while retaining one attention block, learned absolute positions, the residual update and the prediction MLP. Part III and Notebook 7 walk through multi-head attention from scratch. The broader Transformer material is kept in optional Part 2B.

Run the browser demo · Part II slides · Download all notebooks

Run all cells to inspect the saved experiment on CPU. Training is optional. Keep the support files and artifacts from the ZIP beside this notebook.

import math
from pathlib import Path
import torch
from IPython.display import display, HTML
from wordlm import read_json, load_model_npz, count_parameters, tokenize
from run_head_comparison import build_model

ROOT = Path.cwd()
assert (ROOT / 'wordlm.py').exists(), 'Open this notebook beside wordlm.py'
results = read_json(ROOT / 'artifacts/heads/comparison.json')
vocab_data = read_json(ROOT / 'artifacts/vocab.json')
itos = vocab_data['itos']
stoi = {word: i for i, word in enumerate(itos)}
display(HTML('<style>body{font-family:"Avenir Next",sans-serif} .jp-Notebook{max-width:1050px;margin:auto} h1,h2{letter-spacing:-.03em} table{font-variant-numeric:tabular-nums} td,th{padding:10px!important}</style>'))

The comparison

The MLP reads 64 embeddings in order. Attention instead forms a weighted message from the same context. Both attention models use a 64-dimensional representation. One head uses all 64 coordinates for one weight row. Four heads divide the projections into four groups of 16, each with its own weight row.

Four heads do not quadruple the parameter count here: the Q, K, V and output matrices retain shape 64 × 64. Adding heads at fixed per-head width would be a different experiment.

models = {}
for kind in ['mlp', 'attention', 'multihead']:
    model = build_model(kind)
    metadata, _ = load_model_npz(ROOT / f'artifacts/heads/{kind}_seed11.npz', model)
    model.eval()
    models[kind] = model
    print(kind, count_parameters(model), 'parameters; selected step', metadata['selected_step'])
assert count_parameters(models['attention']) == count_parameters(models['multihead'])

mlp 2332832 parameters; selected step 3000
attention 1321120 parameters; selected step 6000
multihead 1321120 parameters; selected step 6000

Implementation

The initializer reuses the Part II embedding tables, projection matrices and prediction MLP from CausalAttentionLM. split_heads makes the head axis explicit. Each head scales by the square root of its own width, 16 here, rather than the total width 64.

forward_details computes every causal query row for inspection. The training loss supervises only the final query, so forward computes that row directly. The two results must agree.

"""One attention block with several heads; the Part II prediction head is unchanged."""
import math
import torch
from torch.nn import functional as F
from wordlm import CausalAttentionLM


class MultiHeadAttentionLM(CausalAttentionLM):
    def __init__(self, vocab_size, context_len, d_model=64, hidden=256, heads=4, pad_id=0):
        if heads < 1 or d_model % heads:
            raise ValueError('d_model must be divisible by the positive head count')
        super().__init__(vocab_size, context_len, d_model, d_model, d_model, hidden, pad_id)
        self.heads = heads
        self.head_width = d_model // heads

    def split_heads(self, rows):
        B, T, _ = rows.shape
        return rows.reshape(B, T, self.heads, self.head_width).transpose(1, 2)

    def forward_details(self, context_ids):
        B, T = context_ids.shape
        positions = torch.arange(T, device=context_ids.device)
        E = self.token_embedding(context_ids) + self.position_embedding(positions)[None]
        Q, K, V = (self.split_heads(layer(E)) for layer in (self.W_Q, self.W_K, self.W_V))
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.head_width)
        future = torch.ones(T, T, dtype=torch.bool, device=E.device).triu(1)
        real = context_ids.ne(self.pad_id)
        mask = future[None] | ((~real[:, None, :]) & real[:, :, None])
        A = F.softmax(scores.masked_fill(mask[:, None], float('-inf')), dim=-1)
        messages = (A @ V).transpose(1, 2).reshape(B, T, self.d_model)
        contextual = E + self.W_O(messages)
        logits = self.vocab_head(F.relu(self.hidden_layer(contextual)))
        return dict(E=E, Q=Q, K=K, V=V, weights=A, messages=messages,
                    contextual=contextual, logits_all=logits, logits=logits[:, -1])

    def forward(self, context_ids):
        B, T = context_ids.shape
        positions = torch.arange(T, device=context_ids.device)
        E = self.token_embedding(context_ids) + self.position_embedding(positions)[None]
        Q = self.split_heads(self.W_Q(E[:, -1:]))
        K, V = self.split_heads(self.W_K(E)), self.split_heads(self.W_V(E))
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.head_width)
        mask = context_ids.eq(self.pad_id)[:, None, None, :]
        A = F.softmax(scores.masked_fill(mask, float('-inf')), dim=-1)
        message = (A @ V).transpose(1, 2).reshape(B, self.d_model)
        updated = E[:, -1] + self.W_O(message)
        return self.vocab_head(F.relu(self.hidden_layer(updated)))

One prompt, four weight rows

Tokenization produces strings, vocabulary lookup produces IDs, and embedding lookup produces vectors. We add one BOS to this new story, crop to 64 IDs and left-pad unused slots. PAD keys receive zero attention weight.

prompt = 'a little dog found a red hat'
tokens = tokenize(prompt)
history = [1] + [stoi.get(t, 3) for t in tokens]
context_ids = history[-64:]
context_ids = [0] * (64 - len(context_ids)) + context_ids
X = torch.tensor([context_ids])
print('Tokens:', tokens)
print('Non-padding IDs:', history)
print('Input shape:', tuple(X.shape))
model = models['multihead']
with torch.inference_mode():
    details = model.forward_details(X)
    torch.testing.assert_close(model(X), details['logits'])
for name in ['E', 'Q', 'K', 'V', 'weights', 'messages', 'contextual', 'logits']:
    print(name, tuple(details[name].shape))

Tokens: ['a', 'little', 'dog', 'found', 'a', 'red', 'hat']
Non-padding IDs: [1, 9, 34, 141, 110, 9, 237, 503]
Input shape: (1, 64)

E (1, 64, 64)
Q (1, 4, 64, 16)
K (1, 4, 64, 16)
V (1, 4, 64, 16)
weights (1, 4, 64, 64)
messages (1, 64, 64)
contextual (1, 64, 64)
logits (1, 4000)

The dimensions are B = 1 example, T = 64 slots, H = 4 heads and d_head = 16. Q/K/V have shape [B,H,T,d_head]. The attention weights have shape [B,H,T,T]. After weighted sums, concatenating the four messages restores width 64.

The following table shows the final query's weights over the real input tokens. Different heads can weight different sources. A weight alone does not tell us which semantic feature a head learned.

from html import escape
real = X[0].ne(0)
words = [itos[int(i)] for i in X[0, real]]
A = details['weights'][0, :, -1, real]
torch.testing.assert_close(A.sum(-1), torch.ones(4))
table = '<table><tr><th>Head</th>' + ''.join('<th>'+escape(w)+'</th>' for w in words) + '</tr>'
for head, weights in enumerate(A):
    table += '<tr><th>'+str(head+1)+'</th>' + ''.join(f'<td>{w:.3f}</td>' for w in weights) + '</tr>'
display(HTML(table+'</table>'))

Head	<BOS>	a	little	dog	found	a	red	hat
1	0.116	0.095	0.103	0.334	0.048	0.103	0.109	0.092
2	0.196	0.053	0.190	0.156	0.075	0.038	0.264	0.029
3	0.414	0.042	0.011	0.006	0.348	0.153	0.025	0.000
4	0.087	0.017	0.015	0.048	0.434	0.175	0.224	0.000

Cross-entropy and perplexity

For each held-out target, read the probability assigned to the observed next token, not the generated token. Its loss is −ln p(target). Average these losses over all targets to obtain cross-entropy in nats per token.

Perplexity is exp(cross-entropy). It is the reciprocal of the geometric mean target probability. If a model assigned equal probability to four alternatives at every step, its perplexity would be four. It is not the literal number of words the model considers, and it is not an accuracy percentage. Compare it only with the same tokenizer and evaluation targets.

target_probabilities = torch.tensor([0.5, 0.25, 0.125], dtype=torch.float64)
loss_per_target = -target_probabilities.log()
cross_entropy = loss_per_target.mean()
perplexity = cross_entropy.exp()
print('Individual losses:', loss_per_target.tolist())
print('Mean cross-entropy:', round(float(cross_entropy), 4))
print('Perplexity:', round(float(perplexity), 4))
assert abs(float(perplexity) - 4) < 1e-10

Individual losses: [0.6931471805599453, 1.3862943611198906, 2.0794415416798357]
Mean cross-entropy: 1.3863
Perplexity: 4.0

Measured results

All runs use 6,000 stories split by document, a training-only vocabulary, a 64-token context, 512 examples per update and a 6,000-update ceiling. The three seeds are 11, 29 and 47. Each run restores its best validation checkpoint before scoring the 120,393 test targets.

The four-head model reuses the single-head learning rate and weight decay. These choices were frozen before any test scoring. This is a fixed-settings head-count comparison, not an exhaustive tuning study. The original two-model benchmark remains in artifacts/benchmark.json; this repeat is stored separately.

for kind, summary in results['aggregate'].items():
    print(kind)
    for metric in ['test_loss', 'test_perplexity', 'runtime_seconds']:
        m = summary[metric]
        print(f"  {metric}: {m['mean']:.4f} ± {m['sample_std']:.4f}")
print('Device:', results['environment'])
print('Per-seed checkpoint steps:')
for kind, runs in results['runs'].items():
    print(kind, [r['selected_step'] for r in runs])

attention
  test_loss: 3.4450 ± 0.0105
  test_perplexity: 31.3445 ± 0.3286
  runtime_seconds: 65.5563 ± 0.7278
mlp
  test_loss: 3.9433 ± 0.0077
  test_perplexity: 51.5909 ± 0.3968
  runtime_seconds: 58.2284 ± 0.1396
multihead
  test_loss: 3.3597 ± 0.0212
  test_perplexity: 28.7846 ± 0.6135
  runtime_seconds: 69.7269 ± 1.2688
Device: {'accelerator': 'Apple Metal (MPS)', 'cuda_available': False, 'cuda_version': None, 'mps_available': True, 'platform': 'Darwin arm64', 'python': '3.11.11', 'torch': '2.14.0'}
Per-seed checkpoint steps:
attention [6000, 5250, 5250]
mlp [3000, 2250, 2250]
multihead [6000, 5250, 6000]

Training the four-head model

Set RUN_TRAINING = True for a short training exercise. prepare_data.py --profile dgx downloads the pinned subset if it is not already cached. The download checks the source revision and checksum. A smoke run verifies the loop; it does not reproduce the published scores.

For the full three-seed comparison, run:

python run_head_comparison.py --data-dir work/data --output-dir work/head-repeat --device auto

The script uses CUDA if available, then Apple MPS, then CPU. It saves each completed run so the same command can resume. Full training replaces no bundled checkpoint when a separate output directory is supplied.

RUN_TRAINING = False
if RUN_TRAINING:
    from wordlm import build_corpus, make_all_windows, TrainConfig, train_model
    corpus = build_corpus(ROOT / 'work/data', 'dgx')
    windows = make_all_windows(corpus, 64)
    fresh_model = build_model('multihead', len(corpus.vocab.itos), seed=11)
    config = TrainConfig(steps=100, batch_size=512, learning_rate=0.003,
                         weight_decay=0.001, seed=11, eval_every=50)
    training_result, _ = train_model(fresh_model, *windows['train'][:2],
                                    *windows['validation'][:2], config)
    print(training_result['validation_loss'])
else:
    print('Training is off. The earlier cells inspect the completed experiment.')

Training is off. The earlier cells inspect the completed experiment.

Generation and timing

Choose a training-story prefix, a held-out story or an authored outside-domain prompt in the browser demo. Read the actual source story when using a dataset prompt. Unknown words become UNK, so an unfamiliar scientific prompt also tests vocabulary coverage.

The browser downloads the real FP32 ONNX checkpoints, warms them up, then computes a fresh forward pass after each chosen token. It shows actual new-token counts, stopping reason, total generation seconds and milliseconds per model call. An EOS decision is a model call but adds no text token. Downloads and warm-up are excluded. Hardware, runtime, context length, output length and decoding settings all affect time.

More heads may improve held-out loss without improving every continuation. Compare the complete test-set result as well as examples. These small models can repeat words, lose the plot and produce incorrect claims.