20  Output Layer & Decoding

TipTL;DR

A language model’s final layer projects each hidden state to logits over the vocabulary; a softmax turns logits into a probability distribution, and the model is trained by categorical cross-entropy against the true next token: \[p = \operatorname{softmax}(h\,W_{\text{out}}), \qquad \mathcal{L} = -\log p_{\text{true}}.\] But a distribution is not text. Decoding is the separate algorithm that turns per-step distributions into a sequence — greedy, beam search, or sampling (temperature, top-k, nucleus). The decoding choice, not just the model, decides whether output is accurate, diverse, or degenerate.

Depends on: BERT & GPT · Information Theory

20.1 Why this matters

Every generative chapter so far ends at “predict the next token” and stops. Two things in that phrase were deferred and are owed a home here:

  1. The output layer itself — how a hidden vector becomes a probability over tens of thousands of tokens, and the softmax + cross-entropy pair that the information-theory chapter introduced for the binary case and promised to generalise to multi-class.
  2. Decoding — given the model’s distribution at each step, how do you actually pick tokens? This is assumed everywhere and explained nowhere, yet it is the difference between fluent text and repetitive nonsense from the same model.

Decoding is pure inference-time algorithm: no parameters, no training. It is the cheapest, most overlooked lever on output quality.

20.2 The mechanism

20.2.1 The output head: logits → softmax → cross-entropy

The final hidden state \(h \in \mathbb{R}^{d_{\text{model}}}\) is projected to a vector of logits over the vocabulary \(V\), then normalised:

\[ z = h\,W_{\text{out}}, \quad W_{\text{out}} \in \mathbb{R}^{d_{\text{model}}\times |V|}, \qquad p_i = \frac{e^{z_i}}{\sum_{j} e^{z_j}}. \]

This is the multi-class softmax — the generalisation of the sigmoid from the binary case. Training minimises categorical cross-entropy, which for a one-hot target reduces to the negative log-probability of the correct token:

\[ \mathcal{L} = -\sum_i y_i \log p_i = -\log p_{\text{true}}. \]

Minimising this is exactly maximising the likelihood the information-theory chapter framed as minimising the KL divergence between the data and the model — now over \(|V|\) classes instead of two. The exponential of the per-token cross-entropy is perplexity, the standard LM metric.

NoteWeight tying

The output projection \(W_{\text{out}}\) and the input embedding matrix have the same shape (\(|V| \times d_{\text{model}}\)) and similar roles. Tying them — using one matrix for both (Press & Wolf, 2017) — saves a large block of parameters (the vocabulary embedding is often the single biggest tensor in a small model) and usually improves quality. Standard in GPT-2 onward.

20.2.2 Decoding: turning distributions into text

At each step the model gives a distribution \(p\) over the next token. A decoding strategy chooses one. The split is deterministic (maximise probability) versus stochastic (sample for diversity).

Deterministic Stochastic (sampling) Greedy — argmax each step fast, repetitive, can dead-end Beam search — keep top-b paths better for MT; bland for open-ended Temperature — sharpen/flatten p τ<1 safer, τ>1 more random Top-k / Top-p — truncate the tail sample from the head only Maximising probability gives safe but dull text; sampling gives diversity — the art is truncating the tail. Same model, different decoder → fluent essay or repetitive loop.
  • Greedy — take the argmax at each step. Fast and deterministic, but myopic: it commits to a locally-best token and often loops (“the the the”).
  • Beam search — keep the \(b\) highest-probability partial sequences and expand all of them, returning the best complete one. Strong for tasks with a single correct answer (translation, summarisation); for open-ended generation it produces bland, repetitive text because high-probability sequences are dull.
  • Temperature — divide logits by \(\tau\) before softmax. \(\tau<1\) sharpens toward the mode (more conservative); \(\tau>1\) flattens (more random); \(\tau\to0\) recovers greedy.
  • Top-k sampling (Fan et al., 2018) — restrict sampling to the \(k\) most likely tokens, renormalise, sample. Cuts off the unreliable tail.
  • Nucleus (top-p) sampling (Holtzman et al., 2020) — restrict to the smallest set of tokens whose cumulative probability exceeds \(p\). Unlike top-k, the candidate set grows and shrinks with the model’s confidence.
WarningPitfall: likelihood maximisation degenerates

The most striking result here: always picking high-probability tokens produces worse text. Holtzman et al. (Holtzman et al., 2020) showed that maximisation (greedy/beam) drives models into bland, repetitive loops — human text is not the highest-probability continuation; it constantly takes lower-probability turns. This is why open-ended generation samples (nucleus) and only closed-ended tasks (translation) use beam search. “Most likely” and “good” diverge.

TipGoing deeper
  • The binary softmax/cross-entropy this generalises — Information Theory.
  • Making autoregressive decoding faster (KV-cache, speculative decoding) — KV-Cache.

20.2.3 Minimal sketch

The output softmax with temperature, plus greedy / top-k / nucleus selection.

import numpy as np

rng = np.random.default_rng(0)

def softmax(z, temp=1.0):
    z = z / temp
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()

logits = rng.normal(size=8)                    # one step's logits over 8 tokens

print("greedy:", int(np.argmax(logits)))

def top_k(p, k):
    idx = np.argsort(p)[-k:]
    q = np.zeros_like(p); q[idx] = p[idx]
    return q / q.sum()

def nucleus(p, top_p):
    order = np.argsort(p)[::-1]
    cum = np.cumsum(p[order])
    keep = order[:np.searchsorted(cum, top_p) + 1]
    q = np.zeros_like(p); q[keep] = p[keep]
    return q / q.sum()

p = softmax(logits, temp=0.8)
print("top-k=3 support :", np.flatnonzero(top_k(p, 3) > 0))
print("nucleus p=.9 sup:", np.flatnonzero(nucleus(p, 0.9) > 0))
greedy: 6
top-k=3 support : [2 6 7]
nucleus p=.9 sup: [0 2 3 5 6 7]

20.3 Application & impact

Choice Best for Failure mode
Greedy quick, deterministic output repetition, dead-ends
Beam search translation, summarisation bland in open-ended tasks
Temperature global diversity dial too high → incoherent
Top-k simple tail-cutting fixed k ignores confidence
Nucleus (top-p) open-ended generation the modern default
Weight tying parameter savings — (essentially free win)
  • Decoding is a product surface: the temperature and top_p knobs every LLM API exposes are exactly this chapter. Tuning them changes behaviour with no retraining.
  • Nucleus sampling is the de facto default for chat and creative generation; beam search survives in machine translation.
  • The output head is also an efficiency target — the vocabulary projection is huge, motivating weight tying and (in some models) factorised softmaxes.
NoteKey takeaway

The model outputs a distribution (logits → softmax, trained by cross-entropy); decoding is the separate algorithm that turns distributions into text. Maximising probability (greedy/beam) is right for closed-ended tasks but degenerates on open-ended ones, where sampling with a truncated tail (nucleus) wins. Same weights, different decoder, completely different output.

→ Next: Scaling Laws & Emergence