Scaled Dot-Product Attention

Implement scaled dot-product attention.

Given Query (Q), Key (K), and Value (V) matrices, compute: $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

where d_k is the dimension of the keys (number of columns in K).

Input:

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

Input:

Output:

[0.7311 0.2689]
[0.2689 0.7311]

Reasoning:

The input provides the sequence length n = 2 and dimension d = 2, along with the Q, K, and V matrices.
We compute the dot product of Q and K^T, which is $\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ , then scale it by $\frac{1}{\sqrt{d}} = \frac{1}{\sqrt{2}}$ .
We apply the softmax function to the scaled dot product: $\text{softmax}\left(\frac{1}{\sqrt{2}}\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}\right) = \begin{bmatrix} \frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}} } + e^{\frac{1}{\sqrt{2}}}} & \frac{e^0}{e^{\frac{1}{\sqrt{2}}} + e^{\frac{1}{\sqrt{2}}}} \\ \frac{e^0}{e^{\frac{1}{\sqrt{2}}} + e^{\frac{1}{\sqrt{2}}}} & \frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + e^{\frac{1}{\sqrt{2}}}} \end{bmatrix} \approx \begin{bmatrix} 0.7311 & 0.2689 \\ 0.2689 & 0.7311 \end{bmatrix}$ .
The final output is obtained by multiplying this softmax result with the V matrix, which yields $\begin{bmatrix} 0.7311 & 0.2689 \ 0.2689 & 0.7311 \end{bmatrix} \begin{bmatrix} 1 & 0 \ 0 & 1 \end{bmatrix} = \begin{bmatrix} 0.