Self Attention Worked Example

Deep Learning interview question · Medium difficulty

Can you walk through a self-attention computation step by step?medium

Status
✓
Type
applied
Topic
self-attention-worked-example
Frequency
common
Tags
self-attention, query-key-value, worked-example, softmax
Answer

Convert tokens to embeddings, project each into Query, Key, and Value vectors using learned weight matrices (Q=XWq, K=XWk, V=XWv), compare each token's Query against every token's Key, apply softmax to get attention weights, then take the weighted sum of Value vectors — that weighted sum is the token's new, context-aware representation.

Explanation

For the sentence "The cat sat": step 1 converts each token to an embedding (The→x1, cat→x2, sat→x3). Step 2 projects each embedding into its own Query, Key, and Value using the same learned weight matrices applied to every token, so "sat" ends up with q_sat, k_sat, v_sat. Step 3 compares the Query of the token being updated against every token's Key — for "sat," that's q_sat·k_The, q_sat·k_cat, and q_sat·k_sat. Step 4 runs softmax over those scores to get normalized attention weights, e.g. The→0.05, cat→0.75, sat→0.20. Step 5 combines the Value vectors using those weights: new representation of "sat" = 0.05·v_The + 0.75·v_cat + 0.20·v_sat. The result is that "sat" now carries information pulled from "cat" because the weights said that connection mattered — and this same five-step process runs independently for every token in the sequence, in parallel, which is what makes self-attention both context-aware and fast to compute.

Follow-upWhy does every token use the same Wq/Wk/Wv weight matrices instead of token-specific ones?
Follow-upWhat would the attention weights look like for a token that has no strong relationship to any other token in the sentence?
Follow-upWhy is this computation naturally parallelizable across tokens, unlike an RNN's hidden state update?