The Attention Mechanism: Mathematical Deep Dive

Open a UPI app's support chat and type: "UPI payment failed but amount deducted, refund please." A triage model reading that sentence has to decide, in effect, what "failed" is even about. Grammatically it could attach to "UPI" (the network failed), to "payment" (the transaction failed), or to nothing in particular. A human reads the whole sentence and instantly knows "failed" describes the payment, not the network — NPCI's rails are up, the debit went through, the credit didn't land. The question this chapter answers is a narrow, precise one: what arithmetic lets a model compute, for every word, how much it should "look at" every other word before deciding what that word means in context? That arithmetic is the attention mechanism, and it is exactly three matrix operations — a dot product, a scale, a softmax — repeated everywhere in every modern language model.

From a fixed bottleneck to dynamic focus

Before attention existed (Bahdanau, Cho and Bengio, 2014), machine translation systems used an encoder-decoder RNN. The encoder read the source sentence one token at a time and compressed everything it had seen into a single fixed-length vector — the final hidden state. The decoder then had to generate the entire translation from that one vector alone. For "UPI payment failed but amount deducted" this means squeezing seven words' worth of meaning into one vector of, say, 512 numbers, and expecting the decoder to reconstruct "deducted" correctly ten words later using only that compressed summary. Performance degraded sharply on long sentences for exactly this reason: the fixed vector is a bottleneck, and information about early words gets diluted by the time the encoder reaches the end.

Attention removes the bottleneck by letting the decoder look back at every encoder hidden state at every decoding step, weighting each one by how relevant it is right now. Instead of one summary vector, the decoder gets a fresh, re-weighted blend of the entire input for each word it produces. Vaswani et al. (2017, "Attention Is All You Need") then took the next step: if attention lets you dynamically re-weight relevant information, why keep the RNN at all? Self-attention lets every token in a sequence attend to every other token in the same sequence, computing contextual representations without any recurrence. That generalization — from decoder-looks-at-encoder to token-looks-at-token — is what this chapter derives in full.

Query, key, value: the formal abstraction

The cleanest way to understand Query, Key and Value is a search analogy you have already used: searching your UPI transaction history for "Swiggy." You type a query ("Swiggy"). The app doesn't compare your query letter-by-letter to the rupee amounts — it compares your query against an indexed key for every transaction (typically the merchant description). Wherever the key matches the query well, the app pulls the corresponding value (the actual transaction record: amount, date, status) and shows it to you, weighted by how well each key matched. A fuzzy match to "Swiggy Instamart" still surfaces, just ranked lower than an exact "Swiggy" match.

Self-attention runs the same three-role search for every token against every token, all at once, and instead of returning a ranked list it returns a soft, weighted blend of all the values. Concretely, every token's embedding is passed through three separate learned linear maps:

Q = X · W_Q     (queries:  what am I looking for?)
K = X · W_K     (keys:    what do I contain, as a label?)
V = X · W_V     (values:  what do I contain, as content?)

where X is the matrix of token embeddings (one row per token), and W_Q, W_K, W_V are learned weight matrices, each of shape d_model × d_k. Note immediately: W_Q and W_K are different matrices. A token's query vector and its key vector are not the same object — a word asks a different kind of question than it advertises about itself. This distinction matters enough that we return to it below.

Given Q, K and V, the full scaled dot-product attention formula is:

Attention(Q, K, V) = softmax( Q·Kᵀ / √d_k ) · V

Read it left to right. Q·Kᵀ is a matrix of raw similarity scores: row i, column j is the dot product of token i's query with token j's key — a single number measuring how well token j answers what token i is looking for. Dividing by √d_k rescales those raw scores (derived below). softmax, applied along each row, turns the row of scores into a probability distribution — non-negative weights that sum to 1, telling token i exactly what fraction of its attention to spend on every other token. Multiplying that weight matrix by V then produces, for every token, a new vector: a weighted average of everyone's value vector, weighted by relevance. That new vector is the token's contextual representation — the thing that actually carries the "failed refers to payment, not UPI" information forward into the next layer.

Why divide by √d_k? A variance argument, not a convention

