Transformer Architecture Explained: Trace a Decoder-Only LLM

Give a language model the prompt “The capital of France is” and it will generate “Paris”. In the inference primer, we followed the outer loop that makes this happen: prefill reads the prompt, decode samples one new token at a time, and the KV cache grows by one column each step. But what happens inside a single forward pass? How does four words of text turn into a ranked list of scores predicting which word should come next?

Before getting lost in matrix shapes, we need four plain-language definitions:

• Transformer: A neural network architecture built around attention. Traditional recurrent networks processed tokens one by one in a sequential chain, forgetting early context. A transformer processes sequence positions in parallel and lets each position pull information directly from any other position.
• Attention: The mechanism that lets tokens exchange information. When the model processes the word “is”, attention allows it to look back at “capital” and “France” so it knows what question is being asked.
• Decoder-only: An architecture where attention is strictly causal—each position is only allowed to look at itself and previous positions, never into the future. Because an autoregressive model must generate text without seeing future words, modern generative language models (from GPT-2 and Llama to modern reasoning models) are decoder-only.
• Logits: The raw, unnormalized score numbers produced for every word in the vocabulary at the end of the pass. A higher logit means higher confidence. Converting logits through a softmax function produces valid probabilities that sum to 1.

In this lesson, we open up one forward pass of a decoder-only transformer. First, we will examine what each internal layer actually does to the representation—why the numbers change even where the tensor shape stays invariant. Then, we will trace a concrete toy batch through every component boundary and verify why only the last position supplies the scores for text generation.

The one-sentence version

A decoder-only transformer keeps a sequence matrix of width d, mixes context across positions with causal self-attention, transforms each position with a feed-forward MLP, and repeats that block; the final hidden state at the last position is projected to vocabulary logits to score the next token.

Everything else in this lesson makes that sentence concrete.

Lab: Run the companion CPU exercise at /labs/track1/02-tiny-decoder/ (stack blocks, tie weights, verify forward pass).

The anatomy of a forward pass: what each layer contributes

If you inspect a transformer purely as a list of tensor shapes, it looks mysterious: the sequence enters a block as [B, T, d], passes through attention, comes out as [B, T, d], enters an MLP, and comes out as [B, T, d]. Why pass data through layers if the shape never changes?

The shape stays constant because of the residual stream, but the information inside that shape is radically transformed at each step:

1. Input representation (Token + Position Embedding): The model maps discrete token IDs into static vectors of length d. In isolation, the vector for “capital” knows about the word’s dictionary definitions, but it has no idea whether this sentence is about financial investments, uppercase letters, or geography. Positional embeddings add a signal indicating whether “capital” appeared at position 0, 1, or 2.

2. Causal Self-Attention (Contextualization across time): Self-attention is the only component in the block that allows tokens to communicate across sequence positions. At position 3 (“is”), attention calculates compatibility scores with all previous tokens (“The”, “capital”, “France”), extracts their relevant features, and combines them into position 3’s vector. The shape remains [B, T, d], but position 3’s vector now encodes a context-dependent concept: “the entity whose capital is sought is France”. The causal mask guarantees that position 1 (“capital”) cannot cheat by looking ahead at “France”.

3. Feed-Forward Network / MLP (Per-position feature transformation and memory): While attention mixes information across sequence positions, the MLP processes each position independently. It projects the representation into a wider intermediate dimension (typically 4 × d or 8/3 × d in SwiGLU), applies a non-linear activation function, and projects back to d. Mechanistic interpretability research indicates that MLPs act like key-value associative memories: if self-attention assembled the concept “capital of France”, the MLP activates the stored factual association “Paris”.

4. Residual Connections (Identity highways): Each sub-layer computes an additive modification Δx, added directly to the residual stream: x_out = x_in + SubLayer(x_in). Instead of forcing each layer to re-encode the entire prompt from scratch, the residual stream acts as a central communication bus. Layers read from the bus and write additive updates to it. This prevents vanishing gradients and allows models to stack dozens or hundreds of blocks deep.

5. Output Projection / Un-embedding (Mapping concepts to vocabulary scores): After the final block, the contextual vector at each position is multiplied by a vocabulary matrix of shape [d, V]. This compares the final representation against every candidate token in the vocabulary, producing a logit score for each.

See the shape trace first: a tiny batch preview

Now that we know what each layer contributes semantically, trace one concrete batch through the complete network to see the exact shape changes at every component boundary:

