Causal Attention Mask

mediumFREEAttention MechanismsLLM 1: Foundations

Implement causal (autoregressive) masked attention.

In causal attention, each position can only attend to itself and previous positions. This is achieved by masking future positions with -infinity before softmax.

Given Q, K, V matrices, compute attention with a causal mask:

Compute scores = QK^T / sqrt(d_k)
Apply causal mask: set scores[i][j] = -inf where j > i
Apply softmax row-wise
Multiply by V

Input: causal_attention(Q, K, V) receives three NumPy arrays of shape $(n, d)$ (queries, keys, values), where $d_k = d$ .

Output: Return the masked attention output as an $(n, d)$ array. The tests print each value to 4 decimal places.

Example

Input

Output

[1.0000 0.0000]
[0.3302 0.6698]

Reasoning

Inputs: $n = 2$ tokens with $d_k = 2$ features each. $Q$ , $K$ and $V$ are all the $2 \times 2$ identity: token 0 is $[1, 0]$ , token 1 is $[0, 1]$ .
Scores: each query–key dot product divided by $\sqrt{2} \approx 1.4142$ . Matching tokens score $1/\sqrt{2} \approx 0.7071$ , different tokens score $0$ : $\begin{bmatrix}0.7071 & 0\\0 & 0.7071\end{bmatrix}$ .
Mask: a query at position $i$ may not look at a key at position $j > i$ . Only row 0, column 1 qualifies, so it becomes $-10^9$ .
Softmax per row: row 0 is $[0.7071, -10^9]$ ; $e^{-10^9}$ is effectively $0$ , so the weights are $[1, 0]$ . Row 1 is $[0, 0.7071]$ : with $e^{0.7071} \approx 2.0281$ , the weights are $\frac{1}{3.0281} \approx 0.3302$ and $\frac{2.0281}{3.0281} \approx 0.6698$ .
Output: a $2 \times 2$ matrix whose rows are those weights mixing the rows of $V$ . Since $V$ is the identity, the output equals the weights: [1.0000 0.0000], [0.3302 0.6698].
Notice that the first token can only attend to itself, so its row is always exactly its own value vector.

Constraints

1 <= n <= 10, 1 <= d <= 10
Use -1e9 as the mask value (approximating -inf)
Round to 4 decimal places

Learn this in

Transformer ArchitectureLLM Study Plan PretrainingLLM Study Plan CNNs & Sequence ModelsML Study Plan Machine TranslationNLP Study Plan

More from LLM 1: Foundations

Transformer Block Forward Passhard Feed-Forward Networkeasy Layer Normalizationeasy KV-Cache Attentionhard Multi-Head Attentionmedium

Concept primer

Background Knowledge

The problem involves implementing causal attention, a type of attention mechanism used in deep learning models, particularly in transformer architectures. Causal attention is also known as autoregressive attention, where each position in the input sequence can only attend to itself and previous positions. This is in contrast to regular attention, where each position can attend to all other positions. The key concept here is the use of a mask to restrict the attention mechanism from attending to future positions.

In the context of attention mechanisms, the query (Q), key (K), and value (V) matrices are used to compute the attention weights. The attention weights are computed by taking the dot product of Q and K, followed by a softmax operation. The causal mask is applied to the attention weights to prevent attending to future positions. The mask sets the attention weights to -infinity for positions where j > i, where i is the current position and j is the position being attended to.

The problem requires a good understanding of linear algebra and matrix operations, as well as the concept of softmax and its application in attention mechanisms. The softmax function is used to normalize the attention weights, ensuring that they sum up to 1. The dimensionality of the input matrices, particularly the embedding dimension d_k, plays a crucial role in computing the attention weights.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Compute the attention weights using the Q, K, and V matrices
Apply the causal mask to the attention weights
Normalize the attention weights using the softmax function
Compute the final output by multiplying the normalized attention weights with the V matrix

This approach involves a combination of matrix multiplications, element-wise operations, and softmax normalization.

Step-by-Step Strategy

To implement the solution, follow these steps:

Read the input matrices Q, K, and V, and store them in suitable data structures.
Compute the attention weights by taking the dot product of Q and K, and scaling by the square root of the embedding dimension d_k.
Apply the causal mask to the attention weights by setting the weights to -infinity where j > i.
Normalize the attention weights using the softmax function, applied row-wise.
Compute the final output by multiplying the normalized attention weights with the V matrix.

Common Pitfalls

When implementing the solution, watch out for the following:

Incorrectly applying the causal mask, which can lead to attending to future positions.
Forgetting to scale the attention weights by the square root of the embedding dimension d_k.
Applying the softmax function incorrectly, or not applying it row-wise.

Time & Space Complexity

The time complexity of the solution is expected to be O(n^2 * d_k), where n is the number of input positions and d_k is the embedding dimension. The space complexity is expected to be O(n^2), as we need to store the attention weights and the output matrix. However, the actual complexity may vary depending on the specific implementation and the libraries used.

solution.py

Python 3.13

Test results