Transformers & Self-Attention: The Architecture That Ate Deep Learning

Canonical lesson. This is the single source of truth for self-attention used by both the Deep Learning track and the NLP & Transformers track (Track 5). If you arrived here from a Track 5 prerequisite link (e.g., from Multi-Head Attention, Positional Encoding, or BERT), you are in the right place — work through this lesson, then return to the NLP-track sequel.

The transformer is not magic. It is three ideas welded together: (1) self-attention as a learned soft lookup over the sequence, (2) residual streams that make depth cheap, and (3) a feed-forward sublayer that does the actual "thinking" in each block. Once you see those three pieces clearly, every modern LLM is just "what if we did that 96 times and trained on the entire internet?"

#Why RNNs Had to Go

Recall where we left off in the previous lesson. The LSTM solved the vanishing-gradient problem with a cell-state highway, and Bahdanau attention solved the fixed-context-vector bottleneck by letting the decoder look back at every encoder state. But seq2seq with attention still had two fundamental problems that scaling could not fix:

1. Sequential compute. To compute hidden state h_50, you must first compute h_1, h_2, ..., h_49 in order. There is no way to parallelize this in the time dimension. A modern GPU has thousands of cores sitting idle while the LSTM crawls forward one token at a time.
2. Effective context still degrades. Even with the cell-state highway, an LSTM's memory is a fixed-size vector that gets overwritten at every step. Information from 500 tokens ago has been multiplied by 500 forget gates, each slightly less than 1. The signal is technically present but practically faded. Long-document understanding (legal contracts, code repositories, scientific papers) demands that every token can reach every other token equally easily.

Bahdanau already showed us that letting the decoder attend to every encoder state worked beautifully. So Vaswani and his co-authors asked the obvious follow-up: what if attention is enough on its own? What if we just drop the recurrence and let attention do all the work — within a sequence (self-attention), across sequences (cross-attention), and even from a token to its own past (causal self-attention)?

#The Core Move: Self-Attention vs Cross-Attention

Self-attention is what makes a transformer feel "contextual": the same word ("bank") gets different output representations depending on its neighbors ("river bank" vs "investment bank") because its query vector matches different keys in each context. RNNs eventually achieved this through their hidden state too, but only after sequentially processing the whole prefix. Self-attention does it in one shot, in parallel, with O(1) hops between any two positions.

#Query, Key, Value: A Soft Differentiable Lookup

The visualization above shows attention scores between every pair of tokens. Each cell (i, j) is how much token i attends to token j. Notice three things: (1) the diagonal is usually bright — tokens often attend to themselves; (2) some columns light up across many rows — those are "popular" tokens that many positions want context from (often nouns or verbs); (3) the pattern is asymmetric — attention from "it" to "animal" is not the same as attention from "animal" to "it".

#The Scaled-Dot-Product Attention Formula

Here is the equation that does the actual work — the one you will see written on every transformer slide, every system-design whiteboard, every blog post about LLMs:

Let us unpack the four operations.

#Multi-Head Attention: Many Lookups in Parallel

Why do heads end up specializing? Because they are initialized with different random weights and trained jointly with the rest of the network. Stochastic gradient descent breaks symmetry — each head finds a different niche in the loss landscape. Interpretability work has found heads that specialize in:

• Positional heads. Attend mostly to immediate neighbors (left or right).
• Syntactic heads. Attend from verbs to their subjects, or from pronouns to their referents.
• Induction heads. Attend from "X" to the position right after a previous occurrence of "X" — the mechanism behind in-context learning in LLMs.
• Punctuation-tracking heads. Attend to matching parentheses, quotes, or section boundaries.

Modern LLMs use 8–96 heads per layer. The original transformer used 8. Claude Sonnet and GPT-4 likely use 64+. The exact count is a hyperparameter; you trade head count for per-head dimension at fixed compute budget.

#Positional Encoding: Attention Is Permutation-Invariant

This is a disaster for language. "Dog bites man" and "Man bites dog" have identical token bags but very different meanings. Without injecting position information somewhere, a transformer would treat them as the same input.

The original transformer added a fixed sinusoidal pattern to each position's embedding:

Modern alternatives:

