Deep Clarification

Scaled Dot-Product Attention: the core operation

Start with a sequence of $n$ token embeddings, each of dimension $d_{model}$, stacked into $X \in \mathbb{R}^{n \times d_{model}}$. Attention is a learned, content-addressable lookup. Three projections turn each token into a query, a key, and a value:

$$Q = XW^Q,\quad K = XW^K,\quad V = XW^V$$

with $W^Q, W^K \in \mathbb{R}^{d_{model} \times d_k}$ and $W^V \in \mathbb{R}^{d_{model} \times d_v}$. The whole operation:

$$\text{Attention}(Q,K,V) = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

Walk it term by term, because each piece is load-bearing.

The score matrix $QK^\top$

$QK^\top \in \mathbb{R}^{n \times n}$. Entry $(i,j)$ is $\langle q_i, k_j \rangle = \sum_{c=1}^{d_k} q_{ic}k_{jc}$ — the dot product of token $i$‘s query with token $j$‘s key. This is an unnormalized similarity in a learned subspace. It is not similarity in the original embedding space; $W^Q$ and $W^K$ rotate/scale tokens into a space where “relevance for this attention head” becomes geometric alignment. Two tokens that are unrelated semantically can have high $q_i \cdot k_j$ if the head’s job is, say, “verb agrees with subject.”

Note the asymmetry: $q_i \cdot k_j \neq q_j \cdot k_i$ in general, because $W^Q \neq W^K$. The relation “token $i$ attends to $j$” is directed. This matters — it’s what lets a pronoun query its antecedent without the antecedent symmetrically querying the pronoun.

Why $\sqrt{d_k}$

Suppose query and key components are independent, mean 0, variance 1. Then $q_i \cdot k_j = \sum_{c=1}^{d_k} q_{ic}k_{jc}$ is a sum of $d_k$ zero-mean unit-variance terms, so it has variance $d_k$ and standard deviation $\sqrt{d_k}$. For $d_k = 128$ that’s logits with magnitude ~11. Feeding logits of that scale into softmax pushes it toward a one-hot distribution, where the gradient of softmax is nearly zero (saturated). Dividing by $\sqrt{d_k}$ renormalizes the score variance back to ~1, keeping softmax in its responsive regime and gradients healthy. It’s variance control, not a cosmetic constant.

The softmax row-normalization

Softmax is applied row-wise. For row $i$:

$$\alpha_{ij} = \frac{\exp(s_{ij})}{\sum_{j’=1}^{n}\exp(s_{ij’})},\qquad s_{ij} = \frac{q_i\cdot k_j}{\sqrt{d_k}}$$

Each row $\alpha_{i\cdot}$ is a probability distribution over all $n$ positions: nonnegative, sums to 1. This is the attention weight — how much token $i$ pulls from each token $j$. The convexity matters: the output for token $i$ is a convex combination of value vectors, so it lives in the convex hull of ${v_j}$. Attention can interpolate among values but cannot extrapolate beyond them — the nonlinearity and range-extension come from the surrounding MLP, not attention.

Weighted value aggregation

$$\text{out}i = \sum{j=1}^{n} \alpha_{ij}, v_j$$

The query/key machinery decides the weights; the values carry the content that actually gets written into the residual stream. This Q/K vs V split is the crux: “which tokens are relevant” ($QK^\top$) is decoupled from “what information moves” ($V$). A head can route based on syntactic position while transporting semantic content.

The Jacobian of softmax (why gradients flow the way they do)

For a single softmax row $\alpha = \text{softmax}(s)$:

$$\frac{\partial \alpha_i}{\partial s_j} = \alpha_i(\delta_{ij} - \alpha_j)$$

In matrix form $J = \text{diag}(\alpha) - \alpha\alpha^\top$. Two consequences:

• When one $\alpha_i \to 1$ (saturated/confident attention), the whole Jacobian $\to 0$ — the head stops learning where to look. This is exactly the failure mode the $\sqrt{d_k}$ scaling defends against.
• It’s coupling, not independence: raising the score on token $j$ necessarily drains weight from all others ($-\alpha_i\alpha_j$ off-diagonals). Attention is zero-sum across positions within a row.

Multi-head attention

A single softmax mixture is a weak averaging operation — one head can only express one routing pattern per position. Multi-head runs $h$ attentions in parallel in disjoint subspaces:

$$\text{head}_m = \text{Attention}(XW^Q_m, XW^K_m, XW^V_m),\quad m=1\dots h$$ $$\text{MHA}(X) = \big[\text{head}_1 ,|, \cdots ,|, \text{head}_h\big],W^O$$

Typically $d_k = d_v = d_{model}/h$, so the concatenation returns to $d_{model}$ and total compute matches one full-width head. $W^O \in \mathbb{R}^{hd_v \times d_{model}}$ mixes the heads back into the residual stream.

Mechanistically it’s cleaner to view $W^O$ as block-partitioned by head: $\text{MHA}(X) = \sum_m (\text{head}_m) W^O_m$, where $W^O_m$ is head $m$‘s slice. Each head independently reads from the residual stream via $W^{QKV}_m$ and writes back via $W^O_m$ — heads communicate only through the residual stream, never directly. The product $W^V_m W^O_m$ is the head’s “OV circuit” (what it writes), and $W^Q_m{}^\top W^K_m$ is its “QK circuit” (where it reads); these are the natural low-rank objects to analyze, not the four matrices separately.

Causal masking

For autoregressive decoding, token $i$ must not see $j > i$. Add a mask $M$ before softmax where $M_{ij} = -\infty$ for $j > i$, else $0$:

$$\alpha = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)$$

$\exp(-\infty)=0$ zeroes future positions after normalization is set up, so the surviving weights still sum to 1 over the visible prefix. In practice $-\infty$ is a large negative constant (e.g. $-10^9$). The score matrix becomes lower-triangular-effective; this is what makes the $n^2$ structure a strict prefix-causal computation and what enables KV-caching at inference (keys/values for past tokens never change, so they’re computed once).

