Construct Causal Attention Maskeasy

Construct Causal Attention Mask

easy Inside an LLM 15 min

NumPy allowed · plain Python earns 10 bonus points

Tokens attend to tokens

Q · Kᵀ / √d → softmax → weighted values, computed in the page.

Open lesson →

01tokens

02Q K V

03scores

04attention

modelslearnrelationships

models32%

learn38%

relationships30%

Q · Kᵀ / √d scores normalized with softmax

Build a causal (autoregressive) attention mask matrix preventing query positions from attending to future key positions during decoding.

Task

Implement causal_attention_mask(seq_len). Return a 2D integer list of size (seq_len, seq_len) where element (i, j) is 1 if j <= i (visible) and 0 if j > i (masked future token).

Function Signature

causal_attention_mask(seq_len: int) -> list[list[int]]

Examples

Example 1: Single Token Sequence
Input: seq_len = 1
Output: [[1]]

Example 2: Three Token Sequence
Input: seq_len = 3
Output: [[1, 0, 0], [1, 1, 0], [1, 1, 1]]

Constraints

seq_len must be a positive integer (`>` 0).
Must raise ValueError if seq_len `<=` 0.
Output matrix must have dimensions exactly seq_len x seq_len.
Diagonal and lower triangle are 1; strictly upper triangle is 0.

Python3Saved ✓

def causal_attention_mask(seq_len):

"""Construct a lower-triangular causal attention mask of shape

(seq_len, seq_len) where mask[i][j] == 1 if j <= i else 0.

Raise ValueError if seq_len <= 0."""

# Your implementation here

pass

Case 1: Single Token Sequence

Input: seq_len = 1

Expected: [[1]]

Case 2: Three Token Sequence

Input: seq_len = 3

Expected: [[1, 0, 0], [1, 1, 0], [1, 1, 1]]