Masked Self-Attention

Recall that standard Scaled Dot-Product Attention computes, for query, key, and value matrices $Q, K, V \in \mathbb{R}^{n \times d}$,

Here $n$ is the sequence length and $d$ the (per-head) hidden dimension. The matrix $QK^\top / \sqrt{d} \in \mathbb{R}^{n \times n}$ contains a raw compatibility score between every pair of positions $(i, j)$: row $i$ says "how much should position $i$ attend to position $j$." In an unmasked (bidirectional) self-attention layer, this is fine — every position may look at every other position, which is exactly what you want for tasks like masked language modeling (BERT-style) or general sequence encoding, where the whole input is visible at once.

Causal masking modifies this computation by forbidding position $i$ from attending to any position $j > i$. Formally, define a mask matrix $M \in \{0, -\infty\}^{n \times n}$ with

and compute

Adding $-\infty$ to a score before the softmax sends its exponential to $0$, so after normalization the attention weight assigned to any "future" position is exactly zero. Row $i$ of the resulting attention-weight matrix is a proper probability distribution supported only on $\{1, \dots, i\}$ — a lower-triangular structure. This is why the mask is often visualized and implemented as a lower-triangular boolean matrix. The terminology worth fixing precisely: a mask that only depends on relative position (like this triangular one) is a structural mask, distinct from a padding mask, which blocks attention to padding tokens inserted to make sequences in a batch equal length. In practice, both are combined by addition (or logical OR of "forbidden" flags) before the softmax. The causal mask alone is what defines autoregressive or decoder-style self-attention, as opposed to the bidirectional self-attention used in encoders.

The mechanism sits inside the same attention computation you already know, with one extra step: build the mask, then add it to the scaled scores before the softmax. Everything else — projecting inputs to $Q, K, V$, scaling by $\sqrt{d}$, applying softmax, multiplying by $V$ — is unchanged from ordinary scaled dot-product attention. There are two regimes worth understanding separately, because they explain why this design is both efficient to train and correct to deploy: 1. Training-time (teacher forcing), full-sequence, parallel form. The entire target sequence (or the entire input sequence in a decoder-only model) is available at once. We compute all $n^2$ pairwise scores in one matrix multiplication, then mask out the upper triangle in a single elementwise operation, then run one softmax. This is fully parallel across all positions — no sequential loop over $i$ is needed at training time, which is the key efficiency win over classical recurrent (RNN) autoregressive models. 2. Inference-time, incremental (autoregressive decoding) form. At generation time, tokens are produced one at a time. When generating token $t+1$, only tokens $1, \dots, t$ exist. The causal mask is trivially satisfied because there simply are no future keys/values yet to attend to — but the training-time mask must have matched this exact conditional structure, otherwise the model would have learned to rely on information that will not exist at generation time. This is precisely the training/inference consistency argument for why the mask must be triangular and not, say, banded or random. ```

# Masked (causal) self-attention — full parallel form (training / prefill) #

# Inputs:

# X : (n, d_model) -- input sequence embeddings (n tokens)

# W_Q, W_K, W_V : (d_model, d_k) -- learned projection matrices

# Output:

# O : (n, d_k) -- context-aware output representations function MaskedSelfAttention(X, W_Q, W_K, W_V): n = number_of_rows(X) d_k = number_of_columns(W_K) Q = X @ W_Q # (n, d_k) K = X @ W_K # (n, d_k) V = X @ W_V # (n, d_k) scores = (Q @ transpose(K)) / sqrt(d_k) # (n, n) raw compatibility

# Build causal mask: forbid j > i (attending to the future) for i in 1..n: for j in 1..n: if j > i: scores[i][j] = -infinity

# else: leave scores[i][j] unchanged (j <= i is allowed) weights = row_wise_softmax(scores) # (n, n), each row sums to 1

# row i has zero weight for j > i O = weights @ V # (n, d_k) return O ``` ```

# Incremental (autoregressive) decoding form — used at generation time #

# Maintains a growing cache of keys/values; each new token only ever

# computes attention over already-generated tokens (mask is automatic). function IncrementalMaskedAttention(x_new, W_Q, W_K, W_V, K_cache, V_cache): q = x_new @ W_Q # (1, d_k), the new token's query k = x_new @ W_K # (1, d_k) v = x_new @ W_V # (1, d_k) K_cache = append(K_cache, k) # grows by one row each step V_cache = append(V_cache, v) scores = (q @ transpose(K_cache)) / sqrt(d_k) # (1, t) — t = current length

# No masking loop needed: K_cache contains only positions 1..t (t = current position) weights = softmax(scores) # (1, t) o = weights @ V_cache # (1, d_k) return o, K_cache, V_cache ``` Two implementation notes that matter in real systems. First, in a real Transformer, this whole block is run once per attention head, and results from multiple heads are concatenated and linearly projected — the mask is identical across heads and is usually computed once per layer (or once per forward pass, since it does not depend on the data, only on sequence length). Second, in practice $-\infty$ is realized as a very large negative finite number (e.g. $-10^9$) to avoid `NaN` values arising from $-\infty - (-\infty)$ during numerically-stabilized softmax computation.

