Lesson 13

Decoder-only language modeling and generation

Why GPT uses next-token prediction

A language model can learn from ordinary text without a human label for every sentence. The text supplies its own targets. For token IDs [h,e,l,l,o], input is [h,e,l,l] and target is [e,l,l,o]. Every position predicts the token one step ahead. This simple objective can teach spelling, syntax, patterns, and some facts when the model and dataset are large.

Predict one row: at input position 1, the model reads [h,e]. Its target is the next token l. The later l,o tokens must stay hidden from that prediction.

GPT is called decoder-only because its transformer stack uses causal self-attention. A position can read itself and earlier positions, but not later positions. During generation there are no later tokens yet. Applying the same restriction during training prevents answer leakage.

Key terms and shapes

Input IDs have shape (B,T). Token and position embeddings produce (B,T,D). Transformer blocks preserve that shape. The language-model head returns (B,T,V). For cross-entropy, flatten batch and time to logits (B×T,V) and targets (B×T).

Worked example: a tiny GPT module

This model uses a PyTorch encoder layer with an explicit causal mask. Architecturally, the layer supplies self-attention and a feed-forward network; the causal mask gives it decoder-only behavior for this small learning example.

import torch
from torch import nn

class MiniGPT(nn.Module):
    def __init__(self, vocabulary_size, maximum_length,
                 model_size=16, heads=2):
        super().__init__()
        self.maximum_length = maximum_length
        self.token_embedding = nn.Embedding(vocabulary_size, model_size)
        self.position_embedding = nn.Embedding(maximum_length, model_size)
        self.block = nn.TransformerEncoderLayer(
            d_model=model_size,
            nhead=heads,
            dim_feedforward=32,
            dropout=0.0,
            batch_first=True,
        )
        self.language_head = nn.Linear(model_size, vocabulary_size)

    def forward(self, token_ids):
        batch_size, length = token_ids.shape
        positions = torch.arange(length, device=token_ids.device)
        x = self.token_embedding(token_ids)
        x = x + self.position_embedding(positions)
        mask = nn.Transformer.generate_square_subsequent_mask(
            length, device=token_ids.device
        )
        x = self.block(x, src_mask=mask)
        return self.language_head(x)

torch.manual_seed(13)
model = MiniGPT(vocabulary_size=3, maximum_length=6)
x = torch.tensor([[0, 1, 2, 0]], dtype=torch.long)
logits = model(x)
print(tuple(logits.shape))  # Expected: (1, 4, 3)

The position tensor has no batch axis, so PyTorch broadcasts its vectors across every sequence in the batch. The maximum length is a hard limit of this learned position table. A prompt longer than that must be shortened or the model design must change.

One training step

sequence = torch.tensor([[0, 1, 2, 0, 1]], dtype=torch.long)
context = sequence[:, :-1]
targets = sequence[:, 1:]
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01)

optimizer.zero_grad()
logits = model(context)
loss = nn.functional.cross_entropy(
    logits.reshape(-1, 3),
    targets.reshape(-1),
)
loss.backward()
optimizer.step()

print(tuple(context.shape))  # Expected: (1, 4)
print(loss.ndim)             # Expected: 0

Every target is an integer class ID. The model receives raw logits, and cross-entropy performs the probability calculation in a numerically stable way. A useful training program repeats this step over many batches and measures validation loss on held-out text.

Generate from a prompt

Generation repeatedly keeps a legal context, predicts the final position, selects one next token, and appends it. Greedy argmax is deterministic. Probability sampling is more varied. Temperature below one sharpens the distribution; above one flattens it.

Training scores every position in parallel because the full target text is available behind a causal mask. Generation has no future text, so it adds one token at a time. Confusing these two modes often causes an off-by-one target or an attempt to generate all new tokens in one forward call.

model.eval()
generated = torch.tensor([[0]], dtype=torch.long)

with torch.no_grad():
    for _ in range(5):
        context = generated[:, -model.maximum_length:]
        next_logits = model(context)[:, -1, :]
        next_id = next_logits.argmax(dim=-1, keepdim=True)
        generated = torch.cat((generated, next_id), dim=1)

print(tuple(generated.shape))  # Expected: (1, 6)

The generated content is not meaningful before adequate training. The code demonstrates the algorithm, not a pretrained model.

Try it

For batch size 4, context length 8, model width 24, and vocabulary size 30, state the shapes after embedding, after the transformer block, and after the language-model head.

Recap

GPT learns by predicting every next token under a causal mask. Token and position embeddings enter transformer blocks, and a language-model head produces vocabulary logits. Training uses shifted targets and cross-entropy. Generation predicts one token, appends it, and repeats. A small implementation teaches the mechanics; useful language behavior requires suitable data, scale, and evaluation.

References: PyTorch transformer encoder layer and causal-mask helper.