Positional information

Raw attention is permutation-equivariant: permute the input rows and the output permutes identically — $QK^\top$ depends only on content, not order. Pure attention literally cannot tell “dog bites man” from “man bites dog.” Position must be injected:

• Sinusoidal/learned absolute: add $P \in \mathbb{R}^{n\times d_{model}}$ to $X$ before projection. Order leaks into $Q,K$ through the addition.
• RoPE (rotary), now dominant: rotate $q_i$ and $k_j$ by position-dependent angles. Pair up dimensions; for pair $c$ with frequency $\theta_c$, multiply by the 2×2 rotation $R(i\theta_c)$. The key identity:

$$\langle R(i\theta)q_i,; R(j\theta)k_j\rangle = \langle q_i,; R((j-i)\theta)k_j\rangle$$

The score depends only on the relative offset $j-i$, not absolute positions, yet it’s implemented by rotating each token independently — no $n^2$ relative-bias table. Different dimension pairs use geometrically spaced $\theta_c$ (a la $10000^{-2c/d}$), giving a spectrum of wavelengths from fast (local) to slow (long-range), which is what gives RoPE its length-extrapolation behavior.

Cost structure

$QK^\top$ is $n \times d_k \times n$ → $O(n^2 d_k)$ time and the score matrix is $O(n^2)$ memory. This quadratic-in-sequence-length term is the entire scaling story of transformers. FlashAttention doesn’t change the FLOPs but tiles the computation so the $n\times n$ matrix is never materialized in HBM — it streams blocks through SRAM, recomputing softmax statistics with an online (running-max, running-sum) normalizer, turning the memory bottleneck from $O(n^2)$ to $O(n)$. That’s why it’s a memory-IO win, not an arithmetic one.

A useful one-line mental model: attention is a differentiable, soft, content-addressable dictionary — queries look up keys, the match scores become a probability-weighted read over values, and the only nonlinearity in the whole operation is the softmax that turns scores into a normalized mixture. Everything else (the heads, RoPE, masking, scaling) is engineering around keeping that lookup expressive, ordered, causal, and trainable.

If you want, I can go a level deeper on any one of these — the residual-stream/circuits view (QK and OV circuits, composition between heads), the FlashAttention online-softmax recurrence in full, or the exact backward pass through $QK^\top V$.

The Transformer architecture relies on the Attention Matrix as its core computational primitive. This operation is not a fixed logic gate but a parametric linear layer followed by a normalization and a score-function based aggregation.

Below is the rigorous mechanical breakdown of the mathematical operations within a Self-Attention layer (specifically the “Scaled Dot-Product” attention described in the original VITS paper).

1. The Input Representation

Let the input sequence at a specific stage (e.g., post-FFN in an encoder) be a matrix $X$: $$ X \in \mathbb{R}^{S \times D} $$ Where:

• $S$: Sequence length (number of tokens).
• $D$: Hidden dimension size ($d_{model}$).

2. Linear Projections (The Q, K, V Mechanism)

Attention does not operate on raw inputs. It operates on three learned projections of $X$. These are learned linear transformations.

$$ Q = X W_Q $$ $$ K = X W_K $$ $$ V = X W_V $$

Where $W_Q, W_K, W_V \in \mathbb{R}^{D \times K}$ are learnable weight matrices, and $K$ is the key vector dimension.

• Typically, dimensions are split as $K = \frac{D}{h}$, where $h$ is the number of “heads”.
• Mathematical implication: This decomposes the sequence into directionally distinct feature subspaces while keeping dimensions compact. However, for the mechanics of a single head, we assume the output dimensions $Q, K, V$ all have size $K$.

3. The Dot-Product Calculation (Similarity Scoring)

Attention is the calculation of “affinity” between the Query and Keys. This is computed via the matrix product:

$$ A_{raw} = Q K^T $$

Dimensional Analysis:

• $Q: (S \times K)$
• $K^T: (K \times S)$
• Result $A_{raw}: (S \times S)$

Mechanical Interpretation: The element at row $i$, column $j$ ($A_{ij}$) represents the dot product of the vector embedding of token $i$‘s Query and token $j$‘s Key: $$ A_{ij} = \mathbf{q}i \cdot \mathbf{k}j = \sum{d=1}^K q{id} k_{jd} $$ Why this specific multiplication? The dot product $\mathbf{q} \cdot \mathbf{k}$ correlates with $\cos(\theta)$ if features are normalized. If $\mathbf{q}$ and $\mathbf{k}$ are high-dimensional vectors representing text embeddings, a high dot product indicates that the semantic information encoded in $\mathbf{q}_i$ is proportional to the semantic “address” $\mathbf{k}_j$.

4. Scaling (Stability)

Before applying Softmax, the scores must be scaled: $$ A_{score} = \frac{A_{raw}}{\sqrt{K}} $$

Mathematical justification:

• Vanishing Gradients: If $K$ is large, the distribution of dot products $\mathbf{q} \cdot \mathbf{k}$ tends toward $\mathcal{N}(0, K)$ (by the Central Limit Theorem, assuming unit variance inputs).
• Softmax Sensitivity: The Softmax function $\sigma(x) = \frac{e^x}{\sum e^x}$ has derivatives that vanish as input variance increases. If the variance of the dot products increases (due to large $K$), the values of $e^x$ become numerically unstable (overflow) or the Softmax becomes too saturated (near 0 or 1), which prevents learning.
• Recalibration: Dividing by $\sqrt{K}$ brakes the variance of the scores back to a standard normal distribution $\mathcal{N}(0, 1)$, making the incoming gradients to the weights $W_Q$ and $W_K$ tractable during backpropagation.

5. Softmax Normalization (The Attention Weights)

$$ A = \text{softmax}S(A{score})_{ij} $$

The Softmax is applied row-wise.

• Input: $A_{score}$ where $A_{score} \in \mathbb{R}^{S \times S}$.
• Operation: For each row $i$, normalize the values such that $\sum_{j=1}^{S} A_{ij} = 1$.