• Small input: A batch of token IDs input_ids of shape [B=2, T=4] drawn from a vocabulary of size V=16, with model width d=8.
• Action: Map IDs through an embedding table [16, 8], add position embeddings [4, 8], pass through a pre-norm decoder block (LayerNorm → multi-head attention with h=2 heads of width d_head=4 → residual addition → LayerNorm → MLP expanding to width m=32 → residual addition), final LayerNorm, and a linear vocabulary projection [8, 16].
• Observable result:
  • Token + position embeddings: [2, 4, 8]
  • Multi-head attention internal scores: [2, 2, 4, 4] (causally masked lower-triangular)
  • Residual stream across the block: strictly invariant at [2, 4, 8]
  • Vocabulary logits: [2, 4, 16]
  • Generation output (last position only): next_token_logits of shape [2, 16]
  • Training loss slice: logits[:, :-1, :] [2, 3, 16] paired with next-token targets [2, 3].

Decoder-only transformer

A small decoder-only transformer traced from token IDs to logits. Each decoder block keeps the [2, 4, 8] shape.

Read this diagram as text

Text equivalent: input token IDs have shape [2, 4]. Token embedding plus position information has shape [2, 4, 8]. A repeated decoder block applies layer normalisation, causal self-attention, a residual connection, layer normalisation, an MLP with hidden width 32, and another residual connection; its output remains [2, 4, 8]. A final layer normalisation keeps [2, 4, 8]. A vocabulary projection produces logits of shape [2, 4, 16]. In this next-token language model, logits at position t score the token that follows position t.

A consistent shape trace

Use these dimensions throughout the lesson:

Symbol	Meaning	Value in this trace
B	Batch size (sequences processed together)	2
T	Sequence length (tokens per sequence)	4
d	Model width (hidden dimension)	8
h	Attention heads	2
d_head	Width per head (d / h)	4
m	MLP hidden width (4 × d in this trace)	32
V	Vocabulary size	16

The numbers are deliberately small. A real model uses the same pattern with larger values; the pattern is what matters.

From token IDs to a sequence matrix

The model begins with token IDs, not text:

input_ids: [B, T] = [2, 4]

An embedding table of shape [V, d] = [16, 8] turns each ID into a vector. The result is a sequence matrix:

token_embeddings: [B, T, d] = [2, 4, 8]

Tokens carry no order information by themselves, so the model adds position information. In this trace that is an additive position embedding broadcast over the batch:

position_embeddings: [T, d] = [4, 8]
hidden_state: [B, T, d] = [2, 4, 8]

The exact position method varies by model: the original Transformer used sinusoidal positional encodings, and GPT-2’s implementation uses learned position embeddings. This trace uses additive embeddings so the tensor shapes stay visible.

One decoder block

A decoder block has two sub-layers, each wrapped in a residual connection. This trace uses pre-normalisation, which is common in modern decoder-only models; the original Transformer used post-normalisation. Placement affects training stability, but it does not change the shape trace.

1. Layer normalisation keeps [2, 4, 8].
2. Causal self-attention also keeps [2, 4, 8]. Internally it forms queries, keys, and values, splits them into h = 2 heads of width d_head = 4, computes an attention score matrix of shape [B, h, T, T] = [2, 2, 4, 4], applies a causal mask so position t cannot attend to positions after t, and merges the heads back.
3. Residual addition adds the attention output to its input, keeping [2, 4, 8].
4. Layer normalisation keeps [2, 4, 8].
5. MLP (feed-forward network) expands each position to an intermediate [2, 4, 32], applies a non-linearity, and projects back to [2, 4, 8]. Positions are transformed independently; the block output matches the residual stream’s shape.
6. Residual addition keeps [2, 4, 8].

The causal mask prevents future-token leakage: the representation at position t can use only positions 0 through t. That masking, together with the next-token training objective and the vocabulary output head, makes this a next-token language model rather than a bidirectional encoder.

This lesson names the attention step and its mask. Deriving queries, keys, values, scaled dot-product attention, and softmax is the subject of the planned self-attention lesson.

Stacking blocks

The block is repeated N times. N changes depth and parameter count, but not the tensor shape: every block receives [B, T, d] and returns [B, T, d]. A two-block toy and a hundred-block model have the same trace shape.

After the last block, a final layer normalisation keeps:

final_hidden_state: [B, T, d] = [2, 4, 8]

From hidden state to logits: generation versus training

The final layer normalisation outputs a contextual hidden representation for every sequence position:

final_hidden_state: [B, T, d] = [2, 4, 8]

To turn these vectors into predictions, the model applies a linear projection from model width d to vocabulary size V. If the embedding and output matrices are tied (as in GPT-2), this projection reuses the embedding table transposed; if untied (as in Llama 2), it is an independent matrix of shape [d, V] = [8, 16]. In both cases:

logits: [B, T, V] = [2, 4, 16]

