Attention Mechanisms

Background

Bahdanau et al. (2014) introduced attention for machine translation as an alignment mechanism for seq2seq models
Vaswani et al. (2017) generalized it to self-attention - the foundation of all modern LLMs
Flash Attention (Dao et al. 2022) made attention memory-efficient for long sequences

The Concept

Attention computes a weighted sum of values, where the weights are determined by compatibility between a query and a set of keys. Self-attention lets each token in a sequence attend to all other tokens, building context-aware representations.

How it Links

Self-attention is the core building block of the Transformer
Cross-attention (queries from one sequence, keys/values from another) is used in encoder-decoder models and multimodal attention
Connects to memory-augmented networks, retrieval augmentation (RAG), and graph attention networks

Concrete Example

Reading comprehension: given a passage and a question, cross-attention lets each question word query the passage for relevant information. The word "where" in the question attends strongly to location words in the passage. The model learns which parts of the passage answer the question.

Prerequisites

RNNs, LSTMs, GRUs, Math you need, Backpropagation & Optimizers

Learning outcomes

Trace the lineage from Bahdanau additive attention through Luong multiplicative variants to scaled dot-product self-attention.
Derive the scaled dot-product formula softmax(QK^T / sqrt(d_k)) V from a kernel-regression viewpoint and explain why the sqrt(d_k) factor matters.
Define multi-head attention, count its parameters and FLOPs, and contrast self, cross, and causal masking.
Explain the O(n^2 d) memory and compute wall, and describe how FlashAttention preserves exact attention while linear attention and Performers approximate it.
Implement scaled dot-product attention from scratch in NumPy and compare with PyTorch nn.MultiheadAttention.

TL;DR

Attention is a learned soft lookup. Given a query $q$ , a set of keys $K$ and values $V$ , attention returns a convex combination $\sum_j \alpha_j v_j$ where $\alpha_j \propto \exp(q \cdot k_j)$ . Bahdanau [1] introduced this for translation in 2014; Luong [2] simplified it to a dot product; Vaswani [3] dropped recurrence entirely and built the Transformer out of scaled dot-product self-attention $\mathrm{softmax}(QK^\top/\sqrt{d_k})V$ , with multiple heads operating in parallel. The $\sqrt{d_k}$ factor keeps the softmax out of saturation. The whole thing is $O(n^2 d)$ in compute and memory, which is both a strength (fully parallel across positions) and a weakness (long context is expensive). FlashAttention [4] recovers exact attention at near-linear memory by tiling the softmax through SRAM; linear attention [5] and Performers [6] approximate the softmax kernel for true $O(n d)$ scaling, at some accuracy cost surveyed in [7].

Intuition

A recurrent encoder reads a sentence and compresses everything it has seen into a single hidden state. By the end of a long sentence the early words are squeezed into a tiny corner of that vector. Attention asks a different question: “which of the encoder hidden states should I look at right now?” and answers it with a learned weighted average. Nothing is compressed; everything stays addressable.

The shape stays the same when self-attention swaps the encoder for the same sequence: every token is allowed to look at every other token, and the model learns who looks at whom. That single substitution removes the recurrence, makes training fully parallel across positions, and is the engineering reason Transformers displaced RNNs [3].

Attention is content-based addressing

Memory in classical computing is address-based: you ask for slot 17, you get whatever is in slot 17. Attention is content-based: you ask “what’s similar to my query $q$ ?” and you get a soft weighted average of every slot, where slots that match $q$ get more weight. The softmax is the differentiable replacement for an argmax-style hard lookup. That makes attention a kind of differentiable hash table, and the gradient through it tells the network how to reshape its keys so that the right values surface for each kind of query.

A short chronology of how the idea matured before becoming the Transformer:

2014Bahdanau additive attentionpaper
Bahdanau, Cho, Bengio remove the seq2seq encoder bottleneck by letting the decoder attend over all encoder hidden states, with a small MLP scoring function.
2015Luong multiplicative attentionpaper
Luong, Pham, Manning simplify the score to a dot product (and a general bilinear form), trading expressivity for speed.
2017Vaswani: scaled dot-product self-attention + multi-headbreakthrough
Attention Is All You Need drops recurrence entirely. The Transformer is built from softmax(Q K^T / sqrt(d_k)) V with h parallel heads.
2020Linear attention (Katharopoulos et al.)efficient
Replace exp(q . k) with phi(q)^T phi(k) for a feature map phi, drop the explicit n x n matrix, get O(n d^2) compute and constant-memory autoregressive decoding.
2020Performers (Choromanski et al.)efficient
Random-feature approximation of the softmax kernel; provably unbiased, scales to long sequences, and remains drop-in for trained Transformers.
2022FlashAttention (Dao et al.)system
Exact attention with O(n) memory by tiling Q, K, V into SRAM blocks and recomputing the softmax on the fly. Practical 2-4x speedups on long sequences.