The scaling factor looks arbitrary until you check what happens without it. Assume, as is roughly true early in training, that each component of q and k is drawn independently with mean 0 and variance 1. The dot product q · k = Σ q_i k_i is a sum of d_k independent terms, each with variance E[q_i²]·E[k_i²] = 1 (since the components are independent and zero-mean). Variance is additive over independent terms, so Var(q · k) = d_k. For a typical single-head dimension of d_k = 64, that means dot products routinely land in the range of roughly ±8 (one standard deviation of √64). Feeding scores of that magnitude into softmax pushes it toward a near one-hot distribution — one weight near 1, the rest near 0 — because softmax is exponential and small differences at large magnitude get exponentially amplified. Gradients through a saturated softmax vanish almost everywhere, so the network stops learning. Dividing every score by √d_k multiplies the variance by (1/√d_k)² = 1/d_k, exactly cancelling the d_k growth and restoring unit variance regardless of how large d_k is. This is not a tuning knob picked by trial and error — it is the unique scalar that keeps the pre-softmax variance constant as dimensionality changes.

Worked example: self-attention on three tokens

Take the sentence fragment "UPI payment failed" as three tokens, and suppose (for hand-traceable arithmetic) that after projection through W_Q, W_K, W_V each token already has these 2-dimensional Q, K, V vectors — small, invented numbers chosen so every step can be checked by hand, not real learned weights:

token     Q        K         V
UPI      [1, 0]   [1,  0]   [1, 0]
payment  [0, 1]   [0,  1]   [0, 2]
failed   [1, 1]   [1, -1]   [1, 1]

Here d_k = 2, so the scale factor is √2 ≈ 1.4142. Trace the computation for the query "failed," Q = [1, 1]:

Step 1 — raw scores (Q · K for each key):

Q·K(UPI)     = [1,1]·[1, 0] = 1·1 + 1·0 = 1
Q·K(payment) = [1,1]·[0, 1] = 1·0 + 1·1 = 1
Q·K(failed)  = [1,1]·[1,-1] = 1·1 + 1·(-1) = 0

Step 2 — scale by 1/√d_k = 1/√2:

[1/1.4142, 1/1.4142, 0/1.4142] = [0.7071, 0.7071, 0.0000]

Step 3 — softmax. Compute e^x for each scaled score, then normalize by the sum:

e^0.7071 = 2.0281   e^0.7071 = 2.0281   e^0 = 1.0000
sum = 2.0281 + 2.0281 + 1.0000 = 5.0562

weight(UPI)     = 2.0281 / 5.0562 = 0.4011
weight(payment) = 2.0281 / 5.0562 = 0.4011
weight(failed)  = 1.0000 / 5.0562 = 0.1978

Step 4 — weighted sum of values:

context(failed) = 0.4011·[1,0] + 0.4011·[0,2] + 0.1978·[1,1]
                = [0.4011, 0] + [0, 0.8022] + [0.1978, 0.1978]
                = [0.5989, 1.0000]

Running the same four steps for the other two queries (verified numerically below) gives the complete 3×3 attention weight matrix and output:

          UPI     payment  failed
UPI     [0.4011, 0.1978, 0.4011]
payment [0.2840, 0.5760, 0.1400]
failed  [0.4011, 0.4011, 0.1978]

context(UPI)     = [0.8022, 0.7967]
context(payment) = [0.4240, 1.2920]
context(failed)  = [0.5989, 1.0000]

Here is the reference implementation, tested against exactly the numbers above:

import numpy as np

def attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    scores = scores - scores.max(axis=-1, keepdims=True)  # numerical stability
    weights = np.exp(scores)
    weights = weights / weights.sum(axis=-1, keepdims=True)
    return weights @ V, weights

Q = np.array([[1,0],[0,1],[1,1]], dtype=float)
K = np.array([[1,0],[0,1],[1,-1]], dtype=float)
V = np.array([[1,0],[0,2],[1,1]], dtype=float)