Mathematical Function: $$ A_{ij} = \frac{e^{A_{score}{ij}}}{\sum{l=1}^{S} e^{A_{score}_{il}}} $$

Graph Interpretation: The Attention Matrix $A$ is a row-stochastic matrix (an adjacency matrix of a Markov process).

• Row $i$: Represents the query for token $i$.
• Row Sum = 1: Represents that the model allocates its “query energy” 100% across all keys in the sequence.
• Result: Token $i$ creates a probability distribution over all tokens $j$, determining what values it should attend to.

6. Weighted Summation (Value Retrieval)

The core goal of attention is to retrieve information from the Values, weighted by the Attention Matrix:

$$ \text{Output} = A V $$

Dimensional Analysis:

• $A: (S \times S)$
• $V: (S \times K)$
• Result: $(S \times K)$

Mechanical Logic: $$ \text{Output}{ik} = \sum{j=1}^{S} A_{ij} v_{jk} $$

This is a linear combination of the value vectors $\mathbf{v}{1}, \dots, \mathbf{v}{S}$.

• If attention $A_{ij}$ is high, key $j$ strongly aggregates into the query for token $i$.
• If all weights are zero, the output is zero.
• Interpretation: The final output vector for position $i$ is a weighted profile of the original Vector Representations in the sequence. It effectively replaces token $i$ with a mixture of context.

7. Multi-Head Attention (Scaling $H$)

The full mechanism runs these operations in parallel $h$ times (number of heads).

1. Linear Decomposition: The hidden dimension $D$ is split into $h$ heads. Each head $m$ has its own $W_Q^m, W_K^m, W_V^m$. $$ Q^m, K^m, V^m \in \mathbb{R}^{S \times (D/h)} $$

2. Parallel Matrices: The $h$ individual attention matrices $A^m$ are computed independently but causally in the same layer.

3. Concatenation: $$ \text{Concat}(A^1 \dots A^h) \otimes D_{concat} $$ Typically, one sums the outputs scaled by $\frac{1}{\sqrt{h}}$: $$ \text{Output}_{heads} = \frac{1}{\sqrt{h}} \text{Concat}(\text{Head}_1, \dots, \text{Head}_h) $$ (Normalization of concatenation prevents gradients from exploding if the concatenation preserves sum-norms).

4. Final Projection: $$ \text{FinalOutput} = \text{Output}_{heads} W_O $$ Where $W_O \in \mathbb{R}^{(D \times h) \times D}$ projects the aggregated head information back to the original model dimension $D$.

8. Causal Attention Masking (Decoder Specifics)

In Decoder layers, the attention mechanism includes a Mask. $$ \text{MaskedScore}{ij} = A{score}{ij} - M{ij} $$

Where $M_{ij} = -\infty$ if $j > i$ and $0$ otherwise.

• Mechanically: In a matrix product, subtracting $-\infty$ results in $-\infty$ (representing double negative infinity with exponent).
• Softmax behavior: $$ e^{-\infty} = 0 $$ This forces the Softmax row-sum to be 0 for terms $j > i$. The token $i$ only attends to positions $j \leq i$. This ensures autoregressive generation (no “peeking” into the future).

9. Why Linear? (The “Dot-Product” Obsession)

There are alternative attention types (e.g., sparse, bilinear). However, Dot-Product Attention is persistent because:

1. Triangular Kernel Momentum: It allows for matrix-vector multiplication in $O(S \cdot D^2)$) for all tokens, which is hardware-optimized (GEMM operations).
2. Variable NCS (Compute Spend): Unlike previous methods where embeddings are fixed, the vector direction can change via $W_Q$ and $W_K$ per head, creating “spotlights” over semantic regions.
3. Probe Mechanism: Since $K$ is relative to $Q$, the Attention Matrix $A$ functions as a probe.
   • If $Q$ is token $i$, $A_{ij}$ tells you “how much information token $j$ must contribute to token $i$.”
   • This structure enables the Transformer to learn “syntactic rules” or “long range dependencies” not by mixing features (CNN) but by mixing positions (Attention).

10. Internal Feasibility Analysis (Gradient Flow)

The total differential of the output w.r.t an input token $x_{i}$ is:

$$ \frac{\partial \text{Output}}{\partial X_i} $$

The chain rule involves four components:

1. Projection Gradient: $\frac{\partial A}{\partial Q} \cdot \frac{\partial Q}{\partial X_i}$.
2. Masked Attention Weight Gradient: $\frac{\partial A}{\partial A_{ij}} \cdot V_j$ (Softmax derivative).
3. Value Gradient: From $W_Q$, $W_K$, $W_V$.
4. Softmax Stability: Because Softmax is differentiable and bounded, gradients do not vanish during backprop. The row-wise normalization keeps gradient norms within a constant range, preventing explosions that standard matrix multiplication would cause during sequence length.

Summary of Mechanics

• Transformation: Input vectors $\to$ Q, K, V spaces via learned linear maps.
• Scoring: Aggregation of semantic affinity via inner product.
• Normalization: Probabilistic mapping to context weights (Softmax).
• Aggregation: Retrieval of values via probabilistic summation.
• Projection: Re-centering vectors into the original space to allow stacking/module interaction.

This is the operational kernel of the Transformer. It is not a hard-coded logic gate but a continuous, learnable matrix operation that optimizes the trade-off between computational cost ($O(S^2)$ for self-attention) and expressiveness ($O(d_{model})$ for parameter sharing).

Surface explanation

Surface: The Transformer attention mechanism computes output through a weighted aggregation of input values. It maps a sequence of input tokens into a representation space where relevance is measured by the alignment between query and key vectors.

Canonical Operation:

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

Level 0 Context:

• Inputs: Sequences of query vectors $Q$, key vectors $K$, and value vectors $V$ (derived from projective weights).
• Action: Compute similarity scores $S = \frac{Q K^T}{\sqrt{d_k}}$, normalize to probabilities via row-wise softmax ($\alpha$), then aggregate values ($\alpha V$).
• Constraint: The operation is globally self-attentive (each position attends to all others).

