The Transformer mental model¶
An LLM is a function that repeatedly answers one narrow question:
Given the tokens visible so far, what probability should I assign to every possible next token?
Everything else - prose, code, tool calls, apparent planning - emerges from training that next-token function and then sampling or selecting from its output. The Transformer is the circuit that turns context tokens into the next distribution.
This chapter builds a reliable map before the later chapters open each box.
The whole path in one picture¶
flowchart LR
textInput[/Input text/]
tokenizer[Tokenizer]
tokenIds[Token IDs]
embeddings[Token vectors]
blockOne[Transformer block 1]
blockMany[Transformer blocks 2 to L]
finalNorm[Final normalization]
vocabulary["Vocabulary logits (one score per token)"]
probabilities[Next-token probabilities]
chosenToken([Chosen token])
textInput --> tokenizer
tokenizer --> tokenIds
tokenIds --> embeddings
embeddings --> blockOne
blockOne --> blockMany
blockMany --> finalNorm
finalNorm --> vocabulary
vocabulary --> probabilities
probabilities --> chosenToken
chosenToken -.->|"Append and repeat"| tokenIds For a decoder-only language model, the major tensors are:
| Stage | Typical shape | Meaning |
|---|---|---|
| Token IDs | (batch, sequence) | Integer vocabulary indexes |
| Hidden states | (batch, sequence, d_model) | One contextual vector per token position |
| Attention queries | (batch, heads, sequence, d_head) | What each position seeks |
| Attention scores | (batch, heads, sequence, sequence) | Pairwise compatibility before/after masking |
| Vocabulary logits | (batch, sequence, vocab_size) | Unnormalized next-token scores |
The exact dimension order varies by implementation. Track the named axes, not just the tuple.
What a Transformer block does¶
The 2017 Transformer paired attention with a position-wise feed-forward network and wrapped both in residual connections and normalization. Modern decoder-only models preserve that division even when they change the details (original paper, Section 3).
flowchart LR
hiddenIn[/Hidden state for every position/]
normAttention[Normalize]
attention["Causal self-attention (mix across positions)"]
addAttention((Add))
normFfn[Normalize]
ffn["Feed-forward network (transform each position)"]
addFfn((Add))
hiddenOut[Updated hidden states]
hiddenIn --> normAttention
normAttention --> attention
attention --> addAttention
hiddenIn --> addAttention
addAttention --> normFfn
normFfn --> ffn
ffn --> addFfn
addAttention --> addFfn
addFfn --> hiddenOut That gives the most useful first distinction:
- Attention moves information between token positions. A pronoun can pull information from a preceding name; a closing brace can attend to earlier structure.
- The FFN transforms each position independently. The same learned function is applied to every position. In a sparse MoE block, a router chooses which subset of FFNs processes each position.
- Residual paths preserve a running stream. Each sublayer writes an update into the existing hidden state instead of replacing it wholesale.
Meta's released Llama 3 inference code makes this modern pre-normalized form especially readable: the entire block forward pass is two residual additions (source lines 222-248).
A concrete one-token trace¶
Suppose the current text is:
The real tokenizer may split leading spaces and words differently, but imagine the conceptual token IDs are:
- The embedding table retrieves a vector for each ID.
- Positional information distinguishes the first token from the fifth.
- In each block, causal attention lets position 5 read positions 1 through 5, never the unseen answer.
- The FFN transforms the newly contextualized vector at each position.
- The final vector at position 5 is projected to one logit per vocabulary token.
- Softmax converts the logits to probabilities. A trained model may assign a high probability to a token representing a leading-space form of
Paris. - Decoding selects a token. The runtime appends it and performs another step.
The model does not retrieve an English sentence from an internal database. It computes the next-token distribution from learned weights and the current context. Retrieval-augmented systems can place database results into the context, but the Transformer still processes them as tokens.
Training and generation use the same causal model differently¶
During pretraining¶
One sequence supplies many supervised positions at once:
A causal mask prevents each position from reading a later target. Because all positions are already known during training, their losses can be calculated in parallel. This is called teacher forcing: the input prefix uses the real previous tokens, not tokens sampled by the model.
During autoregressive generation¶
Only the newest position needs a new prediction. Previously computed keys and values can be retained in a KV cache. Generation is still sequential over newly produced tokens: token t+1 cannot be selected until token t is known.
sequenceDiagram
participant Runtime
participant Model
Runtime->>Model: Prompt tokens
Model-->>Runtime: Logits for next token
Runtime->>Runtime: Select and append token
Runtime->>Model: New token plus KV cache
Model-->>Runtime: Next logits
Runtime->>Runtime: Repeat until stop Encoder, decoder, and decoder-only models¶
The original Transformer is an encoder-decoder model for translation (Vaswani et al.):
- the encoder reads the full source sequence with bidirectional attention;
- the decoder uses causal self-attention over generated target tokens;
- decoder cross-attention reads encoder outputs.
Most general-purpose generative LLMs discussed in this book are decoder-only: they stack causal blocks and predict the next token. Do not look for a separate encoder inside a decoder-only checkpoint just because the word "Transformer" appears in both designs.
Parameters, activations, and context are different things¶
These concepts are often blended together:
| Concept | Exists for how long? | Example |
|---|---|---|
| Parameters | Stored with the model | Attention projection weights |
| Activations | One forward/backward pass | Current token hidden states |
| Optimizer state | Training only | Adam moments |
| KV cache | A generation request | Prior keys and values per layer |
| Context tokens | The request | System message, prompt, generated prefix |
A 7B-parameter model does not have a seven-billion-token context. A longer context increases activation and KV-cache work; it does not add learned parameters.
Five mental models to reject¶
"Attention is a database lookup"¶
Attention produces a weighted mixture of value vectors from visible positions. It can behave like retrieval in some heads, but the weights are continuous, contextual, and recalculated in every layer.
"Each attention head has a permanent English-language job"¶
Heads can show interpretable patterns, but assigning permanent human labels is an analysis claim that requires measurement. The architecture only defines projections and weighted combinations.
"More parameters run for every token"¶
True for a dense layer, but not for a sparse MoE layer. An MoE can store many expert FFNs while activating only a subset per token. Attention, embeddings, normalization, and other shared components still run.
"An MoE expert is a complete mini-LLM"¶
In Mixtral, DeepSeekMoE, Qwen MoE, DBRX, and OLMoE, experts replace FFN sub-blocks. They are not independently prompted agents with their own complete attention stacks.
"The model thinks once, then writes"¶
The model updates its computation after every selected token. Apparent plans may be represented in the context and hidden states, but the external operation remains repeated next-token prediction.
Read released code with this map¶
When opening a model implementation, locate these in order:
- model configuration:
d_model, layer count, heads, vocabulary; - token embedding and final vocabulary projection;
- one block's
forwardmethod; - attention projections and mask;
- dense FFN or MoE router and experts;
- top-level loop over blocks;
- cache update and generation wrapper.
For a compact real trail, see Meta's released Llama 3 model.py. The file contains RMSNorm, RoPE, grouped-query attention, SwiGLU, the block, and the top-level decoder in that order.
Check your understanding¶
You should now be able to answer:
- Which sublayer mixes information across token positions?
- Which axis does the vocabulary projection create?
- Why can training score all positions in parallel while generation cannot choose all future tokens in parallel?
- What survives through a residual connection when a sublayer writes a small update?
- In a standard Transformer MoE, which sub-block is usually replaced by experts?
Next: embeddings and position.