output, weights = attention(Q, K, V)
# weights ≈ [[0.4011,0.1978,0.4011],[0.2840,0.5760,0.1400],[0.4011,0.4011,0.1978]]
# output  ≈ [[0.8022,0.7967],[0.4240,1.2920],[0.5989,1.0000]]

Look closely at the "failed" row: it spends almost equal weight on "UPI" (0.4011) and "payment" (0.4011), and less on itself (0.1978) — even though "UPI" is two positions back and "payment" is immediately adjacent. This is the entire point of attention: relevance is computed from the content of Q and K, not from proximity. An RNN's hidden state naturally decays information from far-back tokens; self-attention has no notion of distance built into the formula at all — every pair of tokens starts on equal footing, and it is purely the learned Q and K projections that decide who matters to whom. (Position is reintroduced separately, via positional encodings added to the input embeddings before Q/K/V are computed — a mechanism outside this chapter's scope.)

Multi-head attention

A single attention computation produces exactly one weight distribution per query token — one way of deciding what matters. But relevance is not one-dimensional: "failed" might need to attend to "payment" for its semantic role and, separately, to "UPI" for the entity being described. Multi-head attention runs several independent scaled dot-product attentions in parallel, each with its own learned W_Q^i, W_K^i, W_V^i projecting into a smaller subspace, then concatenates the results and mixes them with one more learned matrix:

head_i = Attention(Q·W_Q^i, K·W_K^i, V·W_V^i)          i = 1 ... h
MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_O

If the model dimension is d_model = 512 and there are h = 8 heads, each head works in d_k = d_model / h = 64 dimensions, so the total parameter count for Q, K, V projections across all heads matches a single 512-dimensional head — multi-head attention is not "more capacity," it is the same capacity split into h independent softmax distributions instead of one. Each head is free to specialize: empirically, some heads in trained transformers learn to track adjacent-word relationships, others track long-range syntactic dependencies (subject to its verb across a clause), others track coreference. A single head, forced to express all of these as one softmax row per token, could only ever represent one such pattern per token at a time; concatenating heads lets several coexist.

Self-attention, cross-attention, and causal masking

The worked example above is self-attention: Q, K and V all come from the same sequence. The original Bahdanau mechanism this chapter opened with is cross-attention: Q comes from the decoder's current state, while K and V come from the encoder's output sequence — the decoder is "querying" the source sentence. Modern transformer encoder-decoder models (translation, summarization) use both: self-attention within the encoder and within the decoder, plus cross-attention from decoder to encoder.

A decoder generating text left-to-right (as GPT-style models do) must not let a token attend to positions that come after it — that would let it "see the answer" during training. This is enforced by causal masking: before the softmax step, every score at position (i, j) with j > i is set to −∞. Since e^(−∞) = 0, softmax assigns those positions exactly zero weight, without changing the formula at all — masking is implemented as one addition before the existing softmax, not a separate mechanism.

Common misconception: attention weights are not a fixed similarity table

Students who have used cosine similarity between static word embeddings (word2vec, GloVe) often assume attention weights work the same way — a fixed, symmetric, context-independent number that could, in principle, be precomputed once for every word pair and reused everywhere "UPI" and "payment" co-occur. This is wrong on two separate counts, both visible in the numbers already computed above.

First, it isn't symmetric. Compare the weight from "UPI" attending to "payment" against the weight from "payment" attending to "UPI," directly from the matrix: weight(UPI→payment) = 0.1978, but weight(payment→UPI) = 0.2840. These are different numbers computed from the same pair of tokens. That is only possible because W_Q ≠ W_K — a token's query vector (what it looks for) and its key vector (what it advertises) come from different learned projections, so Q_i · K_j and Q_j · K_i have no reason to match. A static cosine-similarity table is symmetric by construction; attention is not, and the asymmetry is doing real work — it lets "failed" search hard for its subject while "UPI," appearing early, doesn't need to search for "failed" nearly as hard.

Second, it isn't fixed. Q and K are computed from each token's current hidden state, not from a static vocabulary lookup. In a deeper network, that hidden state already carries information mixed in from earlier layers' attention, so the same word "payment" produces a different K vector depending on the sentence it appears in, its position, and the layer. There is no single number "how much does UPI attend to payment" — only "how much does UPI attend to payment, in this sentence, at this layer, in this head," recomputed fresh on every forward pass.