Mechanistic clarification — two levels deeper

Level 1 beneath: Numerical Stability via Scaling and Normalization

Atom: Variance Control via Central Limit Theorem (CLT)

Mechanism: The critical barrier at this level is numerical stability in the exponentiation phase. When dot products $y = q \cdot k$ are computed in $d_k$ dimensions, the magnitude of the result grows linearly with dimension if vectors are not normalized.

1. Distributional Assumption: With independent, zero-mean activations of unit variance, the dot product $Q \cdot K^T$ is a sum of $d_k$ terms. By the CLT, the variance of the attention score grows linearly with dimension: $\text{Var}(S) \approx d_k$.
2. Pathology: In high-dimensional spaces ($d_k=64+$), the resulting scores span large ranges. The exponential function $e^x$ rapidly diverges. Without scaling, scores like $8$ ($e^8 \approx 3000$) and $-8$ ($e^{-8} \approx 0.0003$) incur extreme ratios.
3. Resolution: The scaling factor $\frac{1}{\sqrt{d_k}}$ is applied to keep the pre-softmax logits in a standard distribution ($N(0, 1)$). This ensures exponential terms remain within computable bounds (preventing “exploding softmax” in forward pass).

Atom: Softmax Jacobian Structure and Vanishing Derivatives

Mechanism: The second barrier is the interaction between scaling and the gradient magnitude during normalization.

1. Jacobian: For a single row $z$, the diagonal element of the softmax Jacobian is $\frac{\partial a_i}{\partial z_i} = a_i(1-a_i)$. Off-diagonal elements are $-\frac{a_i a_j}{\sum^2}$.
2. Saturation: When softmax outputs approach probabilities of exactly 1 (dominant key), the derivative $a(1-a)$ approaches 0.
3. Propagation: High variance in the scaling-free scores (which the scaling factors attempt to control) creates sharp gradients in the forward pass, but leads to vanishing gradients in the backward pass for the highest-confidence scores. Even if the attention weight is “right”, the gradient for updating the parameters encoding that weight collapses.

Level 2 beneath: Vector Geometry and Gradient Landscape Propagation

Atom: Dot-Product Geometry and Q/K Decoupling Necessity

Mechanism: This level explains why $Q$ and $K$ must be learnable, separate projections ($W_Q, W_K$) rather than identical matrices.

1. Selectivity Loss: If $W_Q = W_K$, the attention score becomes a quadratic form: $H W^T H^T$. This restricts the energy landscape of the attention matrix, forcing all positions to attend based on the projection of the same subspace.
2. Independence: Decoupling allows $Q$ and $K$ to live in different conceptual subspaces. The theory allows $Q_i \cdot K_j$ to act as a selectivity filter independent of initialization constraints.
3. Observation: While $W_Q$ and $W_K$ initialization is arbitrary, the matrix multiplication architecture allows for higher-rank expressivity in the attention heads.

Atom: Gradient Flow in Multi-Query Architecture (GQA/MHA)

Mechanism: The structure of $K$ and $V$ dimensions affects the variance of the incoming gradient.

1. GQA Modification: In Group Query Attention, all heads share the same $K$ dimension ($d_{kv}$), reducing the memory bandwidth required for looking up keys.
2. Gradient Impact: When comparing standard MHA ($d_k$ per head) vs. GQA, the standard fan-out factor in $K$ updates is reduced. This modifies the gradient variance accumulation; fewer unique $K$ columns generates less variance in the backward flow but imposes structural constraints on the attention content.

Epistemic boundary

Epistemic boundary: Hedged claims regarding the primary driver of the scaling factor and empirical thresholds compared to settled derivation of the Jacobian and Decoupling effects.

Claim: “The scaling factor explicitly addresses gradient stability.” Resolution: Hedged.

• Evidence: The 2017 Transformer paper motivates scaling primarily to prevent “exploding softmax” in the forward pass (numerical stability). While gradient stability follows from numerical stability, empirical evidence does not cite “gradient variance control” as the primary theoretical driver in the original text.
• Scope: It is considered a design bylaw ($\sqrt{d_k}$) but the exact value is tunable (“variable scaling” variants exist) without breaking mathematical correctness.

Claim: “Decoupling $W_Q = W_K$ collapses attention to a naive mean.” Resolution: Supported.

• Evidence: Mathematical ablation in standard literature confirms that identical projections (specifically self-similarity) can result in degenerate attention maps in select regimes.

Confidence Matrix:

Practical implications

• Practical implication: If $d_k$ increases, $\sqrt{d_k}$ must increase to preserve variance, or training will experience “vanishing gradients from saturation.”
• Practical implication: If attention mechanisms are observed to have dead weights (saturation $\approx$ 1), the addition of Residual Connections and Layer Normalization is not optional; it is mathematically necessary to buffer the vanishing gradients caused by the softmax Jacobian.
• Practical implication: Using different weight matrices ($W_Q \neq W_K$) is not just an ablation, but a requirement for independent representation of “what I am searching for” (Q) vs “what I am matching against” (K), expanding the geometric subspace of attention interactions.

Surface explanation

Surface: The accessible-level description of Multi-Head Self-Attention is that it takes a sequence of token embeddings, projects them into query, key, and value matrices, computes scaled dot-product attention scores, applies a softmax to turn those scores into a probability weighting, and calculates a weighted sum of the values. This is formalized as $\text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{QK^T}{\sqrt{d_k}}\right)V$. The mechanism then runs $h$ independent copies (heads) of this process in parallel, concatenates their outputs, and mixes them through a final linear projection. The result is a weighted sum of value vectors, where the weights encode the compatibility between queries and keys.

Mechanistic clarification