• Learned absolute (BERT, GPT-2): just learn an embedding per position index. Simple. Fails to extrapolate beyond training length.
• RoPE — Rotary Position Embedding (LLaMA, GPT-J, Mistral, modern Claude): rotates Q and K vectors by an angle proportional to position. Has nice properties (relative-position-aware, extrapolates better, easy to extend context). The current default in open-source LLMs.
• ALiBi (Press et al. 2022): no PE at all — instead, add a linear bias to attention scores that penalizes distant tokens. Very strong extrapolation properties.

The NLP track goes deeper on position encoding. The DL takeaway: attention is permutation-invariant, so position must be injected, and exactly how you inject it matters for long-context behavior.

#Rotary Position Embedding (RoPE): the actual math

RoPE is named in the previous section but the derivation is worth seeing because it explains every "long-context" trick in the modern LLM stack (NTK scaling, YaRN, LongRoPE) none of which make sense as black-box knobs.

This is why RoPE generalizes to lengths the model never saw at training. Sinusoidal-PE-plus-absolute-attention would have to interpolate the unseen positions into the input embedding; RoPE never embeds an absolute position at all — only the relative dependence is exposed to the model, and the relative dependence is the same at sequence position 1000 as at position 100,000.

#Grouped-Query Attention (GQA): the inference-cost optimization

#Toy Implementation: Scaled-Dot-Product Attention from Scratch

The heatmap above shows learned attention patterns from a trained transformer on the sentence "The animal didn't cross the street because it was too tired." Look at the row labeled "it" — you should see a bright cell on the column "animal", not on "street". The model has learned (through pure self-supervision on text) that "it" co-refers with "animal". No one annotated coreference for it. The attention pattern fell out of the language-modeling loss.

#The Causal Mask: Why GPT Cannot See the Future

#Stacking the Block: Add, Norm, FFN

1. Residual connection (a.k.a. skip connection): the input to each sublayer is added back to the sublayer's output. This is the same idea as ResNets — it gives gradients a direct path back to earlier layers and lets you stack many blocks (32, 64, 96+) without vanishing gradients.
2. LayerNorm: normalizes each token's representation across the feature dimension. Stabilizes training; especially important for deep transformer stacks.
3. Feed-forward network (FFN): a two-layer MLP applied independently to each token (no token-mixing — that already happened in attention). Usually expands to 4 * d_model in the hidden layer, then projects back. This is where most of the "thinking" happens — attention moves information between positions, the FFN refines what each position contains.

The block, in modern pre-norm form:

def transformer_block(x):
    # Pre-norm: normalize first, then sublayer, then add residual
    x = x + multi_head_self_attention(layer_norm(x))   # sublayer 1
    x = x + ffn(layer_norm(x))                          # sublayer 2
    return x

def ffn(x):
    # Position-wise FFN: same weights applied to every token independently
    return W_2 @ gelu(W_1 @ x + b_1) + b_2

That is it. The entire transformer is a stack of these blocks — 6 in the original paper, 12 in GPT-1, 24 in BERT-large, 96 in GPT-3, 80+ in GPT-4-scale models. Each block has the same shape; only the learned weights differ.

The architecture diagram above shows the full data flow through one block: embeddings + positional encoding → multi-head self-attention → residual + LayerNorm → FFN → residual + LayerNorm → output. The decoder variant inserts an additional cross-attention sublayer between self-attention and FFN to attend to encoder outputs.

#Pre-norm vs Post-norm: A Tiny Reordering with Huge Consequences

The difference looks trivial, normalize before vs after the sublayer, but it dramatically changes training stability.

#A Tiny End-to-End Transformer Encoder Block

#Three Families: Encoder-Only, Decoder-Only, Encoder-Decoder

Once you have the basic block, you can stack it in three different ways depending on the task:

The encoder-decoder is the most general form and was the original 2017 design. Decoder-only models eventually won the LLM race because: (1) you only need to maintain one set of weights, simpler training; (2) decoder-only handles "input + output" by concatenating them and using prompt engineering, which turns out to be insanely flexible; (3) scaling laws favor a single deep stack over a split encoder-decoder at fixed parameter count.

Encoder-only models still dominate embedding tasks (semantic search, RAG retrieval, classification) because bidirectional context produces much stronger fixed representations than causal models.

