Causal Attention Mask - Transformer Architecture

Causal (Autoregressive) Attention Mask

What you'll build

The additive mask that stops a decoder position from attending to later positions: 0 where attention is allowed, -inf where it is not.

Background

A language model is trained to predict token t+1 from tokens 0..t. Training runs all positions in parallel, so without a mask position t could simply read token t+1 from its input and the loss would be meaningless. Adding -inf to the forbidden scores makes exp(-inf) = 0 after the softmax, so those keys get exactly zero weight. GPT-style decoders, Mini GPT and the KV Cache problem all rely on this.

Formula

mask[i, j] = 0      if j <= i      (key j is at or before query i)
mask[i, j] = -inf   if j > i       (key j is in the future)

attention = softmax(Q @ K.T / sqrt(d_k) + mask) @ V

For seq_len = 4:

[[0, -inf, -inf, -inf],
 [0,    0, -inf, -inf],
 [0,    0,    0, -inf],
 [0,    0,    0,    0]]

Inputs

• seq_len: positive int

Expected Return Format

Float array (seq_len, seq_len) containing only 0.0 and -inf (float('-inf') / -np.inf).

Edge cases and conventions

• The diagonal is 0: every token attends to itself, so no row is fully masked and the softmax never sees an all--inf row (which would give NaN).
• Rows are queries, columns are keys.

Constraints

NumPy only, no loops (np.triu, np.tril or index comparisons).

Follow-up

• Some code uses a large negative number (e.g. -1e9) instead of -inf. When does that matter (hint: fp16)?
• How do you combine this with a padding mask for a batch of left-padded prompts?

Input: create_causal_mask(3)
Output: [[0.0, -inf, -inf], [0.0, 0.0, -inf], [0.0, 0.0, 0.0]]

Input: softmax(create_causal_mask(4))[1]
Output: [0.5, 0.5, 0.0, 0.0]

import numpy as np

def create_causal_mask(seq_len):
    """
    Create an additive causal attention mask.

    Args:
        seq_len: Sequence length

    Returns:
        mask: (seq_len, seq_len) float array,
              0.0 where key j <= query i, -inf where j > i
    """
    # Your code here
    pass