Level 1 beneath: The tensor flow begins by establishing the operands. Queries ($Q$), keys ($K$), and values ($V$) are not native inputs but independent learned linear maps of the identical input matrix $X \in \mathbb{R}^{n \times d}$ (where $n$ is token count and $d$ is model dimension). Specifically, $Q = X W_Q$, $K = X W_K$, and $V = X W_V$, with the weight matrices learned during training. Because $X$ is the exact same matrix feeding all three, the operation is self-attention. For $h$ heads, these projection matrices are factored into $h$ blocks (each of shape $d \times (d/h)$), executed as a single large general matrix multiply followed by a reshape to $(h, n, d/h)$ per tensor. This means multi-head attention allocates the same rank budget differently, rather than adding net parameter capacity. The raw compatibility is computed as the scaled inner product $S^{(i)} = Q^{(i)} K^{(i)\top}/\sqrt{d_k} \in \mathbb{R}^{n \times n}$, yielding a pairwise similarity score for every position combination, not yet a probability. Masking is applied as an additive bias $M$ (where $M_{ab} = -\infty$ for disallowed pairs and $0$ otherwise) before the softmax. The row-wise softmax is computed as $\exp(x - \max x)/\sum \exp(x - \max x)$, preventing overflow when logits are large, and forcing the row into a distribution summing to 1; the $-\infty$ mask ensures $\exp(-\infty) = 0$, yielding exactly zero weight. The output is then $O^{(i)} = A^{(i)} V^{(i)}$, meaning each output row is a convex combination of value vectors weighted by relevance. Finally, concatenation stacks head outputs, and a learned projection matrix $W_O \in \mathbb{R}^{h d_v \times d}$ mixes them. Without $W_O$, heads would feed the residual stream as disjoint, non-interacting signals. The complete layer output is the residual update $X + O$.

Level 2 beneath: The geometric, statistical, and computational reasons this tensor flow takes its specific form become clear at the next depth. The $\sqrt{d_k}$ scaling is a variance stabilization mechanism. If query and key components are independent, zero-mean, unit-variance (as variance-preserving initialization yields), the variance of their dot product grows linearly with $d_k$. Large-variance logits push the softmax toward near one-hot outputs where gradients vanish; dividing by $\sqrt{d_k}$ forces logit variance back to $\sim 1$, keeping the softmax in a gradient-rich regime. Softmax attention operates as differentiable retrieval. Row $A_a$ is a categorical distribution parameterized by $\exp(q_a \cdot k_b/\sqrt{d_k})$, and the output is its expectation $\mathbb{E}{b \sim A_a}[V_b]$. Rather than hard-selecting a single key, attention returns a soft mixture; training learns whether a head sits in a sharp (near-argmax) or broad (near mean-pool) regime. Multi-head architecture provides independent similarity functions, not enforced orthogonality. While the $d_k = d_v = d{model}/h$ split makes multi-head parameter-equivalent to single-head, the architecture allows the model to learn different projection functions per subspace. Specialization (e.g., induction heads, syntactic heads) is emergent from training dynamics and surrounding non-linearities, not architecturally enforced. The mechanism itself is entirely position-blind; it has no inherent notion of order. Positional addressing must be injected before the QKV projections (e.g., via sinusoidal embeddings or Rotary Position Embeddings that twist $Q$/$K$ vectors), cleanly separating content-addressing from positional-addressing. Computationally, the complexity profile is dominated by $n^2$. The $QK^\top$ and $AV$ operations require $O(n^2 d_k)$ FLOPs, and the intermediate $n \times n$ matrix requires $O(n^2)$ memory. Because this matrix must be written to High Bandwidth Memory (HBM) in standard implementations, the operation is memory-bandwidth-bound, not compute-bound, making context length a binding constraint. At inference, autoregressive generation leverages a KV-cache: previous step $K$ and $V$ matrices are cached, so only the new token’s $Q$ is computed, reducing the per-step cost to $O(nd)$ per layer. However, cache memory ($2 \times n \times n_{kv} \times (d_k + d_v) \times b$ bytes) becomes the dominant memory cost, driving engineering focus toward KV-head sharing.

Alternative Level 2 mechanism: Additive attention. Plain-terms statement: A feed-forward network computes compatibility scores instead of a dot product. Epistemic standing: An alternative computational path that predated dot-product attention, largely superseded by the computational efficiency and parallelization benefits of matrix multiplication, though theoretically equivalent in representational capacity.

Alternative Level 2 mechanism: Linear attention. Plain-terms statement: Replaces the softmax kernel with a finite-dimensional feature map $\phi(Q)\phi(K)^\top$, admitting an $O(n)$ associative rewrite (e.g., Performer, FAVOR+). Epistemic standing: An approximation of softmax-as-kernel that sacrifices exact attention weights for linear sequence scaling, trading expressivity for asymptotic efficiency.

Alternative Level 2 mechanism: Multi-Query / Grouped-Query Attention (MQA/GQA). Plain-terms statement: Retains the identical attention math, but shares $K$ and $V$ projections across multiple heads (or all heads). Epistemic standing: Not a competing mathematical framework, but an empirical engineering tradeoff that drastically reduces KV-cache memory and bandwidth at a moderate cost to model quality.

Alternative Level 2 mechanism: State-space models (SSMs). Plain-terms statement: Replaces the $O(n^2)$ attention matrix with a recurrent or selective scan over a continuous-time hidden state. Epistemic standing: A fundamentally different inductive bias that is linear in $n$ and highly efficient for long contexts, but possesses weaker explicit content-based retrieval capabilities compared to softmax attention.

Epistemic boundary

Epistemic boundary: The tensor flow, the $1/\sqrt{d_k}$ scaling as an operation, masking as an additive bias, KV-cache correctness, and the parameter-equivalence geometry are fully settled and production-verified mathematics. However, the theoretical justification for the $\sqrt{d_k}$ scaling relies on the assumption that vector components are independent and zero-mean. In deep, highly non-linear networks, activations exhibit complex correlations, meaning the scaling operates in practice as an empirical heuristic that stabilizes learning dynamics, rather than a strict mathematical guarantee. Furthermore, while multi-head specialization is well-documented empirically, there is no tight theoretical characterization of the specialization-versus-redundancy ratio in trained models. The active research frontier includes tight bounds on attention expressivity, whether softmax is strictly necessary or profitably replaceable (e.g., by linear or sparse attention), and how attention mechanisms optimally compose across depth.