Math

Scaled dot-product attention from first principles

From a single query to the matrix form

Start with one query $q \in \mathbb{R}^{d_k}$ and a memory of $n$ key-value pairs $\{(k_j, v_j)\}_{j=1}^n$ with $k_j \in \mathbb{R}^{d_k}$ and $v_j \in \mathbb{R}^{d_v}$ . The attention output is a convex combination of values, weighted by a softmax over query-key compatibility:

\mathrm{Attn}(q, K, V) = \sum_{j=1}^n \alpha_j v_j, \qquad \alpha_j = \frac{\exp(s(q, k_j))}{\sum_{j'} \exp(s(q, k_{j'}))}.

The score function $s(q, k)$ is what distinguishes attention variants. Bahdanau used $s(q, k) = v^\top \tanh(W_q q + W_k k)$ (additive). Luong used $s(q, k) = q^\top W k$ (multiplicative bilinear). Vaswani used the simplest form, plain dot product, with a critical scale factor:

s(q, k) = \frac{q^\top k}{\sqrt{d_k}}.

Stack $n$ queries into $Q \in \mathbb{R}^{n \times d_k}$ , $n$ keys into $K \in \mathbb{R}^{n \times d_k}$ , and $n$ values into $V \in \mathbb{R}^{n \times d_v}$ . The vectorized form is

\mathrm{Attn}(Q, K, V) = \mathrm{softmax}\!\left( \frac{Q K^\top}{\sqrt{d_k}} \right) V \in \mathbb{R}^{n \times d_v},

with the softmax taken row-wise. Each row of the output is a per-query weighted average of value vectors.

Why divide by sqrt(d_k)? The variance argument

Suppose the entries of $q$ and $k$ are independent random variables with mean $0$ and variance $1$ . Then $q^\top k = \sum_{i=1}^{d_k} q_i k_i$ is a sum of $d_k$ independent products, each with mean $0$ and variance $1$ , so $\mathrm{Var}(q^\top k) = d_k$ . As $d_k$ grows, the unscaled dot products spread over a range of $\pm \sqrt{d_k}$ .

A softmax fed inputs with variance $d_k$ saturates: one entry dominates, the others get nearly zero weight, and the gradient through the softmax collapses (the off-diagonal entries of the softmax Jacobian go to zero [8]). Dividing by $\sqrt{d_k}$ brings the variance back to $1$ , which keeps the softmax in a regime where its gradient is informative. Empirically the un-scaled version is hard to train at $d_k$ greater than about $32$ [3].

Softmax row-wise gradient

Let $a = \mathrm{softmax}(z)$ for a single row. The Jacobian is

\frac{\partial a_i}{\partial z_j} = a_i (\delta_{ij} - a_j),

so $J = \mathrm{diag}(a) - a a^\top$ . Two consequences. First, $J$ is at most rank $n - 1$ (the all-ones vector is in its kernel, which is why softmax is shift-invariant). Second, when one $a_i$ is near $1$ and the rest are near $0$ , $J \approx 0$ except for the dominant entry, so gradients vanish through saturated softmax: this is the engineering reason for the $\sqrt{d_k}$ rescaling.

Kernel-regression view

Fix a bandwidth $h$ and consider Nadaraya-Watson kernel regression with a Gaussian kernel:

\hat{f}(q) = \sum_j \frac{K_h(q, k_j)}{\sum_{j'} K_h(q, k_{j'})} v_j, \qquad K_h(q, k) = \exp\!\left(-\frac{\lVert q - k \rVert^2}{2 h^2}\right).

Expand the squared distance:

\exp\!\left(-\frac{\lVert q - k \rVert^2}{2 h^2}\right) = \exp\!\left(-\frac{\lVert q \rVert^2 + \lVert k \rVert^2 - 2 q^\top k}{2 h^2}\right).

