Transformer

Attention (Transformer)

Terms:

• Self-Attention
• Masked Attention
• Sparse Attention
• Flash Attention
• 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(\frac{Q{K}^T}{\sqrt{d_k}})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, so “it” in a sentence can literally point at “the cat” three tokens back and pull its information forward.

In matrix form

$$\left[\begin{array}{c}\text{softmax}\left(\frac{⟨{\mathbf{q}}_1,{\mathbf{k}}_1⟩}{\sqrt{d}}\right){\mathbf{v}}_1^T+\cdots +\text{softmax}\left(\frac{⟨{\mathbf{q}}_1,{\mathbf{k}}_n⟩}{\sqrt{d}}\right){\mathbf{v}}_n^T\\ ⋮\\ \text{softmax}\left(\frac{⟨{\mathbf{q}}_m,{\mathbf{k}}_1⟩}{\sqrt{d}}\right){\mathbf{v}}_1^T+\cdots +\text{softmax}\left(\frac{⟨{\mathbf{q}}_m,{\mathbf{k}}_n⟩}{\sqrt{d}}\right){\mathbf{v}}_n^T\end{array}\right]\in {\mathbb{R}}^{m\times d}$$

Attention mechanism

• Query roughly speaking is what am I looking for.

• Key is what I represent

• Value is what I actually contain

• Query (Q): “What am I looking for?”

• Key (K): “What do I have?”

• Value (V): “What do I give you if you choose me?”

• 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

Like conceptually, there are no physical properties that allow us to distinguish Q from K. Like for all we known, Q is K and K is Q.

• Query (Q): “What am I looking for?”
• Key (K): “What do I have?”
• Value (V): “What do I give you if you choose me?”

Think of a library:

• Query = what you want to read about
• Keys = the summary on the card catalog for each book
• Values = the full content of each book

Question

in attention, is W_Q, W_K and W_V really necessary? Just use the raw embedding as your Q, K and V. the embedding is going to change too anyways

$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$

• Inputs:
  • Query $Q\in {\mathbb{R}}^{m\times d}$
  • Value $V\in {\mathbb{R}}^{n\times d}$
  • Key $K\in {\mathbb{R}}^{L\times C}$

$\mathbf{V}=\left[\begin{array}{c}{\mathbf{v}}_1^T\\ ⋮\\ {\mathbf{v}}_n^T\end{array}\right]\in {\mathbb{R}}^{n\times d},\mathbf{K}=\left[\begin{array}{c}{\mathbf{k}}_1^T\\ ⋮\\ {\mathbf{k}}_n^T\end{array}\right]\in {\mathbb{R}}^{n\times d},\mathbf{Q}=\left[\begin{array}{c}{\mathbf{q}}_1^T\\ ⋮\\ {\mathbf{q}}_m^T\end{array}\right]\in {\mathbb{R}}^{m\times d}$

Output: an $m\times d$ matrix

The dot product of the query and key tells you how well the key and query are aligned. Large positive dot product = “this key matches what the query is asking about”, so softmax puts most of its mass there and the output is pulled toward that value.

Then (Softmax operation is row-wise, i.e., $\text{softmax}(\mathbf{z}{)}_i=\frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}$):

What is $d_k$? I think that is the number of dimensions

• It’s just a scaling factor
• I asked the professor and he said empirically, it gives the best performance

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), 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 be nudged during training.

Attention Layer

This is cross-attention

Building up Attention from scratch (CS231n 2025 Lec 8)

Bahdanau attention (the original motivation)

Seq2seq (Sutskever 2014) bottlenecked the full input through a single fixed-size vector $c=h_T$. Long inputs (T=1000) can’t compress into one vector. Bahdanau et al. (ICLR 2015) let the decoder look back at all encoder states each step.

For each decoder timestep $t$:

1. Alignment scores: $e_{t,i}=f_{att}(s_{t-1},h_i)$ where $f_{att}$ is a small learned MLP
2. Attention weights: $a_{t,\ast }=\text{softmax}(e_{t,\ast })$ so ${\sum }_i{a}_{t,i}=1$
3. Context vector: $c_t={\sum }_i{a}_{t,i}{h}_i$
4. Decoder step: $s_t=g_U(y_{t-1},s_{t-1},c_t)$

