Heads, norms, residuals, and MLPs
Attention mixes tokens, the MLP mixes channels, and normalized residual paths let both transformations accumulate without losing the original stream.
Domain autocomplete needs a model whose every component the team can explain.
This chapter moves from track width through every path to test the residual contract. We will make the idea concrete in a small code example before placing it back inside a production system with one-device educational build.
Assert unchanged `[batch, time, width]` shape, causal independence from future tokens, finite outputs, and gradient flow through a stack of blocks.
Verify the behavior illustrated in decoder_block.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 transformer block works because several modest mechanisms cooperate: heads mix information across positions, an MLP transforms each position, and residual paths keep the whole stack trainable.
Multiple attention heads project the same sequence into smaller subspaces, allowing different relation patterns to be represented in parallel. Their outputs are concatenated and projected back to model width. The feed-forward MLP then expands and contracts each position independently; attention performs token mixing, while the MLP performs channel mixing.
Residual additions require each sublayer to return the same shape it receives. Layer normalization controls feature scale, and in a pre-norm block it sits before each sublayer, leaving a direct residual path through deep stacks. Dropout and initialization still matter, but shape and mask tests should pass before any optimization experiment.
- QKV projections
- causal masks
- multi-head attention
- pre-norm
Follow the reasoning before touching the code.
- 01
Track width through every path
Verify that head width times head count equals the attention projection width and that each sublayer projects back to model width before residual addition.
- 02
Separate the two mixers
Read attention as communication between positions and the MLP as a shared nonlinear transform at each position. This makes ablations interpretable.
- 03
Test the residual contract
Assert unchanged `[batch, time, width]` shape, causal independence from future tokens, finite outputs, and gradient flow through a stack of blocks.
import torch
class DecoderBlock(torch.nn.Module):
def __init__(self, width=64, heads=4):
super().__init__()
self.norm1 = torch.nn.LayerNorm(width)
self.attn = torch.nn.MultiheadAttention(width, heads, batch_first=True)
self.norm2 = torch.nn.LayerNorm(width)
self.mlp = torch.nn.Sequential(
torch.nn.Linear(width, 4 * width), torch.nn.GELU(),
torch.nn.Linear(4 * width, width))
def forward(self, x):
time = x.shape[1]
mask = torch.triu(torch.ones(time, time, device=x.device,
dtype=torch.bool), diagonal=1)
h = self.norm1(x)
attended, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
x = x + attended
return x + self.mlp(self.norm2(x))A compact pre-norm decoder block keeps both residual paths shape-compatible.
What to notice
- Both attention and MLP return model-width tensors, so each can be added to its residual input.
- Layer normalization receives the current residual stream before the sublayer, which makes this a pre-norm block.
YOUR TURN Pause and predict before revealing the answer.
Change the final MLP layer to output `2 * width`. Which invariant fails first?
Need a hint?
Look at the final addition rather than the attention computation.
Check your reasoning
The residual addition fails because `[batch, time, 2 × width]` cannot be added to `[batch, time, width]`. Residual branches must return exactly the residual stream’s shape.
Lessons 1, 2 in this course.
Connect the chapter to the course’s quantitative spine.
Attention mixes tokens, the MLP mixes channels, and normalized residual paths let both transformations accumulate without losing the original stream. 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: Attention mixes tokens, the MLP mixes channels, and normalized residual paths let both transformations accumulate without losing the original stream. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Track width through every path” into a written invariant and a failing test.
- 02
Promote decoder_block.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Test the residual contract” as the release checkpoint: Assert unchanged `[batch, time, width]` shape, causal independence from future tokens, finite outputs, and gradient flow through a stack of blocks.
- 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.