Building the Transformer Architecture from Mathematical First Principles

A Bengaluru startup runs a cricket-commentary chatbot on a 7-billion-parameter transformer, served in fp16 on 80 GB GPUs. Every question so far in this curriculum has asked "what does attention compute?" Tonight's question is different: the IPL final is two days away, the product team expects 200 concurrent long conversations during the last over, and the on-call engineer has to answer a question no forward-pass diagram can answer: how many GPUs do we book? Get the arithmetic wrong and the service falls over mid-match, or the company burns lakhs on GPUs it didn't need. That number does not come from reading the "Attention Is All You Need" architecture diagram. It comes from three pieces of mathematics that sit underneath the diagram: why the attention scores are divided by a specific constant, how gradients actually flow backward through softmax, and how memory scales with sequence length during autoregressive decoding. This chapter derives all three from scratch, traces a complete numeric example through both the forward and backward pass, and returns to the Bengaluru GPU order at the end with an exact figure.

Scope: derivation, not description

You already know the shape of a transformer block: project into queries, keys and values, compute softmax(QKₜ/√d_k)V, run it through multiple heads, stack the blocks, add positional information. This chapter does not re-derive that shape. It answers four questions the shape alone does not answer, each of which is exactly the kind of question a systems engineer or a research reviewer actually asks: where does the √d_k come from, mathematically, not by convention; what does the gradient look like when you differentiate through a softmax that sits inside a bilinear form; when does the quadratic cost of attention actually start to dominate the rest of the network; and how many bytes does one token of context cost you at inference time, in production, on a GPU you are paying for by the hour.

Deriving the √d_k Scaling Factor from Variance, Not Convention

Take a single query vector q and a single key vector k, both in R^d_k, and assume every component of both is drawn independently with mean 0 and variance 1 — a reasonable approximation immediately after the linear projections that produce Q and K, whose weights are initialized to keep activations close to this regime. The raw attention score is the dot product s = q·k = Σ_i=1^d_k q_i k_i.

Its expectation is E[s] = Σ_i E[q_i]E[k_i] = 0, since q_i and k_i are independent and each has mean 0. Its variance, using independence across the d_k terms of the sum, is Var(s) = Σ_i Var(q_i k_i). For each term, Var(q_i k_i) = E[q_i²k_i²] − (E[q_i k_i])² = E[q_i²]E[k_i²] − 0 = 1·1 = 1, because E[q_i²] equals Var(q_i) when the mean is 0. Summing d_k such terms gives Var(s) = d_k. The standard deviation of a raw, unscaled attention score therefore grows as √d_k, not as a constant, purely as a consequence of summing d_k independent products.

This matters because softmax is scale-sensitive: for two logits that differ by Δ, softmax output ratios scale as e^Δ. If the typical spread of scores grows with √d_k, then for large d_k (64, 128, 256 for individual attention heads in real models) the pre-softmax logits are typically tens of standard deviations wide, softmax saturates to a near one-hot distribution, and the gradient through softmax — derived explicitly in the next section — shrinks toward zero almost everywhere. Dividing every score by exactly √d_k rescales the standard deviation back to 1, independent of d_k: Var(s/√d_k) = Var(s)/d_k = d_k/d_k = 1. This is not a tuned hyperparameter found by grid search; it is the unique normalizing constant that keeps the variance of the pre-softmax logits fixed as the head dimension changes, which is exactly the argument the original Vaswani et al. (2017) paper gives for it.

Numerical Stability: Why Softmax Subtracts the Row Maximum First

