The Transformer Architecture: Attention is All You Need

An Indian conversational-AI startup — the kind now shipping Hindi and Hinglish customer-support bots for banks and telecom operators — rents a single 80 GB Nvidia A100 to serve a 7-billion-parameter transformer. The model's weights, stored in 16-bit precision, occupy a fixed 14 GB. The question the infrastructure team actually loses sleep over is not "how many FLOPs does one forward pass cost?" It is: how many customers can this one card hold in conversation at the same time before it runs out of memory and starts rejecting connections? The answer has almost nothing to do with the 7 billion parameters and everything to do with a structure this chapter dissects from the ground up: the query–key–value mechanism inside self-attention, and the cache it forces every production system to carry. By the end, you will be able to compute that concurrent-user number yourself, and explain exactly why two transformers with identical parameter counts can serve wildly different numbers of simultaneous conversations on the same GPU.

The exact computation behind self-attention

Vaswani et al. (Attention Is All You Need, NeurIPS 2017) defined self-attention with one compact formula:

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

Every symbol in that line is a matrix built from the same input token embeddings, but through three different learned linear projections. If x is a token's embedding vector (a row of the input matrix X), the model computes:

Q = X · W_Q     (what this token is "looking for")
K = X · W_K     (what this token "offers" to others as a match key)
V = X · W_V     (the actual content this token contributes if matched)

W_Q, W_K, and W_V are separate weight matrices, learned by gradient descent, each of shape (d_model, d_k). Because they are separate matrices, the same embedding produces three genuinely different vectors — a token is not compared to other tokens using its raw embedding, but using a projection of it built specifically for asking questions (Q), a different projection built for being matched against (K), and a third built for what to actually hand over once matched (V). This three-way split is the single most load-bearing design decision in the architecture: it is what lets "the same word" play a different role depending on whether it is doing the attending or being attended to.

Once Q and K exist, QK⁠ᵀ is a matrix of raw similarity scores — one row per query token, one column per key token, each entry the dot product between that query and that key. A large dot product means "this key strongly matches what I'm looking for." Softmax turns each row of scores into a probability distribution that sums to 1, and that distribution is then used as the weights in a weighted average of the value vectors V. The output for a given token is therefore a blend of every other token's content, weighted by how relevant each one is to it — computed, not hand-coded.

Why divide by √d_k, and not by d_k, or not at all

The scaling term is easy to wave past as a cosmetic detail. It is not. Assume — as is roughly true after the projections settle during training — that the individual components of a query vector q and a key vector k are independent random variables with mean 0 and variance 1. The dot product q · k = Σ qᵢ kᵢ sums d_k independent terms, each a product of two independent, zero-mean, unit-variance variables. Each such product term itself has mean 0 and variance 1 (the variance of a product of two independent zero-mean unit-variance variables is the product of their variances, which is 1). Summing d_k independent, identically distributed terms adds their variances: the dot product's variance is d_k, so its standard deviation grows as √d_k.

That matters because softmax is exponential. If dot products routinely land at magnitude √d_k and d_k is large — 64, 128, or more, as it is inside every production transformer — the raw scores spread out over a huge range before ever reaching softmax. Softmax on widely-spread inputs saturates: one entry gets driven toward 1 and the rest toward 0, the attention distribution collapses to a near one-hot pick, and the gradient flowing back through softmax vanishes almost everywhere except at that one spike. Training stalls. Dividing every score by √d_k rescales the dot product's standard deviation back down to 1 regardless of how large d_k is, keeping the softmax input in a range where it produces a genuinely soft, differentiable distribution instead of a frozen spike. This is a variance-stabilization step, not a convenience — and you will watch it happen numerically in the worked example below.

A full worked pass, verified end to end

Take a two-token sequence: the tokens "UPI" and "down" (as in "UPI down" — a payments-outage support ticket), with toy embeddings of dimension d_model = 4:

x(UPI)  = [1, 1, 1, 0]
x(down) = [0, 2, 1, 1]

Use a single attention head with d_k = d_v = 2, and fix small integer projection matrices so every step is checkable by hand:

import numpy as np

