Compute the output of scaled dot-product attention for a single query attending to two key/value pairs, and explain why the scores are divided by $\sqrt{d_k}$. Use $q = [1, 0]$, keys $k_1 = [1, 0]$, $k_2 = [0, 1]$, values $v_1 = [1, 0]$, $v_2 = [0, 2]$, with $d_k = 2$.

Attention lets a query pull a weighted blend of values, where the weights come from how well the query matches each key: \begin{equation} \text{Attention}(q, K, V) = \text{softmax}\!\left(\frac{q K^\top}{\sqrt{d_k}}\right) V \end{equation} Step 1: raw scores\begin{equation} q \cdot k_1 = 1\cdot1 + 0\cdot0 = 1, \qquad q \cdot k_2 = 1\cdot0 + 0\cdot1 = 0 \end{equation} Step 2: scale by $\sqrt{d_k} = \sqrt{2} \approx 1.414$\begin{equation} s_1 = \frac{1}{\sqrt{2}} \approx 0.707, \qquad s_2 = \frac{0}{\sqrt{2}} = 0 \end{equation} Step 3: softmax\begin{equation} a_1 = \frac{e^{0.707}}{e^{0.707} + e^{0}} = \frac{2.028}{3.028} \approx 0.670, \qquad a_2 \approx 0.330 \end{equation} Step 4: weighted sum of the values\begin{equation} \text{out} = 0.670\,[1,0] + 0.330\,[0,2] = [0.670,\; 0.660] \end{equation} Why divide by $\sqrt{d_k}$The score is a dot product of two $d_k$-dimensional vectors. If the entries are roughly independent with unit variance, the dot product is a sum of $d_k$ such terms, so its variance grows like $d_k$ and its magnitude like $\sqrt{d_k}$. Large-magnitude scores push the softmax into a saturated region where one weight is near 1 and the rest near 0; there the softmax gradient is almost zero, which stalls learning. Dividing by $\sqrt{d_k}$ rescales the scores back to roughly unit variance, keeping the softmax in a responsive range with healthy gradients regardless of the key dimension. The same mechanism, applied with the queries, keys and values all derived from the same sequence, is self-attention, and it is computed for every position in parallel — the property that let transformers replace recurrence.

Complete Training For Your Trading Interview

Interview Questions Machine Learning Scaled Dot-Product Attention by Hand

Unlock

Sorry, you don't have access to this interview question. Sign up to practice with our wide range of interview questions or use the filters to find free examples!

Python scratchpad

# Run Python right here in your browser
def two_sum(nums, target):
    seen = {}
    for i, n in enumerate(nums):
        if target - n in seen:
            return [seen[target - n], i]
        seen[n] = i
    return []

print(two_sum([2, 7, 11, 15], 9))  # [0, 1]

Run Python in your browser

Write and execute code — numpy, pandas & scikit-learn included — right on the question page.

My Notes

Notes are saved in this browser only. Sign up to keep them in your account.

The Interview Guide for Quants and Traders

Scaled Dot-Product Attention by Hand

What do you think of this page?