PyTorch University Path · Integration capstone

Capstone: build and defend your own transformer

Prerequisites: embeddings, attention, positional information, residual connections, normalization, cross-entropy, train/validation/test splits, and optimization. Review attention, RoPE, transformer blocks, and optimization and generalization.

Download the transformer components, the existing CPU training runner, and the transformer correctness tests. These are the reference implementation and executable evidence, not hidden notebook state.

Quick map

1. Choose sequence reversal or a tiny text language model.
2. Freeze IDs, shapes, masks, and success criteria.
3. Rebuild attention and transformer blocks from the specification.
4. Pass causality, padding, cache, and checkpoint tests.
5. Train, evaluate, ablate, and write an honest report.

Choose one primary track

1. Track A: reverse variable-length sequences. An encoder reads ordinary tokens plus padding. A decoder returns the source tokens in reverse order followed by an end token. Use this track to inspect padding masks, causal masks, cross-attention, and decoding without target leakage.
2. Track B: tiny text language model. A decoder-only transformer predicts the next character or subword. Use a small text file that you have permission to use. Split contiguous text by position, fit the vocabulary on training text only, and report next-token loss and perplexity on later held-out text.

Track B is closer to NLP. Its tiny local corpus still measures only a narrow next-token task.

The supplied runner directly implements Track A and a decoder-only generated cyclic-pattern task. The cyclic task is a mechanics check for causal attention and cached decoding; it is not real text.

Use it before adapting the data path for Track B. The introductory mini-GPT lesson explains shifted targets and generation.

# Existing runnable reference commands
.venv-learning/bin/python examples/pytorch/build_transformer.py \
  seq2seq --task reverse --steps 500

.venv-learning/bin/python examples/pytorch/build_transformer.py \
  decoder-only --steps 500

.venv-learning/bin/python -m unittest tests/test_transformer_builds.py

Write the specification before code

A minimal Track B specification could say: character vocabulary learned from training text only; context length 32; width 64; four heads; two causal blocks; held-out cross-entropy below a unigram baseline; fixed-seed reload produces matching logits. Every phrase maps to a test or tensor choice.

A capstone is not “make a transformer somehow.” Freeze a small specification so every tensor and test has a clear target. The reference quick model uses a vocabulary of 16, model width 32, 4 attention heads, head width 8, feed-forward width 64, one layer, and dropout 0. A full teaching run uses wider or deeper settings but the same interfaces.

B = batch size
S = source length
T = decoder length
D = model width
H = number of heads
Dh = D / H
V = vocabulary size

source IDs:          (B, S)
decoder input IDs:   (B, T)
token embeddings:    (B, sequence length, D)
split Q, K, V:       (B, H, sequence length, Dh)
attention scores:    (B, H, query length, key length)
output logits:       (B, sequence length, V)

State special IDs explicitly: padding, beginning, end, and first ordinary token. Require D to be divisible by H.

Define whether masks use True for allowed positions or blocked positions. The reference uses Boolean masks and rejects rows in which every key is masked, because softmax over no legal choices is undefined.

Component contract

Embedding: map integer token IDs to vectors. Apply RoPE to query and key vectors after splitting heads.

RoPE changes pairwise coordinates according to position while preserving vector norms. Values remain unrotated.

Attention: project the input into query, key, and value tensors. Calculate scaled dot products, apply legal-position masks before softmax, mix values, merge heads, and apply an output projection.

scores = (Q @ transpose(K)) / sqrt(Dh)
weights = softmax(masked scores over the key axis)
attention output = weights @ V

The scale prevents dot products from growing with head width. Softmax is along the key axis because each query chooses a mixture of key positions.

A causal decoder allows key position j only when j <= i for query position i. A padding mask prevents real tokens from attending to padding.

Block: use pre-normalization, residual connections, attention, and a position-wise feed-forward network. A decoder block has masked self-attention and, for Track A, cross-attention over encoder outputs. Cross-attention queries come from the decoder; keys and values come from the encoder.

x = x + self_attention(layer_norm(x))
x = x + cross_attention(layer_norm(x), encoder_memory)  # Track A only
x = x + feed_forward(layer_norm(x))

Output: normalize final hidden values and project each position to V logits. Calculate cross-entropy against shifted targets. Ignore padding targets for variable-length Track A examples.

Track A data and leakage rules

For source [4, 9, 6], the answer is [6, 9, 4]. Decoder input is [BOS, 6, 9, 4].