Before differentiating through softmax, one practical detail has to be correct or the forward pass itself breaks. Softmax of a row of logits s is exp(s_j)/Σ_k exp(s_k). float32 overflows above roughly exp(88) ≈ 1.6×10³&sup8;. If two logits are, say, 1000 and 1002 — entirely plausible before scaling, given the variance result above — exp(1000) overflows to inf and the division produces nan, silently corrupting every downstream gradient. The fix exploits the fact that softmax is invariant to a constant shift of all logits in a row: exp(s_j−m)/Σ_k exp(s_k−m) = exp(s_j)/Σ_k exp(s_k) for any constant m, since the e^−m factor cancels top and bottom. Choosing m = max_j(s_j) guarantees the largest shifted logit is exactly 0, so the largest exponential is exp(0)=1 and every other term is a fraction, with no overflow possible. For the [1000, 1002] example: shifted logits are [−2, 0], giving exp values [0.1353, 1], sum 1.1353, and softmax [0.1192, 0.8808] — mathematically identical to the unstable computation, but representable in float32.

Backpropagation Through Attention: A Fully Worked Example

Now trace gradients through one full scaled dot-product attention block, with every intermediate value computed to four decimal places and cross-checked in code. Use n=2 tokens and d_k=d_v=2, with

Q = [[1, 0],
     [0, 1]]

K = [[1, 0],
     [0, 1]]

V = [[1, 2],
     [3, 4]]

Forward. S = QKₜ = [[1,0],[0,1]] (Q and K are both the identity matrix here, chosen deliberately so every later matrix multiplication reduces to something you can verify by hand). Scaled: S′ = S/√2 = [[0.7071, 0], [0, 0.7071]]. Row-wise softmax of row 1, [0.7071, 0]: exp(0.7071)=2.0281, exp(0)=1, sum=3.0281, giving A row 1 = [0.6698, 0.3302]; row 2 is the mirror image, [0.3302, 0.6698], by the symmetry of S′. The output is O = AV: row 1 = 0.6698×[1,2] + 0.3302×[3,4] = [1.6605, 2.6605]; row 2 = 0.3302×[1,2] + 0.6698×[3,4] = [2.3395, 3.3395].

Backward. Take the simplest possible loss, L = ΣO (sum of every entry of the output), so ∂L/∂O is a 2×2 matrix of ones. Four chain-rule steps carry this gradient back to Q, K and V.

Step 1, through O = AV: since this is a plain matrix product, ∂L/∂A = (∂L/∂O)Vₜ and ∂L/∂V = Aₜ(∂L/∂O). With ∂L/∂O all ones and Vₜ=[[1,3],[2,4]], every row of ∂L/∂A comes out to [3, 7] (each entry is a column sum of Vₜ: 1+2=3 and 3+4=7). And ∂L/∂V = [[1,1],[1,1]], because each row of A sums to 1, so Aₜ applied to a matrix of ones just redistributes the ones.

Step 2, through the softmax, is the step students most often get wrong. Softmax is not element-wise, so its Jacobian is not diagonal: perturbing one score in a row shifts every probability in that row, because they must keep summing to 1. For a softmax row a = [a_1,...,a_m] with upstream gradient da = [∂L/∂a_1,...,∂L/∂a_m], the gradient with respect to the pre-softmax logit s_j is ∂L/∂s_j = a_j·(da_j − Σ_k a_k da_k). For row 1, a=[0.6698,0.3302], da=[3,7], so Σ_k a_k da_k = 0.6698×3 + 0.3302×7 = 4.3210. Then ∂L/∂s′_{1,1} = 0.6698×(3−4.3210) = −0.8847, and ∂L/∂s′_{1,2} = 0.3302×(7−4.3210) = 0.8847. Row 2 works out symmetrically to [−0.8847, 0.8847]. Notice the two entries in each row sum to exactly zero: since Σ_j a_j = 1 is a constant no matter how the logits move, the softmax Jacobian always projects the incoming gradient onto the subspace of vectors that sum to zero. This is a structural property of softmax, not a coincidence of these numbers.

Step 3 undoes the scaling: since S′=S/√d_k is linear, ∂L/∂S = (∂L/∂S′)/√d_k = −0.8847/1.4142 = −0.6256 (and +0.6256 for the positive entries), giving ∂L/∂S = [[−0.6256, 0.6256], [−0.6256, 0.6256]].

