Decoder-only LLMs

Problem

We want a model that writes text. Text is a sequence of tokens, and the most general thing a model can do with a sequence is predict what comes next given everything so far: P(token_t | token₁, …, token_t−1). Learn that well and you can continue any prompt, answer questions phrased as prompts, translate, summarize and write code, because all of those are "text that continues this text".

Two requirements follow. The model must never use the future when predicting a token (otherwise it would just copy the answer), and it must be cheap to train on trillions of tokens, which means scoring every position of a long sequence in one pass instead of one position at a time. The decoder-only Transformer meets both with a single trick, the causal mask.

Intuition

Read a sentence with the next words covered by a card, sliding the card right one word at a time. At each position you guess the covered word. A decoder-only model does exactly that, but with every card position at once: it sees the whole sequence, yet each position's computation is blocked from looking at anything to its right. One forward pass therefore gives T separate predictions for a sequence of T tokens, and the loss is the average of T cross-entropies (see the training loop).

The stack is the Transformer's decoder with the cross-attention half removed: there is no encoder to attend to, so each block has just two sublayers.

Mechanics

Embeddings: tokens and positions

Attention treats its input as an unordered set, so order has to be put in by hand. GPT-2 does it the simplest way: a second lookup table with one learned vector for each position 0 … 1023, added to the token vector. Both tables are nn.Embedding (a row lookup, lines 188-197); by default its weights are drawn from a standard normal (lines 179-181), so GPT-style code overrides that with a small standard deviation such as 0.02. The context length is a hard limit: a learned table has no row for position 1024. See embeddings for what the vectors come to mean.

Causal self-attention

Inside each block, every position makes a query, key and value, scores its query against the keys and mixes the values (see attention). The causal mask adds −∞ to every score where the key's position is later than the query's, so softmax gives those keys weight exactly 0. For four tokens the allowed pattern is a lower triangle:

In PyTorch's float-mask form (0 allowed, -inf blocked) the grid above is torch.triu(torch.full((T, T), -inf), diagonal=1) (lines 29-42, exposed as Transformer.generate_square_subsequent_mask, lines 302-310).

The mask gives a property worth stating exactly: the output at position t depends only on tokens 1 … t. Run causal attention on the four-token example from the attention page (the, cat, sat, down), then run it again on just the first three tokens. The first three output rows agree to within 0e+00. Without the mask, dropping the last token changes those rows by up to 0.36. This is why a decoder can be trained on whole sequences and still be used one token at a time: the answer for the prefix is the same either way. (Stacking layers keeps the property, because layer k+1 at position t only reads layer k at positions ≤ t.)

The block, pre-norm

Each block is x ← x + Attention(LayerNorm(x)), then x ← x + FFN(LayerNorm(x)). This "pre-norm" order is the norm_first=True branch of TransformerEncoderLayer (lines 946-952), and it leaves the residual stream untouched by any normalization, which keeps gradients healthy in deep stacks (see LayerNorm and residuals). Because the stream is never normalized inside the blocks, one extra LayerNorm is applied after the last block. GPT-2's paper describes both changes: layer normalization moved to the input of each sub-block, and an additional one after the final block. The feed-forward network is Linear(768 → 3072) → GELU → Linear(3072 → 768), a width of four times the model dimension.

Why reuse an encoder layer for a decoder? Because "encoder" and "decoder" in PyTorch differ by cross-attention, not by masking. TransformerDecoderLayer (lines 985-1201) takes a memory from an encoder and attends to it. A decoder-only model has none, so what it needs is a self-attention layer plus a feed-forward network: exactly an encoder layer, run with a causal mask.

The LM head and weight tying

After the final LayerNorm each position holds a vector of 768 numbers. A linear map to 50257 numbers gives one logit per vocabulary entry; softmax turns them into the next-token distribution. That map has shape [50257, 768], the same as the token table, and GPT-2 reuses the table's weights for it (weight tying): the vector that represents a token on the way in is the vector it is scored against on the way out. It saves 38,597,376 parameters, 23.7% of what the untied model would have. In PyTorch it is one line, head.weight = tok.weight, because Embedding and Linear store weights with the same layout.

Training versus generation

Training uses teacher forcing. Take a sequence, feed tokens 1 … T−1 as inputs, and use tokens 2 … T as targets. For [a, b, c, d] the inputs are [a, b, c] and the targets [b, c, d]. The model sees the true previous tokens at every position and all 3 predictions are scored in one pass.

Generation is a loop. Feed the prompt, read the logits at the last position, choose a token (the largest logit for greedy decoding, or a sample, see decoding), append it, and repeat. Each new token needs the model run again. Without a cache the work at step i covers the whole prefix; for a 10-token prompt and 20 new tokens that is 390 token positions processed, against 29 when each position's keys and values are remembered (the KV cache).

Worked example

Implementation

Two ways to build the stack. The short one reuses PyTorch's layer, and works because an encoder layer run with a causal mask is a decoder-only block:

import torch
from torch import nn