Practical implications

• Practical implication: Because context length is a memory-bandwidth problem before it is a compute problem (due to the $n \times n$ intermediate matrix and KV-cache dominance), optimization must target memory I/O. This mechanistic reality directly enabled FlashAttention, which restructures the $QK^\top \to \text{softmax} \to AV$ sequence by tiling blocks into fast on-chip SRAM and recomputing the softmax normalization on the fly. This avoids writing the $n \times n$ matrix to HBM, reducing intermediate memory complexity from $O(n^2)$ to $O(n)$ without altering the mathematical output, making the operation fundamentally compute-bound with only a small constant FLOP overhead. It also drives KV quantization, KV eviction, and paged attention.
• Practical implication: Head count and head dimension are coupled but separable knobs. The budget is $h \cdot d_k \le d$. Deciding whether to allocate this budget to many narrow heads or fewer wide heads is a quality-versus-speed tradeoff with no closed-form mathematical answer, and the optimal configuration is empirically regime-dependent (differing across vision, language, and long-context tasks).
• Practical implication: The output projection matrix $W_O$ performs load-bearing work. Initializing $W_O$ to the identity matrix or removing it measurably degrades performance. Heads do not operate as independent, pre-formed features; the model must actively learn to read and aggregate them jointly, reframing $W_O$ from a passive re-mixing step to an active learned aggregator.
• Practical implication: Causal masking is computationally free. Because it is implemented purely as an additive bias on the logits before the softmax, it requires no separate code path or branching logic, allowing a single unified kernel to handle encoder, decoder, and prefix-LM causal behaviors seamlessly.
• Practical implication: Position encoding exists on a separate design surface. Methods like Rotary Position Embeddings (RoPE), ALiBi, and absolute embeddings are not part of the attention math itself. Swapping them only changes the input to the QKV projections, leaving the core attention mechanism intact.

Surface explanation

Surface: Attention is a weighted sum of value vectors, where each weight is the softmax-normalized dot-product similarity between a query and a key — Attention(Q,K,V) = softmax(QK^T/√d_k)·V, run in parallel across h heads and projected back through W^O. The document you pasted documents this level completely and correctly: the operations, shapes, gradients, masking, complexity. It describes what is computed, not why this computation does work. (Confidence: high — settled linear algebra.)

Two of its statements at this level need correcting, and the corrections are not pedantry — each one is the loose end that the mechanism beneath picks up, so they anchor the descent rather than sitting beside it.

• The scaling argument rests on a false premise as written. The document says “two random d_k-dimensional unit vectors” with “variance d_k” — but unit vectors have dot products bounded in [−1, 1], so the variance cannot be d_k. The correct premise (Vaswani et al.’s footnote in “Attention Is All You Need”) is that q and k have independent components with mean 0 and variance 1; then q·k = Σᵢ qᵢkᵢ is a sum of d_k unit-variance terms, giving mean 0, variance d_k, std √d_k — which is why you divide by √d_k. The distinction matters because scaling is variance-control of the pre-softmax logits, the thread the deeper QK-circuit level picks up. (Confidence: high — footnote verified against external sources.)
• The softmax-gradient line is notationally loose. ∂L/∂scaled_scores = softmax'(·)·∂L/∂attention_weights writes softmax’ as a scalar multiplier, but it is a per-row m×m Jacobian. The explicit Jacobian given immediately below it in the document is correct; the summary line above it is not a clean matrix equation. (Confidence: high.)

Mechanistic clarification — three levels deeper

Level 1 beneath: the operation is a differentiable content-addressable lookup — a soft dictionary read.

A hard dictionary does out = store[k] for an exact key match — non-differentiable, because the index is a step function. Attention is the relaxation of exactly that lookup along two axes: (a) it retrieves by similarity rather than exact match — the dot product Q[i]·K[j] is a soft match score (“how much does entry j answer query i?”); (b) it returns a convex blend of all values weighted by match quality rather than one entry. This is genuinely beneath the surface because the formula shows three matrices but does not reveal that the three projections play asymmetric retrieval roles (probe / stored-key / stored-content). (Confidence: high — the associative-memory / soft-retrieval framing is a settled interpretation.)

The softmax is the load-bearing piece, and its causal mechanism is exponential gap-amplification. softmax behaves as a soft argmax because exp turns an additive difference between two scores into a multiplicative ratio between their weights. As the top score pulls ahead, its exponential dominates the normalizer and that weight collapses toward 1 (retrieve the single best match); as scores flatten, the exponentials equalize and weights approach uniform (average over everything). The exp nonlinearity is the actual mechanism that makes the operation interpolate continuously between “look up one thing” and “average over all things” — and which behavior occurs is learned, per query, per head. (Confidence: high.)

This view licenses predictions the recipe cannot. It predicts copying (set V to carry token identity, set the QK match to “find the position I want,” and the output moves that token here). It predicts induction heads — the circuit implementing “I saw AB earlier, I’m now at A, so predict B” by querying for the token that followed the current token last time — readable straight off the retrieval interpretation, unreadable off “weighted sum of values.” It also makes the KV-cache obvious rather than a trick: in autoregressive generation, cached K and V are the stored dictionary entries, so only the new query needs computing. (Confidence: high for the framing; induction heads are a well-replicated empirical finding — Olsson et al. 2022, building on Anthropic’s transformer-circuits work — though calling them the mechanism of in-context learning is a partial, debated claim.)

Level 2 beneath: the head is a pair of low-rank bilinear circuits reading from and writing to a shared bus.

“Query–key similarity” is not a property of tokens; it is a single learned bilinear form. Folding the projections in: scoreᵢⱼ = (xᵢW^Q)(xⱼW^K)^T = xᵢ(W^Q W^K^T)xⱼ^T = xᵢ·M·xⱼ^T. The two learned matrices never appear apart — they collapse into one d_model × d_model interaction matrix M = W_QK = W^Q W^K^T. Q and K are an implementation of M; M is the genuine learned object. (Convention note: Anthropic’s framework writes W_QK = W_Q^T W_K under a left-acting/column-vector convention; this is the same object under the right-acting/row-vector convention your document uses — both correct and internally consistent.) (Confidence: high — substitution.)