At this point, a fundamental question arises: we have logits for all 4 positions in the sequence. Which ones do we actually use?

The answer depends entirely on whether we are generating text or training the model:

1. During generation: only the last position matters

When generating text autoregressively (as described in the inference primer), we want to predict what token comes after the entire prompt.

Because this is a causal model, position 0 was only allowed to attend to token 0. Position 1 only attended to tokens 0 and 1. Only position $T-1$ (the fourth token in our prompt, index 3) was allowed to attend to the complete sequence. Its hidden state is the only one that integrates the full context:

# Extract the unnormalized score vector for the final position:
next_token_logits = logits[:, -1, :]  # Shape: [B, V] = [2, 16]

# Convert logits to probabilities and sample the next token:
next_token_id = torch.argmax(next_token_logits, dim=-1)  # Shape: [B]

The model drops the first $T-1$ logit vectors because they only score continuations for earlier prefixes of the prompt. We append next_token_id to our prompt and run the next forward pass.

2. During training: every position is evaluated in parallel

If we only used the last position during training, pretraining on a 4,096-token document would yield only one learning signal per forward pass. That would be wildly inefficient.

Because of the causal mask, position 0’s representation depends only on token 0. Therefore, the logits produced at position 0 are a valid prediction of what should follow token 0 (which is token 1). Similarly, position 1’s logits predict token 2, and position 2’s logits predict token 3.

Training extracts $T-1$ predictions from a single forward pass by shifting logits and targets:

# Predict the next token for every position up to T-1:
shift_logits = logits[:, :-1, :].contiguous()  # Shape: [B, T-1, V] = [2, 3, 16]
shift_targets = input_ids[:, 1:].contiguous()  # Shape: [B, T-1] = [2, 3]

loss = torch.nn.functional.cross_entropy(
    shift_logits.view(-1, 16), shift_targets.view(-1)
)

The last position in logits ($t=3$) is dropped during training because there is no ground-truth fifth token in this input sequence. We explore this alignment and padding mechanics in detail in Lesson 3.

What shapes reveal, and what they do not

Shapes tell you the dimensions of parallelism: how many sequences, how many positions, how wide each representation is, how many heads, and how many vocabulary scores. They are the first thing to check when debugging a model.

Shapes do not tell you:

• whether the causal mask is actually applied correctly;
• whether normalisation is stable;
• how much memory or time a pass will cost;
• whether the model has learned anything useful.

A correct shape is necessary, not sufficient. It is the floor of debugging, not the ceiling.

Three transformer families

The same building blocks appear in three common arrangements. The difference is which tokens each position is allowed to see, and what the model outputs.

Transformer architecture families

	Encoder-only	Decoder-only	Encoder-decoder
Attention	Bidirectional self-attention	Causal self-attention	Encoder bidirectional; decoder causal plus cross-attention
Mask	Bidirectional; padding mask when needed	Causal mask	Causal mask in decoder; source visible through cross-attention
Typical output	One representation per input token, or a pooled vector	Next-token logits at every position	A generated target sequence
Strong fit	Classification, embeddings, retrieval, reranking	Open-ended generation, chat, code completion	Translation, summarisation, other sequence-to-sequence tasks

