Handout #12 — Masks, Architecture, and Computational Cost

A language model needs more than an attention formula: it needs a precise rule for which information is available to each prediction. That rule determines valid training, padding behavior, architectural variants, and which computations can be reused during generation.

1The causal mask defines an information boundary

Output at position $i$ predicts the next token, so it may use the current input and earlier inputs: $j\le i$. Set future-key scores to negative infinity before softmax. With zero-based indices and four positions, the additive mask is:

The diagonal is allowed: input at $i$ is context for target $i+1$. The mask is a rule about availability, not a learned judgment of relevance. An allowed key can still receive a tiny weight. A dense implementation may calculate all scores before masking, so causality alone does not automatically save half the arithmetic.

1.1 Masking changes the normalization

Return to the three scores $[\log2,0,0]$ and values $[2,0]$, $[0,2]$, $[4,4]$ from Handout #11. At query position 1, only keys 0 and 1 are allowed. Masking the third score gives $[\log2,0,-\infty]$, weights $[2/3,1/3,0]$, and output $[4/3,2/3]$.

If instead you zero the third weight after the original softmax, you get $[1/2,1/4,0]$, sum $3/4$, and output $[1,1/2]$. Dividing the surviving weights by their sum would recover the correct result, but pre-softmax masking normalizes once over exactly the allowed set.

1.2 Test causality, rather than only inspecting the triangle

Run the model on two sequences identical through position $i$, but with different suffixes after $i$. With dropout disabled, the logits through $i$ should agree within floating-point tolerance. This tests the whole computation, including position handling, normalization, and every layer's mask.

The reasoning is inductive. Initial row $i$ depends only on token $i$ and its position. A masked block retrieves rows $j\le i$; if those rows already depend only on their prefixes, no future input enters. Token-wise normalization, residual addition, and FFNs preserve the boundary. Repeating this argument across layers establishes causality.

2Padding needs both an attention rule and a loss rule

Batch tensors need a common length. Padding fills unused slots, but PAD is not ordinary linguistic context. Keep three distinct objects: input IDs, a real-token mask, and target labels. Confusing their roles is a common source of models that “work” while optimizing the wrong objective.

2.1 Shift first, then mark valid targets

Let B be BOS, E be EOS, and P be PAD. For [B, a, b, E, P], form:

EOS is a valid target and must be learned. The final input is real, yet its target is padding and its loss is ignored. Thus a real input mask cannot be copied directly into a loss mask. A padded input row's output should also not be scored.

The attention rule combines causality with valid keys:

Masking PAD keys prevents their values from entering real-token representations. Ignoring PAD targets prevents the loss from rewarding PAD prediction. One does not substitute for the other. Compute the mean over the number of valid targets, not over all padded tensor entries.

2.2 Why left and right padding differ in practice

Right padding keeps real training tokens aligned from column 0. Left padding makes every prompt's last real token sit in the final column, so batched generation can simply read logits[:, -1]. These are convenient conventions, not mathematical restrictions. With right-padded prompts, reading the final column for a short sequence reads a PAD query; instead gather that sequence's last real position.

Position IDs matter too. For left padding, a useful convention numbers real tokens from 0 using the cumulative sum of the real-token mask minus one. Assign a harmless in-range ID to PAD slots. Otherwise the same prompt can receive different learned absolute positions depending on the lengths of its batch companions. The attention mask alone does not fix that change.

3Architectures differ in what each query can read

The same dot-product operation supports three attention patterns. Track where queries, keys, and values come from, then determine which edges are allowed.

An encoder-only model builds bidirectional representations for tasks such as classification or masked-token prediction. A decoder-only LM predicts continuations causally. An encoder–decoder model first represents a source sequence, then generates a target conditional on that source. Architecture alone does not specify the training objective, but the objective and information boundary must agree.

3.1 Cross-attention is rectangular

Let encoder output $E\in\mathbb R^{B\times S\times D}$ have source length $S$, and decoder states $Z\in\mathbb R^{B\times U\times D}$ have target length $U$:

Each target query produces weights over source keys. Softmax is across $S$, not $U$. Cross-attention does not need a source-side triangular mask because the entire source is available before generation begins; source padding still must be masked. The encoder returns a sequence of contextual vectors, not necessarily one compressed vector.

3.2 Teacher forcing in an encoder–decoder

For source the gray chunky cat and target sat on the mat E, the decoder input is B sat on the mat and the labels are sat on the mat E. Every decoder position can see all the source, but only its target prefix. In the standard architecture, each decoder layer's cross-attention reads the final encoder output, with its own projections; it does not read only the encoder layer with the same depth index.