From this, d_k reveals itself as a rank bottleneck, not a “dimension of keys.” Because W_Q, W_K ∈ ℝ^{d_model×d_k} both factor through the shared d_k axis, rank(W_Q W_K^T) ≤ min(rank W_Q, rank W_K) ≤ d_k. So d_k is the rank budget of the bilinear form that scores every token pair — choosing d_k = d_model/h chooses how expressive each head’s where-to-look pattern can be. (Confidence: high.)

The output side factors symmetrically: the QK/OV-circuit decomposition. What a head writes is governed by a second collapsed matrix W_OV = W^V W^O acting on the residual stream, also rank-bounded. A head is therefore two near-independent circuits: a QK circuit (W^Q W^K^T) that decides where attention reads from (the pattern), and an OV circuit (W^V W^O) that decides what gets copied from there to here. (Elhage et al. 2021, “A Mathematical Framework for Transformer Circuits.”) The decomposition is exact, not merely suggestive, because softmax is the only nonlinearity that mixes tokens — everything else in a head is linear in the token representations — so a head factors cleanly into a nonlinear attention pattern (QK) times a linear read-write channel (OV). (Confidence: high on the algebra and the framing being standard in mechanistic-interpretability work; the interpretive claim that the factored matrices are the “natural unit” is well-supported but interpretive.)

The keys, queries, and values are low-rank linear reads of a shared residual stream — the running sum every layer reads from and adds to — and that stream carries position, not just content. The bilinear input xᵢ is not pure token content: it carries positional information too — either an additive sinusoidal/learned positional embedding folded into x, or, in RoPE-style models, a rotation applied to q and k inside the QK product. So xᵢ M xⱼ^T is content-and-position-addressable, which is precisely what lets a QK circuit key on “the token after a token like this one” rather than token identity alone. Induction heads exploit exactly this: a previous-token head writes a “what preceded me” signal into x, and the downstream head’s QK circuit reads that positional signal — closing the gap left at Level 1, where pure content similarity could never by itself select “the next position.” (Confidence: high.)

And this isolates attention’s whole job: it is the only operation that moves information between token positions. MLPs, LayerNorm, and the nonlinearities all act strictly position-wise. (Positional encodings are not a counterexample: they inject position information into each token’s own representation but do not route content from one position to another — the inter-position transport is done entirely by attention.) So mechanistically the entire job of attention is inter-position routing: read from the bus at some positions, write to the bus at others. (Confidence: high — exact algebra from Elhage et al. 2021.)

The scaling correction, reframed at this level: √d_k protects the gradient that learns M. Scaling controls the variance of the bilinear form’s output before softmax, keeping M’s logits in softmax’s sensitive, non-saturated range. Saturate it and the soft lookup hard-collapses to one-hot, and — critically — the softmax Jacobian diag(a) − aaᵀ goes to zero, so gradients to the QK circuit vanish. Scaling is gradient-flow protection for the learning of M, not merely numerical hygiene. (Confidence: high.)

Level 3 beneath: why softmax and why √d_k specifically — a competitive, entropy-maximizing allocation forced to spend all its weight.

(There is a real depth tension worth flagging here: one descent treats the softmax’s thermodynamic structure as a genuine third mechanistic rung beneath the QK/OV factoring; another folds the softmax’s causal mechanism into Level 1 as gap-amplification and reserves “Level 3” for the epistemic-semantic boundary below. Both readings name real structure — the atoms here are vertical mechanism regardless of which rung they are numbered, and the epistemic-boundary material is kept separate from them.)

Softmax is the Boltzmann/Gibbs distribution: the unique distribution maximizing entropy subject to a fixed expected score. The row scores play the role of negative energies and √d_k plays the role of temperature. This is the deeper “why softmax” beneath the soft-argmax behavior. (Confidence: high — settled identity.)

The normalization is a hard zero-sum constraint with no “attend to nothing” option, which forces attention sinks. Every row must sum to 1. When a token’s correct behavior is to retrieve no information, the model must still put its weight somewhere — so trained models learn attention sinks: they dump excess attention mass onto a semantically empty token (very often the first/BOS token). This is a direct consequence of the missing null option in softmax, observed across essentially all large models (Xiao et al. 2023 and subsequent work). The surface formula (rows sum to 1) contains this fact but does not reveal it as a mechanism. (Confidence: high.)

And √d_k is a temperature set at initialization, not a law of the converged model. The variance argument holds at initialization, under the i.i.d.-component assumption. After training, the entries of W_Q/W_K are correlated and their scale is learned, so a head can set its own effective temperature by growing or shrinking W_QK. The √d_k is a prophylactic keeping softmax in its high-entropy, well-gradiented regime early so training can move; it is not the operative temperature of a trained model. (Confidence: medium for the strong “merely an init heuristic” reading — it remains a sensible default and the post-training temperature story is partly empirical; this would resolve with an empirical reference on learned attention-logit scale post-training.)

Alternative Level 2 mechanism — kernel smoother / Nadaraya–Watson. The operation reads as a Nadaraya–Watson estimator where exp(qᵢ·kⱼ) is an (unnormalized, RBF-like) similarity kernel and the output is the kernel-weighted average of values (Tsai et al. 2019). This is the continuous-statistics face. Analogous to: classical kernel regression, where you predict at a point by averaging nearby observations weighted by a distance kernel. The transfer holds because the output is literally a kernel-weighted average; it does not transfer cleanly where the kernel is fixed — classic Nadaraya–Watson uses a fixed kernel, whereas here the kernel is learned (it is the M of Level 2), so it is a learned-kernel smoother. This view makes the linear-attention family possible — replace exp(q·k) with a factorizable feature map φ(q)·φ(k) and reassociate (φ(Q)φ(K)^T)V → φ(Q)(φ(K)^TV), collapsing the n² term to O(n). But this is a lossy approximation, not a costless refactor: dropping the exp/softmax removes the normalization-induced competition among keys (the row-sum-to-1 constraint), changing the function class and empirically tending to degrade quality. It predicts which approximations preserve behavior, and at what cost. Epistemic standing: a co-equal lens on the same operation, not a deeper rung. (Confidence: high for the identity.)