#Why Transformers Cost: The O(N^2) Wall

Self-attention has one well-known weakness: its compute and memory both scale as O(N^2) in sequence length. The QK^T matrix is N-by-N; for a 100k-token context, that is 10 billion entries per head per layer. At 32 layers and 32 heads, you are looking at 10 trillion attention entries to materialize during the forward pass — completely infeasible without engineering tricks.

The modern fixes (covered in depth in the NLP track):

• FlashAttention (Dao et al. 2022, 2023) — still O(N^2) in compute, but tiled into GPU SRAM so you never materialize the full N-by-N matrix. 2-4x speedup, 10x+ memory reduction. Now the default attention kernel in PyTorch.
• Sliding-window / local attention (Mistral, Longformer) — each token only attends to a window of W neighbors. O(N * W) instead of O(N^2). Cheap but loses long-range information; usually combined with a few "global" attention heads.
• Linear / Performer / Linformer. Approximate the softmax attention with a linear-cost surrogate. Faster but consistently slightly worse quality.
• Mamba and state-space models (Gu & Dao 2023) — an entirely different architecture with O(N) inference cost, competitive quality. The first serious challenger to the transformer in years.

Long-context LLMs (1M+ tokens) are a major frontier; expect this list to keep growing.

#Why Transformers Won

Stepping back: why did this one architecture eat all of deep learning?

1. Parallel training. Every position is computed in one shot, every layer is the same shape. GPUs ride this beautifully. RNNs cannot be parallelized in time; CNNs are local. Transformers are the first architecture that lets you spend a million dollars of compute on one training run productively.
2. Long-range dependencies are O(1) hops. Token 1 can attend directly to token 100,000 in one layer. In an LSTM that would require 100,000 sequential steps of memory propagation. In a CNN with kernel size 3, that would require log_3(100,000) ≈ 11 layers.
3. Scales beautifully. Empirically, loss falls as a smooth power law in compute, data, and parameters (Kaplan 2020, Chinchilla / Hoffmann 2022), but only for transformers, at least so far. Other architectures plateau.
4. Universal. The same block works for text (GPT, BERT, T5), images (ViT), audio (Whisper), proteins (AlphaFold Evoformer, ESM-2), code (Codex), video (Sora). One architecture, many modalities. This was not obvious in 2017.
5. Interpretability surface. Attention weights are explicitly visible — you can stare at them, find heads that specialize in coreference or induction, and start to reverse-engineer what the model "does." Anthropic's mechanistic interpretability program lives on top of this fact.

#Common Mistakes

#Production Architecture

#Key Takeaways

• Self-attention is a learned, soft, parallel lookup over the sequence. Each token produces a Query (what I want), Key (what I have), and Value (what to bring back). The score softmax(QK^T / sqrt(d_k)) is the lookup; the multiplication by V is the retrieval. The whole operation is differentiable end-to-end.
• The sqrt(d_k) scaling is not optional. Without it, dot-product magnitudes scale with sqrt(d_k), softmax saturates, gradients vanish, and training collapses. This is the single most common implementation bug in hand-rolled attention.
• Multi-head attention lets different heads specialize. Splitting d_model across H heads gives the model H independent lookup patterns per layer. Interpretability work has found heads that handle coreference, induction, syntax, and punctuation — all emerging from random initialization plus gradient descent.
• Positional encoding is mandatory because attention is permutation-invariant. Without injecting position info (sinusoidal, learned, RoPE, or ALiBi), 'dog bites man' and 'man bites dog' would be literally identical inputs to the network.
• The transformer block = MHA + residual + LayerNorm + FFN + residual + LayerNorm. Pre-norm (LayerNorm before the sublayer) is the modern default; it gives gradients a clean residual highway and lets you train 96+ layer stacks stably.
• Three families, one block. Encoder-only (BERT) uses bidirectional attention. Decoder-only (GPT, Claude, LLaMA) uses causal-masked attention and dominates generative LLMs. Encoder-decoder (T5, original transformer) uses a bidirectional encoder plus a causal decoder with cross-attention.
• Attention is O(N^2) in sequence length. The central engineering challenge of long-context LLMs. FlashAttention, sliding-window attention, and Mamba are the major modern responses.

#Quick Check