The Transformer Architecture
Attention built up from scratch — Query/Key/Value, the scaling factor, multi-head attention, positional encoding, and a full transformer block
The Transformer Architecture
TL;DR
Attention lets every position in a sequence directly relate to every other position, by computing a Query/Key/Value comparison and taking a weighted average of Values. That mechanism has no inherent sense of order, so position must be added explicitly (sinusoidal encoding, or the modern standard, RoPE). Stack attention, a feed-forward layer, normalization, and residual connections into a block, and stack blocks into a transformer.
| Property | Value |
|---|---|
| Level | Intermediate |
| Reading time | ~22 minutes |
| Prerequisites | Training Loops & Optimization |
| You will understand | Attention's mechanics in full, why scaling and positional encoding are necessary, and how a transformer block fits together |
The Motivation
Before transformers, sequence models (RNNs) processed tokens one at a time, carrying information forward in a hidden state — which meant information from early in a long sequence had to survive many sequential steps to influence something much later, and in practice, it often didn't. Attention removes that bottleneck entirely: every position can look directly at every other position, in one step, regardless of distance.
Scaled Dot-Product Attention, in Full
For each token, three vectors are computed by learned linear projections of its embedding:
| Vector | Role |
|---|---|
| Query (Q) | "What is this token looking for?" |
| Key (K) | "What does each token offer, that a query might match against?" |
| Value (V) | "What does each token actually contribute, once it's been attended to?" |
import torch, math
def scaled_dot_product_attention(Q, K, V):
d_k = Q.shape[-1]
scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k) # (seq_len, seq_len)
weights = torch.softmax(scores, dim=-1) # rows sum to 1
return weights @ V # (seq_len, d_v)
seq_len, d_k = 4, 8
Q = torch.randn(seq_len, d_k)
K = torch.randn(seq_len, d_k)
V = torch.randn(seq_len, d_k)
output = scaled_dot_product_attention(Q, K, V) # (4, 8)Why the sqrt(d_k) Scaling Exists
Q @ K.T is a dot product across d_k dimensions — as d_k grows, the magnitude of that dot product tends to grow too (more terms summed), even if each individual component is small. Large values fed into softmax push it toward a near-one-hot distribution, which means most gradients through the softmax become vanishingly small — the network stops getting useful learning signal from attention weights that are already saturated near 0 or 1.
Dividing by sqrt(d_k) keeps the scores in a range where softmax's gradient is still meaningful. Skip this scaling on a model with a large per-head dimension and attention weights become nearly one-hot early in training, well before the network has actually learned anything useful — the softmax saturates for the wrong reason (raw magnitude) rather than a genuinely confident, learned preference.
Multi-Head Attention: Several Smaller Attentions, in Parallel
Instead of one attention computation over the full embedding dimension, multi-head attention splits it into several smaller ones running in parallel, each with its own learned Q/K/V projections:
class MultiHeadAttention(torch.nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_head = d_model // num_heads
self.qkv = torch.nn.Linear(d_model, 3 * d_model)
self.out = torch.nn.Linear(d_model, d_model)
def forward(self, x):
B, T, D = x.shape
qkv = self.qkv(x).view(B, T, 3, self.num_heads, self.d_head)
q, k, v = qkv.unbind(dim=2) # each: (B, T, heads, d_head)
q, k, v = (t.transpose(1, 2) for t in (q, k, v)) # (B, heads, T, d_head)
attn = scaled_dot_product_attention(q, k, v) # per-head attention
attn = attn.transpose(1, 2).reshape(B, T, D)
return self.out(attn)Each head can specialize. One head might learn to attend mostly to the immediately preceding token, another to the subject of the sentence, another to a matching bracket or quote far away — the same way different filters in a convolutional network specialize to different visual patterns. A single large attention head has to represent all of that in one shared space; several smaller heads can divide the work.
Positional Encoding: Why It's Needed at All
Attention, as defined above, is permutation-invariant — swap the order of two input tokens, and the set of Query/Key/Value comparisons made is identical, just relabeled. Nothing about the mechanism itself encodes "this token came before that one." Without adding position explicitly, "the dog bit the man" and "the man bit the dog" would look identical to the attention computation.
Two ways to add position
Sinusoidal positional encoding (the original approach)
A fixed (not learned) pattern of sine/cosine waves at different frequencies, added directly to each token's embedding before the first attention layer — different positions get distinguishable, unique patterns added to their embeddings.
RoPE (Rotary Position Embeddings)
RecommendedThe modern standard. Instead of adding a position signal to the embedding, RoPE rotates the Query and Key vectors by an angle proportional to their position, applied at attention time. The key property: the dot product between a rotated Query and a rotated Key ends up depending mathematically only on their relative distance, not their absolute positions — which generalizes better to sequence lengths not seen during training than a fixed additive encoding does.
The intuition behind RoPE's relative-position property: rotating two vectors by the same angle preserves the angle between them. Rotating Q by an angle proportional to its position and K by an angle proportional to its position means their dot product naturally reflects the difference in those angles — i.e., the relative distance between the two tokens — rather than either token's absolute position.
Assembling a Full Transformer Block
One transformer block
Input
Attention sub-layer
Feed-forward sub-layer
class TransformerBlock(torch.nn.Module):
def __init__(self, d_model, num_heads, d_ff):
super().__init__()
self.attn = MultiHeadAttention(d_model, num_heads)
self.norm1 = torch.nn.LayerNorm(d_model)
self.ff = torch.nn.Sequential(
torch.nn.Linear(d_model, d_ff),
torch.nn.GELU(),
torch.nn.Linear(d_ff, d_model),
)
self.norm2 = torch.nn.LayerNorm(d_model)
def forward(self, x):
x = x + self.attn(self.norm1(x)) # residual connection around attention
x = x + self.ff(self.norm2(x)) # residual connection around the feed-forward layer
return xResidual Connections: Why x + Matters
Each sub-layer's output is added to its input, rather than replacing it. This matters enormously for training very deep stacks: without residual connections, gradients have to flow backward through every single transformation in every layer, and can shrink toward zero over enough layers (vanishing gradients). A residual connection gives the gradient a direct, unimpeded path backward through the addition, alongside the path through the sub-layer — even if a given sub-layer's own gradient is small, the residual path keeps signal flowing.
Layer Norm, Not Batch Norm
Transformers normalize with layer normalization (normalizing across each token's own feature dimension) rather than batch normalization (normalizing across the batch dimension for each feature). Two reasons this matters for sequence models specifically: layer norm's statistics don't depend on other examples in the batch, so behavior at inference time (often batch size 1) matches training exactly; and it handles variable-length sequences naturally, since each token's normalization is entirely self-contained.
Concept Checks
Check yourself
Why does removing the sqrt(d_k) scaling factor from attention tend to hurt training, rather than just changing the numbers slightly?
Because as the per-head dimension grows, unscaled dot products tend to grow large in magnitude, and large inputs to softmax push it toward a near-one-hot output where the gradient is nearly zero almost everywhere. That kills the learning signal through attention early in training — the softmax saturates because of raw input magnitude, not because the network has actually learned a confident, meaningful preference.
Attention as defined is permutation-invariant. What does that mean concretely, and why does it force positional encoding to exist?
It means swapping the order of two tokens in the input produces the same set of Query/Key/Value comparisons, just relabeled — nothing in the raw attention computation distinguishes 'first' from 'second.' Without adding position explicitly, two sentences that are only distinguished by word order would look identical to attention, which is why every transformer adds positional information before or during attention.
Why do residual connections specifically help train very deep transformer stacks?
Because they give the gradient a direct additive path backward through each sub-layer, alongside the path through the sub-layer's own transformation. Even if a particular sub-layer's local gradient is small, the residual path preserves signal flowing back to earlier layers, which prevents the vanishing-gradient problem that would otherwise make very deep stacks difficult or impossible to train effectively.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
| Query / Key / Value | Learned projections representing what a token seeks, offers, and contributes |
sqrt(d_k) scaling | Keeps attention scores in a range where softmax's gradient stays meaningful |
| Multi-head attention | Several smaller, parallel attention computations, each able to specialize |
| Permutation invariance | Attention alone has no sense of order — positional encoding must be added |
| RoPE | Rotates Q/K vectors so their dot product depends on relative, not absolute, position |
| Residual connections | Additive shortcuts that preserve gradient flow through deep stacks |
| Layer norm | Normalizes per-token, independent of batch statistics — suited to variable-length sequences |
| Transformer block | Attention + residual + norm, then feed-forward + residual + norm |
Next
With the core mechanism built, the next page covers keeping a network from simply memorizing its training data: Regularization & Generalization.
Training Loops & Optimization
Backpropagation, the training loop step by step, and how SGD, Adam, and AdamW actually differ — including the bug AdamW specifically fixed
Regularization & Generalization
Overfitting, the train/validation gap as the central diagnostic, and the mechanics of dropout, weight decay, and normalization