Target is [6, 9, 4, EOS]. During training, causal masking stops the decoder position for 6 from reading later answer tokens.

During inference, targets do not exist; generation starts from BOS and appends one prediction at a time.

source:         [4,   9,   6, PAD, PAD]
decoder input: [BOS, 6,   9,   4, PAD, PAD]
target:        [6,   9,   4, EOS, PAD, PAD]

Teacher-forced accuracy measures logits produced while earlier correct answer tokens are present. Greedy exact-match accuracy measures autoregressive generation, where one wrong token can influence later steps.

Report both. A model can have strong teacher-forced token accuracy and weaker generated exact match.

The existing generator samples training and held-out rows independently. Because the token space and lengths are finite, an independently sampled held-out row can still duplicate a training sequence by chance.

State this toy-data limit. Independence of the random draw is not proof that every sequence is unique.

Track B tiny text language model

Use an explicit local corpus path rather than a hidden download. Normalize text with a documented rule.

A character vocabulary is easiest to reproduce: assign one ID to every training character plus special or unknown symbols. A subword tokenizer is more efficient but adds another learned artifact that must be saved and evaluated.

Split text contiguously: earlier text for training, a later section for validation, and the final section for test. Randomly splitting overlapping windows leaks nearly identical contexts across sets.

Create windows only after splitting the raw text. For context length T, input IDs contain positions 0 through T - 1 and targets contain positions 1 through T.

text IDs: [11, 4, 8, 8, 2]
input:    [11, 4, 8, 8]
target:   [4,  8, 8, 2]

# For B windows and vocabulary V:
# input IDs shape: (B, T)
# logits shape:    (B, T, V)
# targets shape:   (B, T)

Report held-out mean cross-entropy and perplexity. Perplexity is the exponential of mean token loss, so it depends on tokenization. Character-level and subword perplexities are not directly comparable.

held-out loss = mean negative log probability of each true next token
perplexity = exp(held-out loss)

Also sample several fixed prompts. Samples help reveal repetition, broken characters, and context use, but they are qualitative examples rather than an aggregate metric.

A tiny model trained on a tiny corpus does not measure broad language understanding, factuality, safety, or useful conversation. Do not call cyclic-pattern accuracy “language accuracy,” and do not compare toy perplexity with a large pretrained model that uses another tokenizer and dataset.

Training and reproduction protocol

Record the source revision, task, seed, model configuration, optimizer, learning rate, weight decay, batch size, number of steps, split procedure, vocabulary, context length, and hardware. Set one CPU thread for this teaching run.

Store the best validation state in memory. Touch the test split only when model choices are complete.

The reference uses AdamW, gradient clipping, and generated mini-batch indices from a seeded generator. Clipping limits the total gradient norm before the update. It protects against rare large updates but does not repair consistently bad optimization.

optimizer.zero_grad(set_to_none=True)
logits = model(inputs)
loss = masked_cross_entropy(logits, targets)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()

For a reproduction, run the exact same command twice in the same environment and compare metrics and checkpoint outputs. Then change the seed deliberately and report whether the main conclusion survives. Reproducibility means another learner can recreate the procedure and understand differences; it does not require every future library and device to produce identical floating-point bits.

Tests required before training claims

Predict each failure signal: a broken causal mask changes an earlier logit when a future token changes. A broken padding mask changes a real-token output when only padding changes. A broken reload changes logits for the same fixed batch.

1. Head split and merge: split (B, T, D) into (B, H, T, Dh), merge it, and require exact equality with the original tensor.
2. Gradient path: run backward and require nonzero gradients for query, key, and value projections.
3. Causality: change a future token and verify earlier decoder logits do not change.
4. Padding: change padded source IDs while keeping the mask and verify valid outputs do not change.
5. RoPE: verify rotation preserves norms and same-position dot products within tolerance.
6. Cache equivalence: with dropout disabled, compare full-prefix logits against one-token cached logits at every position.
7. Checkpoint reload: save state plus configuration only at an explicit path, rebuild the model, reload, and require equal evaluation logits.

The downloadable test suite implements these kinds of checks for the reference components. Run it before a long experiment. A falling training loss cannot prove that a causal mask or cache offset is correct.

Greedy decoding and sampling

Greedy decoding picks the largest next-token logit. It is deterministic and appropriate for exact algorithmic targets such as sequence reversal.