Time complexity. Computing $QK^\top$ costs $O(n^2 d)$ scalar multiply-adds (an $n \times d$ matrix times a $d \times n$ matrix). Applying the mask is an elementwise operation on an $n \times n$ matrix, costing $O(n^2)$ — asymptotically dominated by the score computation. The softmax is likewise $O(n^2)$, and the final multiplication $\text{weights} \cdot V$ costs another $O(n^2 d)$. So the total time complexity per masked self-attention layer is

identical to unmasked self-attention — masking adds no asymptotic overhead; it only zeroes out entries that were already going to be computed. This is an important honesty point: causal masking does not save compute relative to bidirectional attention in the naive dense implementation, since we still materialize the full $n \times n$ score matrix and simply discard (mask) half of it. (Specialized implementations can skip computing the strictly-upper-triangular half, roughly halving the constant factor, but the asymptotic order remains $\Theta(n^2 d)$.)

Space complexity. The dense attention-weight matrix requires $O(n^2)$ memory per head (plus $O(nd)$ for $Q, K, V$ and outputs). This is the well-known quadratic memory bottleneck of Transformer self-attention, and it applies to masked attention exactly as it does to unmasked attention — masking changes values inside the matrix, not its shape. During incremental generation, however, a key-value cache stores only the $O(td)$ keys and values generated so far at step $t$, avoiding recomputation of earlier positions' representations; this is a major practical optimization enabled precisely by the causal structure (past keys/values never need to be revisited or recomputed because future tokens cannot influence them).

Correctness argument. The claim to verify is: training with the causal mask produces a model whose predicted distribution for token $t+1$ depends only on tokens $1, \dots, t$, matching exactly what is available at generation time. Consider the computational dependency graph of the masked attention layer. Row $i$ of the output, $O_i$, is a function only of $\{Q_i\}$ and $\{K_j, V_j : j \le i\}$, because $\text{weights}_{ij} = 0$ for all $j > i$ by construction (softmax of $-\infty$ is exactly $0$, not merely small). Since $Q_i, K_j, V_j$ are themselves computed from $X_i, X_j$ respectively via the same linear projections applied position-wise, output $O_i$ is a function only of $X_1, \dots, X_i$. Stacking $L$ such masked layers preserves this property by induction on depth: if every layer's output at position $i$ depends only on inputs at positions $\le i$, then layer $\ell+1$'s output at position $i$ — being a function of layer $\ell$'s outputs at positions $\le i$ — also depends only on inputs at positions $\le i$. (Feedforward sublayers and layer normalization are applied position-wise and do not break this invariant.) This gives an exact conditional-independence guarantee: the model's final hidden state at position $i$, and hence its predicted next-token distribution $P(x_{i+1} \mid x_1, \dots, x_i)$, is a genuine function of only the causal prefix. This is exactly the factorization used in the chain rule of probability for autoregressive language modeling,

