DAI developerLab· about 60 min· server sandbox

Lab: scaled dot-product attention with a causal mask

Implement softmax, attention(Q, K, V) and a causal mask in NumPy, and verify the properties: weights sum to 1, no token looks ahead.

Teaches: Attention

Requires: Matrices and matrix multiplicationEmbeddings — words as vectors

Theory

Attention(Q,K,V) = softmax(QKᵀ/√d)·V. Causal mask: set the scores for future positions to −∞ before softmax.

Sub-tasks

softmax (row-wise, stable) — softmax(z) along the last axis, numerically stable (subtract the max).
attention — attention(Q, K, V, mask=None) returns (output, weights).
causal mask — causal_mask(n) returns an n×n boolean matrix where True = may be seen.

Passes when: max_abs_diff <= 0.000001

The starter code

runs in an isolated sandbox on the server

import numpy as np


def softmax(z, axis=-1):
    # TODO: stabil softmax
    ...


def causal_mask(n):
    # TODO: True där j <= i (får ses), False ovanför diagonalen
    ...


def attention(Q, K, V, mask=None):
    """Q, K: (n, d); V: (n, dv). mask: (n, n) bool, True = tillåten. Returnerar (out, weights)."""
    # TODO: scores = Q @ K.T / sqrt(d); maska med -inf; softmax; out = weights @ V
    ...

You write the code; tests you cannot see decide whether it holds up. Create a free account to run the lab.

Try the diagnostic Create a free account

Expected results

The weights sum to 1 row-wise; with a causal mask the weights above the diagonal are exactly 0; difference from the reference < 1e-6.

Common mistakes

softmax without subtracting the max → overflow for large scores.
Masks after softmax instead of before (then the rows do not sum to 1).
Forgets √d.