Transformer Architecture from Scratch

During the 2023 IPL, Cricbuzz and ESPNcricinfo both run live commentary bots that turn ball-tracking data into English sentences within a second of the delivery. Somewhere in that pipeline a sentence like this gets generated: "The umpire warned the bowler after he overstepped the crease twice." Any Indian cricket fan resolves "he" to "the bowler" instantly. A 2016-era commentary bot, built on a recurrent neural network (RNN) reading the sentence one word at a time, routinely got this wrong on longer sentences — not because it didn't know cricket, but because of an architectural limitation: by the time an RNN's hidden state reaches the word "he", the influence of "bowler" — nine words back — has been diluted by nine rounds of matrix multiplication and squashed through nine activation functions. The gradient that would teach the network "he = bowler" has to survive that same nine-step chain during backpropagation, and it shrinks at every step. This is the vanishing-gradient problem, and it is the specific, concrete reason the transformer architecture (Vaswani et al., "Attention Is All You Need," 2017) exists. A transformer does not read "he" and then separately recall "bowler" — it lets every word directly examine every other word in a single computation, with no intermediate chain to decay. That direct, all-pairs comparison is called self-attention, and it is the one mechanism this entire architecture is built around. This chapter builds it from the ground up: vectors in, contextualized vectors out, every matrix multiplication shown.

From words to vectors: embeddings and positions

A transformer never sees the word "bowler". It sees a vector — a list of real numbers, typically 512 or 768 of them in practice, that the model has learned to associate with that word during training. This is the embedding layer: a lookup table of shape (vocabulary size × d_model) where d_model is the chosen vector width. Feed in a sentence of n tokens and you get an n × d_model matrix, X, one row per word.

There is an immediate problem. Matrix multiplication and the attention mechanism you're about to see are both permutation-equivariant: if you shuffle the rows of X, the output rows shuffle in exactly the same way, but nothing about the computation itself changes based on which position each row started at. Self-attention, on its own, cannot tell "the bowler warned the umpire" from "the umpire warned the bowler" — it is fundamentally an operation on a set of vectors, not a sequence. The fix is positional encoding: a fixed vector added to each token's embedding, one vector per position, so that the same word carries different information depending on where it sits in the sentence. Vaswani et al. used sinusoids of different frequencies:

PE(pos, 2i)   = sin( pos / 10000^(2i / d_model) )
PE(pos, 2i+1) = cos( pos / 10000^(2i / d_model) )

where pos is the token's position (0, 1, 2, …) and i indexes pairs of dimensions. Low-index dimensions oscillate fast (fine-grained position information), high-index dimensions oscillate slowly (coarse position information) — the same trick a clock uses seconds, minutes, and hours to encode a wide range of time with a few hands. The sinusoidal choice has a specific mathematical payoff: PE(pos+k) can be written as a linear function of PE(pos) for any fixed offset k, which lets the network learn to attend by relative position, and lets it generalize to sentence lengths longer than anything seen during training — something a learned position lookup table cannot do past its trained maximum length. The input to the first transformer layer is simply X + PE, an n × d_model matrix that now carries both content and order.

Self-attention: query, key, value

Take the embedded, position-tagged vector for "he" — call it x. To decide what "he" should "gather" from the rest of the sentence, the transformer computes three different projections of x using three learned weight matrices, W_Q, W_K, W_V (each d_model × d_k):

Q = X W_Q     # "what am I looking for?"
K = X W_K     # "what do I contain, that others might look for?"
V = X W_V     # "what do I actually hand over, if matched?"

Every one of the n tokens produces its own query, key, and value row, giving three n × d_k matrices. For a given token's query vector q, its attention scores against every key are computed as dot products: q · k_j measures how well token j's key aligns with what the current token is querying for. A high dot product means "this is relevant to me." Stack this across all n queries against all n keys in one matrix multiplication, Q Kᵀ, and you get an n × n matrix of raw similarity scores — every token compared against every other token, in one shot, with no sequential chain.

These raw scores are not yet usable as weights — they can be arbitrarily large or negative, and their scale grows with d_k. Assume, as is standard at initialization, that each component of q and k is drawn independently with mean 0 and variance 1. Then:

Var(q · k) = Var(Σ_i=1^d_k q_i k_i) = Σ_i=1^d_k Var(q_i k_i) = Σ_i=1^d_k E[q_i²]E[k_i²] = d_k · 1 · 1 = d_k

