LLM Pretrain: Anatomy of a Decoder-Only TransformerLLM 預訓練：深入解析純解碼器 Transformer

A transformer is often explained as a stack of attention layers, treated as if serving the model is an entirely separate infrastructure problem. In reality, the two are inseparable. A model's mathematical architecture directly determines what data must be saved during text generation—and that data dictates memory capacity, server batching behavior, latency, and ultimately your cloud bill.

This is the first article in a three-part series that follows a decoder-only transformer from language modelling to production serving. Before we can solve the memory bottlenecks (Part 2) or scale the engine to handle thousands of concurrent users (Part 3), we must understand the engine itself. In this post, we will strip away the magic and look at the exact mechanics of how a prompt becomes a predicted token, following the data through tokenization, causal self-attention, and feed-forward layers.

The content is structured as:

1. The mechanics of next-token prediction
2. Why transformers?
3. Decoder architecture
4. Where inference begins

A separate notebook series implements the same ideas in PyTorch and then carries the model into serving experiments.

1. The mechanics of next-token prediction

Consider a short sequence:

The company reported higher revenue.

A language model learns from this sequence by repeatedly answering one question:

Given the tokens seen so far, which token should come next?

The same sentence creates several training examples:

This next-token task is the foundation for both training and generation. For every position, the model learns a probability distribution for the next token:

For example, after the context The company reported higher, the model might assign different probabilities to revenue, earnings, costs, and many other tokens. Training encourages it to assign more probability to the token that actually appears next.

Repeating this task across large amounts of text helps the model learn patterns in language, documents, and code. These patterns are distributed across many internal weights rather than stored as individual sentences in one location.

Next-token prediction has a strict information rule: the prediction at position $t$ may use the current context and earlier tokens, but not later tokens. Regardless of the underlying architecture, the model must be prevented from looking ahead at future information during training. The training sequence supplies both inputs and targets simply by shifting the same tokens by one position:

Input:   The company reported higher revenue
Target:      company reported higher revenue .

For tokens $y_1,\ldots,y_T$, the cross-entropy objective is:

The model compares its predictions with the actual next tokens, adjusting its weights when it guesses incorrectly. The fundamental engineering challenge here is speed: because each position only relies on earlier tokens, the predictions for an entire training sequence could theoretically be calculated simultaneously. How a model actually achieves that parallelism without violating the strict "no looking ahead" rule is what separates modern architectures from the previous generation of neural networks.

2. Why transformers?

Before 2017, sequence modelling was dominated by recurrent neural networks (RNNs) and LSTMs. They pass information through recurrent state updates, meaning they must process text strictly word-by-word. This creates a hard dependency between consecutive positions, which severely limits parallelism during training and forces distant tokens to communicate through a long, easily degraded information path.

The transformer replaced recurrent state updates with self-attention as its primary mechanism for exchanging information across positions. Attention gives each token a direct, instant mathematical path to the earlier context it is allowed to use, completely removing the step-by-step bottleneck. As a result, all known positions in a massive training sequence can be processed in parallel.

The original architecture introduced in Attention Is All You Need paired an encoder stack with a decoder. However, many general-purpose generative LLMs today have dropped the encoder entirely. To maintain the strict "no looking ahead" rule while still processing tokens in parallel, a decoder-only transformer relies on a causal mask— a mathematical filter that actively hides future tokens from earlier positions during self-attention.

Depending on how this visibility is restricted, the transformer architecture dictates its ideal use case:

3. Decoder Architecture

To understand how a model actually generates text, we first need to look at the end-to-end anatomy of the process. Think of this as the roadmap of a single prompt traveling through the model's architecture.

The complete generation path is:

Prompt text
    ↓
Tokenizer and token IDs
    ↓
Input representations and positional information   (Section 3.1)
    ↓
Stack of decoder transformer blocks                (Section 3.2)
    ├── RMSNorm and residual connections         (Section 3.2.3)
    ├── Causal self-attention                    (Section 3.2.1)
    └── Gated feed-forward network               (Section 3.2.2)
    ↓
Final normalisation and vocabulary projection      (Section 3.3)
    ↓
Next-token probabilities                           (Section 3.3)
    ↓
Selected token, appended to the prompt             (Section 3.3)

The following subsections will progressively unpack this pipeline, starting from the initial representations that enter the model and moving all the way through to the final token prediction.

3.1 Input representations and positional information

Text does not enter the model as words. A tokenizer divides it into tokens from a fixed vocabulary and maps each token to an integer ID. Depending on the tokenizer, a token may represent a word, a subword, punctuation, whitespace, or a byte-level fragment.

Each ID acts as an index to retrieve a specific row from a learned vocabulary matrix (often called an embedding table). If the vocabulary has 50,000 tokens, this matrix has the shape [50000, d_model]. This lookup operation maps the integer ID to a continuous embedding vector of width $d_{model}$. For batch size $B$ and sequence length $S$, the initial residual stream has shape:

For example,

Batch size:       2
Sequence length:  128 tokens
Model width:      4,096 values per token

