Self-attention, cross-attention, and causal masking

Same mechanism, three wirings

Everything so far has been one operation: take a query, score it against some keys, spend a budget over the matching values. Nothing in that operation says where the query comes from, or where the keys and values come from. Change only the sources and you get every form of attention in a transformer — there is no second mechanism to learn.

Encoder self-attention reads all three from the same input sequence. The input is fully known, so every position may look at every other, forwards and backwards. This is what we traced by hand.

Decoder self-attention also reads all three from one sequence — but that sequence is the output being written, and the rest of it does not exist yet. Position 3 may look at positions 1, 2 and 3, and must not look at 4.

Cross-attention is the interesting one: the query comes from the decoder, the keys and values from the encoder's output. The decoder asks the question; the input sentence answers. This is the only place the two stacks touch, and it is what makes translation possible — the decoder writing the French can consult the English at every step.

Masking is one line, not a mode

The obvious way to stop a position seeing the future is to generate one token at a time and never show it more. That works, and it is also unbearably slow to train: a 1,000-token sequence would need 1,000 forward passes to learn from.

Instead the whole sequence goes through at once and the illegal scores are set to negative infinity before the softmax. Since $e^{-∞} = 0$, those positions receive exactly zero weight, and because softmax renormalises over what remains, every row still sums to 1 — position 2 spends its full budget across positions 1 and 2, not a partial budget with the rest leaked away.

Read the grid above as four separate readers. Position 1 has only itself, so it is forced to 1.000. Position 2 splits 0.198 / 0.802. By position 4 the budget is spread over all four. The upper triangle is not computed-then-discarded in a trained implementation — it is never usefully computed at all, which is a saving flash-attention kernels take seriously.

The subtle part is that this makes training honest at scale. Every position in a long sequence is simultaneously a training example predicting its own next token, all from one forward pass, and none of them can cheat by reading ahead. That single trick is why next-token prediction over a trillion tokens is affordable.

Self-attention, cross-attention, and causal masking are three configurations of the same underlying mechanism — softmax(QK^T / sqrt(d_k)) * V — that differ in where Q, K, and V come from and whether the score matrix is modified before softmax.

Self-attention

In self-attention, the queries, keys, and values are all derived from the same input sequence. Given a sequence of N token representations X of shape (N, d_model):

Q = X @ W_Q
K = X @ W_K
V = X @ W_V

Every token in the sequence attends to every other token in the same sequence, including itself. The attention matrix is (N, N) — position i computes a compatibility score with every position j, then takes a weighted sum of all value vectors to produce its output.

Self-attention is the mechanism used in both encoder models (BERT) and decoder models (GPT). In BERT, self-attention is bidirectional — every token can attend to every other token, regardless of position. Token 5 can attend to token 50, and token 50 can attend to token 5, with independent, learned weights. This is possible because BERT's training objective (masked language modeling) allows the model to see the full context.

In GPT-style models, self-attention is combined with a causal mask (covered below) that restricts which positions each token can attend to.

Self-attention as a computational graph

Self-attention turns a sequence into a fully-connected graph. Each token is a node. Each attention weight a[i][j] is a directed, weighted edge from node j to node i (information flows from j into the update of i). The edge weight is not static — it is computed dynamically from the content of the nodes via the query-key dot product. This makes the Transformer a form of message-passing neural network where the message function is content-dependent.

A single self-attention layer can propagate information between any two tokens in one step. Two layers can compose patterns: layer 1 might move information from token A to token B, and layer 2 can then use B's updated representation (which now contains A's information) to inform token C. This composition across layers is why deep Transformers can implement complex reasoning chains — each layer adds a hop in the information routing.

Cross-attention

In cross-attention, the queries come from one sequence and the keys and values come from a different sequence. This enables one representation to read from another.

The encoder-decoder Transformer (used in the original Vaswani et al. 2017 architecture, T5, and BART) has three types of attention:

• Encoder self-attention — the source sequence attends to itself. Q, K, V all from the encoder.
• Decoder self-attention — the target sequence attends to itself (with causal masking). Q, K, V all from the decoder.
• Encoder-decoder cross-attention — the decoder attends to the encoder's output. Q from the decoder, K and V from the encoder.

import torch
import torch.nn as nn
import torch.nn.functional as F

class CrossAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x_query, x_context):
        """
        x_query:   (B, N_q, d_model) — the sequence that asks
        x_context: (B, N_kv, d_model) — the sequence that answers
        """
        B, N_q, d_model = x_query.shape
        N_kv = x_context.shape[1]

        Q = self.W_Q(x_query).view(B, N_q, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_K(x_context).view(B, N_kv, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_V(x_context).view(B, N_kv, self.n_heads, self.d_head).transpose(1, 2)

        # Attention: (B, H, N_q, N_kv) — rectangular, not square
        scores = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5)
        weights = F.softmax(scores, dim=-1)
        out = weights @ V  # (B, H, N_q, d_head)

        out = out.transpose(1, 2).contiguous().view(B, N_q, d_model)
        return self.W_O(out)


d_model = 512
n_heads = 8

cross_attn = CrossAttention(d_model, n_heads)

decoder_states = torch.randn(1, 15, d_model)   # 15-token target
encoder_output = torch.randn(1, 30, d_model)   # 30-token source

output = cross_attn(decoder_states, encoder_output)
print(f"Decoder query shape:  {decoder_states.shape}")  # (1, 15, 512)
print(f"Encoder context shape: {encoder_output.shape}")  # (1, 30, 512)
print(f"Cross-attention output: {output.shape}")         # (1, 15, 512)

The attention matrix in cross-attention is rectangular: (N_query, N_context). Each decoder position distributes its attention across all encoder positions, but encoder positions do not attend to each other (that was handled by the encoder's self-attention layers). The output sequence has the same length as the query sequence (N_query), not the context sequence.

Cross-attention in multimodal models

Cross-attention is the standard mechanism for connecting modalities. In Flamingo (Alayrac et al. 2022), text tokens generate queries that attend to visual features (keys and values extracted from image patches via a vision encoder). The text stream can selectively read from the image — the query "what color is the car?" produces attention weights that concentrate on the image region containing the car.

The same pattern appears in text-to-image generation (Stable Diffusion uses cross-attention from image latents to text embeddings), speech recognition (Whisper uses cross-attention from decoder text tokens to audio encoder features), and retrieval-augmented generation (some architectures use cross-attention from the language model to retrieved document representations, rather than concatenating them into the context).

Causal masking

Autoregressive language models generate text left to right: at each step, the model predicts the next token given all previous tokens. During generation, this is naturally enforced — token 5 has not been generated yet when the model is predicting token 4. But during training, the model sees the entire sequence at once (teacher forcing). The model processes "The cat sat on the mat" as a single forward pass, computing all positions simultaneously. Without intervention, the self-attention at position 3 ("sat") would attend to positions 4, 5, 6 — future tokens that it should not have access to.

Causal masking enforces the autoregressive constraint during training. Before the softmax, the upper-triangular entries of the score matrix are set to negative infinity:

import torch
import torch.nn.functional as F
import numpy as np

def causal_scaled_dot_product_attention(Q, K, V):
    """Scaled dot-product attention with causal mask."""
    N = Q.shape[0]
    d_k = K.shape[-1]

    scores = Q @ K.T / np.sqrt(d_k)  # (N, N)

    # Causal mask: position i can attend to positions 0..i only
    mask = np.triu(np.ones((N, N)), k=1)  # upper triangle (above diagonal)
    scores = np.where(mask == 1, -1e9, scores)

    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = weights / weights.sum(axis=-1, keepdims=True)

    output = weights @ V
    return output, weights


np.random.seed(42)
N, d = 6, 64
Q = np.random.randn(N, d) * 0.1
K = np.random.randn(N, d) * 0.1
V = np.random.randn(N, d) * 0.1

output, weights = causal_scaled_dot_product_attention(Q, K, V)

print("Attention weights after causal masking:")
print("(Each row can only attend to positions up to and including itself)\n")
for i in range(N):
    row = " ".join(f"{w:.3f}" for w in weights[i])
    print(f"  token {i}: [{row}]")

The output shows the triangular structure clearly. Token 0 puts all its weight on itself (it has nothing else to attend to). Token 1 distributes weight between positions 0 and 1. Token 5 can attend to all six positions. Positions in the upper triangle have weight exactly 0.000.

The mask is deterministic and fixed — a static lower-triangular matrix of ones. It has no learnable parameters. In practice, it is precomputed once and broadcast across all heads and all layers.

Why masking enables parallel training

Without causal masking, autoregressive training would require sequential computation: compute token 0's representation, then use that to compute token 1's, and so on — exactly the RNN bottleneck that attention was designed to eliminate. The mask allows the model to process all positions in parallel while ensuring that each position's output depends only on positions to its left. The loss at position i is cross_entropy(predicted[i], actual[i+1]), and the gradient for that loss flows only through positions 0 through i — the mask guarantees this.

This is what makes Transformer training vastly more efficient than RNN training. A GPT-3-scale model processes a 2048-token sequence in a single forward pass, computing 2048 next-token predictions simultaneously. An equivalent RNN would require 2048 sequential steps per sequence.

Bidirectional versus causal attention

The distinction between bidirectional and causal attention defines two families of models:

Bidirectional (no mask): BERT, RoBERTa, ELECTRA, DeBERTa. Every token attends to every other token. This produces contextualized representations where each token's output incorporates information from the full sequence. Bidirectional models excel at understanding tasks — classification, named entity recognition, extractive QA, semantic similarity — because they can use both left and right context.

Bidirectional models cannot generate text autoregressively. Because every position's representation depends on every other position (including future tokens), you cannot compute a partial sequence and extend it. BERT predicts masked tokens — given "The [MASK] sat on the mat," it predicts the masked position using both left and right context — but this is a fill-in-the-blank operation, not sequential generation.

Causal (lower-triangular mask): GPT, Llama, Mistral, Claude, Gemini. Position i attends only to positions 0 through i. This enables autoregressive generation: at inference time, the model computes each new token using only the tokens that precede it, which is exactly the pattern the causal mask enforces during training.

The tradeoff is that causal models see less context per position. Position 10 in a causal model uses only 11 tokens of context (positions 0–10), while the same position in a bidirectional model uses the full sequence. This is why BERT outperforms GPT-2 on many NLU benchmarks despite having fewer parameters — bidirectional context is strictly more informative for understanding tasks. But causal models can generate coherent text, and at sufficient scale (GPT-3 and beyond), they match or exceed bidirectional models even on understanding tasks by leveraging the massive increase in training data and parameters.

Prefix attention — a hybrid

Some models use a hybrid: bidirectional attention over a prefix (the input/prompt), then causal attention for the generated tokens. T5 and UL2 use this for encoder-decoder tasks. The prefix sees full bidirectional context, giving the model maximum understanding of the input. The generated suffix uses causal attention, enabling autoregressive generation. This combines the strengths of both approaches at the cost of a more complex attention mask:

def prefix_causal_mask(prefix_len, total_len):
    """
    Prefix tokens attend to all prefix tokens (bidirectional).
    Generated tokens attend to all prefix tokens + previous generated tokens (causal).
    """
    mask = np.zeros((total_len, total_len))

    # Prefix: fully bidirectional
    mask[:prefix_len, :prefix_len] = 1

    # Generated: attend to all prefix + causal within generated
    for i in range(prefix_len, total_len):
        mask[i, :prefix_len] = 1        # can see all prefix tokens
        mask[i, prefix_len:i + 1] = 1   # can see previous generated tokens

    return mask

mask = prefix_causal_mask(prefix_len=4, total_len=8)
print("Prefix-causal mask (4 prefix tokens, 4 generated tokens):")
for i in range(8):
    row = " ".join(f"{int(v)}" for v in mask[i])
    label = "prefix " if i < 4 else "gen    "
    print(f"  {label}{i}: [{row}]")