x_UPI  = np.array([1, 1, 1, 0], dtype=float)
x_down = np.array([0, 2, 1, 1], dtype=float)
X = np.stack([x_UPI, x_down])          # shape (2, 4)

W_Q = np.array([[1,0],[0,1],[1,0],[0,1]], dtype=float)
W_K = np.array([[1,0],[0,1],[0,1],[1,0]], dtype=float)
W_V = np.array([[1,1],[0,1],[1,0],[0,0]], dtype=float)

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

d_k = 2
scores = Q @ K.T
scaled = scores / np.sqrt(d_k)

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

attn = softmax(scaled)
output = attn @ V

print(Q, K, V, scores, scaled, attn, output)

Multiplying row by row: Q(UPI) = x(UPI)·W_Q = [2, 1] and Q(down) = [1, 3]; K(UPI) = [1, 2] and K(down) = [1, 3]; V(UPI) = [2, 2] and V(down) = [1, 2]. This code was executed (not hand-waved) to confirm every downstream number. For the query token "UPI", the raw scores against the two keys are Q(UPI)·K(UPI) = 2·1 + 1·2 = 4 and Q(UPI)·K(down) = 2·1 + 1·3 = 5. Scaling by √2 ≈ 1.41421 gives 2.8284 and 3.5355. Softmax over those two numbers produces attention weights [0.3302, 0.6698] — "UPI" attends 33% to itself and 67% to "down", which is sensible: "down" is the token carrying the outage signal. The output for "UPI" is the weighted blend 0.3302·[2,2] + 0.6698·[1,2] = [1.3302, 2.0000], again confirmed by execution, not estimation.

Multi-head attention: parallel subspaces, not parallel guesses

A second head, built from an independently-initialized triple W_Q2, W_K2, W_V2 over the same inputs, produces its own attention distribution for "UPI". Its projection matrices, fixed to small integers exactly as head 1's were:

W_Q2 = np.array([[1,0],[1,0],[1,1],[0,0]], dtype=float)
W_K2 = np.array([[1,3],[0,0],[0,0],[1,1]], dtype=float)
W_V2 = np.array([[1,2],[0,0],[0,0],[3,1]], dtype=float)

Q2 = X @ W_Q2
K2 = X @ W_K2
V2 = X @ W_V2

scores2 = Q2 @ K2.T
scaled2 = scores2 / np.sqrt(d_k)
attn2 = softmax(scaled2)
output2 = attn2 @ V2

Multiplying row by row exactly as with head 1: Q2(UPI) = x(UPI)·W_Q2 = [3, 1]; K2(UPI) = x(UPI)·W_K2 = [1, 3] and K2(down) = x(down)·W_K2 = [1, 1]; V2(UPI) = x(UPI)·W_V2 = [1, 2] and V2(down) = x(down)·W_V2 = [3, 1]. For the query token "UPI", the raw scores are Q2(UPI)·K2(UPI) = 3·1 + 1·3 = 6 and Q2(UPI)·K2(down) = 3·1 + 1·1 = 4, scaling to [4.2426, 2.8284], softmax weights [0.8044, 0.1956], and output 0.8044·[1,2] + 0.1956·[3,1] = [1.3911, 1.8044] — every number in this second head hand-checkable the same way head 1's were, and this code block was executed (not hand-waved) to confirm it. Notice this head reaches an almost opposite conclusion from head one — it weights "UPI" toward itself far more strongly. That divergence is the entire point of splitting attention into multiple heads with independent projections rather than using one large head with d_k = d_model: each head's W_Q/W_K pair carves out a different similarity geometry over the same embedding space, so one head can specialize in a local, self-referential relationship while another specializes in a longer-range dependency. Empirically, Voita, Talbot, Moiseev, Sennrich, and Titov (Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned, ACL 2019) confirmed this is not just a theoretical possibility: trained transformer heads specialize — some track positional/local patterns, some track rare, syntactically-important tokens — and a substantial fraction of heads can be pruned after training with only minor quality loss, because their specialized job turned out to be redundant with another head's.

The two head outputs are concatenated — [1.3302, 2.0000, 1.3911, 1.8044] — and passed through one more learned matrix W_O (the identity matrix in this toy example, so the concatenation is also the final output) to produce the token's final self-attention representation, which then flows into the feed-forward sublayer.