Residual-stream shape: [2, 128, 4,096]

Self-attention by itself has no notion of token order. The original transformer added fixed sinusoidal positional encodings to token embeddings. Many current decoders instead use Rotary Position Embeddings (or RoPE). RoPE rotates pairs of query and key coordinates according to token position, making relative position affect their dot product without adding a separate vector to the residual stream.

RoPE is normally a deterministic transformation rather than a learned projection matrix. Its configuration still matters: base frequencies (where different vector dimensions rotate at progressively slower speeds to encode both fine-grained and coarse-grained distances), scaling method, and maximum context assumptions affect long-context behaviour.

Token representation
        ↓
Q and K projections
        ↓
Position-dependent rotation
        ↓
Attention scores

The matching token is 1 position earlier.
The matching token is 100 positions earlier.

3.2 Inside the transformer block

The token representations pass through a sequence of transformer blocks:

Initial representations X₀
        ↓
Transformer block 0
        ↓
Transformer block 1
        ↓
       ...
        ↓
Transformer block L−1
        ↓
Final representations Xᴸ

Every block shares the same structural blueprint but contains its own unique learned weights. Each block reads the current residual stream, injects new contextual information, and passes the updated representations up to the next block. Through this sequence, the model's understanding of the text is progressively refined across the stack.

A modern pre-norm decoder block relies on three fundamental mechanisms:

1. Causal self-attention: which routes and mixes information across different token positions.
2. Feed-forward network (MLP): which transforms the representation at each position independently.
3. RMSNorm and residual connections: which stabilize the learning process and ensure information flows cleanly through the deep network without degrading.

The block receives a residual stream $x_\ell$ and produces the input to the next block, $x_{\ell+1}$. By stacking many blocks with different learned parameters, the model builds deep, complex representations.

The mathematical flow through layer $\ell$ ties these three mechanisms together:

xₗ
 ↓ RMSNorm → Attention update → Residual addition
hₗ
 ↓ RMSNorm → MLP update → Residual addition
xₗ₊₁

Here, $x_\ell$ is the block input, $h_\ell$ is the intermediate representation after the attention update, and $x_{\ell+1}$ is the final block output passed to the next layer in the stack.

3.2.1 Causal self-attention

Self-attention is the mechanism that decides which tokens in the context should influence the current token. To achieve this, it relies on three elements: Queries, Keys, and Values.

Queries, Keys, and Values

For a sequence of hidden states $X$, the attention layer creates queries, keys, and values through learned linear projections:

These three representations serve distinct roles:

• A query describes what the current token is looking for;
• A key describes what a token offers for matching; and
• A value carries the information returned when that token receives attention weight.

The core operation is scaled dot-product attention:

The query-key dot product ($QK^\top$) measures the compatibility between tokens. Dividing this score by $\sqrt{d_k}$ prevents the magnitude of the dot product from growing too large with the head width. Without this scaling, the values would push the softmax function into an extremely saturated regime, leading to vanishing gradients during training.

The Causal Mask

In the formula above, $M$ represents the causal mask. To enforce the strict autoregressive rule (no looking into the future), entries in $M$ representing future positions are set to $-\infty$ before softmax is applied.

Because the softmax of $-\infty$ is zero, token $i$ can attend to itself and earlier positions ($j \le i$), but its attention weight for any future position ($j>i$) is zeroed out. This mask is what allows the transformer to process all positions in a training sequence simultaneously while preserving the integrity of next-token prediction.

Note: Attention is best understood as contextual information routing, not a strict database lookup. A high attention score does not act as definitive proof for why the model produced a specific answer. A useful representation emerges across many heads and layers, and it should not be assumed that one attention head corresponds cleanly to one human-interpretable function.

Multi-head and grouped-query attention

Instead of applying one massive attention operation across the full model width, multi-head attention splits the queries, keys, and values into smaller, parallel "heads." This allows the model to attend to different types of relationships simultaneously—for example, one head might track grammatical structure while another resolves pronoun references.

Each head performs attention in its own learned subspace. The results are then concatenated and passed through a final projection matrix $W_O$:

Concatenation places the head outputs side by side. The learned matrix $W_O$ then mixes the information across those distinct head dimensions and maps the combined result back to the original model width ($d_{model}$). This step is crucial: the attention output must match the width of the residual stream before the two can be added together. Even if the concatenated width already equals $d_{model}$, $W_O$ ensures the model learns how to optimally combine the different heads rather than leaving them as separate groups of features.

Not every modern decoder uses a strict 1:1 ratio of query heads to key-value heads. Grouped-query attention (GQA) is an architectural variant that allows several query heads to share a single key-value head.

8 query heads
2 key/value heads

Query heads 1–4 → share KV head 1
Query heads 5–8 → share KV head 2

This is a critical optimization choice that has major consequences for production serving, which will be explored in depth later in this series.

3.2.2 Gated feed-forward network (MLP)