The $\lVert q \rVert^2$ term cancels in the row-wise softmax (it is constant in $j$ ). If we further assume $\lVert k_j \rVert$ is approximately constant across keys (true for layer-normalized inputs, true on average for random init), the $\lVert k_j \rVert^2$ terms are also approximately constant in $j$ and cancel. What remains is

\hat{f}(q) \approx \mathrm{softmax}_j(q^\top k_j / h^2)\, v_j.

Setting $h^2 = \sqrt{d_k}$ recovers exactly scaled dot-product attention. So softmax attention is kernel regression with a learned key/value memory and the bandwidth absorbed into the scale factor. This view also makes every approximation strategy a choice of kernel: linear attention picks $\phi(q)^\top \phi(k)$ , Performers pick a random-feature kernel, etc.

Multi-head attention

A single scaled dot-product attention head produces a $d_v$ -dim output per query. Multi-head attention [3] runs $h$ heads in parallel, each with its own $(W_Q^{(i)}, W_K^{(i)}, W_V^{(i)})$ , and concatenates their outputs:

\mathrm{MHA}(X) = \mathrm{Concat}\!\left( \mathrm{head}_1, \dots, \mathrm{head}_h \right) W_O,

with

\mathrm{head}_i = \mathrm{Attn}(X W_Q^{(i)},\, X W_K^{(i)},\, X W_V^{(i)}).

If the model dimension is $d$ and there are $h$ heads, each head uses $d_k = d_v = d / h$ , so the total parameter count is the same as a single head with model dimension $d$ . The point of using $h$ heads is not capacity but diversity: each head can specialize on a different relation (positional, syntactic, content), and the concatenation followed by $W_O$ lets the network mix those specializations.

Self vs cross vs causal: the three attention modes

The same scaled dot-product formula is reused in three different roles by routing different tensors into $Q$ , $K$ , $V$ , and by applying different masks to the score matrix.

Self-attention. $Q$ , $K$ , $V$ all come from the same sequence $X$ : $Q = X W_Q$ , $K = X W_K$ , $V = X W_V$ . Used in encoder layers and inside masked decoder layers.
Cross-attention. $Q$ comes from the decoder, $K$ and $V$ from the encoder output. Used in encoder-decoder Transformers (the original NMT setup) and in the second sub-layer of every decoder block.
Causal (masked) self-attention. Self-attention with a triangular mask that sets $s_{ij} = -\infty$ for $j > i$ . After softmax those positions get weight zero, so token $i$ only attends to tokens $1, \dots, i$ . Required for autoregressive language models so that prediction of token $t+1$ does not peek at token $t+1$ during training.

The mask is added pre-softmax, not post-softmax, so masked positions must use $-\infty$ rather than $0$ . A common bug is to apply the mask after softmax and renormalize: this still works for forward pass but breaks the gradient because the renormalization is non-smooth across the mask boundary.

Mode	Q from	K, V from	Mask	Typical use
Self	sequence	same sequence	none	encoder, BERT
Causal self	sequence	same sequence	upper triangular -inf	GPT, decoder pre-cross-attn
Cross	decoder	encoder output	none	NMT, encoder-decoder

Variant	Mask	Use case	Example architecture
Self-attention (no mask)	None	Bidirectional contextualization, full-sentence understanding	BERT encoder, ViT, T5 encoder
Causal self-attention	Upper-triangular -inf (token i can only see tokens 1..i)	Autoregressive language modelling, next-token prediction	GPT family, decoder-only LLMs
Cross-attention	None on K, V; sometimes a padding mask	Conditional generation: decoder Q reads encoder K, V	Original NMT Transformer, T5 decoder, Whisper, Flamingo

Complexity: feature, then bug

For a sequence of length $n$ and head dimension $d$ , scaled dot-product attention does:

Score matrix $QK^\top$ : shape $(n, n)$ , cost $O(n^2 d)$ .
Softmax: cost $O(n^2)$ .
Multiply by $V$ : cost $O(n^2 d)$ .

Total: $O(n^2 d)$ in compute and $O(n^2)$ in attention-matrix memory. Per-head FLOPs across an entire sequence are roughly $4 n^2 d$ for one attention layer, dominated by the two matrix multiplies (one with $K^\top$ , one with $V$ ).

The feature side: those matrix multiplies are the most parallel operation in the architecture. There is no time-step dependency to break, so a Transformer can train at full GPU utilization across the entire sequence dimension. RNNs cannot; that is why Transformers train an order of magnitude faster on the same hardware [3].