Step 4, through S=QKₜ: ∂L/∂Q = (∂L/∂S)K and ∂L/∂K = (∂L/∂S)ₜQ. Because both Q and K equal the identity matrix in this example, both simplify to ∂L/∂Q = ∂L/∂S = [[−0.6256,0.6256],[−0.6256,0.6256]] and ∂L/∂K = (∂L/∂S)ₜ = [[−0.6256,−0.6256],[0.6256,0.6256]].

The code below implements exactly these four steps and reproduces every number above:

import numpy as np

Q = np.array([[1.0, 0.0], [0.0, 1.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0]])
V = np.array([[1.0, 2.0], [3.0, 4.0]])
d_k = Q.shape[1]

S = Q @ K.T
S_scaled = S / np.sqrt(d_k)
shifted = S_scaled - S_scaled.max(axis=1, keepdims=True)
expS = np.exp(shifted)
A = expS / expS.sum(axis=1, keepdims=True)
O = A @ V

dO = np.ones_like(O)               # dL/dO for L = sum(O)
dA = dO @ V.T
dV = A.T @ dO
dS_scaled = np.zeros_like(A)
for i in range(A.shape[0]):
    a, da = A[i], dA[i]
    dS_scaled[i] = a * (da - np.sum(a * da))
dS = dS_scaled / np.sqrt(d_k)
dQ = dS @ K
dK = dS.T @ Q

print(np.round(A, 4))          # [[0.6698 0.3302] [0.3302 0.6698]]
print(np.round(O, 4))          # [[1.6605 2.6605] [2.3395 3.3395]]
print(np.round(dS_scaled, 4))  # [[-0.8847 0.8847] [-0.8847 0.8847]]
print(np.round(dS, 4))         # [[-0.6256 0.6256] [-0.6256 0.6256]]

Run it and the printed arrays match the hand derivation to rounding. The diagram below lays out this exact forward-then-backward path as a computational graph, with the value flowing down in blue and the gradient flowing up in red, labeled at each stage with the formula just derived.

Compute Cost: When Does Attention Actually Dominate the Feed-Forward Block?

A transformer layer spends FLOPs in two very different places: the attention sub-layer and the position-wise feed-forward sub-layer. It is common to hear that "attention is the expensive part," but that is only true past a specific, derivable sequence length. Attention computes S=QKₜ (an n×d_model matrix times a d_model×n matrix, costing on the order of n²d_model multiply-adds) and then AV (another n²d_model), for a total of roughly 2n²d_model FLOPs per layer. The feed-forward block, with its standard 4× hidden expansion, applies two linear layers per token: d_model→4d_model and 4d_model→d_model, each costing n·d_model·4d_model multiply-adds, for a total of roughly 8n·d_model² FLOPs per layer.

Setting the two equal to find the crossover: 2n²d_model = 8n·d_model², which simplifies to n = 4·d_model. For a model with d_model=4096 (a realistic mid-size figure), the crossover sits at n = 16,384 tokens. Below that context length, the feed-forward block is the FLOP bottleneck, not attention — the opposite of most students' intuition, which fixates on attention's quadratic term without checking where the constants put the actual crossover for realistic dimensions. Above roughly 16k tokens, attention's n² term overtakes the feed-forward's fixed-per-token cost, which is precisely the regime where engineering techniques that specifically target the attention computation (rather than the whole network) start to earn their complexity.

KV-Cache Memory: Sizing the Bengaluru GPU Order

During autoregressive generation, a transformer would naively recompute K and V for every previous token at every new decoding step, an O(n²) waste. Production inference engines instead cache K and V for every layer once and reuse them, appending only the new token's K and V at each step. The cost of this optimization is memory, and that memory is exactly what constrains how many concurrent conversations one GPU can serve.

