Train a causal language model
Causal LM training is next-token alignment plus a no-looking-ahead proof; everything else scales that contract.
Domain autocomplete needs a model whose every component the team can explain.
This chapter moves from shift once, visibly to earn the large run. We will make the idea concrete in a small code example before placing it back inside a production system with one-device educational build.
Overfit a tiny batch, inspect loss against a uniform-vocabulary baseline, then add validation, gradient, checkpoint, and throughput instrumentation.
Verify the behavior illustrated in causal_lm_step.py, then measure decision utility, slice reliability, and system latency.
Without the causal mask, each position sees its target and validation loss becomes impossibly good. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
A causal language model learns one repeated game: from each prefix, assign high probability to the token that actually came next.
One token sequence supplies many training pairs. The inputs omit the final token and targets omit the first, creating a one-position shift. A causal mask ensures the hidden state at position i cannot inspect later input positions; without it, training leaks the answer and produces an impossibly low loss.
Cross-entropy averages negative log probability over valid target tokens. Padding targets must be ignored, and document packing must respect declared boundaries so one document does not accidentally condition on unrelated text unless that is intentional. Before a full run, overfit a tiny batch and verify that changing a future token cannot change earlier logits.
- multi-head attention
- pre-norm
- next-token loss
- KV cache
Follow the reasoning before touching the code.
- 01
Shift once, visibly
Print an input row and its target row. At every position, confirm the target is exactly the next token, including how end-of-sequence is represented.
- 02
Prove causal isolation
Change a suffix while preserving a prefix and compare prefix logits. They should match within numerical tolerance when dropout is disabled.
- 03
Earn the large run
Overfit a tiny batch, inspect loss against a uniform-vocabulary baseline, then add validation, gradient, checkpoint, and throughput instrumentation.
import torch
vocab, width = 32, 16
tokens = torch.tensor([[2, 7, 9, 4, 3], [2, 5, 8, 6, 3]])
inputs, targets = tokens[:, :-1], tokens[:, 1:]
embedding = torch.nn.Embedding(vocab, width)
position = torch.nn.Embedding(16, width)
block = torch.nn.TransformerEncoderLayer(width, 4, batch_first=True)
head = torch.nn.Linear(width, vocab)
places = torch.arange(inputs.shape[1])
hidden = embedding(inputs) + position(places)
mask = torch.triu(torch.ones(4, 4, dtype=torch.bool), diagonal=1)
logits = head(block(hidden, src_mask=mask))
loss = torch.nn.functional.cross_entropy(logits.reshape(-1, vocab), targets.reshape(-1))
loss.backward()
print(round(loss.item(), 3), logits.shape)Shift targets and mask a tiny transformer for next-token training.
What to notice
- Eight input positions produce eight next-token targets and logits shaped `[2, 4, 32]`.
- The upper-triangular mask prevents each position from reading tokens to its right during the transformer block.
YOUR TURN Pause and predict before revealing the answer.
For a vocabulary of 32 tokens, what loss do uniform predictions give, and why is it a useful smoke-test baseline?
Need a hint?
Uniform probability is 1/32, and negative log likelihood uses the natural logarithm.
Check your reasoning
The loss is −log(1/32) = log(32) ≈ 3.47. A randomly initialized model should be near that scale; a far lower untrained loss can reveal leakage or a target-shift bug.
Lessons 1, 2, 3 in this course.
Connect the chapter to the course’s quantitative spine.
Causal LM training is next-token alignment plus a no-looking-ahead proof; everything else scales that contract. The scaled dot-product attention below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
Scaled dot-product attention
| Symbol | Meaning / shape / unit |
|---|---|
Q, K, V | query, key, value matrices |
d_k | head width |
M | causal or padding mask |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: Q, K, V, d_k, M.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c14(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
Attention temperature lab
Change score temperature and inspect the exact softmax distribution.
Assumption: Fixed scores [1.2, 0.3, −0.4, 1.8].
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| token 1 | 1.2 | 0.2915 |
| token 2 | 0.3 | 0.1185 |
| token 3 | -0.4 | 0.0589 |
| token 4 | 1.8 | 0.5311 |
Trace the complete operating path.
Carry the chapter’s central idea into the system: Causal LM training is next-token alignment plus a no-looking-ahead proof; everything else scales that contract. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Shift once, visibly” into a written invariant and a failing test.
- 02
Promote causal_lm_step.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Earn the large run” as the release checkpoint: Overfit a tiny batch, inspect loss against a uniform-vocabulary baseline, then add validation, gradient, checkpoint, and throughput instrumentation.
- 04
Serve or materialize behind explicit one-device educational build.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
Without the causal mask, each position sees its target and validation loss becomes impossibly good. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
Decoder-only transformer from scratch
Specify tokenizer, attention blocks, causal training, unit tests, sampling, and KV-cache parity without a transformer module.
- Shape/mask tests
- Tiny-batch overfit
- Held-out evaluation
- Cached/uncached parity
Read primary material with a purpose.
Return to the opening failure.
Version tokenizer, vocabulary, special tokens, architecture, weights, prompt template, and generation policy together.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce tokenizer, attention, rope, residual stack, training loop, generation, and kv cache.