📘

Causal Attention Mask

MediumAttention Mechanisms 📚 LLM 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:

Line 1: n d
Next n lines: Q matrix
Next n lines: K matrix
Next n lines: V matrix

Output: The masked attention output, values rounded to 4 decimal places.

Example:

Input:

Output:

[1.0000 0.0000]
[0.2689 0.7311]

Reasoning:

We start by computing the scores using the given Q and K matrices: $scores = QK^T / \sqrt{d_k} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}^T / \sqrt{2} = \begin{bmatrix} 1/\sqrt{2} & 0 \\ 0 & 1/\sqrt{2} \end{bmatrix} / \sqrt{2} = \begin{bmatrix} 1/2 & 0 \\ 0 & 1/2 \end{bmatrix}$
Then, we apply the causal mask to the scores: since $j > i$ only when considering the second row and first column, we set $scores[1][0] = -\inf$ , resulting in $scores = \begin{bmatrix} 1/2 & 0 \\ -\inf & 1/2 \end{bmatrix}$
Next, we apply softmax row-wise to the masked scores: for the first row, $softmax([1/2, 0]) = [1, 0]$ ; for the second row, $softmax([-inf, 1/2]) = [0, 1]$ normalized to $[0.2689, 0.7311]$ due to the finite representation of $-\inf$
The final output is obtained by multiplying the softmax scores with the V matrix: $\begin{bmatrix} 1 & 0 \\ 0.2689 & 0.7311 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 0 \\ 0.2689 & 0.7311 \end{bmatrix}$

Constraints:

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

More from LLM 1: Foundations

Transformer Block Forward PassHard Feed-Forward NetworkEasy Layer NormalizationEasy KV-Cache AttentionHard Multi-Head AttentionMedium

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

0/0

Run code to see test results.