For a model with L layers and hidden size d_model, running in a format that uses p bytes per number, one token of cached K and V across every layer costs 2·L·d_model·p bytes (the factor of 2 is for storing both K and V; d_model appears because summing d_head over all attention heads recovers d_model regardless of how many heads split it). For the Bengaluru startup's 7B-class model, L=32 and d_model=4096, running in fp16 (p=2 bytes): 2×32×4096×2 = 524,288 bytes per token, i.e. 512 KiB per token. At a 2048-token context, that is 512 KiB × 2048 = 1,048,576 KiB = exactly 1 GiB of cache per active conversation. For 200 concurrent conversations during the match, the KV-cache alone needs 200 GiB. Model weights add another 7×10&sup9;×2 bytes ≈ 13.04 GiB, fixed regardless of how many users are connected. Total: roughly 213 GiB, which will not fit on two 80 GB GPUs (160 GB) and needs at least three, with real deployments typically adding a further margin for activation memory and request batching overhead. That is the number the on-call engineer takes to the GPU order form, derived rather than guessed.

Assembling the Block: How the Derived Pieces Compose into One Layer

The four quantities derived above are not four isolated facts; they are the load-bearing constants and dynamics inside every transformer block actually built. One encoder or decoder block is LayerNorm(x + MultiHead(x)), followed by LayerNorm(x′ + FFN(x′)). MultiHead splits the model dimension into h heads, each of size d_k = d_model/h, runs the scaled dot-product attention derived in full above independently per head — using the exact √d_k normalization from the first section, computed with the numerically stable, max-subtracted softmax from the second section — concatenates the h head outputs back to width d_model, and adds the result to the block's input x through a residual connection before LayerNorm rescales the sum. The residual path itself costs nothing in FLOPs, and LayerNorm costs only O(n·d_model) per layer, negligible next to either sub-layer's matrix multiplies; the two components that actually dominate a block's compute, and the ones whose crossover was derived above, are exactly the attention matrix multiplies and the FFN matrix multiplies inside that same block. Stack N such blocks and apply the KV-cache accounting above per layer, and the result is the full encoder or decoder stack. Building the architecture, in this sense, is exactly assembling these four derived pieces in this order: nothing about the composition itself (where the residual attaches, where LayerNorm sits, how heads split and recombine) requires mathematics beyond what the four sections above already established.

Correcting a Common Misconception

Students who have only seen the attention formula written down, without deriving it, commonly treat the √d_k term as an empirically discovered hyperparameter — something like a learning rate, found by trying a few values and keeping whichever worked best in the original experiments. This is incorrect, and the derivation above shows exactly why: √d_k is not a value that was searched over, it is the unique constant, derived analytically from the variance of a sum of d_k independent products, that keeps the pre-softmax logit variance at exactly 1 regardless of head dimension. Nothing about it was tuned. The consequence of getting this wrong in practice is concrete, not cosmetic: build a transformer with, say, d_k=1024 per head and skip the scaling, and Var(s)=1024, giving a typical logit spread of roughly 32 standard deviations before softmax. Softmax with inputs that spread out saturates almost everywhere into a near one-hot distribution, and by the Jacobian formula derived above, the softmax gradient a·(da−Σa·da) shrinks toward zero whenever a itself is close to a one-hot vector (a≈0 or a≈1 makes the product tiny). The model does not merely train a little worse; large blocks of it stop receiving gradient signal at all, and this failure mode scales directly with head dimension, which is exactly why production architectures pin per-head d_k near 64-128 and grow capacity by adding more heads instead of widening each one — d_model climbs into the thousands, but d_k does not.

Active Recall

Attempt each question before reading the worked answer beneath it.