(the middle step uses independence of q_i and k_i, and E[q_i k_i]=0 so Var(q_i k_i)=E[q_i²k_i²]). So the standard deviation of the raw dot product grows as √d_k. As d_model (and therefore typically d_k) gets larger, the scores spread out further, and feeding a wide-spread vector into a softmax pushes it toward a near one-hot distribution — almost all the probability mass on the single largest score, everything else pinned near zero. That kills learning: the gradient of softmax with respect to its inputs is smallest exactly where the output is saturated near 0 or 1. Vaswani et al.'s fix is to divide every score by √d_k before the softmax, which — per the derivation above — renormalizes the standard deviation back to 1 regardless of d_k. This is exactly why the mechanism is named scaled dot-product attention:

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

The softmax is applied row-wise, turning each token's n raw scores into a probability distribution that sums to 1 — the attention weights. Multiplying this n × n weight matrix by V (n × d_v) produces the output: each token's new representation is a weighted average of every token's value vector, weighted by how relevant that token's key was to this token's query. This is the entire self-attention operation, and it is nothing more than two matrix multiplications and a softmax.

Worked example: three tokens, computed by hand

Take the truncated phrase "Bumrah bowled it" — three tokens, embedded in d_model = 4 for illustration (real models use 512+). The embedding vectors below are deliberately simple integers chosen so the arithmetic is checkable by hand; they are illustrative, not the output of a trained model:

X = [[1, 0, 1, 0],   # Bumrah
     [0, 1, 0, 1],   # bowled
     [1, 1, 0, 0]]   # it

Project into Q, K, V using fixed 4×3 weight matrices (d_k = d_v = 3):

import numpy as np

X  = np.array([[1,0,1,0], [0,1,0,1], [1,1,0,0]], dtype=float)
Wq = np.array([[1,0,0], [0,1,0], [1,0,1], [0,1,0]], dtype=float)
Wk = np.array([[0,1,0], [1,0,1], [0,0,1], [1,1,0]], dtype=float)
Wv = np.array([[1,0,0], [0,1,0], [0,0,1], [1,1,1]], dtype=float)

Q, K, V = X @ Wq, X @ Wk, X @ Wv
d_k = Q.shape[1]                         # 3
scores  = (Q @ K.T) / np.sqrt(d_k)

def softmax(row):
    e = np.exp(row - row.max())
    return e / e.sum()

weights = np.array([softmax(r) for r in scores])
output  = weights @ V

Working through the matrix multiplications by hand: row 1 of X, [1,0,1,0], dotted against each column of W_Q gives q₁ = [2, 0, 1]; the same for rows 2 and 3 gives q₂ = [0, 2, 0] and q₃ = [1, 1, 0]. Repeating with W_K gives k₁ = [0, 1, 1], k₂ = [2, 1, 1], k₃ = [1, 1, 1]. With W_V: v₁ = [1, 0, 1], v₂ = [1, 2, 1], v₃ = [1, 1, 0].

Now the scaled scores for token 1's query against all three keys, dividing by √3 ≈ 1.732: q₁·k₁ = (2)(0)+(0)(1)+(1)(1) = 1 → 0.577; q₁·k₂ = (2)(2)+(0)(1)+(1)(1) = 5 → 2.887; q₁·k₃ = (2)(1)+(0)(1)+(1)(1) = 3 → 1.732. Softmax of [0.577, 2.887, 1.732] gives weights [0.0702, 0.7070, 0.2228] — the executed code confirms this exactly. Token "Bumrah" ends up attending most heavily (70.7%) to token 2, "bowled". Its output row is the weights blended with V: 0.0702·v₁ + 0.7070·v₂ + 0.2228·v₃ = [1.000, 1.637, 0.777].

Token 2's row is more interesting. q₂ = [0, 2, 0] has a nonzero value only in its second coordinate. Look at the second coordinate of every key: k₁, k₂, k₃ all have second coordinate exactly 1. So q₂ · k_j = 2 × 1 = 2 for every single j — the three raw scores are tied before scaling, and stay tied after. A tied score set softmaxes to a perfectly uniform distribution: weights₂ = [0.3333, 0.3333, 0.3333], exactly a plain average of the three value vectors, [1.000, 1.000, 0.667]. This is not a coincidence to wave away — it is the mechanism working exactly as designed: query₂ happens to probe a subspace of the keys where all three tokens are indistinguishable, so attention correctly reports "I have no preference here" via a flat distribution. Row 3 (token "it") computes to weights [0.1679, 0.5329, 0.2992] and output [1.000, 1.365, 0.701], following the identical procedure. Every one of these nine decimal values above was produced by running the code, not asserted — you should re-run it and check.

