Attention, masks, and next-token prediction

An interactive field guide

A Transformer mixes information between positions. A mask decides which information is available. It does not decide whether the values are discrete or continuous.

Before this article: Probability and tensors

What was already possible

A convolution uses a fixed local neighborhood. Attention computes a content-dependent weighted sum. Each position projects its features into a query Q, a key K, and a value V. Dot products between queries and keys produce attention scores; softmax converts allowed scores to weights.

What changes, and why

For autoregressive training, position i must not see the target it predicts. In a shifted decoder, token i can attend through i and predicts token i+1. A prefix can be fully visible if it is supplied context and excluded from target prediction. Bidirectional attention is also valid over known tokens in masked generation.

Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.

Project Q, K, V

Split feature channels into heads, not sequence positions.

q,k,v = project(x).reshape(B,N,3,heads,d).unbind(2)

Mask scores

Blocked positions must have zero attention after softmax.

scores = q @ transpose(k) / sqrt(d)
scores = mask_blocked_with_minus_inf(scores)

Mix values

A residual connection preserves the old state alongside the new mixture.

weights = softmax(scores, dim=-1)
y = x + output_projection(weights @ v)

Inspect the pinned source

\operatorname{Attention}(Q,K,V)=\operatorname{softmax}(QK^\top/\sqrt{d}+M)V

Work through the mechanism

Select a query row in the matrix. Causal mode allows only its past and present columns. Prefix mode uses three known prefix tokens. Masked mode observes even positions and hides odd ones. The matrix shows permission, not learned attention weight; allowed weights still depend on Q and K.

Change a parameter, inspect the computed values, then compare this small example with the model described above.

Analytical teaching model · controls change the example, not a trained research model.

Query position i 30 causal · 1 prefix · 2 masked 0

Rows are queries; columns are keys. In masked mode, only even-indexed positions contain observed values.

Computed values for the current example
Quantity	Value
Selected row	3
Visible keys	4

Training is not sampling

Training can evaluate all shifted targets in one masked forward pass because the true history is available. Sampling has to produce the history first. A causal KV cache saves old keys and values. With bidirectional context, adding a token can change earlier states, so that cache is not directly reusable.

What this does not establish

Position encoding tells attention about location; a causal mask tells it about access. Neither replaces the other. One can build a continuous causal model, a discrete masked model, or a bidirectional image encoder.

Open the derivation

Q and K have width d per head. Their dot products give scores for each query/key pair. The additive mask uses zero for allowed pairs and negative infinity for forbidden pairs. Softmax is applied over keys.

\operatorname{Attention}(Q,K,V)=\operatorname{softmax}(QK^T/\sqrt d+M)V

Check your understanding

Is a row of visible cells a probability distribution?

No. It is an access mask. Softmax over content-dependent allowed scores produces attention weights.

Sources and next steps

Return to the learning map

Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.