Misconception: "attention is basically free once the sequence is in memory"

Students who have understood that a transformer sees the whole sequence at once, unlike an RNN's strictly sequential pass, often conclude that this parallelism makes attention computationally cheap at any length — that because there's "no loop," cost barely depends on sequence length the way an RNN's does. This is backwards. Computing QK⁠ᵀ for a sequence of length n produces an n × n matrix of scores: attention is quadratic — O(n²) — in both compute and memory, per layer, per head. An RNN pays a constant cost per step and takes n steps, so its total cost is O(n); a transformer pays O(n²) total because every one of the n tokens compares itself against all n others. Doubling the context window quadruples the attention cost. This is precisely why long-context serving is hard and why the next section — the KV cache — is the dominant memory cost in production inference, not a footnote to it.

The real bottleneck: KV-cache economics in production

During autoregressive generation, a transformer produces tokens one at a time, and each new token needs to attend back over every token generated so far. Recomputing K and V for the entire prefix at every single new token would be wasteful, since those vectors for earlier tokens never change — so production systems cache them: the KV cache stores every previous token's key and value vectors, for every layer and every attention head, and simply appends to it as generation proceeds. The cache's size is what actually consumes GPU memory during a live conversation, separately from the fixed cost of the model's weights, and it grows with every token the user and the model exchange.

The per-sequence KV-cache size, in bytes, is:

KV_bytes = 2 × n_layers × n_kv_heads × head_dim × seq_len × bytes_per_value

(the leading 2 accounts for storing both K and V). The n_kv_heads term is where architecture choices bite hardest. Vaswani et al.'s original design uses one key/value head per query head (multi-head attention, MHA). Shazeer (Fast Transformer Decoding: One Write-Head Is All You Need, 2019) proposed multi-query attention (MQA): every query head still exists, but all of them share a single key/value head, cutting the cache by a factor equal to the head count. Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebrón, and Sanghai (GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023) generalized this to grouped-query attention (GQA): query heads are split into a handful of groups, each group sharing one key/value head — a tunable middle ground between MHA's quality and MQA's memory savings. Meta's Llama 2 70B (Touvron et al., 2023) ships with GQA: 8192 hidden size, 80 layers, 64 query heads, but only 8 key/value heads.

Plug Llama 2 70B's real configuration into the formula for a 4096-token conversation in 16-bit precision: 2 × 80 × 8 × 128 × 4096 × 2 bytes = 1,342,177,280 bytes ≈ 1.34 GB per conversation. Had the same model instead used ordinary multi-head attention with 64 key/value heads (one per query head), the identical formula with n_kv_heads = 64 gives 10.74 GB per conversation — an 8× increase, for zero change in the model's ability to represent relationships between tokens, purely from how key/value heads are shared. That difference alone can be the gap between a GPU serving one user's long conversation and running out of memory before serving anyone.

Return to the opening scenario. Mistral 7B (Jiang et al., 2023) uses GQA with hidden size 4096, 32 layers, 32 query heads, and 8 key/value heads. Its weights in 16-bit precision occupy 7 × 10⁹ × 2 bytes = 14 GB. On an 80 GB card, reserving roughly 10% headroom for activations and runtime scratch space leaves about 58 GB for KV caches. Each 4096-token conversation costs 2 × 32 × 8 × 128 × 4096 × 2 ≈ 0.537 GB, so 58 / 0.537 ≈ 108 conversations can run concurrently. Swap that same model's attention design for MQA (a single shared key/value head instead of 8) and each conversation's cache shrinks to 0.0671 GB, pushing concurrent capacity to roughly 864 — an 8× jump in how many customers one GPU serves, with no change to the number of parameters, purely from an attention-head-sharing decision made at model-design time. This is the calculation the infrastructure team from the opening paragraph actually runs, and it is why GQA/MQA are now near-universal in models built for high-concurrency serving rather than a niche optimization.

Active recall

Attempt every question before reading its answer.

