Who is allowed to see what?

Attention tells us how to mix a set of values. It does not by itself decide which values a query should be permitted to use. That decision depends on the task.

A model classifying a complete sentence can read both ends of it. A model predicting the next output token cannot consult the answer it is being asked to predict. A model translating a source may read all of that source while seeing only the existing target prefix. We will express each case with two questions: where do the vectors come from, and which query–key pairs are visible?

1. Self-attention uses one sequence in three roles

In self-attention, the same sequence of hidden states supplies the inputs to the Q, K, and V projections. The projections remain different learned matrices. “Self” describes the source sequence, not equality between Q and K and V.

With an input of length S, the score matrix is (S, S). In bidirectional self-attention, every non-padding source position is available to every query. A representation at the first position can incorporate information from the last position in one attention layer.

That is appropriate when the whole sequence is available for the task. It becomes a problem if those later tokens are the targets we are pretending to predict.

2. A causal mask prevents access to future positions

Number target-prefix positions from zero. A query at position i can use keys j with \(j\leq i\). It sees its current input token and all previous ones, but not later inputs.

The mask is applied before softmax. Exponentiating a negative-infinite masked logit contributes zero. The permitted entries are still normalized together. Zeroing entries after softmax without renormalizing would produce a different, generally non-normalized operation.

For a target [4, 7, 9, EOS], the decoder receives [BOS, 4, 7, 9]. Its query at position zero can see BOS and predicts 4. At position one, it sees BOS, 4 and predicts 7.

We need both the target shift and the mask. A causal mask does not help if we feed the desired answer at the same position and train the model to reproduce it there.

3. Parallel training, sequential generation

During training, the correct target sequence is known. We can create all shifted inputs and compute all target-position losses in parallel, provided the causal mask prevents future leakage. This use of correct previous tokens is called teacher forcing.

During generation, later output tokens are unknown. We must predict one, append it to the prefix, and run the next step. The difference is availability of the target sequence, not a change from non-causal training to causal inference.

Our code compares full-prefix and incremental-prefix predictions. It also changes a future token and checks that earlier logits remain unchanged. Those are stronger checks than merely asserting a triangular mask has the expected shape.

4. Padding is a separate visibility problem

Batches contain sequences of different lengths. Padding fills unused positions so tensors can be rectangular. Those placeholders should not contribute source information.

A source-key mask of shape (B, 1, 1, S) broadcasts across heads and queries. In target self-attention, a target-key mask combines with the causal condition to produce (B, 1, T, T).

def attention_masks(source_ids: IntArray, target_ids: IntArray, pad: int = 0) -> tuple[BoolArray, BoolArray]:
    """IDs: (batch, length). Return source-key and causal-target masks."""
    source_visible = (source_ids != pad)[:, None, None, :]
    target_visible = (target_ids != pad)[:, None, None, :]
    length = target_ids.shape[1]
    causal = np.arange(length)[None, :] <= np.arange(length)[:, None]
    target_visible = target_visible & causal[None, None, :, :]
    return source_visible, target_visible

def attention_masks(source_ids: Tensor, target_ids: Tensor, pad: int = 0) -> tuple[Tensor, Tensor]:
    source_visible = (source_ids != pad)[:, None, None, :]
    target_visible = (target_ids != pad)[:, None, None, :]
    length = target_ids.shape[1]
    index = torch.arange(length, device=target_ids.device)
    causal = index[None, :] <= index[:, None]
    target_visible = target_visible & causal[None, None, :, :]
    return source_visible, target_visible

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

Padding keys are masked out. Padding query positions can still produce numerical outputs in our implementation, but those outputs have no task meaning. We ignore their training losses. Masking every key for a padding query would create an all-masked softmax row; our convention avoids that undefined operation while still excluding padding from the objective.

Adding positional vectors can make a padding representation nonzero even when its token embedding is zero. Masking is necessary; a zero PAD lookup row alone does not solve the problem.

5. Cross-attention reads a different sequence

In an encoder–decoder model, the encoder reads a source sequence and produces memory states. The decoder constructs target-side states from the prefix it has seen. Cross-attention lets those target states retrieve source information:

Y has target length T and the source memory has length S. Scores have shape (T, S) and output states retain target length T. Source and target do not need equal lengths.

In ordinary offline translation, all source positions are available. Cross-attention therefore excludes source padding without imposing a target-style triangle on the source. Simultaneous or streaming tasks can require different source-availability rules. The task determines visibility.

6. One convention, with an API trap

Every custom function in this course uses True = visible. This matches the boolean mask convention of PyTorch’s scaled_dot_product_attention ↗ .

Boolean masks passed to nn.MultiheadAttention ↗ , however, use True = masked out. Moving a boolean mask between these APIs without checking the convention can invert the information flow. The shape can be correct while the computation is wrong.

Our explicit function returns weights for inspection and rejects all-masked rows. A fused kernel may have its own edge-case behavior. Read an API’s mask semantics before swapping implementations.

Carry this forward

We have one attention operation, multiple choices of input stream, and explicit visibility masks. Next we will run several learned attention projections in parallel. Multiple heads do not change the visibility rules; they provide different learned views under those rules.