The [original Transformer](https://arxiv.org/abs/1706.03762) was an encoder-decoder model. [BERT](https://arxiv.org/abs/1810.04805) popularised bidirectional encoder pre-training with a masked language-model objective, and the [GPT-2 report](https://cdn.openai.com/better-language-models/language-models.pdf) popularised decoder-only language modelling. These are common fits, not requirements: decoder-only models can also be prompted for classification, embedding, translation, and summarisation.

A BERT-style masked-language-model encoder is a poor natural generator because its objective does not train it to predict the next token from a prefix. Decoder-only models can be prompted to classify or embed, but they spend capacity on generation machinery that a pure encoder does not need. Encoder-decoder models pay for two stacks and cross-attention, which is worthwhile when source and target are genuinely different sequences.

Architecture decision example

A support workflow needs two capabilities: route tickets by intent, and draft replies. If the product only needed routing, an encoder-only model would be the more direct fit. If it needs open-ended replies, a decoder-only model is the natural choice. If tickets arrive in one language and replies must be drafted in another, an encoder-decoder model maps the source sequence to a target sequence explicitly. The decision is driven by the output the product must produce, not by which architecture is newest.

Where this trace stops

This trace is one small, consistent design. Real models vary in ways that do not change the shape story:

• activation function (GELU, SwiGLU, and others);
• normalisation placement (pre-LN versus post-LN) and normalisation type (LayerNorm, RMSNorm);
• position method (sinusoidal, learned, or rotary);
• attention variant (multi-head, multi-query, grouped-query, or latent);
• depth, width, and MLP expansion ratio;
• whether embedding and output weights are tied.

The width-preserving residual stream transfers; block details and ordering vary by model. Tokenization algorithms and the derivation of attention are covered in their own planned lessons.

Interview questions

• Why does the residual stream keep the shape [B, T, d] through every block? Answer: Each sub-layer is designed to return a tensor of the same width, so it can be added back to its input. That residual path is what lets depth increase without changing the representation size.
• What does the causal mask prevent? Answer: It prevents position t from using information from positions after t, so the model can be trained to predict the next token from the prefix alone.
• Why is an encoder-only model a poor fit for open-ended generation? Answer: It sees the whole input bidirectionally and is not trained to produce a next-token distribution at every position in the same way a decoder is.

What comes next: Where did the input numbers come from?

In this trace, we began with integer IDs like [2, 4] and looked up rows in an embedding table of shape [16, 8]. But where did those numbers come from? How does raw text like “The capital of France is” become clean integer IDs, why are spaces often absorbed into word tokens, and how does the output projection layer physically turn a 1D hidden state vector back into vocabulary logits?

That is the subject of the next lesson: Tokenisation and Embeddings: From Text to Tensors.

Key takeaways

• A decoder-only transformer maintains a residual stream of shape [B, T, d] across repeated blocks.
• Causal self-attention contextualizes representations across time; the feed-forward MLP transforms features and recalls associations at each position independently.
• The residual stream acts as an additive highway (x + SubLayer(x)), allowing deep networks to preserve identity and propagate gradients cleanly.
• Logits at the last position ($T-1$) score the next continuation for autoregressive generation. Training extracts $T-1$ next-token targets across the entire prompt in parallel using a loss-side shift.

Module exit check: tensor shapes and boundaries

1. 01For the baseline setup [B=2, T=4, d=8, h=2, V=16], recompute all internal shapes if attention heads increase to h=4 while model width remains d=8. Which tensor shapes change, and which stay invariant?

   Show a hint

   Per-head dimension d_head = d / h = 8 / 4 = 2. The Q, K, V head projections become [B, h, T, d_head] = [2, 4, 4, 2], and the causal attention score matrix becomes [B, h, T, T] = [2, 4, 4, 4]. Every other shape is strictly invariant: input_ids [2, 4], token embeddings [2, 4, 8], residual stream [2, 4, 8], MLP intermediate [2, 4, 32], and final logits [2, 4, 16].

2. 02A colleague claims that for sequence length T=4, the training loss function compares logits [2, 4, 16] directly against targets [2, 4]. Using the shape trace from this lesson, explain why this is incorrect and write the exact shapes passed to the loss.

   Show a hint

   Causal next-token prediction requires predicting token t+1 from prefix 0..t. The last position (t=3) has no subsequent token to predict in this sequence. The loss function slices logits[:, :-1, :] to shape [2, 3, 16] and compares against labels[:, 1:] of shape [2, 3]. Direct unshifted comparison would evaluate predicting the current token rather than the next token.

3. 03In generation mode, you feed a batch of size B=1 with 4 tokens ["The", "capital", "of", "France"] into this model. From logits of shape [1, 4, 16], which slice provides the scores for the next token? If the top-scoring token ID is sampled and appended to the prompt, what is the exact shape of input_ids passed into the next forward pass (assuming no KV cache is used)?

   Show a hint

   The slice logits[:, -1, :] (or logits[0, 3, :]) of shape [1, 16] provides the scores, because position 3 is the only position that attended to all four prompt tokens. Appending the sampled token yields an input sequence of length 5, so the next un-cached forward pass receives input_ids of shape [1, 5].

Related lessons

• LLM Internals & Architecture · Inference foundationsHow LLM inference works9 min
• LLM Internals & Architecture · Inference foundationsUnderstanding the KV Cache12 min

Sources and scope

• Vaswani et al., “Attention Is All You Need” (arXiv:1706.03762) introduced the Transformer as an encoder-decoder architecture.
• Radford et al., “Language Models are Unsupervised Multitask Learners” (GPT-2 report) describes decoder-only transformer language modelling.
• Devlin et al., “BERT” (arXiv:1810.04805) introduced bidirectional encoder-only pre-training.
• Raffel et al., “T5” (arXiv:1910.10683) frames NLP tasks as text-to-text with an encoder-decoder model.
• PyTorch CrossEntropyLoss documentation defines the loss used for next-token training.

This trace is an illustrative small model. It is not a specification for any production architecture. Attention derivation, tokenization, positional encodings, and training objectives are covered in later lessons.