Self Attention With Weights
Why cannot self-attention just use raw input vectors, and what are Queries, Keys, and Values?
How this lesson unfolds.
Without trainable projection weights, self-attention cannot differentiate between searching for information and broadcasting information. Learn how Query, Key, and Value matrices decouple these roles and how scaled dot products stabilize training.
Why do raw input embeddings need to be projected through trainable matrices W_Q, W_K, W_V into specialized subspaces?
Imagine a brilliant multi-talented scholar attending a world conference. If the scholar walks into the lecture hall, they cannot wear their lab coat, their evening tuxedo, and their running shoes all at once. When asking a question (Query), they wear their investigator hat. When advertising their specialty to others on a badge (Key), they list their credentials. When delivering a keynote speech (Value), they share their actual knowledge. Trainable projection matrices W_Q, W_K, W_V are the specialized wardrobe: they project the same underlying person into three distinct functional roles.
- Learning goalExplain why trainable linear projections W_Q, W_K, W_V are required to break dot-product symmetry and decouple token embeddings into specialized functional subspaces.
- Helpful foundationsQKV triad fundamentals · Linear transformations and matrix rank
- Common trapAssuming that self-attention dot products are symmetric. Because qᵢ = xᵢ W_Q and kⱼ = xⱼ W_K use different weight matrices, qᵢ · kⱼ ≠ qⱼ · kᵢ, making attention intrinsically directed and asymmetric.
- Next connectionNext, examine the batch matrix transformations, tensor shapes, and parameter counts in Part B.
More context & mechanics
Directly using raw token embeddings x as queries, keys, and values imposes severe limitations:
- The Symmetry Trap: If q = x and k = x, then compatibility s_(i, j) = xᵢ · xⱼ = xⱼ · xᵢ = s_(j, i). Attention would be forced to be completely symmetric! Token A would attend to Token B with the exact same strength that Token B attends to Token A.
- Asymmetric Linguistic Reality: In language, relationships are fundamentally asymmetric. A verb like "ate" urgently needs to attend to its direct object "pizza", but "pizza" does not need to attend back to "ate" with equal weight.
- Linear Subspace Decoupling: Trainable projection matrices W_Q, W_K, W_V ∈ R^d_(model) × dₖ break this symmetry:
q = x W_Q, k = x W_K, v = x W_V
Now, qᵢ · kⱼ ≠ qⱼ · kᵢ, enabling rich, directed semantic routing.
In transformer blocks, linear projection layers account for a large portion of computational FLOPs. In GPT-2 (d_(model) = 768), each projection matrix has 768 × 768 = 589,824 parameters. Together, W_Q, W_K, W_V contain 1.77M trainable weights per layer, learning specialized grammatical, factual, and stylistic routing patterns.
Vaswani et al. (2017) introduced trainable linear projections to enable multi-head attention. Earlier attention mechanisms (e.g., Luong 2015, Bahdanau 2014) used single additive or multiplicative matrices; projecting into three independent subspaces revolutionized transformer expressivity.
- Decouple the roles: One vector cannot seek, advertise, and deliver at the same time.
- Asymmetric routing requires asymmetric weights: W_Q and W_K break the symmetric dot-product trap.
How do batched projection tensors transform across dimensional contracts, and how are parameters accounted for in W_Q, W_K, W_V?
Think of an airport luggage scanner. A continuous conveyor belt carries luggage trays (Batch B, Sequence T). Each tray holds an item described by 768 physical measurements (d_(model)). The scanner has three specialized optical sensors (W_Q, W_K, W_V), each configured with 768 lenses. As the trays pass through, all three sensors flash simultaneously, outputting three parallel streams of processed digital scans without pausing the belt.
- $$Input sequence batch $X \in \mathbb{R}^{B \times T \times d_{model}}$ is projected via batched matrix multiplication into Query, Key, and Value tensors of shape $\mathbb{R}^{B \times T \times d_k}$.$$
- Learning goalCalculate tensor shape transitions, parameter counts, and memory footprints for batched Query, Key, and Value linear projections.
- Helpful foundationsBatched matrix multiplication (GEMM) · PyTorch tensor reshaping and slicing
- Common trapAssuming that projection matrices alter sequence length T. Linear projections operate strictly on the embedding dimension d_(model), leaving sequence length T completely invariant.
- Next connectionNext, discover why attention dot products must be scaled by √(dₖ) to prevent softmax saturation in Screen 2.
More context & mechanics
#### 1. Batched Projection Tensor Contracts
Let the input batch of token representations be X ∈ R^B × T × d_(model), where:
- B is the batch size,
- T is the sequence length (context window),
- d_(model) is the embedding dimension.
The trainable projection weights are:
W_Q ∈ R^d_(model) × dₖ, W_K ∈ R^d_(model) × dₖ, W_V ∈ R^d_(model) × d_v
The batched projections execute as:
Q = X W_Q ∈ R^(B × T × dₖ)
K = X W_K ∈ R^(B × T × dₖ)
V = X W_V ∈ R^(B × T × d_v)
(In standard architectures, dₖ = d_v = d_(model)).
In production implementations, rather than launching three separate matrix multiplications, frameworks fuse the three weight matrices into a single monolithic weight tensor:
W_(QKV) ∈ R^d_(model) × (3 · d_(model))
A single batched GEMM computes:
X_(QKV) = X W_(QKV) ∈ R^B × T × (3 · d_(model))
This result is then sliced in memory into Q, K, V, maximizing GPU compute utilization and minimizing kernel launch overhead.
For a single attention layer with d_(model) = 768:
Params(W_Q) = 768 × 768 = 589,824
Params(W_K) = 768 × 768 = 589,824
Params(W_V) = 768 × 768 = 589,824
Total QKV Parameters = 3 × 589,824 = 1,769,472 parameters ( ≈ 1.77M)
In PyTorch, fused QKV projection is written as `self.qkv = nn.Linear(d_model, 3 * d_model, bias=qkv_bias)`. During inference, fusing these matrix operations achieves up to 1.8× speedup on NVIDIA Tensor Cores compared to three sequential linear calls.
The fused projection trick was first popularized by Fast-Transformer (NVIDIA) and Hugging Face in 2019 to maximize tensor core pipeline occupancy. Today, fused QKV projection is an industry standard across all open-source transformer libraries.
- [B, T, d] × [d, d] = [B, T, d]: The batch and sequence dimensions remain untouched; only the feature dimension transforms.
- Fuse three into one: One giant GEMM is faster than three small ones.
Why do unscaled dot products explode with high dimensionality, and how does scaling by √(dₖ) prevent catastrophic gradient vanishing?
Imagine adjusting the sensitivity dial on a sound amplifier. If you hook up 4 microphones, the total noise adds up moderately, and you can easily distinguish between soft voices and loud shouts. But if you connect 512 microphones together without turning down the master volume, the cumulative signal instantly redlines into deafening screeching feedback. The amplifier saturates at maximum volume, making a whisper and a scream sound completely identical: blown-out static. Softmax saturation is that blown-out static: when inputs get huge, softmax outputs hard 1s and 0s, and all training gradients die.
- Learning goalExplain the statistical relationship between high hidden dimensionality and softmax saturation, deriving the vanishing gradient trap of unscaled dot products.
- Helpful foundationsSoftmax derivative calculus · Variance and standard deviation of independent random variables
- Common trapThinking that layer normalization eliminates the need for attention scaling. LayerNorm normalizes vectors before the projection; the dot product itself multiplies two normalized vectors, creating a new sum of variance dₖ that must be scaled separately.
- Next connectionNext, derive the formal variance proof and step-by-step statistical expectation in Part B.
More context & mechanics
The necessity of dividing attention scores by √(dₖ) is rooted in the interplay between vector geometry and the derivative of the softmax function:
- The Unscaled Tragedy:
The inner product q · k = ∑_(i=1)^(dₖ) qᵢ kᵢ sums dₖ independent random terms. If each component has variance 1, the sum has variance dₖ, and standard deviation σ = √(dₖ).
- For dₖ = 64: σ = 8.0.
- For dₖ = 512: σ ≈ 22.6.
- The Softmax Plateau:
When input logits differ by values greater than 10 or 20, softmax saturates: the largest element approaches 1.0, while all other elements collapse to 0.0.
- Vanishing Gradient Catastrophe:
The derivative of softmax with respect to input logit zᵢ is:
(∂ softmax(z)ᵢ) / (∂ zⱼ) = αᵢ (δ_(i, j) - αⱼ)
When αᵢ → 1.0 (or αⱼ → 0.0), the derivative collapses:
αᵢ (1 - αᵢ) → 1.0 × (1 - 1.0) = 0.0
The gradient vanishes completely, halting backpropagation through the entire attention block!
Without the (1) / (√(dₖ)) scaling factor, transformer models with standard embedding sizes (dₖ = 64 to 128) fail to train entirely; the loss gets trapped in flat plateaus on the very first optimization step.
In spring 2017, while developing the original Transformer architecture, the Google Brain team noticed that for large values of dₖ, dot-product attention performed significantly worse than additive attention (Bahdanau). Vaswani et al. suspected that large dot products were pushing softmax into regions with vanishingly small gradients, and introduced (1) / (√(dₖ)) scaling to reconcile the speed of dot-products with the stability of additive attention.
- Large dimensions blow out softmax: Summing d components sums d variances.
- Scale by root d to keep gradients alive: Division by √(dₖ) restores unit variance and prevents the saturation freeze.
How do we formally prove that dividing the dot product by √(dₖ) rescales its statistical variance back to exactly 1.0?
Think of tossing 64 fair coins where heads is +1 and tails is -1. If you toss 1 coin, your variance is 1. If you sum 64 independent tosses, your expected sum is still 0, but your typical distance from 0 has spread out by a factor of √(64) = 8. If you want your 64-coin sum to have the exact same spread as a single coin toss, you must divide the sum by 8. Dividing by the square root of the number of coins resets the spread back to 1.
- $$Assuming independent zero-mean unit-variance components, the variance of their inner product is $d_k$; dividing by $\sqrt{d_k}$ scales variance by $(\frac{1}{\sqrt{d_k}})^2 = \frac{1}{d_k}$, restoring $\operatorname{Var} = 1.0$.$$
- Learning goalDerive the formal mathematical proof that Var(q · k) = dₖ and demonstrate that scaling by √(dₖ) enforces unit variance.
- Helpful foundationsProbability theory (expectation, variance of sums) · Properties of scalar scaling on variance
- Common trapAssuming that dividing by dₖ would be better than √(dₖ). Dividing by dₖ would shrink variance to 1/dₖ → 0, compressing all logits into near-zero and destroying attention sharpness into uniform noise.
- Next connectionNext, learn how causal attention masks enforce temporal directionality in autoregressive models in Chapter 16 Screen 2.
More context & mechanics
#### 1. The Mathematical Variance Proof
Assume that components of Query vector q and Key vector k are independent random variables with zero mean and unit variance:
E[qᵢ] = E[kᵢ] = 0, Var(qᵢ) = Var(kᵢ) = 1 (∀ i ∈ 1, ṡ, dₖ)
Consider the product of a single component pair Xᵢ = qᵢ kᵢ:
- Expectation of Xᵢ:
E[Xᵢ] = E[qᵢ kᵢ] = E[qᵢ] E[kᵢ] = 0 × 0 = 0
- Variance of Xᵢ:
Var(Xᵢ) = E[Xᵢ²] - (E[Xᵢ])² = E[qᵢ² kᵢ²] - 0 = E[qᵢ²] E[kᵢ²]
Since Var(qᵢ) = E[qᵢ²] - (E[qᵢ])² = 1 ⟹ E[qᵢ²] = 1, we have:
Var(Xᵢ) = 1 × 1 = 1
Now consider the unscaled dot product S = q · k = ∑_(i=1)^(dₖ) qᵢ kᵢ:
- Because components are mutually independent, the variance of the sum is the sum of the variances:
Var(S) = Var( ∑_(i=1)^(dₖ) Xᵢ ) = ∑_(i=1)^(dₖ) Var(Xᵢ) = ∑_(i=1)^(dₖ) 1 = dₖ
The standard deviation is σ_S = √(dₖ).
Let scaled score be S̃ = (S) / (√(dₖ)). Using the variance property Var(c Y) = c² Var(Y):
Var(S̃) = Var( (S) / (√(dₖ)) ) = ( (1) / (√(dₖ)) )² Var(S) = 1/dₖ × dₖ = 1.0
The scaled dot product S̃ has mean 0 and variance 1.0, regardless of whether dₖ = 16 or dₖ = 4096! \blacksquare
Compare dₖ = 64 with unscaled vs. scaled scores:
- Let components average ≈ 1.0. Unscaled dot product S = ∑_(i=1)^(64) (1.0 × 1.0) = 64.0.
Softmax with input 64.0:
e^(64) ≈ 6.235 × 10^(27) ⟹ Floating point overflow / hard saturation!
- Now scale by √(64) = 8.0:
S̃ = (64.0) / (8.0) = 8.0
e^(8.0) ≈ 2,980.96 ⟹ Stable, smooth, non-overflowing softmax!
In FlashAttention and Triton kernels, the scaling factor (1) / (√(dₖ)) is premultiplied into the Query tensor before GEMM computation: Q_(scaled) = Q × (1) / (√(dₖ)). This saves an extra tensor element-wise scaling pass across HBM memory.
The algebraic law that the variance of the sum of independent random variables equals the sum of their variances was formulated by Pafnuty Chebyshev in 1867. Its application to scale attention dot products in 2017 stands as one of the most elegant examples of classic probability theory solving modern deep learning engineering bottlenecks.
- Var(q · k) = dₖ: Unscaled variance matches the dimension.
- Divide by root d to restore unit variance: ((1) / (√(d)))² · d = 1.0.
How do you calculate scaled dot products and row-wise softmax step-by-step for a 3-token sequence?
Put on your math boots and walk through the numbers. Suppose you have 3 tokens and head dimension dₖ = 4, so sqrt(dₖ) = 2. We compute the 9 pairwise dot products of Q and K. If the raw dot products for token 1 against tokens 1, 2, 3 are [4.0, 2.0, 0.0], we divide by 2 to get scaled scores [2.0, 1.0, 0.0]. Exponentiating gives [e², e¹, e⁰] = [7.389, 2.718, 1.000], with sum 11.107. Normalizing gives attention weights [0.665, 0.245, 0.090]. Token 1 devotes 66.5% of its attention to itself, 24.5% to token 2, and 9.0% to token 3.
- $$Trace the exact numerical execution of S = Q K^T / sqrt(d_k) and its row-wise Softmax transformation into attention weights A.$$
- Learning goalExecute the hand calculation of matrix multiplication Q K^T, scale by 1/sqrt(dₖ), apply numerical-stability subtraction, and compute row-wise Softmax weights.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · Matrix Multiplication by Hand · Numerical Stability in Softmax (x - max(x))
- Common trapForgetting that Softmax must be computed along rows (axis=-1) rather than columns. Each row corresponds to a single query token deciding how to distribute its attention budget over all available keys.
- Next connectionOnce the attention probability matrix A is computed, we multiply it by the Value matrix V to synthesize the contextual output vectors.
More context & mechanics
Let T=3 and dₖ=4 (so √(dₖ) = 2). Suppose the computed Query and Key matrices produce the raw score matrix Q K^\top:
Q K^\top = 4.0 & 2.0 & 0.0 \ 1.0 & 3.0 & -1.0 \ 2.0 & 0.0 & 4.0
Step 1: Scale by (1) / (√(dₖ)) = 0.5:
S = (Q K^\top) / (2) = 2.0 & 1.0 & 0.0 \ 0.5 & 1.5 & -0.5 \ 1.0 & 0.0 & 2.0
Step 2: Compute Softmax along Row 1:
- Exponentiate: e^(2.0) ≈ 7.389, e^(1.0) ≈ 2.718, e^(0.0) = 1.000.
- Row sum: ∑ = 7.389 + 2.718 + 1.000 = 11.107.
- Normalize:
α_(11) = (7.389) / (11.107) ≈ 0.665, α_(12) = (2.718) / (11.107) ≈ 0.245, α_(13) = (1.000) / (11.107) ≈ 0.090
Row 1 sums to 0.665 + 0.245 + 0.090 = 1.000.
Repeating for rows 2 and 3 yields the full attention weight matrix A ∈ R^(3 × 3) where every row is a probability distribution.
This exact calculation is optimized in modern GPU attention kernels like FlashAttention (Dao et al., 2022), which tiles the Q K^\top product into SRAM to perform online Softmax without materializing the full T × T matrix in High Bandwidth Memory (HBM).
Computing dot products by hand was the standard pedagogical technique used in early neural network courses. In 2022, Tri Dao et al. analyzed these exact numerical steps to invent FlashAttention, proving that fused SRAM tiling could speed up attention computation by 3-4x.
- Hand-Calculation Checklist:
- Q × K^\top (Raw dot products)
- Divide by √(dₖ) (Scaled scores)
- Exponentiate row entries
- Divide by row sum (Attention weights A)
How does multiplying the attention weight matrix by the Value matrix assemble final contextual vectors?
Now that token 1 has assigned percentages to the other tokens (0.665 to itself, 0.245 to token 2, 0.090 to token 3), it goes shopping. It takes 66.5% of Value vector 1, adds 24.5% of Value vector 2, and adds 9.0% of Value vector 3. The result is a single blended vector z₁ containing information from all three tokens, customized precisely according to how relevant each was to query 1.
- $$The final context matrix Z = A V computes a linear combination of value vectors, where each row is a customized weighted blend of all tokens.$$
- Learning goalCalculate the context matrix Z = A V by performing the linear combination of Value row vectors weighted by the attention matrix.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · Linear Combinations of Vectors · Row-Vector Matrix Multiplication
- Common trapThinking the output vectors Z are the attention weights themselves. The attention weights A are scalar coefficients (summing to 1 per row); the output Z consists of new feature vectors residing in value space R^(d_v).
- Next connectionUnderstanding how weighted values assemble allows us to analyze what attention weights actually represent: semantic relational evidence versus incidental token proximity.
More context & mechanics
Given attention weights A ∈ R^(T × T) and the Value matrix V ∈ R^(T × d_v), the output context matrix is:
Z = A V ∈ R^(T × d_v)
For each individual token i, its contextualized output vector zᵢ is the linear combination of all value row vectors v₁, ṡ, v_T:
zᵢ = ∑_(j=1)^T A_(ij) vⱼ
Worked Numerical Example:
Using the Row 1 weights computed in Screen 3 (A_(1) = [0.665, 0.245, 0.090]) and 2D values:
v₁ = [1.0, 2.0]^\top, v₂ = [3.0, 0.0]^\top, v₃ = [0.0, 4.0]^\top
The blended vector z₁ is:
z₁ = 0.665 1.0 \ 2.0 + 0.245 3.0 \ 0.0 + 0.090 0.0 \ 4.0 = 0.665 + 0.735 + 0.000 \ 1.330 + 0.000 + 0.360 = 1.400 \ 1.690
Token 1 now holds a rich representation carrying information from tokens 1, 2, and 3.
In decoder transformers like GPT, this operation is how tokens gather factual information from prior tokens in the prompt. In multimodal transformers (CLIP, Flamingo), visual tokens and text tokens blend through this exact linear combination of Values.
The formulation of attention as a convex sum over value vectors was derived from Kernel Density Estimation and Nadaraya-Watson kernel regression (1964), where non-parametric regression estimates values by taking distance-weighted averages of observed target values.
- Golden Axiom: "Attention weights are the recipe; Values are the ingredients."
- Each row of A contains the blending percentages; multiplying by V produces the finished contextual dish.
How do attention weights differentiate between genuine semantic affinity and incidental token proximity?
In a sentence like "The animal did not cross the street because it was too tired", the pronoun "it" is adjacent to "was" and "too", but its true semantic referent is "animal" six words earlier. A naive proximity filter would connect "it" to "street" or "tired". But self-attention's query for "it" matches strongly with the key for "animal", directing high attention weight across the grammatical distance. Attention discovers relational evidence, breaking free from the tyranny of adjacency.
- Learning goalAnalyze how trained Query-Key projections overcome positional proximity bias to capture long-range semantic and syntactic dependencies.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · What Does BERT Look At? An Analysis of Attention (Clark et al., 2019) · Coreference Resolution & Dependency Parsing
- Common trapEquating high attention weight directly with causal importance or absolute explanation. Attention weights indicate which inputs the model selected for representation blending, but nonlinear downstream transformations also modulate feature importance.
- Next connectionIn free-form text generation, bidirectional attention creates an illegal information leak: tokens can attend to future answers. To build generative LLMs, we must enforce Causal Masking.
More context & mechanics
Sequential models like RNNs or n-gram Markov chains suffer from locality bias: mutual information decays exponentially with sequence distance |i - j|. In self-attention, the Query-Key projection maps tokens into an abstract feature space where geometric distance reflects task-specific relevance, not token order:
Affinity(i, j) ∝ qᵢ^\top kⱼ
If token i is a pronoun seeking its antecedent, W_Q projects token i into an antecedent-seeking direction, while W_K projects candidate noun phrases into that same direction. Because the dot product is independent of position index |i - j|, two tokens 500 positions apart can achieve an attention weight as high as adjacent tokens.
Furthermore, attention matrices are intrinsically asymmetric (A_(ij) ≠ A_(ji)):
A_(ij) = (exp(qᵢ^\top kⱼ / √(dₖ))) / (∑ₘ exp(qᵢ^\top kₘ / √(dₖ))) ≠ (exp(qⱼ^\top kᵢ / √(dₖ))) / (∑ₘ exp(qⱼ^\top kₘ / √(dₖ))) = A_(ji)
A dependent clause attends heavily to its governing verb, while the verb distributes its attention across multiple arguments.
This capability enables transformers to resolve coreference, track code variable declarations across hundreds of lines of python script, and perform multi-hop mathematical reasoning.
In 2019, Clark, Khandelwal, Levy, and Manning published What Does BERT Look At?, demonstrating that individual attention heads autonomously specialize into classical NLP pipeline components: some heads track direct objects, others track coreferent pronouns, and others track sentence-level punctuation.
- Golden Axiom: "Proximity is geography; Attention is meaning."
- Attention connects tokens by semantic evidence, completely independent of how many words separate them in the raw sequence.
Why does the dot product variance scale linearly with dimension dₖ, requiring division by sqrt(dₖ) to prevent softmax saturation?
Imagine rolling dₖ six-sided dice and summing their scores. If you roll 4 dice, the sum varies around 14 by a couple of points. But if you roll 1,024 dice, the sum can swing wildly by hundreds of points. When you take the dot product of two 64-dimensional vectors, you are summing 64 independent products: the total variance is 64, so the typical score swings between -16 and +16. When passed to softmax, e^(16) ≈ 8,800,000, which completely saturates the probabilities into a hard 1.0 and freezes all gradients to zero. Dividing by √(64) = 8 pulls the scores back into a gentle [-2, +2] range where learning thrives.
- $$Prove that the variance of the sum of d_k independent zero-mean unit-variance components equals d_k, causing logits to explode unless scaled by 1/sqrt(d_k).$$
- Learning goalDerive the mathematical proof that Var(q · k) = dₖ and prove that scaling by 1/sqrt(dₖ) enforces unit variance.
- Helpful foundationsVariance of independent sums · Attention Is All You Need (Vaswani et al., 2017)
- Common trapBelieving 1/sqrt(dₖ) is an arbitrary heuristic; it is the exact analytic normalizer required to preserve unit variance under independent Gaussian distributions.
- Next connectionHow do multiple scaled attention heads partition the global embedding space? We explore direct sum subspaces.
More context & mechanics
Assume the components of q, k ∈ R^(dₖ) are independent random variables with zero mean and unit variance:
E[qᵢ] = E[kᵢ] = 0, Var(qᵢ) = Var(kᵢ) = 1
Let S = q · k = ∑_(i=1)^(dₖ) qᵢ kᵢ.
- Compute the expectation of S:
E[S] = ∑_(i=1)^(dₖ) E[qᵢ kᵢ] = ∑_(i=1)^(dₖ) E[qᵢ] E[kᵢ] = 0
- Compute the variance of S:
Var(qᵢ kᵢ) = E[(qᵢ kᵢ)²] - (E[qᵢ kᵢ])² = E[qᵢ²] E[kᵢ²] - 0 = 1 × 1 = 1
By the Bienaymé formula for sums of independent variables:
Var(S) = Var( ∑_(i=1)^(dₖ) qᵢ kᵢ ) = ∑_(i=1)^(dₖ) Var(qᵢ kᵢ) = ∑_(i=1)^(dₖ) 1 = dₖ
The standard deviation of the raw dot product is σ_S = √(dₖ).
To restore unit variance (Var = 1):
S_(scaled) = (q · k) / (√(dₖ)) ⟹ Var( (S) / (√(dₖ)) ) = 1/dₖ Var(S) = dₖ/dₖ = 1
This exact division `scores = (q @ k.transpose(-2, -1)) / math.sqrt(dₖ)` sits at the center of every implementation of Scaled Dot-Product Attention in PyTorch, JAX, and FlashAttention.
Vaswani et al. explicitly proved this scaling necessity in Section 3.2.1 of 'Attention Is All You Need' (2017), noting that without (1) / (√(dₖ)), large values of dₖ push the softmax function into regions with vanishingly small gradients.
- Variance adds up to dₖ; divide by square root of dₖ to tame the variance back to 1.
- Without scaling, softmax saturates and gradients die.
What is the mathematical foundation and geometric intuition of FlashAttention: IO-Aware Tiling on GPU SRAM?
Understand FlashAttention: IO-Aware Tiling on GPU SRAM visually: Dao et al. (2022): Eliminating the O(T²) HBM memory bottleneck via online softmax.. Every transformation preserves core structural invariants while mapping coordinate spaces.
- $$(2022): Eliminating the O(T^2) HBM memory bottleneck via online softmax.$$
- Learning goalMaster the algebraic and geometric derivation of FlashAttention: IO-Aware Tiling on GPU SRAM.
- Helpful foundationsLinear Algebra Fundamentals · Matrix Transformations · Vector Spaces
- Common trapAssuming FlashAttention: IO-Aware Tiling on GPU SRAM is merely a computational trick rather than an invariant spatial law.
- Next connectionConnects to upstream spectral theory and downstream tensor operations.
More context & mechanics
A hardware memory architecture diagram on natural warm linen (#f5f4e9) with a 1px slate border.
Top section: Title in bold sans-serif: "FlashAttention: IO-Aware Tiling on GPU SRAM". Subtitle: "Dao et al. (2022): Eliminating the O(T²) HBM memory bottleneck via online softmax."
Middle section: A hardware memory hierarchy comparison:
Top Layer (Slow High-Bandwidth Memory - HBM): 80 GB capacity, slow 2.0 TB/s bandwidth.
Bottom Layer (Fast On-Chip SRAM): 192 KB per streaming multiprocessor, ultra-fast 19 TB/s bandwidth.
Show standard attention failing: writes the full T x T attention matrix to HBM and reads it back (O(T²) memory bandwidth wall!).
Show FlashAttention: tiles Q, K, V into small blocks of size B_r x B_c that fit entirely in fast SRAM. Computes local attention and rescales dynamically using online softmax without ever writing T x T to HBM.
Bottom section: Performance card: "Reduces memory footprint from O(T²) to O(T); accelerates attention runtime by 2x to 4x on NVIDIA A100/H100."
Crisp hardware diagram, balanced layout, uncluttered design.
This principle is foundational across deep learning optimization, high-dimensional representation, and modern LLM architectures.
"FlashAttention: IO-Aware Tiling on GPU SRAM: Invariance and structure determine stability under spatial and numerical transformations."
How does FlashAttention: IO-Aware Tiling on GPU SRAM operate mathematically at scale?
Visual GPU memory hierarchy comparison:
- Learning goalMaster the algebraic and geometric derivation of FlashAttention: IO-Aware Tiling on GPU SRAM.
- Helpful foundationsFoundations of Build Llm From Scratch · Coordinate Invariance & Geometric Projections · Vaswani et al. (2017), Attention Is All You Need, NeurIPS
- Common trapAssuming FlashAttention: IO-Aware Tiling on GPU SRAM is an arbitrary computational trick rather than an invariant spatial law.
- Next connectionConnects upstream foundational lemmas to modern high-dimensional computational systems.
More context & mechanics
A hardware memory hierarchy diagram in square 1:1 format on dark slate (#0f172a) with clean hardware blocks.
Top section: Title in white sans-serif: "FlashAttention: IO-Aware Tiling on GPU SRAM". Subtitle in muted cyan: "Dao et al. (2022): Eliminating high-latency HBM memory read/writes via online softmax."
Middle section: Visual GPU memory hierarchy comparison:
Top Tier (Slow HBM - 2 TB/s bandwidth, 80 GB capacity): High-Bandwidth Memory where Q, K, V live.
Bottom Tier (Fast SRAM - 19 TB/s bandwidth, 20 MB capacity per GPU): Ultra-fast on-chip cache.
Contrast:
Standard Attention (Disaster): Computes full L x L attention matrix S = Q K^T in HBM, applies softmax in HBM, stores L x L matrix P in HBM (O(L²) memory footprint; memory bottlenecked!).
FlashAttention (Brilliant): Tiles Q, K, V into small blocks (e.g., 64x64) that fit entirely inside SRAM. Computes attention and softmax incrementally using the online softmax trick, writing only the final output O back to HBM!
Bottom section: Computational gain: "Reduces memory traffic by 5x-10x and shrinks attention memory complexity from O(L²) to O(L)."
Clear memory hierarchy blocks, dynamic data transfer arrows, high contrast dark theme.
Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.
Formal derivation and peer-reviewed analysis:
"FlashAttention: IO-Aware Tiling on GPU SRAM: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."