LLM Architecture
How are Transformer blocks composed, stacked, and projected to predict next tokens?
How this lesson unfolds.
A single attention head is only one instrument; an LLM is a full symphony. Follow this visual module to see how Transformer blocks combine attention and feed-forward networks, how depth creates linguistic abstraction, and how contextual vectors project to vocabulary logits.
How do attention and feed-forward sublayers divide labor between token-to-token communication and intra-token computation?
Imagine a team of software engineers working on a major codebase. The workday consists of two alternating phases: standup meetings and private coding time. During the standup meeting (Multi-Head Attention), engineers talk across desks, exchange status updates, and coordinate dependencies between tasks. During private coding time (Feed-Forward Network), every engineer sits quietly at their own desk, opens their personal IDE, and writes code based on what they learned. No talking occurs during coding time; no coding occurs during standup. A transformer block repeats this exact cycle: coordinate across tokens, then think deeply within each token.
- Learning goalExplain the architectural division of labor between attention (token-to-token communication) and feed-forward networks (intra-token non-linear processing).
- Helpful foundationsMulti-head attention mechanisms · Feed-forward neural networks and residual connections
- Common trapBelieving that feed-forward layers allow tokens to communicate with each other. The feed-forward network is strictly position-wise: token i's MLP computation is mathematically independent of token j's MLP computation.
- Next connectionNext, derive the exact mathematical equations, residual invariants, and tensor contract ledgers in Part B.
More context & mechanics
A transformer block is organized into two fundamentally distinct computational engines:
- Engine 1: Multi-Head Attention (Token Communication):
- Operation: Tokens exchange information along the sequence dimension T.
- Scope: Cross-token, context-aware. Token i reads from and writes to tokens j ≤ i.
- Contract: Computes dynamic routing weights on the probability simplex.
- Engine 2: Feed-Forward Network (Token Computation):
- Operation: Applies identical non-linear transformations along the feature dimension d_(model).
- Scope: Intra-token, position-wise independent. Token i is processed in complete isolation from token j.
- Role: Acts as a key-value associative memory that stores factual knowledge and refines representations.
- The Residual Stream Backbone:
Both sublayers wrap around an uninterrupted additive highway (x + f(x)), ensuring gradients flow directly across dozens of stacked blocks without decay.
Modern LLMs stack between 12 (GPT-2 124M) and 128 (Llama-3-405B) identical transformer blocks. Research by Geva et al. (2021) showed that the feed-forward networks act as key-value stores for factual knowledge (e.g. associating "Paris" with "capital of France"), while attention layers act as routing buses that move facts to the active prediction slot.
The transformer block was introduced by Vaswani et al. in 2017 to replace recurrent and convolutional neural networks. By decoupling cross-token interaction (attention) from deep non-linear processing (FFN), Vaswani created an architecture that scales near-linearly with compute on modern GPU clusters.
- Attention talks; MLP thinks: Attention passes notes between tokens; the MLP calculates answers in private.
- Communicate, then compute: The alternating heartbeat of all modern foundation models.
How are the mathematical forward flow and tensor shape contracts defined across a Pre-LN transformer block?
Think of a main water pipeline passing through an industrial purification plant. The main clean water pipe flows straight through the factory from left to right. At Station 1, a small side pipe siphons off water, cleanses and tests it, and injects fresh minerals back into the main pipe. At Station 2, another side pipe siphons water, carbonizes it, and feeds it back into the main pipe. The main pipeline is never blocked or severed; it simply accumulates clean enhancements at every stage.
- $$A Pre-LN transformer block applies LayerNorm before each sublayer and adds outputs directly to the residual stream: $x^{(1)} = x + \text{MHA}(\text{LN}(x))$ and $x^{(2)} = x^{(1)} + \text{FFN}(\text{LN}(x^{(1)}))$.$$
- Learning goalDerive the mathematical equations and tensor shape contracts for Pre-LN transformer blocks, proving that the additive residual path preserves gradient flow.
- Helpful foundationsLayer normalization concept · Residual connection calculus and chain rule
- Common trapConfusing Pre-LN with Post-LN. In Post-LN, x = LN(x + f(x)), which normalizes the residual stream itself and causes gradients to shrink exponentially in deep networks.
- Next connectionNext, explore how LayerNorm mathematically stabilizes channel distributions in Chapter 20 Screen 3.
More context & mechanics
#### 1. Mathematical Forward Flow Equations
In the modern Pre-LN transformer block, normalization occurs immediately before each transformation:
- Sublayer 1 (Communication):
x^((0.5)) = LayerNorm(x)
x^((1)) = x + MultiHeadAttention(x^((0.5)))
- Sublayer 2 (Computation):
x^((1.5)) = LayerNorm(x^((1)))
x^((2)) = x^((1)) + FeedForward(x^((1.5)))
Where the final output x^((2)) serves as the input x to the subsequent block.
Across every single operation inside the transformer block, the batch and sequence dimensions remain completely invariant:
- Input x: [B, T, d_(model)]
- Normalized x^((0.5)): [B, T, d_(model)]
- Attention output: [B, T, d_(model)]
- Residual sum x^((1)): [B, T, d_(model)]
- Normalized x^((1.5)): [B, T, d_(model)]
- FFN expanded state: [B, T, 4 · d_(model)]
- FFN projected output: [B, T, d_(model)]
- Block output x^((2)): [B, T, d_(model)]
Because the residual connection is pure addition, the gradient of the loss L with respect to input x is:
(∂ L) / (∂ x) = (∂ L) / (∂ x^((2))) ( I + (∂ FFN) / (∂ x^((1))) ) ( I + (∂ MHA) / (∂ x) )
The identity matrix I provides an unbroken gradient highway, allowing signals to propagate backward without vanishing.",
### Real-World Usage
Every frontier model (GPT-3, GPT-4, Llama-2/3, Mistral) uses the Pre-LN transformer block structure. The unbroken identity path enables stable FP16/BF16 training without requiring the complex learning rate warm-up schedules that plagued early Post-LN architectures.
While Vaswani et al. (2017) originally placed LayerNorm after the residual addition (Post-LN), Ruibin Xiong et al. (2020) and Baevski & Auli (2018) proved that Pre-LN eliminates gradient vanishing at initialization, establishing Pre-LN as the universal industry standard.
- x + f(LN(x)): Normalize before the function; add directly to the stream.
- The identity highway: Addition preserves gradients; never break the residual spine.
How does deep stacking identical transformer blocks create progressive linguistic abstraction?
Think of processing a book in an editorial house. The junior proofreaders on the ground floor check spelling, punctuation, and basic grammar (layers 1–4). The section editors on the middle floors track recurring character arcs, plot continuity, and thematic motifs (layers 5–16). The executive editor-in-chief on the top floor evaluates overarching philosophical consistency and dictates the final concluding verdict (layers 17–32). Even though every floor uses the exact same architectural floorplan, depth naturally creates higher levels of abstraction.
- Learning goalAnalyze how depth in stacked transformer layers transforms raw token representations into high-level semantic and reasoning abstractions.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · Language Models are Unsupervised Multitask Learners / GPT-2 (Radford et al., 2019) · Hierarchical Representation Learning in Deep Networks
- Common trapAssuming deeper layers require different mathematical code or different layers types. The blocks are structurally identical; functional specialization emerges entirely from learned weight parameters during gradient descent.
- Next connectionWe trace this progressive abstraction concretely by following how a raw sentence transforms into rich contextual vectors.
More context & mechanics
Let H^((0)) = X ∈ R^T × d_(model) be the input token plus positional embeddings. Passing through L stacked blocks:
H^((l)) = TransformerBlock_l(H^((l-1))) for l = 1, 2, ṡ, L
Empirical probing across layers demonstrates clear functional partitioning:
- Layers 1 ≤ l ≤ L/4 (Surface Syntax): Focus on subword morphology, local n-grams, and punctuation boundaries. Attention patterns are predominantly local diagonal stripes.
- Layers L/4 < l ≤ 3L/4 (Relational Semantics): Resolve coreference (linking pronouns to antecedents), dependency trees, and clause structures.
- Layers 3L/4 < l ≤ L (High-Level Reasoning & Next-Token Synthesis): Perform multi-hop factual retrieval, in-context deduction, and format anticipation, concentrating information into the final token position.
This depth hierarchy explains why scaling model depth (e.g. 12 layers in GPT-2 Small, 32 in LLaMA-3 8B, 96 in GPT-3 175B) yields dramatically superior reasoning and complex instruction following.
In 2019, Tenney et al. published BERT Rediscovers the Classical NLP Pipeline, showing that deep transformer stacks rediscover the traditional NLP pipeline (part-of-speech tagging → parsing → NER → semantic roles → coreference) in sequential order as layer depth increases.
- Golden Axiom: "Syntax at the base, semantics in the middle, reasoning at the summit."
- Identical blocks at every layer produce non-identical depths of understanding.
How does a sequence of raw token vectors progressively absorb multi-hop context across layers?
Imagine whispering in a line of people. In step 1, you can only hear your immediate neighbor. In step 2, your neighbor can tell you what their neighbor said. In step 3, information from three people away reaches you. In a Transformer, because every layer allows direct connections across the whole sequence, layer 1 computes direct relationships (A is friends with B; B lives in Paris). Layer 2 combines these hops (A is connected to Paris). By layer 12, every token vector is a sponge soaked in the full holistic narrative of the prompt.
- Learning goalTrace the contextual accumulation across layers and prove how multi-layer attention resolves multi-hop relational dependencies.
- Helpful foundationsAttention Is All You Need (Vaswani et al., 2017) · Graph Neural Networks & Message Passing · Receptive Fields in Deep Networks
- Common trapAssuming a single attention layer can solve multi-hop reasoning. A single attention layer can only perform a single direct lookup (1 hop). Connecting indirect facts (Fact A -> Fact B -> Conclusion C) requires composing multiple attention layers in sequence.
- Next connectionOnce the final token vector has accumulated all necessary context from the prompt, it must be projected to output vocabulary logits to predict the next word.
More context & mechanics
Consider a prompt: "Alice was born in France. Her mother tongue is..."
- At Layer 0 (Input Embedding):
The vector h_T^((0)) for "is" contains only the isolated word representation of "is" plus its position index. It knows nothing about France or Alice.
- At Layer 1:
Attention connects "tongue" to "mother" and "is" to "tongue". The vector h_T^((1)) encodes the query concept "mother tongue".
- At Layer 2:
The vector h_T^((2)) attends to "Alice", integrating the subject. Simultaneously, earlier tokens integrated "Alice" with "France".
- At Layer 3:
The vector h_T^((3)) resolves the multi-hop connection: Alice → France → French. The final position vector now points strongly toward the lexical subspace of French.
Each layer performs one round of differentiable message passing, transforming the final token from a dumb lexical placeholder into a dense contextual summary.
Multi-hop context integration enables LLMs to answer complex questions that require synthesizing clues scattered across hundreds of pages of context.
This information routing phenomenon was formalized by Geva et al. (2021) in Transformer Feed-Forward Layers Are Key-Value Memories and Meng et al. (2022) in ROME (Locating and Editing Factual Associations in GPT), demonstrating that factual recall flows through specific attention-FFN layer cascades.
- Golden Axiom: "Tokens are information accumulators."
- As depth increases, tokens stop representing words and begin representing whole situations.
How is the final token's contextual vector mapped into next-token vocabulary logits?
Imagine a courtroom judge delivering a verdict. After listening to all the testimony, the judge has a finalized mental summary (the final token vector in R⁷68). To announce the verdict in English, the judge compares this internal thought against an official dictionary of 50,257 legal terms (the vocabulary projection matrix W_vocab). Words that closely match the judge's thought receive high scores (logits). The highest-scoring word—or a word sampled probabilistically via Softmax—becomes the chosen next token.
- Learning goalCalculate next-token logits u = h_T W_vocab^T in R^V and apply Softmax with temperature scaling to compute output token probabilities.
- Helpful foundationsLanguage Models are Unsupervised Multitask Learners / GPT-2 (Radford et al., 2019) · Linear Classifiers & Dot Product Readout · Softmax with Temperature Scaling
- Common trapAssuming the model computes predictions for every token in the sequence during generation. In autoregressive generation, only the very last token position h_T is projected to predict token T+1; earlier positions have already been generated.
- Next connectionBefore examining the full assembly of a Transformer block, we must master its primary stabilizing engine: Layer Normalization.
More context & mechanics
Let h_T^((L)) ∈ R^d_(model) be the output vector of the final Transformer block at the final sequence position T. The Language Model (LM) Head produces next-token predictions through three mathematical steps:
Step 1: Final Layer Normalization:
z_T = LayerNorm(h_T^((L))) ∈ R^d_(model)
Step 2: Linear Vocabulary Projection:
Multiply by the LM head weight matrix W_(vocab) ∈ R^V × d_(model) (where V is vocabulary size, e.g. 50,257):
u = z_T W_(vocab)^\top ∈ R^V
Each logit u_v is the dot product z_T · w_v, measuring how well the context vector aligns with candidate vocabulary word v. In many architectures (including GPT-2), W_(vocab) shares weights directly with the token embedding matrix W_E (Weight Tying).
Step 3: Softmax Probability Distribution:
Convert logits into probabilities with temperature τ > 0:
P(w_(T+1) = v) = (exp(u_v / τ)) / (∑_(j=1)^V exp(uⱼ / τ))
This is the final inference step of every LLM. Sampling strategies (top-k, top-p nucleus sampling, temperature scaling) operate directly on this probability vector P.
Press and Wolf (2017) demonstrated in Using the Output Embedding to Improve Language Models that tying the input embedding weights to the output projection weights (W_(vocab) = W_E) reduces parameter count by tens of millions of weights while improving generalization.
- Golden Axiom: "Normalize, project, Softmax."
- Vector h_T ∈ R^d → logits u ∈ R^V → probabilities P ∈ [0, 1]^V.
Why does forcing an entire sentence through a fixed-size recurrent hidden state vector cause catastrophic forgetting in RNNs?
Imagine attending a 4-hour university lecture where you are strictly forbidden from taking notes on paper. You have only one small index card (the hidden state vector h). After minute 1, you write a brief summary on the card. After minute 2, you erase part of the card to write the new point. By hour 4, your index card has been erased and overwritten 240 times! The memory of what the professor said in minute 1 is completely wiped out. This is the Recurrent Bottleneck: an RNN forces an entire book into a tiny fixed-size vector.
- $$Show that the recursive update h_t = tanh(W_h h_{t-1} + W_x x_t) forces a sequence of arbitrary length T into a fixed d-dimensional vector, losing early context exponentially.$$
- Learning goalDerive the Jacobian contraction of recurrent hidden state updates and quantify the information capacity limits of fixed-size memory vectors.
- Helpful foundationsRecurrent neural networks · Neural Machine Translation by Jointly Learning to Align and Translate (Bahdanau et al., 2014)
- Common trapThinking increasing hidden state size d solves the RNN bottleneck; doubling d increases parameter count quadratically (d²) without fixing the exponential vanishing gradient decay across long temporal horizons.
- Next connectionHow did LSTMs attempt to heal this recurrent wound? We inspect the LSTM dual-highway and conveyor belt gating.
More context & mechanics
In a standard Recurrent Neural Network (RNN):
h_t = \tanh(W_(hh) h_(t-1) + W_(xh) x_t + b_h)
Unrolling through time:
h_T = f(x_T, f(x_(T-1), ṡ, f(x₁, h₀)))
Theoretical compression limits:
- Information Capacity: The final hidden state h_T ∈ R^d has fixed dimensionality d. By Shannon's channel capacity theorem, h_T can hold at most O(d) bits of mutual information about the input history x_(1:T).
- Exponential Memory Decay: The Jacobian of the hidden state across k steps is:
(∂ h_t) / (∂ h_(t-k)) = ∏_(j=1)^k diag(1 - h_(t-j+1)²) W_(hh)^T
If the spectral norm \|W_(hh)\| < 1, gradient sensitivity decays exponentially:
\| (∂ h_t) / (∂ h_(t-k)) \| ≤ σ_(max)^k → 0 as k → ∞
The network has an effective context horizon of fewer than 50 tokens.
This fatal bottleneck is what motivated the development of LSTMs, Attention mechanisms (Bahdanau et al., 2014), and ultimately the modern Transformer, which gives every token direct O(1) access to all previous tokens.
David Rumelhart, Geoffrey Hinton, and Ronald Williams introduced Elman/Jordan recurrent networks in 1986. Dzmitry Bahdanau et al. (2014) formally broke the bottleneck with neural attention.
- You cannot squeeze an encyclopedia into an index card.
- Recurrent compression destroys long-range memory.
How do forget, input, and output gates regulate information flow along the constant error carousel cell state conveyor belt?
Instead of a single erasing index card, an LSTM installs a factory conveyor belt (the cell state C_t) running through the entire sequence. Above the conveyor belt sit three robotic gatekeepers: the Forget Gate decides what old junk to push off the belt; the Input Gate decides what fresh ideas to load onto the belt; and the Output Gate decides what summary to read off the belt for the current step. Because items can ride undisturbed along the belt with simple addition, memories can travel 1,000 steps without fading.
- $$Protect the linear cell state C_t = f_t * C_{t-1} + i_t * C_tilde_t via multiplicative sigmoid gates, preventing exponential gradient decay.$$
- Learning goalFormulate the six equations of the LSTM cell and trace the uninterrupted gradient highway of the cell state.
- Helpful foundationsSigmoid and tanh functions · Long Short-Term Memory (Hochreiter & Schmidhuber, 1997)
- Common trapConfusing cell state C_t with hidden state h_t; C_t is the internal linear memory conveyor belt, while h_t is the non-linearly squashed and gated external representation.
- Next connectionIn modern generative transformers, how does autoregressive generation generate text one token at a time? We examine the sliding window generation loop.
More context & mechanics
An LSTM cell processes input x_t and previous hidden state h_(t-1) using three gating mechanisms:
- Forget Gate (f_t ∈ (0, 1)^d):
f_t = σ(W_f [h_(t-1), x_t] + b_f)
- Input Gate (i_t ∈ (0, 1)^d) and Candidate (C̃_t):
i_t = σ(Wᵢ [h_(t-1), x_t] + bᵢ), C̃_t = \tanh(W_c [h_(t-1), x_t] + b_c)
- Linear Cell State Update (The Conveyor Belt):
C_t = f_t \odot C_(t-1) + i_t \odot C̃_t
Notice the additive update: (∂ C_t) / (∂ C_(t-1)) = f_t. If f_t = 1, gradients propagate across arbitrary time steps with zero attenuation!
- Output Gate (o_t) and Hidden State (h_t):
o_t = σ(W_o [h_(t-1), x_t] + b_o), h_t = o_t \odot \tanh(C_t)
LSTMs powered Google Translate, Siri voice recognition, and handwriting generation from 2014 until Transformers became dominant in 2018-2020.
Sepp Hochreiter and Jürgen Schmidhuber invented the Long Short-Term Memory (LSTM) architecture in 1997, solving the vanishing gradient problem in recurrent networks.
- Cell state is the conveyor belt; gates are the valves.
- Forget old junk, add new ideas, output the filtered thought.