Next-Token Prediction and the Transformer

Ask a large language model anything and, underneath, it does one thing over and over: it reads the tokens so far and outputs a probability for every token that could come next. Choose one, append it, repeat. Training, perplexity, attention and the KV cache all follow from taking that loop seriously. The GPT-2 report states the starting point in one line.

One classification problem per position

The chain rule of probability is exact for any sequence of tokens:

$p(x_1,\ldots,x_T)=\prod_{t=1}^{T}p(x_t\mid x_{<t}).$

No Markov assumption hides here. An $n$-gram model truncates the history; a transformer conditions on the whole prefix that fits in its context window. Language modelling is therefore $T$ classification problems over the vocabulary $\mathcal V$, one per position. From the prefix $x_{<t}$ the network computes logits $z_t\in\mathbb R^{|\mathcal V|}$, and the softmax turns them into a distribution:

$p_\theta(x_t=i\mid x_{<t})=\operatorname{softmax}(z_t)_i=\frac{e^{z_{t,i}}}{\sum_{j}e^{z_{t,j}}}.$

In code, the output at position $t-1$ is the prediction for token $t$, which is why the targets are the inputs shifted left by one. Adding the same constant $c$ to every logit changes nothing, because $e^{z_i+c}/\sum_je^{z_j+c}=e^{c}e^{z_i}/(e^{c}\sum_je^{z_j})$. Only differences between logits matter, so implementations subtract the largest logit before exponentiating and nothing can overflow:

def log_softmax(z):
    m = z.max(axis=-1, keepdims=True)  # shift invariance: subtract the max
    return z - m - np.log(np.exp(z - m).sum(axis=-1, keepdims=True))

Cross-entropy and teacher forcing

Training maximises the likelihood of real text, which is the same as minimising the mean negative log-likelihood, the cross-entropy loss:

$\mathcal L(\theta)=-\frac1T\sum_{t=1}^{T}\log p_\theta(x_t\mid x_{<t}).$

For one position with target $y$, the loss is $-z_y+\log\sum_je^{z_j}$, and its gradient is $\partial\ell/\partial z=\operatorname{softmax}(z)-e_y$: predicted probabilities minus the one-hot target. Each wrong token's logit is pushed down in proportion to the probability it took.

During training the model conditions on the true prefix, never on its own guesses. This is teacher forcing, and together with the causal mask below it lets one forward pass score all $T$ positions in parallel. At generation time the model must condition on its own samples instead, so early mistakes can compound (sometimes called exposure bias).

A useful check: a freshly initialised model predicts nearly uniformly, so its loss should start near $\ln|\mathcal V|$. A first logged loss far above that points at the initialisation or the loss code.

Perplexity, bits and compression

Perplexity is the exponential of the cross-entropy, the reciprocal of the geometric-mean probability of the observed tokens:

$\operatorname{PPL}=e^{\mathcal L}=\Big(\prod_{t=1}^{T}p_\theta(x_t\mid x_{<t})\Big)^{-1/T}.$

A model that gives every token probability $1/|\mathcal V|$ has $\mathcal L=\ln|\mathcal V|$, so its perplexity is exactly $|\mathcal V|$: perplexity is the effective number of tokens the model is choosing between, and a perfect model scores 1. Dividing $\mathcal L$ by $\ln2$ gives bits per token, and the bits are literal. An arithmetic coder driven by the model's probabilities encodes a text in about $\sum_t-\log_2p_\theta(x_t\mid x_{<t})$ bits, within a couple of bits for the whole message, so a better predictor is a better compressor (Delétang et al., 2024).

Per-token perplexity depends on the tokeniser: splitting the same text into fewer, longer tokens raises it without changing how well the text is predicted. Compare models on the same text in bits per byte, with the same context length, and remember that perplexity measures prediction, not helpfulness or truth.

The decoder-only transformer

In most current LLMs, the network that computes $z_t$ is a stack of identical blocks:

tokens x_1 … x_T
  → embedding lookup: one row of a |V| × d table per token (chamber 1)
  → positional information (learned vectors, or RoPE inside attention)
  → L blocks, each:
        h ← h + Attention(Norm(h))   mixes information across positions
        h ← h + MLP(Norm(h))         processes each position on its own
  → final Norm → unembedding (d × |V|) → logits z_t

Attention is the only place where positions exchange information. Residual connections let each block add a correction to the stream rather than replace it, and normalising before each sublayer (pre-norm, as in GPT-2; many newer models use RMSNorm) keeps deep stacks trainable. Apart from the mask, attention scores ignore token order, so position is injected explicitly: GPT-2 adds a learned vector per position, while rotary embeddings, RoPE (Su et al., 2021), rotate each query and key by an angle proportional to its position so that $q\cdot k$ depends on the relative offset. Open-weight families such as Llama, Mistral and Qwen use RoPE.

Each block holds about $4d^2$ parameters in attention and $8d^2$ in an MLP with a $4d$ hidden layer, so $N\approx12Ld^2$ without embeddings (Kaplan et al., 2020). For GPT-2 small, $L=12$, $d=768$ and $|\mathcal V|=50{,}257$ give $12\cdot12\cdot768^2+50{,}257\cdot768\approx123.5$ million; position embeddings, biases and norms bring it to the 124 million of the released checkpoint.