The bug side: $O(n^2)$ memory means doubling sequence length quadruples the attention activation memory. At $n = 2048$ , attention is the dominant cost; at $n = 32{,}768$ , it pushes most consumer GPUs out of memory. Long-context models cannot afford the naive implementation.

Why attention is permutation-equivariant without positional encodings

Pure self-attention with no positional encoding is permutation-equivariant: if you shuffle the input tokens by a permutation $P$ , the output rows are shuffled by the same $P$ but the values in each row are unchanged. The reason is that every operation in attention (matrix multiplies, row-wise softmax) treats the sequence dimension as a set: there is no convolution over time, no recurrence. Permuting rows of $X$ permutes rows of $Q$ , $K$ , $V$ identically, permutes $QK^\top$ by $P$ on rows and columns, and the row-wise softmax sees the same multiset of scores.

This is why Transformers must inject positional information explicitly: sinusoidal encodings, learned embeddings, RoPE, ALiBi. Without them, the model genuinely cannot distinguish “the cat sat on the mat” from “mat the on sat cat the” [3].

Efficient variants

The $O(n^2)$ wall has spawned a small zoo of “efficient transformers”, surveyed in detail by Tay et al. [7]. Three families are worth knowing.

FlashAttention: tiled exact softmax

FlashAttention [4] keeps exact attention (same outputs as naive scaled dot-product, bit for bit) but cuts the memory footprint from $O(n^2)$ to $O(n)$ . The key observations:

GPUs are bandwidth-bound, not flop-bound, on attention. A modern GPU has $\sim 100$ TFLOPs of compute and $\sim 1$ TB/s of HBM bandwidth. Reading and writing the $(n, n)$ attention matrix is the bottleneck, not the matrix multiplies.
The softmax is associative if you carry running statistics. For two blocks of scores $z^{(1)}, z^{(2)}$ with running max $m$ and running sum $\ell$ , you can update $(m, \ell)$ as you stream blocks through SRAM, and at the end produce the same softmax output as one giant softmax over the concatenation.

The algorithm splits $Q$ into row blocks $Q_i$ of size $B_r$ and $K, V$ into column blocks $K_j, V_j$ of size $B_c$ , both small enough to fit in on-chip SRAM. For each $(i, j)$ tile:

Load $Q_i, K_j, V_j$ from HBM into SRAM (one read each per tile).
Compute the tile of scores $S_{ij} = Q_i K_j^\top / \sqrt{d_k}$ in SRAM.
Update running max $m_i \leftarrow \max(m_i, \mathrm{rowmax}(S_{ij}))$ .
Rescale the running output and sum: $O_i \leftarrow \mathrm{diag}(e^{m_i^{\text{old}} - m_i^{\text{new}}}) O_i + e^{S_{ij} - m_i^{\text{new}}} V_j$ , $\ell_i \leftarrow e^{m_i^{\text{old}} - m_i^{\text{new}}} \ell_i + \mathrm{rowsum}(e^{S_{ij} - m_i^{\text{new}}})$ .
After all tiles, $O_i \leftarrow O_i / \ell_i$ .

The $(n, n)$ score matrix is never materialized in HBM. Memory drops from $O(n^2)$ to $O(n)$ , and wall-clock attention time drops 2-4x at $n = 4096$ in the original paper. The backward pass is the same idea: recompute the score tile on the fly during backward instead of storing it. FlashAttention-2 and FlashAttention-3 push this further but the algorithmic core is the tiled exact softmax above.

Linear attention and Performers

Linear attention [5] replaces $\exp(q^\top k)$ with $\phi(q)^\top \phi(k)$ for some feature map $\phi: \mathbb{R}^{d_k} \to \mathbb{R}^{d_\phi}$ . Because the kernel is now a dot product, you can rewrite

\sum_j \phi(q)^\top \phi(k_j) v_j = \phi(q)^\top \left( \sum_j \phi(k_j) v_j^\top \right),

and the inner sum is independent of $q$ . Compute it once in $O(n d_\phi d_v)$ , store an $(d_\phi, d_v)$ matrix, and then every query costs $O(d_\phi d_v)$ . Total: $O(n d_\phi d_v)$ , linear in $n$ . Autoregressive decoding becomes a recurrence over a fixed-size state, recovering an RNN-like inference profile from a Transformer-like training profile.