Multi-head attention

A single attention computation forces every token to blend all of its "relevance" reasoning into one n × n weight matrix. But a word like "it" might simultaneously need to track a syntactic relationship (it's the object of "bowled") and a coreference relationship (it refers to the ball). One weight matrix per token can't cleanly represent two different kinds of relevance at once — averaging them together muddies both signals. Multi-head attention runs h independent copies of the entire Q/K/V/softmax pipeline in parallel, each with its own learned W_Q, W_K, W_V, projecting into a smaller d_k = d_model / h per head (in the original paper, d_model = 512, h = 8, d_k = d_v = 64). Each head produces its own n × d_v output; concatenating all h outputs back into an n × d_model matrix and passing it through one more learned projection, W_O, produces the final multi-head output. The parameter count is deliberately similar to a single large head, because h × d_k = d_model — the split doesn't cost extra capacity, it reallocates the same capacity into h independent attention patterns that the model is free to specialize differently during training.

Residual connections, layer normalization, the feed-forward block

The attention sub-layer's output is not used directly — it is added back to its own input: x + Attention(x). This residual (skip) connection matters for the same reason it mattered in ResNets: it gives gradients a direct path back to earlier layers that doesn't require passing through the attention computation, which keeps deep stacks of transformer layers trainable. The sum is then normalized with layer normalization — for each token independently (not across the batch, unlike batch normalization), subtract the mean and divide by the standard deviation across that token's d_model features, then apply a learned scale γ and shift β: LayerNorm(x) = γ · (x − μ) / σ + β. This keeps activations in a stable numeric range as they pass through dozens of stacked layers.

After Add & Norm, each token's vector (independently — no cross-token interaction here) passes through a two-layer feed-forward network: FFN(x) = max(0, x W₁ + b₁) W₂ + b₂. W₁ typically expands d_model up to 4× its width (512 → 2048) before W₂ projects it back down, giving the model extra nonlinear capacity to transform each token's attention-gathered representation. The output goes through a second residual connection and layer norm. One full encoder layer is: self-attention → add & norm → feed-forward → add & norm. Real transformers stack six, twelve, or many more of these layers, each with its own independently learned weights, so representations get progressively more contextualized with depth.

The decoder: masked self-attention

Everything above describes an encoder layer, which is free to let every token see every other token — appropriate for building a representation of a complete input sentence. A decoder, generating the commentary sentence one word at a time, cannot be allowed to do this: at the moment it is producing the word "he", it must not be allowed to peek at "overstepped" or "twice", which haven't been generated yet, or training would let it cheat by looking at the answer. The fix is a causal mask applied to the raw scores before the softmax: set every score at position (i, j) where j > i to −∞, so that after exponentiating, e^(−∞) = 0 and the softmax assigns exactly zero weight to future positions, with no probability mass to renormalize away. For a 4-token decoding sequence the mask added to the raw score matrix looks like:

[[  0, -inf, -inf, -inf],
 [  0,   0,  -inf, -inf],
 [  0,   0,    0,  -inf],
 [  0,   0,    0,    0]]

Masking must happen before the softmax, not after — if you zeroed out the future-position weights after softmax instead, the remaining weights would no longer sum to 1, and you'd need an extra renormalization step that complicates the gradient and doesn't behave identically. The original encoder-decoder transformer additionally gives the decoder a cross-attention sub-layer, where the decoder's queries attend over the encoder's keys and values — this is how a translation model lets the output sentence pull information from the input sentence. Decoder-only architectures (the GPT family, and most large language models you'll encounter later in this course) drop the encoder and cross-attention entirely, keeping only stacked masked self-attention and feed-forward blocks — the same mechanism proven here, used to predict the next token from everything generated so far.

End-to-end architecture

Common misconception: "attention already knows word order"

Because self-attention lets every token look at every other token simultaneously, students often conclude that the transformer must therefore also understand the order of those tokens — after all, it can "see everything." This is false, and it is worth being precise about why. Self-attention is a set operation: strip out positional encoding and feed the model the bag of vectors {embed("Bumrah"), embed("bowled"), embed("it")} in any order, and the Q Kᵀ computation produces the same set of pairwise scores regardless of which row you call "first" — permute the input rows and the output rows permute identically, with no change to which token attends to which. The model would compute numerically identical attention behaviour for "Bumrah bowled it" and, absurdly, "it bowled Bumrah". Nothing internal to the attention formula distinguishes subject from object, first from last, or before from after — that entire notion of sequence is injected purely through the additive positional encoding vectors before the first attention layer ever runs. If you ever have to debug a transformer-based model that seems to ignore word order, positional encoding — whether it's present, whether it's being added at the right stage, whether it's been truncated for sequences longer than it was built for — is the first place to look, not the attention weights themselves.

Active recall

Attempt each question before reading its answer.

1. Why do we divide Q·Kᵀ by √d_k before applying softmax? What specifically goes wrong in training if we skip this?
2. In the worked example, token 2 ("bowled") produced a perfectly uniform attention distribution over all three tokens. Explain why, using the actual key vectors — not just "that's what the numbers gave."
3. What would happen to a sentence's meaning, from the model's point of view, if you removed positional encoding entirely and shuffled the word order?
4. Self-attention costs O(n²·d) per layer; a recurrent layer costs O(n·d²) per layer. For a 50-token sentence with d_model = 512, which is cheaper, and by roughly what factor?
5. Why must the causal mask be applied to the scores before the softmax, rather than zeroing out the same entries in the attention weights after the softmax?
6. With d_model = 512 and h = 8 heads, each head uses d_k = 64. Why not just use one head with d_k = 512 — what is actually gained by splitting?

Answers

1. If q and k have components with mean 0 and variance 1, Var(q·k) = Σ Var(q_i k_i) = d_k, so the standard deviation of the raw score grows as √d_k. Larger d_model (hence larger d_k) produces wider-spread scores, which push softmax toward a near one-hot output. Softmax's gradient vanishes exactly where its output saturates near 0 or 1, so unscaled scores stall learning as d_k grows. Dividing by √d_k renormalizes the standard deviation back to 1 regardless of d_k, keeping softmax in its well-behaved range.
2. q₂ = [0, 2, 0] — its only nonzero component is the second coordinate. The second coordinate of every key in this example is exactly 1 (k₁, k₂, k₃ all have middle value 1), so q₂·k_j = 2×1 = 2 for every j, before and after scaling. Tied scores softmax to a perfectly flat distribution — the query happens to probe a direction in key-space where all three tokens are indistinguishable, so the model correctly reports no preference.
3. Nothing — and that is the point. Self-attention alone is permutation-equivariant: shuffle the input rows and the output rows shuffle identically, with the same pairwise scores attached to the same token pairs. Without positional encoding, "Bumrah bowled it" and "it bowled Bumrah" would produce numerically identical attention computations, because the model has no representation of order at all — only positional encoding, added to the embeddings before the first layer, breaks that symmetry.
4. n²d = 50² × 512 = 1,280,000. nd² = 50 × 512² = 13,107,200. Self-attention is roughly 10× cheaper per layer in raw operations for a sentence this length (n < d_model), and unlike the recurrent layer it computes in one parallel matrix multiply rather than 50 forced-sequential steps — the parallelism, not just the FLOP count, is what actually makes GPU training fast.
5. exp(−∞) = 0 exactly, so masking before softmax gives future positions exactly zero weight while the remaining weights still sum to 1 automatically, with no extra step. Zeroing entries after the softmax would leave the remaining (already-normalized) weights summing to less than 1, requiring a separate renormalization pass, and would also alter how gradients flow back through the masked positions during backpropagation.
6. Splitting costs essentially nothing in parameters, since h × d_k = d_model — the same total capacity as one large head. What's gained is that each of the 8 heads learns its own independent W_Q, W_K, W_V and can specialize in a different kind of relationship (one head might track subject-verb agreement, another coreference, another adjacency) in parallel, within the same layer. A single 512-dimensional head is forced to blend every kind of relevance into one attention pattern per token, which empirically performs worse than several smaller, specialized patterns concatenated together.

Think About It

Think about this: How would you explain transformer architecture from scratch 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 transformer architecture from scratch, 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.