Scaled dot-product attention, step by step

Let $h_i\in\mathbb R^d$ be the input at position $i$. One attention head projects it three ways: a query $q_i=W_Qh_i$ (what this position looks for), a key $k_j=W_Kh_j$ (what position $j$ can be found by) and a value $v_j=W_Vh_j$ (what position $j$ hands over), with $q,k\in\mathbb R^{d_k}$.

1. Score each position: $s_{ij}=q_i\cdot k_j/\sqrt{d_k}$.
2. Mask the future: add $M_{ij}=-\infty$ for $j>i$ (and 0 otherwise), so position $i$ cannot see the tokens it must predict.
3. Normalise: $a_{ij}=e^{s_{ij}+M_{ij}}/\sum_le^{s_{il}+M_{il}}$. Masked entries get $e^{-\infty}=0$, exactly.
4. Average the values: $o_i=\sum_{j\le i}a_{ij}v_j$.

Stacking the rows into matrices gives the whole layer in one line:

$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\Big(\frac{QK^\top}{\sqrt{d_k}}+M\Big)V.$

def causal_attention(Q, K, V):
    T, d_k = Q.shape
    S = Q @ K.T / np.sqrt(d_k)              # (T, T) scores
    S[np.triu_indices(T, k=1)] = -np.inf    # mask j > i
    A = np.exp(S - S.max(axis=1, keepdims=True))
    A /= A.sum(axis=1, keepdims=True)       # each row sums to 1
    return A @ V

Three facts are worth knowing cold. The weights in each row are nonnegative and sum to 1, so every output lies in the convex hull of the visible values: attention blends, it never extrapolates. If the components of $q$ and $k$ are independent with mean 0 and variance 1, then $q\cdot k$ has variance $d_k$; without the division, wide heads would push the softmax into nearly one-hot regions with tiny gradients. And multiplying a query by $c$ multiplies every score by $c$, which is a softmax temperature of $1/c$ in disguise. Temperature returns, applied to the output logits, in the chamber on sampling.

Multi-head attention runs $h$ heads in parallel, each with its own projections to $d_k=d/h$ dimensions, concatenates their outputs and mixes them with $W_O$, at about the cost of one full-width head. Heads specialise: one may attend to the previous token, another may find an earlier copy of the current token and look at what followed it (an induction head). Grouped-query attention (Ainslie et al., 2023) lets several query heads share one key–value head, which shrinks the cache below.

Generation: prefill, decode and the KV cache

Generation runs the opening loop: compute logits at the last position, choose a token (greedily, or by sampling with temperature, top-k or top-p, which is the subject of the sampling chamber), append it, repeat. It has two phases. Prefill processes the whole prompt in one parallel pass. Decode then produces one token per forward pass, and each new query must attend to the keys and values of every earlier position.

Because of the causal mask, earlier positions never see later tokens, so their keys and values never change. The model therefore stores them. The KV cache holds $k_j$ and $v_j$ for every layer, head and position so far; a decode step computes $q,k,v$ for the new token only, appends its $k,v$ and attends over the cache. Without it, every step would reprocess the whole prefix: generating $G$ tokens after a prompt of $P$ would take about $GP+G^2/2$ token passes instead of $P+G$.

The $2N$ comes from one multiply and one add per parameter; Kaplan et al. write the whole forward pass as $C_{\text{forward}}\approx2N+2n_{\text{layer}}n_{\text{ctx}}d_{\text{model}}$ FLOPs per token. Over a context of length $T$ the score matrices have $T^2$ entries per head and layer, so attention's total cost is quadratic in $T$. Kernels such as FlashAttention (Dao et al., 2022) avoid storing those matrices but still do the quadratic arithmetic. The cache costs $2\cdot L\cdot n_{\text{kv}}\cdot d_{\text{head}}$ numbers per token; the chamber on running models turns that into gigabytes.

Read beyond

Read the equation in context

Section 3.2.1 defines Equation (1) with queries, keys and values packed into matrices, so all outputs come from two matrix products. Read three passages around it: the footnote that justifies $\sqrt{d_k}$ with the variance argument above; Section 3.2.2, where $h=8$ heads of $d_k=d_{\text{model}}/h=64$ dimensions replace one wide head; and Section 3.2.3, where the decoder prevents leftward information flow by setting the scores of illegal connections to $-\infty$ inside the softmax. The paper's model is an encoder–decoder for translation. GPT-style models keep only the masked decoder stack, drop the cross-attention to an encoder, and train it with the factorisation from the opening glimpse.

For the whole pipeline written and trained in code, from a bigram baseline to a small GPT with causal self-attention:

Your turn

Six toy tokens have fixed 2-D keys and values. Steer one query until a single token takes almost all the attention, then shorten or turn it until the attention spreads out, and watch the output stay inside the hull of the visible values. The panel underneath shows what goes wrong without the $\sqrt{d}$.

Coding problems