The catch is that $\phi$ must approximate the softmax kernel well. Katharopoulos et al. use $\phi(x) = \mathrm{elu}(x) + 1$ which is cheap but coarse; Performers [6] use random Fourier features for the Gaussian kernel and prove unbiasedness with controllable variance. Both lose accuracy on tasks where exact attention captures sharp lookups, but recover most of the quality on language modelling at long context.

Sparse and local attention

Another family restricts which $(i, j)$ pairs are computed at all: local windows around the diagonal, strided patterns, dilated windows, learned routing (Reformer, Longformer, BigBird). The Tay survey [7] tabulates dozens of variants; the unifying theme is “drop $O(n^2)$ entries by hand or by hashing, keep the rest exact”. Modern long-context models (Llama 3, Mistral) typically combine local sliding-window attention with periodic global attention sinks for streaming long-context decoding, an interpolation between dense and sparse rather than a clean either-or.

Interactive: attention heatmap

The visual demo first. The heatmap shows three precomputed heads on the same 9-token sentence: a diagonal head (each token attends mostly to itself), a previous-token head (each token attends to its left neighbor, a real pattern observed in early-layer GPT heads), and a content-word head (every token attends to the content nouns and verbs in the sentence, a real pattern observed in middle-layer BERT heads). Hover or click any row to highlight that query’s distribution.

Attention heatmap (precomputed, 9 x 9)

Rows are query tokens, columns are key tokens, and each row sums to 1. Hover a row label or cell to highlight that query's attention distribution; click to pin it. Cell color tracks attention weight from light (near 0) to deep blue (near 1).

Implementations

A from-scratch NumPy implementation that mirrors the math 1:1, then PyTorch’s production-grade nn.MultiheadAttention.

python

import numpy as np

def softmax_rows(x):
  """Numerically stable row-wise softmax."""
  x = x - x.max(axis=-1, keepdims=True)
  e = np.exp(x)
  return e / e.sum(axis=-1, keepdims=True)

def scaled_dot_product_attention(Q, K, V, mask=None):
  """Q: (n, d_k), K: (m, d_k), V: (m, d_v), mask: (n, m) of {0, -inf}."""
  d_k = Q.shape[-1]
  scores = Q @ K.T / np.sqrt(d_k)        # (n, m)
  if mask is not None:
      scores = scores + mask              # broadcast -inf into masked entries
  weights = softmax_rows(scores)          # (n, m)
  return weights @ V, weights              # (n, d_v), (n, m)

def multi_head_attention(X, W_q, W_k, W_v, W_o, num_heads, mask=None):
  """X: (n, d_model). W_*: (d_model, d_model). Returns (n, d_model)."""
  n, d_model = X.shape
  d_k = d_model // num_heads
  Q = (X @ W_q).reshape(n, num_heads, d_k).transpose(1, 0, 2)   # (h, n, d_k)
  K = (X @ W_k).reshape(n, num_heads, d_k).transpose(1, 0, 2)
  V = (X @ W_v).reshape(n, num_heads, d_k).transpose(1, 0, 2)
  head_outs = []
  for i in range(num_heads):
      out, _ = scaled_dot_product_attention(Q[i], K[i], V[i], mask)
      head_outs.append(out)
  cat = np.concatenate(head_outs, axis=-1)                       # (n, d_model)
  return cat @ W_o                                               # (n, d_model)

# Sanity check on a tiny sequence.
np.random.seed(0)
n, d, h = 5, 16, 4
X = np.random.randn(n, d)
W_q = np.random.randn(d, d) / np.sqrt(d)
W_k = np.random.randn(d, d) / np.sqrt(d)
W_v = np.random.randn(d, d) / np.sqrt(d)
W_o = np.random.randn(d, d) / np.sqrt(d)
Y = multi_head_attention(X, W_q, W_k, W_v, W_o, num_heads=h)
print(Y.shape)   # (5, 16)

The four matmuls (Q, K, V projections + output projection) plus the n x n softmax are the entire layer. Production code uses einsum/batched matmul instead of a Python loop over heads.

Run a tiny self-attention layer live in the browser:

Idle (packages: numpy)

(no output yet)

Pitfalls

Forgetting the causal mask in a decoder

