Causal Attention Mask

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

Example:

Constraints:

Test Results