Intuition, translating “we see the sky” → “vediamo il cielo”. At the first decoder step “vediamo” = “we see”, so the model learns weights like $a_{1,1}=a_{1,2}=0.45$ (attend to “we”, “see”), $a_{1,3}=a_{1,4}=0.05$. All differentiable, no supervision on weights. It learns alignment as a byproduct of caption quality.

The fixed bottleneck vector $c=h_T$ was the whole translation, compressed. Attention lets the decoder reach back and reread the input as it writes each word, so no single vector has to carry everything.

Generalized Attention Layer (inputs decoupled from scores)

Throw away the RNN; keep the primitive. Inputs:

• Query vectors $Q\in {\mathbb{R}}^{N_Q\times D_Q}$
• Data vectors $X\in {\mathbb{R}}^{N_X\times D_X}$
• Key projection $W_K\in {\mathbb{R}}^{D_X\times D_Q}$ and Value projection $W_V\in {\mathbb{R}}^{D_X\times D_V}$

$K=X{W}_K,V=X{W}_V$ $E=\frac{Q{K}^T}{\sqrt{D_Q}}[N_Q\times N_X]$ $A=\text{softmax}(E,\text{dim}=1),Y=AV[N_Q\times D_V]$

Two upgrades over Bahdanau:

• Scaled dot product replaces the MLP $f_{att}$: $e_{ij}=q_i\cdot k_j/\sqrt{D_Q}$. The $\sqrt{D_Q}$ scaling keeps pre-softmax logits in the “not saturated” regime; without it, large $D_Q$ makes softmax collapse to one-hot and gradients vanish
• Separate key & value matrices, the query compares against $K$ but sums up $V$. Decouples “what to match on” from “what to pass through”. Library analogy: the catalog card (key) advertises the book’s topic; the book itself (value) is what you actually read. You don’t want to be forced to score books using their full contents

Self-Attention = Cross-Attention where Q, K, V all come from X

Each token plays all three roles against every other token: it asks its own question, offers its own advertisement, and provides its own content. The layer is just “every token mixes in a weighted blend of every other token’s value, weighted by query-key similarity.” That’s it.

$Q=X{W}_Q,K=X{W}_K,V=X{W}_V(\text{all }[N\times D_{out}])$

Everything else is the same. The QKV projections are often fused into one matmul: $[QKV]=X[W_Q{W}_K{W}_V]$ with shape $[N\times 3D_{out}]$.

Is self-attention permutation-equivariant?

Yes, $F(\sigma (X))=\sigma (F(X))$. Permuting the inputs permutes $Q$, $K$, $V$ and thus the outputs in the same way. Self-attention works on sets of vectors, not sequences. This is why Transformers need positional encoding (see Positional Encoding); the layer itself has no notion of order.

Masked self-attention

For autoregressive language modeling, token $i$ must not see tokens $j>i$. Implement by setting $E_{ij}=-\infty$ for $j>i$ before the softmax, which forces $A_{ij}=0$. No changes to the forward code path, just a mask.

$e^{-\infty }=0$ in the softmax, so those positions contribute zero weight. The same forward pass trains all $N$ prefixes in parallel without the token at position $i$ ever seeing its own future, which is what makes teacher-forced training tractable on GPUs.

Multi-head self-attention

Run $H$ independent self-attention heads in parallel, each with head dim $D_H=D/H$, then concat and project:

Intuition

Ask multiple questions in parallel. One head can track syntax, another semantics, another long-range co-reference. Each head has its own $(W_Q,W_K,W_V)$ so it learns its own notion of “what to match on” and “what to pass through.” Concatenating lets the next layer read the sequence from several perspectives at once. The total param count is the same as one big head of width $D$ since each head is $D/H$ wide.

name	shape
Queries	$Q=X{W}_Q$	$[H\times N\times D_H]$
Keys	$K=X{W}_K$	$[H\times N\times D_H]$
Values	$V=X{W}_V$	$[H\times N\times D_H]$
Similarities	$E=Q{K}^T/\sqrt{D_Q}$	$[H\times N\times N]$
Attention	$A=\text{softmax}(E,\text{dim}=2)$	$[H\times N\times N]$
Head outputs	$Y=AV$ reshape	$[N\times H{D}_H]$
Output	$O=Y{W}_O$	$[N\times D]$

In practice all $H$ heads run as one batched matmul.

Self-attention is four matmuls

Everything above collapses to four matrix multiplies against $X$:

1. QKV projection: $[N\times D][D\times 3H{D}_H]\Rightarrow [N\times 3H{D}_H]$, split into Q/K/V
2. QK similarity: $[H\times N\times D_H][H\times D_H\times N]\Rightarrow [H\times N\times N]$
3. V-weighting: $[H\times N\times N][H\times N\times D_H]\Rightarrow [H\times N\times D_H]$
4. Output projection: $[N\times H{D}_H][H{D}_H\times D]\Rightarrow [N\times D]$

Compute and memory are both $O(N^2)$ from steps 2+3. At $N=100k$, $H=64$, the $H\times N\times N$ attention matrix alone is 1.19 TB, no GPU fits that. Flash Attention fuses steps 2+3 tile-by-tile, avoiding the full materialization and reducing memory to $O(N)$.

The $N^2$ cost is the price of “every token sees every token.” You never actually need the whole matrix in memory: each row only needs its own softmax denominator, which Flash Attention streams through SRAM.

Spatio-Temporal Self-Attention / Nonlocal Block (CS231n 2025 Lec 10)

The Nonlocal Block (Wang et al CVPR 2018) drops self-attention into a 3D CNN as a residual sub-module, letting any spatio-temporal location attend to any other. Pipeline starting from a 3D-CNN feature map $f\in {\mathbb{R}}^{C\times T\times H\times W}$:

1. Three $1\times 1\times 1$ convs project to $Q,K,V\in {\mathbb{R}}^{C^{\prime }\times T\times H\times W}$ (typically $C^{\prime }=C/2$)
2. Reshape $Q,K$ to $C^{\prime }\times THW$, compute $A=\text{softmax}(Q^{\top }K)\in {\mathbb{R}}^{(THW)\times (THW)}$; every spatio-temporal voxel attends to every other
3. Apply $A$ to $V$ (reshape to $C^{\prime }\times THW$) → $C^{\prime }\times T\times H\times W$
4. $1\times 1\times 1$ conv back to $C$ channels
5. Residual add to the input feature map

Plug it anywhere into an existing 3D CNN (e.g. I3D, SlowFast) for “global temporal context for free”; SlowFast + Nonlocal hits 79.8 on Kinetics-400, vs SlowFast alone trailing.

This is structurally identical to ViT-style self-attention, just over $(T,H,W)$ tokens with 3D-conv projections instead of linear ones. Foreshadows the Video Transformer wave (ViViT, MViTv2, VideoMAE).

Context Parallelism, parallelizing attention across the sequence (CS231n 2025 Lec 11)

When sequence length $N$ gets long (131k for Llama3-405B stage 2), even FSDP+TP can’t fit activations. Context Parallelism (CP) splits the $N\times D$ tokens across GPUs along the sequence axis. Most Transformer sublayers are easy:

• RMSNorm / residual: elementwise, no comm needed
• MLP / QKV projection: matmul on a slice of tokens, same weights everywhere; treat like DP on the sliced tokens
• Self-attention: hard, because every query needs every key & value; the $Q{K}^{\top }$ score matrix is $N\times N$ and spans GPUs

Two implementations:

Option 1, Ring Attention (Liu et al., arXiv 2023). Outer loop over $Q$ tiles (stays local to each GPU); inner loop rotates $K,V$ tiles around a ring of GPUs. Each GPU processes its $Q$ against every $(K,V)$ shard by receiving from its left neighbor and sending to its right. Scales to arbitrarily long sequences (no single GPU ever holds the full $K,V$), at the cost of $N$ peer-to-peer comms per attention layer.

Option 2, DeepSpeed Ulysses (Jacobs et al., arXiv 2023). Parallelize over heads instead of sequence: each GPU gets all tokens for a subset of heads. Max parallelism = $H$ (number of heads). Simpler than Ring Attention but caps at $H$; fine for Llama3 (128 heads) but doesn’t scale past that.

Llama3-405B usage:

• Stage 1 (seq 8192): no CP, TP=8, PP=16, DP=64/128 → 43% MFU
• Stage 2 (seq 131072): CP=16 (so 8192 tokens/GPU), TP=8, PP=16, DP=8 → 38% MFU

CP is Flash Attention-friendly; both avoid materializing the full $N\times N$ attention matrix, just at different granularities (tile-within-GPU vs GPU-across-ring).

Related

• KV Cache
• Transformer
• Positional Encoding
• Flash Attention
• Seq2Seq
• Action Classification