so training a masked Transformer with cross-entropy loss on next-token prediction is a consistent way to fit this factorization — the model literally cannot cheat by looking ahead, because the mask makes the forward computation graph itself acyclic and causally ordered.

Tradeoffs. Masking is cheap to implement and adds no fundamentally new complexity class, but it does not reduce the $O(n^2)$ memory and compute footprint of full self-attention — for long sequences this remains the dominant bottleneck, motivating sparse, linear, or blockwise-recurrent attention variants (a separate topic from masking itself). Masked attention also gives up the bidirectional context that helps encoder-style models (like those used for classification or masked-language-modeling pretraining); this is a genuine representational tradeoff, not a bug — it is the price paid for making the model usable for open-ended generation.

Consider a toy sequence of 4 tokens: `"the cat sat down"`, with $d_k = 2$ for simplicity. Suppose the (already projected) query and key vectors are:

Raw dot products $q_i \cdot k_j$ (before scaling by $\sqrt{d_k} = \sqrt{2} \approx 1.414$), for the row $i=3$ ("sat") against all four keys:

Scaled: $1/1.414 \approx 0.707$ for $j=1,2,3$ and $2/1.414 \approx 1.414$ for $j = 4$.

Without masking, row 3's softmax input is $(0.707, 0.707, 0.707, 1.414)$. Computing softmax:

Sum $\approx 2.03+2.03+2.03+4.11 = 10.2$, giving weights approximately $(0.199, 0.199, 0.199, 0.403)$ — note that "sat" would attend most strongly (weight $0.403$) to "down", a word that comes after it. That is precisely the information leak we must prevent.

With the causal mask, position $j=4$ is forbidden for query $i = 3$, since $4 > 3$. The mask sets that score to $-\infty$:

Softmax now only normalizes over $j = 1, 2, 3$:

Sum $= 2.03 \times 3 = 6.09$, giving weights

Position 3 ("sat") now attends equally to "the", "cat", and itself, and exactly zero weight to "down" — the future token is fully excluded, not merely down-weighted. This row is precisely what will be available when the model must, at generation time, predict what comes after "the cat sat" without yet knowing the word "down" exists.

This table exhibits the triangular support pattern directly: row $i$'s softmax is supported exactly on $\{1, \dots, i\}$, matching the earlier correctness argument concretely rather than abstractly.

Masked self-attention was introduced as part of the original Transformer architecture, designed to replace recurrent and convolutional sequence models with a fully attention-based design. The insight that a single triangular mask could convert a bidirectional attention mechanism into an autoregressive one — while preserving full parallelism across the sequence during training — was central to making self-attention practical for sequence generation tasks like machine translation, and later the foundation for large-scale generative language modeling.

Masked self-attention is the computational core of every GPT-family model and, more broadly, every decoder-only large language model used for open-ended text generation, chat assistants, and code completion. It is also used inside the decoder half of encoder-decoder ("sequence-to-sequence") Transformer models used for machine translation and summarization, where the decoder must generate output tokens autoregressively while separately attending to the (unmasked) encoder output. Production inference systems build entire serving stacks — key-value caching, speculative decoding, continuous batching — around the specific dependency structure that causal masking guarantees: that past keys and values are frozen and reusable once computed, since no future token can ever alter them. Any system claiming to do streaming or incremental text generation with a Transformer — search engine query completion, code autocompletion, real-time chat — relies on this exact guarantee at its core. Understanding masked self-attention is the direct prerequisite for understanding the full Transformer Decoder Stack and how it powers GPT & Autoregressive Language Modeling.

Ask the tutor: How does causal masking interact with multi-head attention and key-value caching when generating text one token at a time?