Neural Mastery
← Practice
Vectorized Causal Triangular Attention Maskinghard

Vectorized Causal Triangular Attention Masking

hard NumPy & Data Wrangling 20-25 min

NumPy allowed · plain Python earns 10 bonus points

Attention step 1/4

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

Prevent autoregressive models from attending to future sequence tokens via upper-triangular masking.

Related Skills

Task

Implement apply_causal_mask(scores, mask_val): for a square attention score matrix scores of shape (S, S), replace all upper-triangular elements scores[i][j] where j > i with mask_val. Keep elements where j <= i intact. Round all elements to 4 decimals.

Function Signature

apply_causal_mask(scores: list[list[float]], mask_val: float = -10000.0) -> list[list[float]]

Examples

Example 1: 3x3 Attention Scores
Input: scores = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] mask_val = -10000
Output: [[1, -10000, -10000], [4, 5, -10000], [7, 8, 9]]
Explanation: Mask upper triangular future positions.
Example 2: 2x2 Minimal Sequence
Input: scores = [[0.5, 0.8], [0.2, 0.4]] mask_val = -1000000000
Output: [[0.5, -1000000000], [0.2, 0.4]]
Explanation: Causal mask on 2-token sequence.

Constraints

  • 1 <= S <= 2048
  • mask_val is large negative float (e.g. -1e4)
Python3Saved ✓
def apply_causal_mask(scores: list[list[float]], mask_val: float = -10000.0) -> list[list[float]]:
"""Mask upper triangular future positions (j > i) with large negative mask_val.
Round output matrix entries to 4 decimals."""
# Implement solution logic here
pass

Case 1: 3x3 Attention Scores
Input: scores = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] mask_val = -10000
Expected: [[1, -10000, -10000], [4, 5, -10000], [7, 8, 9]]
Case 2: 2x2 Minimal Sequence
Input: scores = [[0.5, 0.8], [0.2, 0.4]] mask_val = -1000000000
Expected: [[0.5, -1000000000], [0.2, 0.4]]