If you forget the upper-triangular mask in a causal decoder, training looks fine: the model can peek at future tokens, so per-token loss drops fast. Inference quality, however, collapses: the model has never seen a context where the future was unavailable, and it produces fluent nonsense. Always build the causal mask once, cache it, and assert its shape on every batch. The mask must be added to the score logits before the softmax, with $-\infty$ on masked positions.

Numerical stability of softmax

The softmax $\exp(z_i) / \sum_j \exp(z_j)$ overflows for $z_i$ above about $88$ in fp32 ( $e^{88} \approx 10^{38}$ ). Standard fix: subtract the row max before exponentiating, $\mathrm{softmax}(z)_i = e^{z_i - \max_j z_j} / \sum_{j} e^{z_j - \max_{j'} z_{j'}}$ . Every framework does this internally, but if you implement attention from scratch you must too. FlashAttention’s running-max state machine is exactly the streaming version of this fix.

Attention dropout placement

The original Transformer applies dropout to the attention probabilities after the softmax, so each row is renormalized post-dropout (or not, depending on the implementation). This works but has a subtle bug: if you drop the highest-weight position, the gradient routes weirdly. Modern implementations (HuggingFace, FlashAttention) drop after the matmul with $V$ , which is equivalent in expectation and avoids the renormalization question. Pick one convention and stick with it across pretrain and fine-tune; mixing is a silent failure mode.

Float32 logits, even if the rest is bf16

Mixed-precision training (bf16 / fp16) is standard, but the attention softmax should be done in fp32 because the exponentiation is sensitive to small differences between large logits. PyTorch’s scaled_dot_product_attention does this automatically. If you write attention by hand, cast scores to fp32 before the softmax and back to bf16 after.

Pop quiz

Check yourself

Score: 0 / 5(0 answered)

1. Why does scaled dot-product attention divide by sqrt(d_k)?
2. What does the causal mask in masked self-attention do?
3. What is the time and memory complexity of vanilla scaled dot-product attention for sequence length n and head dimension d?
4. Pure self-attention without positional encodings is:
5. What does FlashAttention achieve compared to naive softmax attention?

Interview Q&A

Derive softmax attention from kernel regression on a learned key-value memory.onsite

Start with Nadaraya-Watson kernel regression. Given a query $q$ and a memory $\{(k_j, v_j)\}$ and a kernel $K_h$ , the kernel regression estimate is

\hat{f}(q) = \frac{\sum_j K_h(q, k_j) v_j}{\sum_{j'} K_h(q, k_{j'})}.

Pick the Gaussian kernel $K_h(q, k) = \exp(-\lVert q - k \rVert^2 / (2 h^2))$ and expand the squared distance:

\exp\!\left(-\frac{\lVert q \rVert^2 + \lVert k \rVert^2 - 2 q^\top k}{2 h^2}\right).

The $\lVert q \rVert^2$ term is constant in $j$ and cancels in the ratio (the row-wise softmax). If keys are layer-normalized so $\lVert k_j \rVert$ is approximately constant, the $\lVert k_j \rVert^2$ term also cancels. What remains is

\hat{f}(q) \approx \mathrm{softmax}_j\!\left(\frac{q^\top k_j}{h^2}\right) v_j.

Setting $h^2 = \sqrt{d_k}$ recovers scaled dot-product attention. So the formula is non-parametric kernel regression with a learned key-value memory and a temperature absorbed into the scale factor. This view makes the design knobs explicit: the kernel choice, the bandwidth, and the implicit assumption that keys are approximately normalized.

Why divide by sqrt(d_k)? Make the variance argument precise.phone

Suppose $q, k \in \mathbb{R}^{d_k}$ with independent components, each with mean $0$ and variance $1$ . Then

\mathrm{Var}(q^\top k) = \mathrm{Var}\!\left(\sum_{i=1}^{d_k} q_i k_i\right) = \sum_{i=1}^{d_k} \mathrm{Var}(q_i k_i) = d_k,

using independence of the $(q_i, k_i)$ products and $\mathrm{Var}(q_i k_i) = \mathrm{E}[q_i^2] \mathrm{E}[k_i^2] = 1$ for centered unit-variance factors. So the unscaled dot product has standard deviation $\sqrt{d_k}$ .

A softmax fed inputs with this scale collapses: one entry dominates and the rest get exponentially small weight. The Jacobian of the softmax in saturated regimes is approximately rank-1, so gradients through every position except the dominant one vanish. Dividing by $\sqrt{d_k}$ rescales to unit variance, putting the softmax back in a regime where its Jacobian is informative. Empirically this is the difference between a Transformer that trains and one that does not at $d_k$ above $\sim 32$ [3].