Bidirectional source attention is legitimate because the source is observed conditioning information. Bidirectional target attention would reveal answers. “The encoder can look right” and “the decoder cannot look right” refer to different sequences.

4Count arithmetic and stored tensors separately

4.1 Derive cost from matrix shapes

Multiplying $A\in\mathbb R^{n\times m}$ and $B\in\mathbb R^{m\times k}$ produces $nk$ entries, each a length-$m$ dot product. Counting each multiply and addition as one operation gives $nk(2m-1)\approx2nmk$ FLOPs. FLOPs counts operations; FLOP/s is a rate. Neither alone determines elapsed time.

For one ordinary MHA block with $D_{ff}=4D$, ignoring biases, dropout, normalization, activations, and the vocabulary head:

The scaling is $O(BTD^2+BT^2D)$ per layer. For these constants, the two terms match at $T=6D$. The attention pairwise term grows quadratically with length, but projections and FFNs can dominate at shorter lengths. Backward adds work; this table is a forward-pass approximation, not total training FLOPs.

4.2 Quadratic storage can matter before quadratic arithmetic dominates

A straightforward attention implementation stores scores or probabilities of shape [B,H,T,T]. At $B=2,H=8,T=1024$, one such tensor has 16,777,216 entries: 32 MiB at two bytes per entry. Doubling $T$ makes that one tensor 128 MiB. Training may store several intermediate tensors and gradients; this is not a whole-model memory estimate.

At fixed $D$, adding heads leaves the leading pairwise arithmetic proportional to $BT^2D$, but increases explicit weight-matrix storage proportional to $BHT^2$. Efficient attention kernels can avoid materializing the full matrix in device memory. That changes storage and IO behavior without changing the dense attention function or its quadratic pairwise arithmetic.

4.3 Why an operation count is not a stopwatch

Reading each input once and writing each output once would move $nm+mk+nk$ elements for matrix multiplication. Real device traffic depends on tiling, cache reuse, kernel fusion, memory hierarchy, and intermediate writes. Treat the slide IO counts as a simplified accounting model, not a guaranteed count of high-bandwidth-memory transfers.

Transformers can parallelize positions within a full-sequence layer; an RNN's state at time $t$ depends on the previous step. More arithmetic can therefore run faster when it maps well to parallel hardware. Transformer layers remain sequential in depth, and autoregressive generation remains sequential in generated tokens.

5KV caching reuses the unchanged past

In a causal decoder, appending a token does not change earlier representations. With model weights and position assignments fixed, their keys and values are unchanged too. Store those tensors for each layer; process only the newly supplied token at the next step.

The query comes from the token currently being processed, not from the unknown token it will predict. Previous queries are unnecessary for the next step because their outputs were already computed. Each layer needs its own cache, since it projects different contextual representations.

5.1 Cached masking uses absolute offsets

If the cache contains positions 0–3 and a new query has position 4, it may read keys 0–4. A fresh $1\times5$ triangle aligned to row 0 would incorrectly allow only key 0. Compare key indices to the query's offset position instead. For one unpadded new token with a cache containing only valid past tokens and itself, all those keys are allowed. Multi-token chunks and padded batches need the general offset-aware rule.

Learned absolute position IDs also continue from the cache length; restarting the new token at position 0 changes the model. Our workshop explicitly tracks this offset and compares cached logits with full-sequence logits. Its cache assumes a growing, unpadded prefix with fixed positions. Cropping a window and reassigning old positions can invalidate the stored tensors.

5.2 Cache memory grows linearly with context

For standard MHA, keys and values each have total width $D$ per token per layer. At $s$ bytes per cached element:

For $B=1,L=24,T=2048,D=1024,s=2$, this is 201,326,592 bytes, or 192 MiB. Batch size 4 makes it 768 MiB. These are cache tensors only, excluding parameters, allocator overhead, and temporary buffers. Cache dtype can differ from parameter dtype. Variants that share fewer K/V heads use their actual total K/V width instead of $D$.

5.3 Linear per step is not linear for the entire generation

With a cache, projections and FFNs for one new token cost $O(BD^2)$ per layer. Its query scores and combines $T$ values, costing $O(BTD)$. Thus one decode step costs $O(BD^2+BTD)$, plus the output vocabulary head.

For prompt length $P$ and $G$ generated-token processing steps, summing context lengths $P+1,\ldots,P+G$ gives $GP+G(G+1)/2$. The aggregate decode work is therefore:

in addition to prefill $O(BLPD^2+BLP^2D)$. The first emitted token comes from prefill; counting its prediction separately changes only endpoint constants. Caching removes repeated past computation, not the need to read an increasingly long context. Small-batch decoding can be limited by reading weights and caches rather than arithmetic throughput.

