Inside the model · Unembed & Sampling

Unembed & Sampling

LM head · tied to WE · real run

shown
real GPT-2
dmodel
768 8 drawn
vocab
50,257
LM head
WEᵀ, tied

All steps

  1. Final LayerNorm

    Only the last position predicts the next token. Its vector, after all 12 blocks, goes through the final LayerNorm lnf. The numbers are from a real GPT-2 small run.

    [1 × 768]

  2. Logits

    The vector is dotted with every row of the embedding matrix: the unembedding WU, or LM head, is WEᵀ in GPT-2. That gives one logit per token. The grid draws 8 of the 768 dimensions for the 6 likeliest tokens; each logit sums all 768.

    [1×768]·[768×50,257] · 77M FLOPs

  3. Temperature

    Logits are divided by the temperature T before softmax. Only the gaps between them matter: lower T stretches the gaps and sharpens the distribution, higher T flattens it.

    z / T

  4. Softmax

    Softmax turns the scaled logits into probabilities over all 50,257 tokens, summing to 1. At low T the top few take most of the mass; raising T shifts it into the long tail.

    p = softmax(z / T)

  5. Sample

    A strategy picks the next token: greedy takes the likeliest, sampling lays the probabilities end to end and a random u picks one, and top-k or top-p drop the tail first. Generation repeats until <|endoftext|> or a length limit.

    sample · greedy · top-k · top-p

Code

x = self.transformer.ln_f(x)
logits = self.lm_head(x[:, [-1], :])      # lm_head.weight is wte.weight (tied)
logits = logits[:, -1, :] / temperature
if top_k is not None:
    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
    logits[logits < v[:, [-1]]] = -float("Inf")
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)   # greedy: probs.argmax(-1)
idx = torch.cat((idx, idx_next), dim=1)              # append, run again

Go deeper