Next Token Prediction — Module 11.1

How This Lesson Fits the Module & Volume

You know a language model emits a probability distribution over the vocabulary. Next-token prediction is the concrete learning problem and the generation loop: given tokens so far, predict the following token—then append it and repeat.

This objective is why causal masking from Vol. 10 matters, why teacher forcing works in training, and why inference is autoregressive. Later lectures on sampling and search only change how you pick from the predicted distribution.

Learning Objectives

By the end of this lesson, students should be able to:

Formulate next-token prediction as supervised classification over the vocabulary at each position.
Explain teacher forcing: train on gold prefixes, not on the model’s own samples.
Shift labels by one position (input_ids vs labels) in a causal LM batch.
Implement a tiny PyTorch training step with cross-entropy on shifted logits.
Describe the generate loop: predict → choose token → append → stop.
Relate context length limits to the upcoming context window lecture.

Definition

Next-token prediction is the task of estimating P(x_t | x₁,…,x_t-1) for each position t. During training, the model sees the true prefix and is scored against the true x_t. During generation, the chosen token becomes part of the next prefix.

Training vs Generation

Train

Full sequence in parallel (causal mask).

Loss

CE at every position vs gold next token.

Generate

One new token at a time.

Stop

EOS, max length, or stop string.

Aspect	Training (teacher forcing)	Generation (inference)
Prefix source	Ground-truth tokens	Model’s own previous outputs
Parallelism	All positions in one forward (masked)	Mostly sequential (KV cache helps)
Objective	Minimize cross-entropy	Produce a coherent continuation
Error compounding	No (gold prefixes)	Yes—mistakes enter the context

Label Shift

For input tokens [x₁, x₂, x₃, x₄], the model at positions 1–3 predicts x₂, x₃, x₄. In code this is usually logits[..., :-1, :] vs labels[..., 1:], with ignore index on padding.

Why It Scales

Unlimited self-supervised text.
No hand-labeled classes needed.
One head: vocab-sized classifier.

Why It’s Hard

Huge vocabulary (|V| ~ 32k–256k).
Long-range dependencies.
Exposure bias at generation time.

Decoding Choices

Greedy / beam search.
Temperature sampling.
Top-k / top-p truncation.

Code: Toy Causal LM Step

import torch
from torch import nn
import torch.nn.functional as F

class TinyLM(nn.Module):
    def __init__(self, vocab_size=100, d=64, n_layers=2):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, d)
        layer = nn.TransformerEncoderLayer(d, nhead=4, batch_first=True)
        self.stack = nn.TransformerEncoder(layer, num_layers=n_layers)
        self.lm_head = nn.Linear(d, vocab_size, bias=False)

    def forward(self, idx):
        # Causal mask: True = ignore (PyTorch additive/bool conventions vary by API)
        T = idx.size(1)
        causal = torch.triu(torch.ones(T, T, dtype=torch.bool, device=idx.device), 1)
        x = self.emb(idx)
        x = self.stack(x, mask=causal)
        return self.lm_head(x)

model = TinyLM()
batch = torch.randint(0, 100, (2, 16))          # (B, T)
logits = model(batch)                           # (B, T, V)
loss = F.cross_entropy(
    logits[:, :-1].reshape(-1, logits.size(-1)),
    batch[:, 1:].reshape(-1),
)
print(loss.item())

# Greedy one-step generation from a prefix
prefix = batch[:1, :8]
with torch.no_grad():
    next_id = model(prefix)[:, -1].argmax(-1)
print("next token id:", int(next_id))

Strengths

Simple, universal objective.
Dense supervision (every position).
Transfers to many downstream tasks.

Tradeoffs

Teacher forcing ≠ free-running generation.
Left-to-right bias; bidirectional tasks need other LMs.
Compute grows with context length.

Common Misconception

“At training time the model generates token-by-token like ChatGPT.” Training uses teacher forcing: the full gold sequence is fed under a causal mask, and losses at all positions are computed in one forward/backward. Autoregressive generation is an inference procedure. Confusing the two makes KV-cache and latency discussions impossible to follow.

Knowledge Check

Short Answer: What does the model predict at position t during training? Answer: The distribution for token x_t+1 (the next token) given x_≤t.
True/False: Teacher forcing feeds the model its own previous predictions during training. Answer: False—it feeds ground-truth tokens.
Multiple Choice: Typical label alignment uses: (a) logits[:-1] vs labels[1:], (b) logits vs labels identical indices with no shift, (c) only the last token. Answer: (a) (with causal LM conventions).
Short Answer: Name one stopping criterion for generation. Answer: EOS token, max new tokens, or a stop string (any one).
True/False: Causal masking lets every position attend to future tokens in training. Answer: False.
Multiple Choice: Exposure bias refers to: (a) dropout, (b) train-on-gold vs generate-on-model prefixes mismatch, (c) FP16 underflow. Answer: (b).
Short Answer: Why is next-token prediction self-supervised? Answer: Targets are the next tokens already present in raw text—no external labels required.
Short Answer: What loss is standard for this task? Answer: Cross-entropy (negative log-likelihood) over the vocabulary.
Multiple Choice: Generation is primarily: (a) fully parallel over future tokens, (b) sequential token-by-token, (c) a single softmax over sentences. Answer: (b).
True/False: Dense supervision means every non-padded position contributes a loss term. Answer: True.

Key Takeaways

Next-token prediction is vocabulary-sized classification at each position under a causal mask.
Training uses teacher forcing and parallel CE; generation appends tokens autoregressively.
Label shift (logits[:-1] vs labels[1:]) is the standard wiring.
Decoding strategies only change how we pick from P(next token | context).
Next: Context Window—how much prefix the model can see.

Trainer’s Guide

Hands-on idea: On a whiteboard, write a 5-token sentence and have students fill the (input, target) pairs for each position.

Discussion prompt: How does exposure bias show up in long story generation? (Early wrong entity name poisons later pronouns.)

Recap: Next-token prediction is the LM’s train and generate objective—classify the next vocab id, then loop at inference. Continue with Context Window.