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.

[1.0000 0.0000]
[0.2689 0.7311]

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}$

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

Editor

Python 3.13.1

Test Results

0/0

Run code to see test results.

📘

Causal Attention Mask

MediumAttention Mechanisms

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

Editor

Python 3.13.1

Test Results

0/0

Run code to see test results.