1. Derive Var(q·k) for query and key vectors in R^d_k with independent, mean-0, variance-1 components, and explain why the correct normalizer is √d_k rather than d_k itself.
2. In the worked backward pass, why do the two entries of dL/dS′ in each row always sum to exactly zero, for any values of Q, K and V?
3. The startup's context length doubles from 2048 to 4096 tokens, with batch size fixed at 200 concurrent sequences and precision fixed at fp16. State what happens, in exact numbers, to: (a) KV-cache memory, (b) attention compute per layer, (c) feed-forward compute per layer, (d) model weight memory.
4. For d_model=4096, at what sequence length n do attention FLOPs and feed-forward FLOPs become equal? Show the equation you solved.
5. Raw attention logits before scaling are [1000, 1002]. Explain, with the actual numbers, why computing softmax directly on these is numerically dangerous, and show the stable computation.
6. The startup switches the KV-cache (only the cache, not the model weights) from fp16 to int8 quantization, and simultaneously doubles the context length from 2048 to 4096 tokens. Does total KV-cache memory for 200 sequences go up, down, or stay the same? Give the exact figure.

Worked Answers

1. With independence, E[q·k]=0 and Var(q·k)=Σ_i=1^d_k Var(q_i k_i)=d_k·1=d_k, since each Var(q_i k_i)=E[q_i²]E[k_i²]=1. Dividing by d_k would give Var(s/d_k)=d_k/d_k²=1/d_k, which shrinks to 0 as d_k grows, over-flattening the softmax input and destroying the model's ability to distinguish tokens. Dividing by √d_k gives Var(s/√d_k)=d_k/d_k=1, exactly independent of d_k, which is the property actually needed.

2. Because softmax rows sum to 1 no matter what the logits are, Σ_j a_j is a constant. The Jacobian identity ∂L/∂s_j=a_j(da_j−Σ_k a_k da_k), summed over j, gives Σ_j a_j da_j − (Σ_k a_k da_k)(Σ_j a_j) = Σ_j a_j da_j − Σ_k a_k da_k = 0, since Σ_j a_j=1. This is a structural consequence of softmax's normalization constraint, true for any Q, K, V, not a coincidence of the chosen numbers.

3. (a) KV-cache: 512 KiB/token × 4096 = 2 GiB/sequence × 200 = 400 GiB, exactly double the 200 GiB at 2048 tokens — linear in context length. (b) Attention compute per layer ∝ n², so it quadruples. (c) Feed-forward compute per layer ∝ n (fixed cost per token, times number of tokens), so it only doubles, not quadruples — the most commonly missed part of this question, since students who correctly remember attention is "quadratic" often assume everything downstream inherits that quadratic scaling. (d) Model weight memory (13.04 GiB) is completely independent of context length and does not change at all.

4. Solve 2n²d_model = 8n·d_model² for n: divide both sides by 2n·d_model to get n = 4d_model = 4×4096 = 16,384 tokens.

5. exp(1000) exceeds float32's representable range (overflow near exp(88)) and returns inf, and inf/inf in the softmax division returns nan, silently poisoning the rest of the computation. The stable path subtracts the row maximum first: shifted logits become [1000−1002, 1002−1002] = [−2, 0], giving exp values [0.1353, 1], sum 1.1353, softmax [0.1192, 0.8808] — the exact same mathematical answer, computed without overflow, because the shift cancels in the ratio.

6. int8 halves bytes-per-token from 512 KiB to 256 KiB. At the new context of 4096 tokens: 256 KiB×4096 = 1 GiB/sequence, ×200 sequences = 200 GiB — identical to the original fp16-at-2048-tokens figure. The two changes exactly offset: doubling context length doubles memory, and halving precision halves it, so total KV-cache memory stays at 200 GiB even though the model now serves twice the context per user. This is precisely the tradeoff production inference engines exploit: int8 or int4 KV-cache quantization is what makes longer context windows affordable without a proportional GPU budget increase.

Think About It

Think about this: How would you explain building the transformer architecture from mathematical first principles 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.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind building the transformer architecture from mathematical first principles, 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.