Prove that pure self-attention is permutation-equivariant, and explain why positional encodings are necessary.onsite

Let $P \in \{0, 1\}^{n \times n}$ be a permutation matrix. Apply it to the input: $X' = P X$ . Self-attention computes $Q' = X' W_Q = P (X W_Q) = P Q$ , similarly $K' = P K$ , $V' = P V$ . The score matrix is

Q' K'^\top = (P Q)(P K)^\top = P (Q K^\top) P^\top.

Row-wise softmax is permutation-equivariant on the rows and the values inside each row form the same multiset on both sides:

\mathrm{softmax}\!\left(Q' K'^\top / \sqrt{d_k}\right) = P\, \mathrm{softmax}\!\left(Q K^\top / \sqrt{d_k}\right) P^\top.

Multiply by $V' = P V$ :

\mathrm{Attn}(Q', K', V') = P\, \mathrm{softmax}\!\left(Q K^\top / \sqrt{d_k}\right) P^\top P V = P \cdot \mathrm{Attn}(Q, K, V),

since $P^\top P = I$ for any permutation. The output is permuted by the same $P$ , never reshuffled. So pure self-attention treats the sequence as a set: no operation distinguishes position $i$ from position $j$ .

This is fatal for any task where order matters (translation, language modelling). The fix is to add or fold in positional encodings at the input: sinusoidal (Vaswani 2017), learned absolute, RoPE (rotary), or ALiBi (linear bias). Each breaks permutation equivariance in a different way; see Transformers for the trade-offs.

Sketch the FlashAttention tiled-softmax algorithm and explain why memory drops from O(n^2) to O(n).research

Setup. Naive attention computes the full $n \times n$ score matrix $S = QK^\top / \sqrt{d_k}$ , materializes it in HBM, applies row-wise softmax, multiplies by $V$ . The $S$ matrix and the softmax probabilities each take $O(n^2)$ memory.

Insight 1: GPUs are HBM-bandwidth-bound on attention. Reading and writing $S$ between HBM and the SM is the dominant cost. If we never materialize $S$ in HBM, we save both memory and time.

Insight 2: softmax has a streaming online form. For a sequence of score blocks $S^{(1)}, S^{(2)}, \dots$ , we can keep a running row-max $m$ and a running denominator $\ell = \sum e^{S - m}$ . When a new block arrives with row-max $m'$ , update

m^{\text{new}} = \max(m, m'), \quad \ell^{\text{new}} = e^{m - m^{\text{new}}} \ell + \mathrm{rowsum}(e^{S - m^{\text{new}}}).

The rescale factor $e^{m - m^{\text{new}}}$ propagates into the running output too. After streaming all blocks, divide the accumulated $V$ -weighted sum by $\ell^{\text{final}}$ to recover the exact softmax-weighted average.

Algorithm. Split $Q$ into row blocks $Q_i$ of size $B_r$ and $K, V$ into column blocks $K_j, V_j$ of size $B_c$ , with $B_r, B_c$ chosen so $Q_i, K_j, V_j$ all fit in SRAM. For each output tile $(i)$ :

Initialize $O_i = 0$ , $\ell_i = 0$ , $m_i = -\infty$ in SRAM.
For each $j$ $j$ :
- Load $K_j, V_j$ from HBM to SRAM.
- Compute $S_{ij} = Q_i K_j^\top / \sqrt{d_k}$ in SRAM.
- $m_i^{\text{new}} = \max(m_i, \mathrm{rowmax}(S_{ij}))$ .
- $\ell_i \leftarrow e^{m_i - m_i^{\text{new}}} \ell_i + \mathrm{rowsum}(e^{S_{ij} - m_i^{\text{new}}})$ .
- $O_i \leftarrow \mathrm{diag}(e^{m_i - m_i^{\text{new}}}) O_i + e^{S_{ij} - m_i^{\text{new}}} V_j$ .
- $m_i \leftarrow m_i^{\text{new}}$ .
After all $j$ : $O_i \leftarrow O_i / \ell_i$ , write $O_i$ back to HBM.

Complexity. The $n \times n$ matrix $S$ is never written to HBM; only the $n \times d$ output and the per-row running statistics ( $n$ scalars each for $m$ and $\ell$ ) leave the SM. Memory drops from $O(n^2)$ to $O(n)$ . Compute is the same $O(n^2 d)$ FLOPs but with much better arithmetic intensity, giving 2-4x wall-clock speedup at $n = 4096$ in the original paper [4]. The backward pass uses the same streaming idea with $O(n)$ extra storage (recompute $S_{ij}$ on the fly during backward).

The result is bit-exact equivalence to naive attention, so FlashAttention is a drop-in replacement, no retraining needed. Modern PyTorch’s scaled_dot_product_attention dispatches to it automatically when shapes and dtypes are supported.

Linear attention promises O(n d^2) compute by replacing exp(q . k) with phi(q)^T phi(k). What does it lose, and when is the trade-off worth it?research

The trick. Replace $\exp(q^\top k)$ with $\phi(q)^\top \phi(k)$ for some $\phi: \mathbb{R}^{d_k} \to \mathbb{R}^{d_\phi}$ . Then

\sum_j \phi(q)^\top \phi(k_j) v_j = \phi(q)^\top \sum_j \phi(k_j) v_j^\top.

The inner sum is independent of $q$ . Compute it once in $O(n d_\phi d_v)$ , store the resulting $(d_\phi, d_v)$ matrix, and answer each query in $O(d_\phi d_v)$ . Total cost over a sequence of length $n$ : $O(n d_\phi d_v)$ , linear in $n$ [5].

For autoregressive decoding the inner sum becomes a running tensor that updates incrementally, giving constant-memory inference (RNN-style), unlike vanilla attention which grows the KV cache as $O(n)$ .

What it loses.

Sharpness. Softmax can produce near-one-hot distributions (“focus on this one token”). A polynomial $\phi$ has bounded entropy in its kernel and tends to produce smoother distributions. On tasks that need sharp lookup (named entity, copy, in-context retrieval), this hurts.
Approximation error. $\phi$ has finite dimension, so it cannot exactly represent the softmax kernel. Performers [6] use random Fourier features so the approximation is unbiased with controllable variance, but at the cost of $d_\phi \gg d_k$ in practice.
Optimization. Linear-attention models often need different learning-rate schedules, longer warmup, and special initialization. They are not drop-in for trained Transformer weights.

When it is worth it. Long-context tasks where $n$ is the binding constraint and exact softmax is unaffordable: long-document classification, very long retrieval-augmented contexts, streaming inference at constant memory. As of 2024, the practical landscape is that FlashAttention covers most of the gap up to $n \sim 32{,}000$ without giving up exactness, so linear attention has been pushed to the very long-context regime ( $n \gg 32{,}000$ ) and to memory-constrained inference. State-space models (S4, Mamba) borrow the same constant-memory recurrence story and are arguably the more competitive long-context alternative today.

Build it yourself

A weekend project that takes attention from formula to working LM head:

Implement the NumPy scaled_dot_product_attention above. Verify it matches torch.nn.functional.scaled_dot_product_attention to within fp32 noise on random inputs.
Build a single-layer Transformer encoder with $h = 4$ heads on a tiny copy task: input is a random sequence of integers, target is the same sequence reversed. With $d = 64$ and $n = 8$ , training on $10^4$ samples reaches near-perfect copy in a minute on CPU.
Replace the encoder with a causal decoder and train a character-level next-token predictor on the first 1000 lines of Hamlet. Read off the attention weights of head 0 layer 0; you will see something close to the previous-token head shown in <AttentionViz />.
Stretch: swap your softmax for an elu(.) + 1 linear-attention kernel and observe the quality drop on copy. Then scale $n$ up to $4096$ and observe the memory advantage.

References

Bahdanau, D., Cho, K., and Bengio, Y. (2014). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015 (arXiv preprint 2014). arXiv:1409.0473
Luong, M.-T., Pham, H., and Manning, C. D. (2015). Effective Approaches to Attention-based Neural Machine Translation. EMNLP 2015. arXiv:1508.04025
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020. arXiv:2006.16236
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., Belanger, D., Colwell, L., and Weller, A. (2020). Rethinking Attention with Performers. ICLR 2021 (arXiv preprint 2020). arXiv:2009.14794
Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. (2022). Efficient Transformers: A Survey. ACM Computing Surveys. arXiv:2009.06732
Goodfellow, I., Bengio, Y., and Courville, A. (2016). Deep Learning. MIT Press. https://www.deeplearningbook.org

Practical Projects Transformers