class GPT(nn.Module):
    def __init__(self, vocab=50257, context=1024, d=768, layers=12, heads=12):
        super().__init__()
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(context, d)
        block = nn.TransformerEncoderLayer(
            d, heads, dim_feedforward=4 * d, dropout=0.1,
            activation="gelu", batch_first=True, norm_first=True,
        )
        self.blocks = nn.TransformerEncoder(
            block, layers, norm=nn.LayerNorm(d), enable_nested_tensor=False
        )
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok.weight          # weight tying
        nn.init.normal_(self.tok.weight, std=0.02)
        nn.init.normal_(self.pos.weight, std=0.02)

    def forward(self, ids):                          # ids: [B, T]
        T = ids.size(1)
        x = self.tok(ids) + self.pos(torch.arange(T, device=ids.device))
        mask = nn.Transformer.generate_square_subsequent_mask(T, device=ids.device)
        x = self.blocks(x, mask=mask, is_causal=True)
        return self.head(x)                          # [B, T, vocab]

model = GPT()
print(sum(p.numel() for p in model.parameters()))  # tied weights are counted once

ids = torch.randint(0, 50257, (2, 16))
logits = model(ids)
loss = nn.functional.cross_entropy(
    logits[:, :-1].reshape(-1, 50257), ids[:, 1:].reshape(-1)
)

@torch.no_grad()
def generate(model, ids, new_tokens):
    model.eval()
    for _ in range(new_tokens):
        logits = model(ids[:, -1024:])               # crop to the context
        ids = torch.cat([ids, logits[:, -1].argmax(-1, keepdim=True)], dim=1)
    return ids

The details that matter: norm_first=True selects the pre-norm branch; batch_first=True makes inputs [B, T, d]; the encoder passes the mask to every layer and, if you give is_causal=None, works out for itself whether the mask is causal (_detect_is_causal_mask), while a wrong is_causal=True hint on a non-causal mask is documented as unpredictable. enable_nested_tensor=False silences a warning: the nested-tensor fast path is disabled for norm_first anyway (line 378). The generate loop recomputes the whole prefix each step, which is correct but wasteful; real implementations cache keys and values.

The longer way builds the block from primitives, replacing the layer with a few lines you control:

class Block(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
        self.heads = heads

    def forward(self, x):                            # x: [B, T, d]
        B, T, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=-1)
        q, k, v = (t.view(B, T, self.heads, -1).transpose(1, 2) for t in (q, k, v))
        a = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, d))
        return x + self.ff(self.ln2(x))

Here scaled_dot_product_attention applies the causal mask itself (the reference code in its docstring builds the lower-triangular mask, adds −∞ elsewhere, softmaxes and multiplies by V), and raises an error if you also pass attn_mask.

Tradeoffs

• Decoder-only versus encoder-decoder or encoder-only. A single causal stack handles any task phrased as text continuation, trains on raw text with no labels, and scales simply. Bidirectional models (encoder-decoder, encoder-only) see the whole input and often win at classification and translation for a given size, but need paired data or a masking objective and can't generate freely.
• The context window is a budget. Attention compares every pair of positions, so cost grows with T²; a learned position table can't go past 1024. Later models replace learned positions with rotary or relative schemes and use grouped-query attention to shrink the KV cache; the Llama 3 paper describes one such recipe (RMSNorm, SwiGLU, rotary positions, grouped-query attention). Those variants are planned for the KV-cache page.
• Tying the head. It removes 23.7% of the untied parameters here and often costs little accuracy, but it forces the input and output spaces to agree. With a small model and a big vocabulary it is a large saving; with a huge model it is a rounding error, which is why larger models often untie it.
• Generation is sequential. Training parallelizes over T; generating T tokens takes T forward passes however many GPUs you own. This is why serving cost is dominated by memory traffic, not arithmetic, and why caching, batching and speculative decoding exist.
• Pitfalls. Exposure bias: the model is trained on true prefixes but generates from its own, so small errors compound. Forgetting to shift targets by one trains the model to copy its input. Padding tokens must be excluded from the loss with ignore_index. And a mask that points the wrong way silently lets the model read the future, giving a suspiciously low training loss and useless generation.

Check yourself

Connections

References

• Radford et al. Language Models are Unsupervised Multitask Learners (GPT-2). OpenAI, 2019 (PDF).
• Radford, Narasimhan, Salimans and Sutskever. Improving Language Understanding by Generative Pre-Training (GPT-1). OpenAI, 2018.
• Brown et al. Language Models are Few-Shot Learners (GPT-3). NeurIPS, 2020 (arXiv 2005.14165).
• Vaswani et al. Attention Is All You Need. NeurIPS, 2017 (arXiv 1706.03762).
• Press and Wolf. Using the Output Embedding to Improve Language Models. EACL, 2017 (arXiv 1608.05859; weight tying).
• Xiong et al. On Layer Normalization in the Transformer Architecture. ICML, 2020 (arXiv 2002.04745; pre-norm).
• Llama Team. The Llama 3 Herd of Models. 2024 (arXiv 2407.21783).
• PyTorch v2.15.0: transformer.py, functional.py, sparse.py, loss.py.