Self-Attention

Self-attention derives queries, keys, and values from the same sequence, then computes softmax(QKᵀ / √d) V. The dot-product matrix contains every query–key compatibility score; dividing by √d keeps their scale manageable before softmax. Padding masks prevent attention to fake tokens added only to make batches rectangular, while causal masks prevent a next-token model from seeing future tokens. Multi-head attention repeats this with several smaller learned projections, then combines the resulting contexts.

The mask is part of the model definition. Without the causal mask, a language model can cheat during training by reading the answer it is asked to predict, producing impressive training loss and a useless generator. Standard attention also creates an n×n matrix, so context length affects memory and compute quadratically; efficient-attention methods are responses to that concrete bottleneck, not merely architectural fashion.

The mechanism is identical; only the source of the three tensors differs. In cross-attention — module 16's original encoder-decoder case — the queries come from one sequence and the keys and values from another, so the two can have different lengths and the attention matrix is rectangular: (target_len × source_len). In self-attention all three come from the same sequence, so the matrix is square (n × n) and its diagonal is a position attending to itself. The shape difference is the quickest way to tell which one a piece of code implements, and the reason nn.MultiheadAttention(q, k, v) takes three arguments rather than one.

Real Transformers use both, often in the same block. A translation decoder layer has self-attention over the partial output followed by cross-attention into the encoded source, and the two have different masks — the self-attention is causal, the cross-attention is not, because the whole source is legitimately available. Confusing them produces either a model that cannot see the source or one that cheats on its own output.

An RNN's timestep t cannot start until t − 1 finishes, so a 512-token sequence is 512 sequential operations however much hardware is available. Self-attention computes QKᵀ for all positions at once — one matrix multiplication — and every position's output is then a weighted sum over the same value tensor. There is no loop over the sequence during training. Two properties follow: the number of sequential steps per layer is 1 rather than n, and the path between any two positions is 1 rather than up to n, so gradients between distant tokens do not traverse the long chain that module 15's vanishing-gradient problem lived in.

This is the whole reason Transformers displaced recurrence at scale. It is not that attention represents language better in principle; it is that the same compute budget buys vastly more training throughput, and the architectures that could consume internet-scale corpora in reasonable wall-clock time were the ones that got large. Generation is the exception and remains sequential — one token at a time, whatever the architecture — which is why inference optimisation is a different problem from training throughput.

A next-token model must not let position i attend to any position after i, or it would have access to the token it is being asked to predict. The mask is applied to the scores *before* the softmax, by setting the disallowed entries to −inf: exp(−inf) = 0, so those positions get exactly zero weight and the remaining weights still sum to 1. The allowed pattern is lower-triangular including the diagonal, which is why the mask is built with torch.triu or torch.tril. Setting the weights to zero *after* the softmax is not equivalent — the row no longer sums to 1, and the gradient still flows through the masked scores.

This is the single most consequential mask in deep learning, and getting it wrong produces the most convincing-looking failure in the field: training loss drops to near zero because the model is reading the answer, and generation is incoherent because that information does not exist at inference. The loss curve looks like a triumph. It also explains a structural fact about training — one forward pass computes the loss for every position at once, which is only sound because the mask makes each position's prediction depend on its own prefix alone.

Self-attention has no notion of order. Permute the input tokens and the set of attention scores is the same set, permuted identically — so 'dog bites man' and 'man bites dog' produce the same representations, reordered. An RNN got order for free from the sequential update; attention must be told. The fix is to add a position-dependent vector to each token embedding before the first layer. The original sinusoidal scheme uses sines and cosines at geometrically spaced frequencies, which needs no parameters and extends to lengths never seen in training; learned absolute embeddings are a simple table lookup; and rotary embeddings (RoPE), now the common choice, rotate the queries and keys by a position-dependent angle so that the dot product depends on *relative* position.

It is not a refinement but a requirement — without it the model is a bag-of-tokens processor, and the failure is total rather than gradual. The scheme also decides how far the model extrapolates: absolute learned embeddings have nothing to look up past their table size, which is a large part of why extending a model's context window is more involved than raising a number in a config.

On a three-token toy sequence with d = 2, the whole operation is four small steps and worth doing once. Project each token to a query, key and value; compute every query-key dot product and divide by √d; softmax each row; and take the weighted sum of the value vectors per row. The point of doing it numerically is to see two things concretely: the output for a position is a blend of *other positions'* values, so the representation genuinely changes based on context, and the arithmetic is nothing but two matrix multiplications with a softmax between them.

Self-attention is usually met as a formula and remains abstract until a row has been computed by hand. It also makes the shapes unforgettable — (n × d) @ (d × n) → (n × n), then (n × n) @ (n × d) → (n × d) — which is the chain every shape bug in a Transformer is measured against.

A self-attention sublayer is never used bare. In practice it is wrapped: x = x + Attention(LayerNorm(x)), and the feedforward sublayer beside it is wrapped the same way. The residual addition is module 10's gradient highway, which is what makes a 96-layer stack trainable at all; the LayerNorm is module 11's scale control, chosen over BatchNorm because it works at batch size 1 and never couples examples. The placement matters: pre-norm — normalizing the sublayer's *input*, as written above — leaves the residual path completely clean and trains stably without careful warmup, which is why modern Transformers use it in preference to the original post-norm arrangement.

It explains why self-attention alone is not a layer you can stack. Without the residual, a deep stack hits exactly module 10's vanishing gradients; without the norm, the activation scale drifts as the residual stream accumulates contributions from every block. The wrapping is also the structural unit module 18 assembles, so recognising it here means the transformer block introduces nothing new.

The failure is worth stating as a scenario, because it is so easy to produce and so convincing while it happens. Training a next-token model with is_causal=False or with no attn_mask, every position attends to the whole sequence including the token it is predicting. Training loss falls to a fraction of what any honest model achieves; validation loss, computed the same unmasked way, falls with it. Nothing looks wrong until generation, where each step genuinely has no future context and the model — which learned to read the answer rather than to predict it — produces fluent-looking nonsense.

It is the clearest example of an evaluation that confirms a bug instead of catching it, because the leak is present in both splits. The lesson generalises past masks: when training and validation share a shortcut, the validation number cannot detect it, and only a test of the *inference procedure* — actually generating, or evaluating with the mask the model will run under — can.