6Encoder–decoder generation has three length terms

Let source length be $S$ and target length be $U$. For fixed layer counts, batch size, and standard attention width, the forward work includes:

$S^2D$ is encoder self-attention, $U^2D$ is the aggregate of causal decoder self-attention, and $SUD$ is cross-attention. The source encoder runs once. Each decoder layer can also precompute its cross-attention keys and values from the final encoder output; these remain fixed throughout target generation. Per-step cross-attention still scores and combines source positions.

The decoder cache grows with $U$; its cross-attention cache is fixed at source length $S$. In standard equal-width MHA their tensor sizes are approximately $2BL_{dec}UDs$ and $2BL_{dec}SDs$. The encoder output itself and parameters require additional memory.

7Assigned at-home reading: assembling a Transformer

Read slides 112–125, “Writing our own Transformer,” at home. They assemble an encoder–decoder implementation. Complete #11 §9.4’s component worksheet and read §7.1–§7.2 below to connect the implementation to masks, positions, and parameter sharing.

The companion script and code-modification experiments are optional ways to test this assigned reading. To first test the shared building blocks in a smaller setting, download 12_causal_decoder.py. It implements a two-layer pre-norm decoder with learned absolute positions, ordinary MHA, a GELU FFN, and a vocabulary head. Run python 12_causal_decoder.py with PyTorch. No dataset, GPU, or external download is needed.

The script checks four meaningful properties: changing future tokens leaves prefix logits fixed; token-by-token cached logits equal full-sequence logits; real outputs in a right-padded batch match separate unpadded execution; and a shifted, PAD-excluding loss backpropagates. The toy batch has six valid targets and initial loss about 2.8076 nats on PyTorch 2.12.0. Random initialization tests wiring, not language quality.

7.1 Read the offset-aware mask

Original workshop · Block.forward · excerpt

key_pos = torch.arange(past + t, device=x.device)
query_pos = past + torch.arange(t, device=x.device)
allowed = (key_pos[None, :] <= query_pos[:, None])[None, None]
allowed = allowed & real[:, None, None, :].bool()
scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)
assert allowed.any(-1).all(), 'each query needs an allowed key'
weights = scores.masked_fill(~allowed, float('-inf')).softmax(-1)

past is cached sequence length; t is the new chunk length. The causal rule has shape [1,1,t,past+t] before combining with batch-specific valid keys. Broadcasting applies the same rule to all heads. The cache has one pair of K/V tensors per block, each [B,H,past,d_h].

7.2 Match the implementation to the slide components

• Cloning layers: ModuleList([Block(d), Block(d)]) constructs independent parameters. [block] * 2 repeats the same object and ties weights; it is not equivalent.
• Embedding and position: the token table maps IDs to vectors; the position table indexes past + arange(t). To try fixed sinusoids, replace the position lookup by a nontrainable buffer while keeping the offset.
• Attention and FFN: fused QKV retrieves context; the expansion/GELU/contraction transforms each row. The residual branches preserve width.
• Prediction head: the final linear layer returns logits. Cross-entropy takes logits; generation applies softmax only when probabilities are needed.
• Encoder extension: remove source causality, preserve source padding masks, and return final source rows. Add decoder cross-attention with decoder queries and projected encoder K/V between self-attention and FFN. Test rectangular shapes before training.
• Initialization: this small script uses PyTorch module defaults. Large-model initialization and training stability require deliberate choices; see Handout #8.

8Optional code reading: masks and loss in nanoGPT

8.1 Recognize the manual attention path

model.py · CausalSelfAttention.forward · verbatim excerpt

            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
            att = F.softmax(att, dim=-1)

The buffer named bias is a lower-triangular availability mask, not a trainable affine bias. The code masks before normalization, just as §1 requires. A second path calls scaled_dot_product_attention with is_causal=True; the helper's existence alone does not prove a particular hardware kernel is used.

8.2 Find where the shift happens

model.py · GPT.forward · verbatim excerpt

            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

GPT.forward does not shift labels: its caller must supply next-token-aligned targets. Flattening merges batch and time, leaving vocabulary as the class axis. Here the ignored label is -1; the workshop uses -100. Match the label sentinel to the loss argument, and never put either negative label into the input embedding lookup.

At generation, the same file projects only x[:, [-1], :] to the vocabulary. This is appropriate for its unpadded or correctly arranged inputs; simply right-padding arbitrary prompts and selecting the final column would introduce §2's error.

Excerpts: nanoGPT, Copyright (c) 2022 Andrej Karpathy, MIT License. License and permission notice.

9Practice set: put the pieces together