Text generation often uses probability sampling. Temperature divides logits before softmax.

Values below one sharpen the distribution; values above one flatten it.

sampling probabilities = softmax(logits / temperature)

Top-k sampling keeps only the k highest-scoring tokens. Top-p sampling keeps the smallest high-probability set whose cumulative probability reaches p.

Record temperature, top-k or top-p, prompt, random seed, and maximum new tokens. Compare methods using fixed prompts.

Do not select only attractive samples.

Track A should use greedy decoding for exact-match evaluation. Track B should report held-out loss independently of sampling settings. Sampling changes generated text, not the teacher-forced test loss.

Controlled ablations

An ablation removes or changes one component to test its contribution. Run a baseline first.

Then change one factor while holding data, seed, training budget, and evaluation fixed. Useful capstone ablations include:

• remove RoPE or replace it with learned absolute position embeddings;
• remove the sqrt(Dh) attention scale;
• change one layer to two while keeping model width fixed;
• compare dropout 0 with modest dropout on a text corpus;
• compare full-prefix generation time with cached one-token generation;
• compare greedy text output with temperature sampling while keeping weights fixed.

Some ablations intentionally break correctness. Removing the causal mask leaks future targets and can make training loss look excellent.

Label that as an invalid model, not an improvement. A fair ablation table includes configuration, parameter count, training budget, validation metric, final test metric if selected, and a short interpretation.

Checkpoint and artifact contract

An inference checkpoint needs model state and enough configuration to recreate the architecture. A training-resume checkpoint also needs optimizer state, current step, scheduler state if present, and random-number states. A text model also needs the exact vocabulary or tokenizer, normalization rule, and special IDs.

# Existing runner writes only when --checkpoint is present.
.venv-learning/bin/python examples/pytorch/build_transformer.py \
  seq2seq --task reverse --steps 500 \
  --checkpoint /tmp/reverse-transformer.pt

Load only checkpoints you trust. A state dictionary is preferable to saving an entire Python model object because the architecture remains explicit in code. Reload in evaluation mode and compare logits on a fixed batch.

Capstone rubric

• 15% specification: all IDs, dimensions, masks, data splits, and success criteria are written before results.
• 20% implementation: attention is implemented from projections through value mixing; no hidden nn.Transformer shortcut.
• 20% correctness: shape, gradient, causality, padding, cache, and checkpoint tests pass.
• 15% training: the complete CPU command, seed, baseline, loss curve, and configuration are recorded.
• 15% evaluation: Track A reports token and exact autoregressive metrics; Track B reports held-out loss, perplexity, and representative fixed-prompt samples.
• 10% ablation: at least two controlled comparisons isolate one factor each.
• 5% honesty: limitations, duplicate risk, dataset rights, compute budget, and unsupported real-world claims are stated.

A strong submission is reproducible and careful, not merely high scoring. Include failed experiments when they explain a design choice.

Failure drill: impossible training success

Check target alignment and the causal mask. The decoder may be reading the token it must predict or a later target.

Modify one future input token and compare all earlier logits. They must remain unchanged.

Confirm decoder input begins with BOS while target ends with EOS. During generation, do not pass the correct answer tokens.

If full-prefix output is correct but cached output differs, compare one position at a time with dropout disabled. Inspect RoPE offsets, cache concatenation order, and mask dimensions. The new token at absolute position 7 must use position 7, not position 0.

Practice checks with worked answers

Quick recap

Choose one track and freeze its contract before coding. Make shapes, masks, data splits, and metrics explicit.

Pass focused correctness tests before trusting training loss. Then reproduce the run, compare controlled ablations, evaluate held-out data, and state what the result does not prove.

Final report template

State the chosen track and intended use. List dataset creation, split boundaries, tokenizer or generated task, tensor shapes, model configuration, parameter count, optimizer, seed, and exact commands.

Report the baseline, validation selection rule, final held-out metrics, and confusion or generation failures. Add correctness-test results, checkpoint reload evidence, and ablation results.

For text, include several fixed prompts and sampling settings. End with limits and one next experiment.

The reverse task proves that the implementation can learn a clean algorithm under controlled conditions. The cyclic decoder task proves next-token and cache mechanics.

A tiny text model measures next-token prediction only on its held-out corpus. None of these establishes broad real-world language performance.

That distinction is part of the capstone, not a footnote.

Reference solution: complete transformer training build. Efficiency extension: KV cache and cached decoding.