Transformer

Attention (Transformer)

Attention mixes information by using query-key similarity to take a weighted sum of values.

Terms:

• Self-Attention
• Cross-Attention
• Multi-Head Attention
• Masked Attention
• Sparse Attention
• FlashAttention
• Paged Attention (used for faster compute)
• Multi-Head Latent Attention
• Multi-Head Attention
• Multi-Query Attention (modern)
• Grouped Query Attention

Confusion:

• https://stats.stackexchange.com/questions/421935/what-exactly-are-keys-queries-and-values-in-attention-mechanisms

$Attention(Q,K,V)=softmax\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$

Intuition

Soft dictionary lookup. Each query asks “who’s relevant to me?”, each key advertises “here’s what I’m about,” each value carries the actual content. The softmax over $Q{K}^T$ picks a weighted blend of values.

Shapes

• Query (Q): what this position is looking for
• Key (K): what each position advertises
• Value (V): what each position contributes if selected

Input:

$$X\in {\mathbb{R}}^{n\times d_{model}}$$

Project into queries, keys, values:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$$ $$W_Q,W_K\in {\mathbb{R}}^{d_{model}\times d_k},W_V\in {\mathbb{R}}^{d_{model}\times d_v}$$

So:

$$Q,K\in {\mathbb{R}}^{n\times d_k},V\in {\mathbb{R}}^{n\times d_v}$$

Scores:

$$\frac{Q{K}^T}{\sqrt{d_k}}\in {\mathbb{R}}^{n\times n}$$

Apply softmax row-wise over the last dimension:

$$A=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}},\text{dim}=-1\right)$$

Then mix values:

$$Y=AV\in {\mathbb{R}}^{n\times d_v}$$

Are W_Q, W_K, and W_V really necessary? Why not use the raw embedding?

Yes. The projections let the model learn separate subspaces for matching and for passing content. Otherwise the same representation would have to serve all three roles.

Why not make queries and keys the same matrix?

Because “what I’m looking for” and “what I offer” are different roles. Dating analogy: the traits you want in a partner are not the same as the traits you advertise about yourself. Separate $W_Q$ and $W_K$ let attention model that asymmetry; tying them would make matching less expressive.

Why \sqrt{d_k}?

Without it, dot products grow with dimension ($q\cdot k$ has variance proportional to $d_k$ for random unit-ish vectors). Then softmax saturates into a near one-hot and gradients vanish.

Dividing by $\sqrt{d_k}$ keeps logit variance around 1, so softmax stays in the informative regime where it can still be nudged during training.

• https://www.youtube.com/watch?v=bCz4OMemCcA

Another vid

• https://www.youtube.com/watch?v=QCJQG4DuHT0

I was really confused about QKV. https://www.reddit.com/r/MachineLearning/comments/19ewfm9/d_attention_mystery_which_is_which_q_k_or_v/ https://stats.stackexchange.com/questions/421935/what-exactly-are-keys-queries-and-values-in-attention-mechanisms

Related

• Bahdanau Attention
• Nonlocal Block
• Context Parallelism
• KV Cache
• Transformer
• Positional Encoding
• FlashAttention
• Seq2Seq
• Action Classification