Computational cost

For a sequence of n tokens with key/query dimension d_k, computing Q·Kᵀ multiplies an n × d_k matrix by a d_k × n matrix, producing an n × n result — that is n × n × d_k multiply-add operations, i.e. O(n²d_k). The subsequent weighted sum with V costs another O(n²d_v). Both terms are quadratic in sequence length. Doubling the input length quadruples the attention compute; a document of 10,000 tokens costs roughly (10,000/1,000)² = 100× the raw attention arithmetic of a 1,000-token document, holding d_k fixed. This quadratic wall — not the model's parameter count — is the specific reason long-context language models need sparse, windowed, or linear-attention approximations rather than plain scaled dot-product attention scaled up naively.

Active recall

Attempt each question before reading its answer.

1. If q and k each have d_k = 64 independent components with mean 0 and variance 1, what is the variance of the raw dot product q · k, and what does dividing by √d_k do to it?
2. Using this chapter's worked-example vectors, compute the raw (unscaled) dot-product score matrix Q·Kᵀ for all three tokens.
3. For the query "payment" (Q = [0, 1]), the scaled scores were [0, 0.7071, −0.7071]. Show the arithmetic that produces the middle softmax weight, 0.5760.
4. If d_model = 512 and a transformer layer uses h = 8 attention heads, what is d_k per head, and why use 8 smaller heads instead of 1 head at the full 512 dimensions?
5. True or false: the post-softmax attention weight matrix is always symmetric (A_ij = A_ji). Justify using this chapter's numbers.
6. A causal decoder processes a 4-token sequence. Before softmax, which entries of the 4×4 score matrix get masked to −∞, and why does that make e^(score) equal to zero at those positions rather than requiring a separate zeroing step?

Answers

1. Each product term q_i k_i has variance E[q_i²]E[k_i²] = 1 (independence, zero mean). Summing 64 independent such terms gives Var(q·k) = 64. Dividing by √64 = 8 scales the variance by (1/8)² = 1/64, restoring variance 1 regardless of d_k — this is why the divisor is √d_k and not d_k itself.
2. Q·Kᵀ — row UPI: [1·1+0·0, 1·0+0·1, 1·1+0·(−1)] = [1, 0, 1]. Row payment: [0·1+1·0, 0·0+1·1, 0·1+1·(−1)] = [0, 1, −1]. Row failed: [1·1+1·0, 1·0+1·1, 1·1+1·(−1)] = [1, 1, 0]. Full matrix: [[1,0,1],[0,1,−1],[1,1,0]].
3. Scaled scores [0, 0.7071, −0.7071] exponentiate to [e⁰, e^0.7071, e^−0.7071] = [1.0000, 2.0281, 0.4931]. Sum = 3.5212. Middle weight = 2.0281 / 3.5212 = 0.5760.
4. d_k = 512 / 8 = 64 per head. Eight smaller heads let the model learn eight independent attention patterns in parallel (e.g. one head can specialize in local adjacency, another in long-range subject-verb dependency) at the same total parameter cost as one 512-dimensional head, which could only express a single weighting pattern per token per layer.
5. False. From the computed matrix, weight(UPI→payment) = 0.1978 while weight(payment→UPI) = 0.2840 — different values for the same token pair. Symmetry would require Q_i·K_j = Q_j·K_i for all pairs, which holds only if W_Q = W_K; in general they are separate learned matrices, so the score matrix, and hence the softmax output, is asymmetric.
6. Positions (i, j) with j > i — token i is not allowed to see tokens after it. Masked entries are set to −∞ before the softmax's exponentiation step; since e^(−∞) = 0, they automatically receive zero weight after normalization, with no separate zeroing operation needed — the mask reuses the same softmax that already runs over every row.

Think About It

Think about this: How would you explain the attention mechanism: mathematical deep dive to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where the attention mechanism: mathematical deep dive is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting the attention mechanism: mathematical deep dive to at least 3 other topics you have studied.