While self-attention routes information between different positions (inter-token), the feed-forward network (or MLP), transforms the representation at each position completely independently (intra-token) using the same shared learned weights.

The original transformer used two linear projections separated by a standard ReLU activation. However, most modern decoders have adopted a "gated" architecture, such as SwiGLU:

In this design, each token enters as a vector of $d_{model}$ values. The MLP temporarily expands it into a much wider feature space, $d_{ff}$, processes it, and compresses it back down:

Input (x):                 d_model values
                                ↓
Up projection (W_up):      d_ff candidate features
Gate projection (W_gate):  d_ff gate values
                                ↓
Activation (SiLU):         Applies non-linearity to the gate
                                ↓
Multiplication (⊙):        Element-wise multiplication
                                ↓
Down projection (W_down):  d_model values

Notice that the network creates two expanded vectors of the exact same size ($d_{ff}$):

• The up projection ($W_{up}x$) creates the raw candidate features.
• The gate projection ($W_{gate}x$) decides how strongly each candidate feature should pass through to the next stage.

The SiLU (Sigmoid Linear Unit) activation transforms the gate projection into a set of learned "dimmer switches." These two wide vectors are then multiplied together one value at a time ($\odot$):

The dimmer switches strengthen useful features and suppress less useful ones. Finally, the down projection ($W_{down}$) compresses these selected features back to the original $d_{model}$ width, allowing the result to be cleanly added back into the residual stream. (Note: this "expansion" strictly refers to the number of features per token, not the number of tokens in the sequence).

Because of this massive internal expansion space, MLPs frequently contain the largest share of a dense transformer's total parameters—though the exact percentage depends on the model's width, vocabulary size, and whether it uses sparse experts (MoE). Research often associates these dense MLP layers with the model's factual recall and associative memory, though it is important to remember that model knowledge is holistic and does not live exclusively in one isolated component.

3.2.3 RMSNorm and residual connections

Modern decoder families commonly use a "pre-normalization" architecture: the input is normalized immediately before passing through a transformation, while an unnormalized residual path bypasses it entirely.

Using RMSNorm, the mathematical flow for one complete block can be written as:

RMSNorm rescales the input to keep activation magnitudes stable before the attention or MLP layers process them. The residual connection (the $+ x_\ell$ and $+ h_\ell$) acts as an information highway. It preserves the existing representation, allowing the attention and MLP layers to add their new results as updates, rather than completely replacing the original information. (This is exactly why both attention and the MLP must project their final outputs back to the model width $d_{model}$—the dimensions must match so they can be added directly back into the residual stream).

While older models like the original transformer used LayerNorm, modern decoders largely prefer RMSNorm. LayerNorm centers (subtracts the mean) and rescales activations; RMSNorm only rescales them. Dropping the mean-centering step provides a noticeable boost to computational speed and efficiency during training and inference without degrading the model's performance.

3.3 Vocabulary projection and token selection

After the sequence passes through the final transformer block ($L-1$), the residual stream holds a deeply contextualized representation for every token in the prompt. However, because next-token prediction operates strictly left-to-right, the model only looks at the representation of the very last token in the sequence to determine what comes next.

Final normalization and projection

This final vector first passes through one last RMSNorm layer to stabilize its values. Then, it hits the language modeling head (or un-embedding matrix). This learned projection matrix maps the vector from the internal model width back out to the full vocabulary size.

Final token representation:  [d_model]
                                ↓
Vocabulary projection:       [vocab_size]

If the model has a vocabulary of 50,000 tokens, this projection outputs a vector of 50,000 raw, unnormalized scores known as logits.

Next-token probabilities

Because logits can be any real number (positive or negative), they must be converted into a valid probability distribution. The model applies a softmax function:

This forces all 50,000 scores into a range between 0 and 1, ensuring they sum perfectly to 1.0 (or 100%). At this moment, the model has completed its forward pass. It has assigned a probability to every possible word, subword, and punctuation mark it knows.

Selection and the autoregressive loop

The model itself does not strictly choose the next token; it only provides the probabilities. The actual selection is dictated by the decoding strategy:

• Greedy decoding: Simply selects the token with the highest probability (the $\operatorname{argmax}$).
• Sampling (Temperature, Top-P, Top-K): Adjusts the logits to make the distribution flatter (more creative) or sharper (more focused), and then randomly selects a token based on those weighted probabilities.

Once a token is chosen, the loop closes: the new token is appended to the original prompt, forming a new, slightly longer sequence. This new sequence is fed directly back into the tokenizer at Step 1, and the entire massive architectural process repeats to predict the token after that.

4. Where inference begins

The architecture we just explored explains exactly what must happen to execute a single forward pass. But production serving introduces an entirely different challenge:

How can a system generate many tokens, for many concurrent requests, without repeatedly recomputing the mathematical history of every prompt?

The next article, LLM Inference: How KV Caching and PagedAttention Save LLM Inference, separates the generation process into prefill and decode phases, exposing exactly how the self-attention state transforms from a theoretical mechanism into a massive production memory bottleneck.