Preporato

The additive causal mask

Easy8 min10 hidden tests

A code-completion model predicts each token from the tokens before it, so attention must never look ahead. Write causal_mask(n, past=0) that returns the additive mask the model adds to its attention scores before the softmax: 0.0 where a query may look at a key and -inf where it may not. The n queries are the newest tokens, and past tokens already sit in a cache before them, so the keys are all past + n tokens. Query i is token past + i and may see keys 0 through past + i, itself included. For n = 3 and no past, the rows are [0, -inf, -inf], [0, 0, -inf] and [0, 0, 0].

Tap a row: what can that token see?

shape (n, past + n) = (3, 3)

Row 2 is the token w, at position 2. It attends to keys 0 to 2: y = w. Nothing comes after it, so the whole row is 0.

Row 2, token w, sees 3 of 3 keys.

Examples

Each row sees one more token than the row above

causal_mask(3)

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

Two new tokens after three cached ones

causal_mask(2, 3)

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

Requirements

  • n is an integer of 1 or more and past an integer of 0 or more.
  • Return a float NumPy array of shape (n, past + n).
  • Entry (i, j) is 0.0 when j <= past + i and -inf (-np.inf) when j > past + i.
The idea behind it

A decoder-only transformer such as GPT is trained to predict every next token at once. To keep that honest, each position may attend only to itself and earlier positions. The usual way is an additive mask: add 0 to the scores the query may use and -inf to the rest, then take the softmax. Since exp(-inf) = 0, blocked keys get a weight of exactly 0, and the other weights still sum to 1. For one sequence the mask is square and upper triangular. With a key-value cache (stored keys and values of earlier tokens), new queries arrive after past cached tokens, so the mask has past + n columns and the boundary shifts right by past. Implementations that work in half precision often use a large negative number in place of -inf, but the idea is the same.

pythonCtrl+Enter runs
import numpy as np

def causal_mask(n, past=0):
"""Additive causal attention mask.

n: number of new query tokens (>= 1)
past: number of earlier tokens already in the cache (>= 0)
Returns an (n, past + n) float array: 0.0 where key j <= past + i,
and -inf where key j is in the future of query i.
"""
raise NotImplementedError

Stuck? Reveal one rung at a time

Worked solution

Give it a real try first. Then read the solution and why each line is there.