Alternative Level 2 mechanism — modern Hopfield / associative memory. Softmax-attention is exactly one update step of a continuous (“modern”) Hopfield network (Ramsauer et al. 2020, “Hopfield Networks is All You Need”): the keys are stored patterns, the query is a probe, and the softmax-weighted readout is one step of energy-minimizing retrieval toward the nearest stored pattern. Its distinctive payoff is capacity — a modern Hopfield net stores exponentially many patterns in the embedding dimension, setting a predictive ceiling on how many distinct things one attention layer can reliably retrieve in-context before stored patterns interfere and retrieval blurs. It predicts retrieval-capacity limits — which neither the lookup nor the kernel view foregrounds. Epistemic standing: co-equal lens, identity is settled; the capacity bound’s tightness in trained models is a theoretical result less directly established empirically. (Confidence: high for the identity.)

The standing relation among the three lenses: the lookup view predicts discrete copying/induction; the kernel view predicts behavior-preserving (and behavior-degrading) approximations; the Hopfield view predicts retrieval-capacity ceilings. Co-equal lenses on one operation, not competitors — the field has not picked one.

To anchor the whole descent in one image: the recipe is how a hash map is implemented; Level 1 is “oh, it’s a dictionary lookup”; Level 2 is “the dictionary’s matching rule is a single low-rank bilinear form over content-and-position, and the read and write paths are separable channels”; the frontier boundary below is “what does this particular trained dictionary actually store and retrieve in practice?” — and that last one we can only partially answer.

Epistemic boundary

Everything in the mechanism descent — the lookup relaxation, the QK/OV factorization, the rank bottleneck, the residual-stream channel, the softmax thermodynamics, the kernel and Hopfield identities — is settled mathematics, derivable from the definitions; there is no open question about what the operation computes. What is frontier, not settled, is what trained heads compute semantically. We have crisp, verified circuits for narrow behaviors (induction heads; “name-mover”/IOI heads; copy-suppression heads), almost all established in small or mid-size models. We do not have a complete circuit-level account of attention in frontier models; most heads in a large model have no clean human-legible function identified. The central obstacle is superposition / polysemanticity — heads and directions represent more features than they have dimensions, so a single head rarely does one human-legible thing. The clean QK/OV math (settled) and the semantic attribution of what those circuits compute (current-best-understanding, frequently contested, often partial) are different epistemic objects. Any confident claim “this head does X” about a frontier model is a current-best hypothesis, not a theorem. (Confidence: medium-low and explicitly frontier — the hedge is deliberate because the field is hedging.)

Two interpretive judgments inside the descent are worth naming as judgments rather than facts. Whether the Level-1 semantic reinterpretation — “attention is a soft dictionary read” — counts as a genuine mechanistic level beneath the surface or a metaphor alongside it is a judgment about mode intent; naming the exp gap-amplification as the causal mechanism moves it toward the causal-process reading, and a definitive call would come from an ML-interpretability domain ruling. Separately, judging Level 2’s rank-bottleneck/residual-stream channel as genuinely vertical beneath Level 1 rather than an adjacent algebraic restatement leans on transformer-internals expertise; the read here is that it is vertical — it specifies what the learnable parts of the lookup actually are.

Practical implications

• Practical implication: never interpret Q or K (or W_Q/W_K) in isolation. Only W_QK = W_Q W_K^T is invariant; the individual factors are gauge-arbitrary (insert any invertible M and M^{-1} between them with no change to the product). When probing or pruning attention, operate on the combined QK and OV matrices, or you are reading noise. Likewise, interpretability merges the matrices: to read what a head does, study W^Q W^K^T and W^V W^O directly, never Q/K/V separately. This is not predictable from the surface formula.

• Practical implication: don’t evict the sink token from the KV cache. Because softmax forces non-zero mass onto a sink, naive sliding-window or “drop the oldest tokens” truncation silently deletes the token the model parks its idle attention on, and quality collapses. Keeping a few initial tokens (the StreamingLLM fix — Xiao et al. 2023, four initial tokens plus a rolling recent-token window) follows directly from the sink mechanism and is unpredictable from the surface formula.

• Practical implication: the real per-head expressivity knob is d_k as a rank, not a “size.” If a head needs richer where-to-look logic, raising d_k (the rank of W_QK) is the lever; adding heads gives more independent patterns, not richer individual ones. This reframes the head-count vs. head-dim trade-off. Relatedly, Grouped-Query / Multi-Query Attention (sharing K,V across heads) are sane precisely because the expressive object is low-rank — you are trading rank for KV-cache size, and the memory view tells you exactly what you are spending.

• Practical implication: if a model fails to combine information across positions, the bug is in attention, not the MLP. Since attention is the sole inter-position mover, a failure to integrate two facts stated in different tokens is an attention-routing failure by elimination — a concrete debugging prior the position-wise/cross-position split hands you for free. (Head composition across layers is the same mechanism: whether one head’s OV-circuit writes into the subspace another head’s QK-circuit reads is what enables multi-step circuits, including the previous-token → induction-head composition.)

• Practical implication: depth has a practical floor that depends on your goal. For using attention — wiring a model, reasoning about the KV-cache, choosing GQA, applying LoRA to attention projections — the QK/OV factoring (Level 2) is the actionable floor; the frontier-semantics boundary changes almost nothing operational. The semantic-circuit frontier matters only if the goal is interpreting or steering model internals (safety, circuit analysis, editing): for a shipping engineer, descending past the QK/OV factoring is intellectually real but practically inert; for interpretability work, it is the whole job.