Looking Inside Large Language Models
Adapted from: Jay Alammar & Maarten Grootendorst, Hands-On Large Language Models (O'Reilly, 2024), Chapter 3 — rewritten and condensed for teaching, with code examples adapted from the original.
Lesson 2 covered how text becomes tokens, and how tokens become embeddings. This lesson goes one level deeper: what actually happens inside the model, between "here's a sequence of embeddings" and "here's the next word"? We'll focus specifically on generative, text-producing transformer models — the kind behind ChatGPT-style assistants — rather than representation models like BERT.
An Overview of Transformer Models
Inputs, outputs, and the generation loop
From the outside, a transformer LLM looks simple: text goes in, text comes out. But that framing hides an important detail — the model doesn't produce its entire response in one shot. It generates one token at a time, and each single-token generation step is one full forward pass through the network (the input flowing through every computation in the model, once, to produce one output).
After each forward pass, the newly generated token gets appended to the end of the input, and that longer sequence is fed back in for the next forward pass. The model is, at its core, just repeatedly answering the question "given everything so far, what's the single most-likely-useful next token?" — with a loop around it that keeps feeding its own output back in as input.
This pattern — a model consuming its own earlier predictions to produce later ones — has a name: autoregressive. It's why text-generation LLMs are called autoregressive models, a label that specifically distinguishes them from representation models like BERT, which don't generate text token-by-token this way at all.
Here's what setting up that loop looks like in code, using Hugging Face's pipeline helper:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
model = AutoModelForCausalLM.from_pretrained(
"microsoft/Phi-3-mini-4k-instruct",
device_map="cuda",
torch_dtype="auto",
trust_remote_code=True,
)
# Create a pipeline
generator = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
return_full_text=False,
max_new_tokens=50,
do_sample=False,
)
prompt = "Write an email apologizing to Sarah for the tragic gardening mishap. Explain how it happened."
output = generator(prompt)
print(output[0]['generated_text'])
The output cuts off abruptly mid-email — not because anything went wrong, but because we capped generation at max_new_tokens=50. Raise that limit and the loop just keeps running, appending one token at a time, until the email is actually finished.
Check: the generation loop
Autoregressive generation
1 / 2How does a transformer LLM actually produce a multi-token response?
Pick an answer to see if you are right.
🧪 A minimal autoregressive loop, illustrated
The components of the forward pass
Every forward pass flows through three major stages, in order:
- The tokenizer — turns text into token IDs (covered in Lesson 2).
- A stack of transformer blocks — the bulk of the model's processing.
- The LM head — a final layer that converts the last block's output into a probability score for every token in the vocabulary.
The tokenizer's vocabulary and the model's token embeddings (recall: one starting vector per vocabulary entry) are directly linked — the embedding matrix's row count equals the vocabulary size.
You can see this structure directly by printing a loaded model. Here's a condensed view of Phi-3's layer stack:
Phi3ForCausalLM(
(model): Phi3Model(
(embed_tokens): Embedding(32064, 3072, padding_idx=32000)
(layers): ModuleList(
(0-31): 32 x Phi3DecoderLayer(
(self_attn): Phi3Attention(
(o_proj): Linear(in_features=3072, out_features=3072, bias=False)
(qkv_proj): Linear(in_features=3072, out_features=9216, bias=False)
(rotary_emb): Phi3RotaryEmbedding()
)
(mlp): Phi3MLP(
(gate_up_proj): Linear(in_features=3072, out_features=16384, bias=False)
(down_proj): Linear(in_features=8192, out_features=3072, bias=False)
(activation_fn): SiLU()
)
)
)
(norm): Phi3RMSNorm()
)
(lm_head): Linear(in_features=3072, out_features=32064, bias=False)
)
A few things worth reading directly off this dump:
embed_tokensconfirms the embedding matrix: 32,064 vocabulary entries, each a 3,072-dimensional vector.- 32 stacked
Phi3DecoderLayerblocks do the actual processing — each one bundling a self-attention sub-layer and an MLP (feedforward) sub-layer. lm_headis a single linear layer mapping a 3,072-dimensional vector to a 32,064-dimensional output — one score per vocabulary token.
Check: forward pass components
Components of the forward pass
1 / 2What are the three major stages a prompt passes through during a forward pass, in order?
Pick an answer to see if you are right.
Choosing a token from the probability distribution
Once the LM head produces a score for every vocabulary token, something has to pick one — that choice is the decoding strategy. The simplest option is greedy decoding: always take the single highest-scoring token. In practice this tends to produce noticeably worse output than adding some controlled randomness — sampling proportionally to each token's score, so a token with a 40% probability gets picked roughly 40% of the time, not automatically because it's on top. (Lesson 1 covers this territory — temperature, top-k, top-p — in much more depth; the point here is just to see where in the forward pass this decision actually happens.)
Greedy decoding is exactly what you get by setting temperature to zero.
Here's the forward pass broken into its raw pieces, bypassing the convenience of .generate():
prompt = "The capital of France is"
# Tokenize the input prompt
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to("cuda")
# Get the output of the model before the lm_head
model_output = model.model(input_ids)
# Get the output of the lm_head
lm_head_output = model.lm_head(model_output[0])
lm_head_output comes out with shape [1, 6, 32064]: batch size 1, 6 input tokens, and a 32,064-length score vector for each token position. Only the last position matters for predicting the next token:
token_id = lm_head_output[0, -1].argmax(-1)
tokenizer.decode(token_id)
# 'Paris'
Check: decoding
Choosing a token from the output distribution
1 / 2What does greedy decoding do, specifically?
Pick an answer to see if you are right.
Parallel token processing and context size
One of the transformer's biggest practical advantages over earlier architectures (like RNNs) is that input tokens can be processed in parallel rather than one-by-one. Each input token flows through its own computation "stream," and — for text generation — only the last stream's final output actually gets used to predict the next token.
That raises an obvious question: if only the last stream's output is used, why bother computing all the others? Because the intermediate results from every stream are still needed — they get pulled into the last stream's computation via the attention mechanism at every transformer block along the way. Nothing is wasted; it's just that only one stream's final output feeds the LM head.
The number of parallel streams a model can maintain at once is its context length — a model with a 4K context length can process (and "remember," in a sense) 4K tokens at once, no more.
Tying this back to the shapes from the code above: the input to lm_head was [1, 6, 3072] — a batch of 1, containing 6 tokens, each represented by a 3,072-dimensional vector (the model's hidden dimension, after flowing through the full stack of transformer blocks):
model_output[0].shape
# torch.Size([1, 6, 3072])
lm_head_output.shape
# torch.Size([1, 6, 32064])
Speeding up generation with a KV cache
Naively, generating the second token means re-running the entire forward pass on the (now one-token-longer) sequence — including recomputing everything for tokens we already fully processed in the previous step. That's wasteful: those earlier tokens haven't changed.
The fix is the KV cache (short for "keys and values cache," named after the key and value vectors central to attention — covered next). By caching those vectors from earlier steps, only the newest token's stream needs to be computed on each step — everything else is reused. This is a genuinely large speedup: in one benchmark, generating 100 tokens took 4.5 seconds with caching enabled versus 21.8 seconds with it disabled — roughly a 5x difference, from one flag.
# Cached (fast)
generation_output = model.generate(input_ids=input_ids, max_new_tokens=100, use_cache=True)
# Uncached (slow)
generation_output = model.generate(input_ids=input_ids, max_new_tokens=100, use_cache=False)
Even 4.5 seconds is a noticeable wait from a user's perspective — which is exactly why production LLM APIs stream tokens to the user as they're generated, rather than waiting for the entire response to finish before showing anything.
Check: context size and caching
Context length and the KV cache
1 / 2What does a model's "context length" actually limit?
Pick an answer to see if you are right.
🧪 Recomputation avoided: cached vs. uncached step counts
Inside the Transformer Block
Each transformer block — and there can be anywhere from a handful (the original 2017 paper) to well over a hundred (in today's largest LLMs) — is made of two sequential pieces:
- An attention layer, which pulls in relevant information from other token positions.
- A feedforward layer (an MLP), which holds the bulk of the model's raw processing/storage capacity.
The feedforward network, intuitively
Think about feeding a language model the partial phrase "The Shawshank" and expecting it to predict "Redemption" next. That completion has to come from somewhere — and it comes from the feedforward network. When the model trained on a large enough text archive that mentioned "The Shawshank Redemption" often enough, its feedforward layers absorbed and stored that association.
This is genuine memorization — but memorization alone doesn't explain why LLMs generalize well to inputs they've never seen. The same machinery that stores specific facts also learns to interpolate between related patterns, which is what lets a model perform reasonably on genuinely novel inputs, not just regurgitate memorized text.
A raw, purely pre-trained language model given "The Shawshank" would just continue the sentence. A modern chat assistant given the same input instead answers as if asked about the film — because it went through instruction-tuning and preference finetuning after pre-training (see Lesson 1 on post-training) specifically to behave that way, rather than defaulting to raw text continuation.
The attention layer, intuitively
Feedforward memorization and interpolation only gets you so far — language depends heavily on context. Consider: "The dog chased the squirrel because it..." — to predict what comes after "it," the model needs to resolve what "it" actually refers to. The dog? The squirrel? Nothing in "it" alone tells you; only the surrounding context does.
Attention is the mechanism that pulls relevant context into a token's representation before the model has to make that call — it's specifically what lets the representation of "it" absorb information about "squirrel" (or "dog," depending on what the rest of the sentence implies).
Attention, mechanically
Attention breaks down into exactly two steps, applied at the position currently being processed:
- Score the relevance of every earlier token to the current one.
- Combine information, weighted by those relevance scores, into a single output vector.
Zooming into how that actually gets computed (this builds on the Q/K/V setup from Lesson 1 — a quick refresher: every token position gets a query, key, and value vector, produced by multiplying the token's input vector by three learned projection matrices):
Step 1 — relevance scoring. Take the current position's query vector and multiply it against the keys matrix (every earlier position's key vector, stacked together). That produces one raw relevance score per earlier position. Pass those scores through softmax so they sum to 1 — now they're a proper probability-like weighting.
Step 2 — combining information. Multiply each earlier position's value vector by its relevance score from step 1, then sum all of those weighted value vectors together. The result is the attention layer's output for the current position — a single vector that blends in exactly the right amount of information from exactly the right earlier tokens.
Just like in Lesson 1, this whole two-step process is duplicated across multiple attention heads running in parallel, each with its own independent Q/K/V projection matrices — letting different heads specialize in tracking different kinds of relationships simultaneously. The heads' outputs get concatenated and projected back down at the end.
Check: the attention mechanism
How self-attention is computed
1 / 3In the sentence "The dog chased the squirrel because it fled," what problem is attention specifically solving when the model processes "it"?
Pick an answer to see if you are right.
🧪 A single attention step, computed by hand
Recent Improvements to the Transformer Architecture
The core ideas from the original 2017 transformer paper still hold up — but a handful of architectural refinements have become standard in modern LLMs, largely aimed at making attention (the most computationally expensive part of the model) cheaper without sacrificing much quality.
More efficient attention
- Local / sparse attention limits how many previous positions a token is even allowed to attend to, trading some long-range context for speed. GPT-3 used this selectively — alternating full-attention blocks with sparse-attention blocks rather than making every block sparse, since restricting every block this way would meaningfully hurt output quality.
- Multi-query attention shares a single keys matrix and a single values matrix across all attention heads (only the query projections stay separate per-head) — a big efficiency win, but one that can cost noticeable quality at larger scale.
- Grouped-query attention (used by Llama 2 and 3) is the middle ground: instead of one shared K/V pair for the whole model, several groups of heads each share their own K/V pair. That recovers most of multi-query's efficiency while giving up much less quality.
- Flash Attention isn't a different attention formula at all — it's an implementation-level optimization that speeds up the same math by carefully managing which values move between a GPU's fast, small on-chip memory (SRAM) and its slower, larger main memory (HBM).
Check: efficient attention variants
Local, multi-query, grouped-query, and Flash Attention
1 / 3Why did GPT-3 alternate between full-attention and sparse-attention blocks instead of using sparse attention everywhere?
Pick an answer to see if you are right.
Modern transformer block tweaks
A handful of small changes have become the norm in recent (2023-2024-era) architectures like Llama 3, on top of the original attention + feedforward structure:
- Pre-normalization — normalizing before the attention and feedforward sub-layers (rather than after, as in the original transformer paper) has been shown to reduce required training time.
- RMSNorm instead of the original LayerNorm — a simpler, more computationally efficient normalization method.
- SwiGLU instead of the original ReLU activation function inside the feedforward layer — one of several newer activation variants shown to improve performance.
Positional embeddings: from absolute to rotary (RoPE)
Attention alone has no inherent sense of token order — without extra information, "the dog chased the cat" and "the cat chased the dog" would attend identically. Positional embeddings are what inject that order information.
The original transformer (and early successors) used absolute positional embeddings: token 1 is tagged "position 1," token 2 is tagged "position 2," and so on — either as a fixed geometric pattern or as learned vectors. This runs into a practical problem at scale: training documents are often much shorter than a model's full context length, so for efficiency, multiple short documents get packed together into a single training context. If a second document starts partway through that packed context, naively numbering its first token "position 50" (say) would incorrectly imply 49 tokens of relevant preceding context, when really the preceding tokens belong to a completely unrelated document.
Rotary positional embeddings (RoPE) sidestep this by encoding position as a rotation applied to vectors in embedding space, capturing both absolute and relative position information — and, importantly, they're applied later in the pipeline than absolute embeddings were: specifically to the query and key vectors, right before the relevance-scoring step of attention, rather than added once at the very start of the forward pass.
Check: modern architecture tweaks
Pre-normalization, RMSNorm, SwiGLU, and RoPE
1 / 2Why can naive absolute positional embeddings cause problems when multiple short documents are packed into one training context?
Pick an answer to see if you are right.
Chapter Summary
- A transformer LLM generates one token at a time, autoregressively — each new token is appended to the input and fed back in for the next forward pass.
- A forward pass has three stages: tokenizer → stack of transformer blocks → LM head, with the LM head producing one probability score per vocabulary token; a decoding strategy (greedy, or sampling — see Lesson 1) picks the actual next token.
- Input tokens are processed in parallel streams, capped by the model's context length — but only the last stream's final output feeds the LM head, since earlier streams' information is folded in via attention along the way.
- Caching keys and values (the KV cache) avoids redundant recomputation across generation steps, producing a large real-world speedup.
- Each transformer block pairs an attention layer (pulls in relevant context) with a feedforward layer (stores memorized patterns and interpolates between them); attention itself is two steps — score relevance, then combine information — run in parallel across multiple heads.
- Modern architectures refine the original design with local/sparse, multi-query, and grouped-query attention and Flash Attention for efficiency, plus pre-normalization, RMSNorm, SwiGLU, and rotary positional embeddings (RoPE) for quality and scalability.
The next step is moving from how a model works to how to actually use one — starting with practical, task-focused applications built on top of everything covered so far.
Exercises
Running an actual transformer (Phi-3, Llama) needs a GPU, gigabytes of weights, and network access to download them — none of which the sandbox running these exercises has. So these exercises implement the exact same math (attention, decoding, normalization) directly with plain Python, on tiny hand-sized inputs — the real formulas, mocked inputs.
Read a dimension d, a token count n, then n key vectors, n value vectors, and finally one query vector (all d-dimensional). Compute the single-head self-attention output for that query: score each key with a dot product against the query, softmax the scores, then return the weighted sum of the value vectors.