Scaled Dot-Product Attention

hardPRONLP FundamentalsMachine Learning 2

Implement Scaled Dot-Product Attention from the Transformer architecture.

Given query matrix $Q$ , key matrix $K$ , and value matrix $V$ :

$\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).

Steps:

Return the attention output matrix, rounded to 4 decimal places.

Input

Q = [[1, 0]]
K = [[1, 0], [0, 1]]
V = [[1, 2], [3, 4]]

Output

[[1.6605, 2.6605]]

Reasoning

We start by computing the matrix product $QK^T$ , which is $[[1, 0]] \cdot [[1, 0], [0, 1]]^T = [[1, 0]] \cdot [[1, 0], [0, 1]] = [[1, 0]]$ .
Then, we scale this result by $\frac{1}{\sqrt{d_k}} = \frac{1}{\sqrt{2}}$ , yielding $[[\frac{1}{\sqrt{2}}, 0]]$ .
Next, we apply the softmax function row-wise: since $[[\frac{1}{\sqrt{2}}, 0]]$ is a single row, this results in $[[\frac{1}{1+e^{-\frac{1}{\sqrt{2}}}}, \frac{1}{1+e^{0}}]] \approx [[0.7311, 0.2689]]$ after normalization, but because we apply softmax to $[\frac{1}{\sqrt{2}}, 0]$ we actually calculate $\text{softmax}([[\frac{1}{\sqrt{2}}, 0]]) = [[\frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + e^{0}}, \frac{e^{0}}{e^{\frac{1}{\sqrt{2}}} + e^{0}}]] = [[\frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + 1}, \frac{1}{e^{\frac{1}{\sqrt{2}}} + 1}]]$ , which then gets multiplied by V.
Finally, multiplying this result by $V$ yields $[[\frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + 1}, \frac{1}{e^{\frac{1}{\sqrt{2}}} + 1}]] \cdot [[1, 2], [3, 4]] = [[\frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + 1} \cdot 1 + \frac

Background Knowledge