1. Why does dividing the dot product by √d_k (rather than by d_k, or not scaling at all) keep softmax well-behaved as d_k grows?
2. In the worked example, suppose x(down) changes from [0, 2, 1, 1] to [0, 4, 1, 1] (only its second component doubles). Trace the full effect on head 1's computation for the query "UPI": which of Q(UPI), K(UPI), K(down), V(down) change, and what happens to the attention weights and the output?
3. A team deploys a GQA transformer with 40 layers, 4 key/value heads, head_dim = 128, serving 8192-token conversations in fp16. Compute the KV-cache memory per conversation.
4. Why can two transformers with the exact same total parameter count require drastically different amounts of GPU memory to serve the same conversation?
5. Why does using more than one attention head tend to outperform one large head with d_k = d_model, even when both have comparable total projection parameters?
6. True or false, and why: "Because a transformer processes the whole sequence in parallel rather than one token at a time like an RNN, its computational cost barely grows with sequence length."

Answers.

1. If the components of q and k are independent, mean-0, unit-variance, each of the d_k product terms in the dot product has variance 1, and summing d_k independent terms sums their variances, giving the dot product a variance of d_k and a standard deviation of √d_k. Dividing by √d_k restores unit standard deviation regardless of how large d_k gets. Without this, larger d_k produces larger-magnitude scores, which push softmax toward a near one-hot output and collapse its gradient almost everywhere — dividing by the raw d_k would overcorrect and shrink the scores' variance to 1/d_k, making the distribution too flat to express confident attention.

2. Q(UPI) and K(UPI) are built only from x(UPI), so they are completely unchanged: Q(UPI) = [2, 1], K(UPI) = [1, 2] still. But K(down) and V(down) are built from x(down), so both shift: K(down) goes from [1, 3] to [1, 5], and V(down) goes from [1, 2] to [1, 4]. The score against "UPI" itself is unchanged at 4 (since neither Q(UPI) nor K(UPI) moved), but the score against "down" rises from 2·1+1·3=5 to 2·1+1·5=7. After scaling by √2, the softmax weights shift from the fairly balanced [0.330, 0.670] to the sharply skewed [0.107, 0.893] — a single embedding change in one token pulls the query's attention almost entirely onto it. The output for "UPI" moves from [1.330, 2.000] to [1.107, 3.786], pulled toward the new, larger V(down). This also demonstrates the saturation risk from the scaling discussion directly: the larger dot product pushes the distribution measurably closer to one-hot.

3. 2 × 40 × 4 × 128 × 8192 × 2 bytes = 671,088,640 bytes ≈ 0.671 GB per conversation.

4. The bulk of inference-time memory beyond the fixed weights is the KV cache, whose size depends on n_kv_heads — a design choice independent of total parameter count. A model using standard multi-head attention (one key/value head per query head) can need 4–8× more KV-cache memory per conversation than a same-sized model using grouped-query or multi-query attention, because the latter share key/value projections across groups of query heads. Two models can match on parameters and differ by many gigabytes per active conversation purely from this structural choice.

5. Different heads specialize on different projection subspaces of the same embeddings, as demonstrated directly in the worked example (head 1 and head 2 reached almost opposite attention distributions for the same query from the same input) and confirmed empirically by Voita et al. (2019), who found individual trained heads specialize in different roles. One large head with d_k = d_model is forced to compute a single softmax distribution mixing every kind of token relationship into one weighting; it cannot simultaneously express "attend locally" and "attend to the semantically related distant token" as two separate, cleanly separable distributions the way multiple smaller heads can.

6. False. Parallelism removes the sequential dependency of an RNN, but the compute and memory cost of forming the n × n attention score matrix is O(n²) in sequence length, per layer and head — worse asymptotic scaling than an RNN's O(n), not better. This is exactly why the KV cache (which adds a fixed amount per new token, accumulating to O(n) total per sequence, but multiplied across every layer and every kv-head) becomes the dominant memory cost as context windows grow, and why halving or doubling a deployed model's context length has an outsized effect on how many users one GPU can serve.

Think About It

Think about this: How would you explain the transformer architecture: attention is all you need 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 transformer architecture: attention is all you need 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 transformer architecture: attention is all you need to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind the transformer architecture: attention is all you need, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.