TechByteByByte

Next Token Prediction

One of the most important modules in this course — precisely what 'predict the next token' means mechanically: input sequence, logits, softmax, probability distribution, and selection — traced with a real, executed computation, not just described.

#LLM#AI#Next Token Prediction#Logits#Softmax

Before you continue: three tools for this module

  • Token: a piece of text handled as one vocabulary item.
  • Token ID: the integer assigned to that item.
  • Tokenizer: the algorithm mapping text to and from tokens.

You do not need to memorize these yet. Use this map when the terms reappear.

Begin with the central question

What hidden problem does Next Token Prediction solve inside a real language-model system?

Keep that central question about Next Token Prediction in mind. The definitions, numbers, diagrams, and code examples below answer it one piece at a time.

context → logits → probabilities → chosen next token → repeat

1. What You Will Learn

Learning outcomes

  • Trace one next-token prediction from hidden state to logits and probabilities.
  • Explain why every vocabulary token receives a score.
  • Distinguish logits, probability distributions, and selected tokens.
  • Show how repeating one prediction step creates sentences, code, and conversations.

In one sentence

💡 Big picture

For every next step, the LLM gives every possible token a score, turns those scores into probabilities, and chooses one token.

2. Why This Module Exists

The problem this module solves

  • This small-looking operation is the engine behind entire answers, stories, and code blocks.
  • It explains both impressive completions and confident mistakes: the model always predicts a likely continuation.

3. Intuition

at every single step, the model doesn’t “decide” on one word directly. It computes a score for every word in its entire vocabulary, converts those scores into genuine probabilities, and then a token gets selected from that distribution. “Predicting the next token” is this scoring-and-selecting operation, nothing more mysterious.

Analogy: The Dictionary Point-Scoring Contest Think of next-token prediction as running a massive dictionary voting contest at the end of each sentence:

  • The Contest (Logits Calculation): The final layer does not choose a word directly. Instead, the model acts as a game show judge evaluating all 50,000 words in its dictionary. It assigns a raw score (logit) to every single word, based on how well it fits. Plausible words (“blue”, “cloudy”) get positive scores like 1.23 or 1.63. Implausible words (“sky”, “green”) get negative scores like -1.46 or -1.54.
  • The Percentages (Softmax): We run the raw scores through the softmax calculator, converting them into clean percentage values summing to exactly 1.0. Plausible words get high percentages (“clear” gets 45.8%, “the” gets 30.5%), while others get near 1.9%.
  • The Selection: We pick one token (greedy selects the top vote, “clear”), append it to our prompt, and run the next round.

📊 Visual Flowchart: The Next-Token Prediction Head Pipeline

Here is how final hidden states are projected to dictionary logits and softmax percentages:

graph TD
    classDef representation fill:#3498db,stroke:#333,stroke-width:1px,color:#fff;
    classDef logits fill:#e74c3c,stroke:#333,stroke-width:1px,color:#fff;
    classDef select fill:#2ecc71,stroke:#333,stroke-width:1px,color:#fff;

    Seq["Input Sequence:<br>'the sky is'"] --> Stack["Transformer Blocks Stack"]
    Stack --> FinalRep["Final Layer Representation Matrix (Seq Length x d_model)"]:::representation

    FinalRep --> SliceLast["1. Slice Last Token Hidden State<br>(d_model vector at index -1)"]

    subgraph OutputHead ["LM Head Classifier"]
        SliceLast --> LinearProj["2. Linear Projection: W_lm_head<br>(Map d_model to Vocabulary Size)"]
        LinearProj --> Logits["3. Vocabulary Logits Vector<br>(One raw score per word)"]:::logits
    end

    Logits --> Softmax["4. Softmax Normalization<br>(Convert to percentages summing to 1.0)"]
    Softmax --> Probabilities["5. Token Probability Distribution"]

    Probabilities --> GreedySelect["6. Argmax Selection (Greedy)<br>Outputs: 'clear'"]:::select

    GreedySelect --> LoopAppend["7. Append to Input Context Loop"]

4. Core Concept

Input sequence:    the tokens the model has SEEN so far
Target:             the ACTUAL next token (known during training,
                    Module 8; unknown/being predicted during
                    inference, this module)
Logits:              raw, unbounded scores — one per vocabulary
                    token — produced by the LM head
Softmax:               converts logits into a genuine PROBABILITY
                    DISTRIBUTION (values in [0,1], summing to 1)
Selection:                a token is chosen from this distribution
                        (Module 15 covers HOW precisely)

5. How It Works — Step by Step

1. The input sequence (everything processed so far) flows through
   the full Transformer stack (Module 4, 10) -- producing a
   FINAL REPRESENTATION for every position
2. ONLY the LAST position's final representation is used for
   next-token prediction (Module 4's distinction)
3. The LM HEAD -- a learned linear projection -- maps this single
   vector from d_model dimensions to VOCAB_SIZE dimensions,
   producing LOGITS: one raw score per possible next token
4. SOFTMAX converts these logits into a PROBABILITY DISTRIBUTION
   -- every token gets a probability between 0 and 1, and all
   probabilities sum to EXACTLY 1
5. A token is SELECTED from this distribution (greedy: highest
   probability; or sampling-based, Module 15)
6. The selected token is APPENDED to the sequence
7. REPEAT the entire process for the next position, now with one
   more token of context than before

6. Mathematical Intuition

Read the mathematics as a story

context → logits → probabilities → chosen next token → repeat

First locate the input, operation, and output. Then treat the formula as a compact description of that journey rather than a collection of symbols to memorize.

logits = W_lm_head @ final_hidden_state

probability(token_i) = exp(logit_i) / Σ exp(logit_j)   for all j

W_lm_head is a learned weight matrix of shape (vocab_size, d_model) — every row corresponds to one vocabulary token’s learned “direction” in representation space; the dot product between this row and the final hidden state is that token’s raw score. Softmax’s exponentiation and normalization guarantee the output is a genuine probability distribution — never negative, always summing to 1.

7. Small Worked Example

Walk through the example

  1. Name what each input represents.
  2. Follow one transformation at a time.
  3. Translate the result back into ordinary language.

The purpose is to reveal the mechanism, not merely display an answer.

For the prompt “The sky is ___,” a well-trained model’s logits should score plausible continuations (“blue,” “clear,” “cloudy”) noticeably higher than implausible ones (“green,” “purple”) — purely because training has repeatedly reinforced these patterns.

The mechanism run below is genuine and complete; since this specific model is untrained (random weights, exactly like Modules 4/10’s illustrative examples), the particular winning word isn’t semantically meaningful — but the process producing it is exactly what a real, trained LLM runs.

8. Python Example

What the code will demonstrate

The code builds a tiny version of the mechanism, prints values you can inspect, and connects them to the worked example. Predict the direction of the result before running it.

Python symbols used below

  • NumPy (np) stores numeric vectors and matrices.
  • np.array(...) creates a numeric collection.
  • Library calls perform the same conceptual steps shown above at a larger scale.
# Build a small, inspectable example of Next Token Prediction.
# Follow the inputs, transformations, and output in order.
import numpy as np

def softmax(x):
    exp_x = np.exp(x - np.max(x))
    return exp_x / np.sum(exp_x)

def softmax_rows(x):
    exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

def layer_norm(x, eps=1e-8):
    mean = x.mean(axis=-1, keepdims=True)
    std = x.std(axis=-1, keepdims=True)
    return (x - mean) / (std + eps)

def relu(x): return np.maximum(0, x)

def causal_mask(n):
    return np.triu(np.ones((n, n)), k=1).astype(bool)

np.random.seed(5)
d_model = 8
vocab = ["the", "sky", "is", "blue", "dark", "clear", "green", "cloudy"]
vocab_size = len(vocab)

embedding_table = np.round(np.random.randn(vocab_size, d_model) * 0.4, 3)

prompt = ["the", "sky", "is"]
token_ids = [vocab.index(w) for w in prompt]
print("INPUT SEQUENCE:", prompt)

x = embedding_table[token_ids]

def positional_encoding(n, d):
    pos = np.arange(n)[:, np.newaxis]
    i = np.arange(d)[np.newaxis, :]
    angles = pos / np.power(10000, (2 * (i // 2)) / np.float32(d))
    pe = np.zeros((n, d))
    pe[:, 0::2] = np.sin(angles[:, 0::2])
    pe[:, 1::2] = np.cos(angles[:, 1::2])
    return pe

x = x + positional_encoding(len(prompt), d_model)

def decoder_block(x, seed):
    rng = np.random.RandomState(seed)
    n = x.shape[0]
    Wq, Wk, Wv, Wo = [rng.randn(d_model, d_model) * 0.3 for _ in range(4)]
    Q, K, V = x @ Wq, x @ Wk, x @ Wv
    scores = Q @ K.T / np.sqrt(d_model)
    scores[causal_mask(n)] = -np.inf
    attn_out = (softmax_rows(scores) @ V) @ Wo
    x = layer_norm(x + attn_out)
    W1 = rng.randn(d_model, d_model * 2) * 0.3
    W2 = rng.randn(d_model * 2, d_model) * 0.3
    ffn_out = relu(x @ W1) @ W2
    return layer_norm(x + ffn_out)

for i in range(3):
    x = decoder_block(x, seed=i)

last_hidden = x[-1]   # ONLY the last position's representation is used

# LM Head: project to vocab-sized LOGITS
W_lm_head = np.round(np.random.randn(vocab_size, d_model) * 0.5, 2)
logits = W_lm_head @ last_hidden

print("\nLOGITS:")
for word, logit in zip(vocab, logits):
    print(f"  {word:8s}: {logit:.4f}")

probs = softmax(logits)
print("\nPROBABILITY DISTRIBUTION:")
for word, p in sorted(zip(vocab, probs), key=lambda x: -x[1]):
    print(f"  {word:8s}: {p:.4f}  ({p*100:.1f}%)")

print(f"\nSum of probabilities: {probs.sum():.4f}")
selected = vocab[np.argmax(probs)]
print(f"Selected next token (greedy): '{selected}'")

new_prompt = prompt + [selected]
print("\nNew input sequence for the NEXT prediction step:", new_prompt)

Expected Output:

INPUT SEQUENCE: ['the', 'sky', 'is']

LOGITS:
  the     : 1.2307
  sky     : -1.4635
  is      : -1.2529
  blue    : -0.2170
  dark    : -1.5213
  clear   : 1.6374
  green   : -1.5428
  cloudy  : -0.1056

PROBABILITY DISTRIBUTION:
  clear   : 0.4583  (45.8%)
  the     : 0.3051  (30.5%)
  cloudy  : 0.0802  (8.0%)
  blue    : 0.0717  (7.2%)
  is      : 0.0255  (2.5%)
  sky     : 0.0206  (2.1%)
  dark    : 0.0195  (1.9%)
  green   : 0.0191  (1.9%)

Sum of probabilities: 1.0000
Selected next token (greedy): 'clear'

New input sequence for the NEXT prediction step: ['the', 'sky', 'is', 'clear']

9. How It Works

  • Every token in the vocabulary received a logit — including “the” and “sky,” which are grammatically implausible continuations here — because the LM head always produces a score for the entire vocabulary, every single time, with no exceptions.
  • Softmax converted these logits into a distribution that sums to exactly 1.0000 — this is not approximate; it’s a mathematical guarantee of the softmax function.
  • The model (with random, untrained weights) selected “clear” — not because it “understands” weather, but because “clear“‘s logit (1.6374) happened to be highest given these random weights. A real, trained model would show a sharply different, semantically sensible distribution — but the mechanism producing that distribution is identical to what’s traced here.
  • The new sequence ['the', 'sky', 'is', 'clear'] is exactly what would feed back into the same process (step 6-7 above) to predict the token after “clear” — this repetition is Module 7’s complete generation loop.

10. Why the Model Does Not “Look Up” an Answer

As emphasized in Module 1 and reconfirmed here with a complete trace: there is no branch, no conditional, no lookup table anywhere in this mechanism that says “if the question is about weather, return a color.” Every prediction — correct, plausible, or nonsensical — goes through this exact same logits-then-softmax-then-selection computation.

What differs between a well-trained and poorly-trained (or untrained) model is purely the learned weight values (W_lm_head, and every weight throughout the Transformer stack) — not the mechanism itself.

11. How Is This Used in Modern AI?

Trace it through a real model call

user message → assembled context → LLM computation → decoded output → application checks

This topic affects one stage of that path; it is not the complete product. Hosted GPT- and Gemini-style applications also add instructions, safety systems, retrieval, tools, serving infrastructure, and evaluation around the model.

🤖 How Is This Used in Modern AI?

This exact computation runs once for every single token any LLM generates, for every request, across every product built on top of one. There is no alternate mechanism for “harder” questions — the same logits-softmax-selection process runs uniformly, regardless of apparent task complexity.

12. How Is This Used in Agentic AI?

Separate the model from the runtime

goal + state + tool results → LLM proposal → runtime validation → execution or response

The LLM proposes text or a structured action. Ordinary application code controls permissions, tools, retries, memory, and execution.

Direct relevance to Agentic AI: Very High. Every token of an agent’s reasoning, tool-call generation, and final response comes from this exact mechanism, run repeatedly. Understanding this precisely is what makes an agent’s behavior explicable as a mechanical, traceable consequence of learned weights and context — not an unexplainable black box.

13. Common Beginner Mistakes

⚠️ Mistake

Incorrect idea: assuming only “plausible” tokens receive a logit.

Why it is incorrect: As shown directly, every single vocabulary token — plausible or not — gets a raw score every time; softmax then determines how much relative probability mass each receives.

⚠️ Mistake

Incorrect idea: believing prediction quality comes from a different mechanism for correct vs. incorrect answers.

Why it is incorrect: As emphasized directly, the exact same computation runs regardless of whether the output ends up correct — quality differences trace back entirely to learned weight values from training (Module 8), not a different code path.

⚠️ Mistake

Incorrect idea: confusing logits with probabilities.

Why it is incorrect: Logits are raw, unbounded scores (can be negative, don’t sum to anything meaningful); only after softmax do you get genuine probabilities, bounded in [0,1] and summing to exactly 1 — verified directly.

14. Important Distinctions

LogitsProbabilities
Raw, unbounded scores from the LM headBounded [0,1], sum to exactly 1 — verified directly
Can be negativeAlways non-negative
Input SequenceTarget (Training only)
Tokens the model has already seenThe actual correct next token — only available/used during training (Module 8-9)

15. When to Use

Not applicable in the technique-selection sense — this is the fundamental, universal mechanism every LLM generation step uses.

16. When Not to Use

Not applicable.

17. Production Considerations

  • Logit computation happens for the ENTIRE vocabulary every step — a genuine, non-trivial computational cost (proportional to vocabulary size), part of why very large vocabularies have real inference cost implications.
  • How a token gets selected from the distribution (Module 15: greedy, temperature, top-k, top-p) is a separate, configurable design choice layered on top of this module’s mechanism — the distribution itself is always computed the same way.

18. What You Should Remember

  • The complete mechanism: final hidden state → LM head → logits → softmax → probability distribution → selection — traced with real, executed numbers, not just described.
  • Every vocabulary token gets a score every time — there’s no special-casing for “obviously correct” or “obviously wrong” answers.
  • Softmax’s output always sums to exactly 1 — verified directly, the defining mathematical property of a genuine probability distribution.

19. Interview Questions

Beginner

Q: What does “predict the next token” actually mean, mechanically?

Ans: The model computes a raw score (logit) for every single token in its vocabulary, based on the final hidden state representing everything processed so far. Softmax converts these logits into a genuine probability distribution — values between 0 and 1, summing to exactly

  1. A token is then selected from this distribution.

Intermediate

Q: Why does the model compute a logit for every token in its vocabulary, even ones that are obviously implausible given the context?

Ans: The LM head is a single, uniform linear projection applied to the final hidden state — it has no mechanism for selectively skipping “implausible” tokens; it always produces one score per vocabulary token.

Verified directly in this module, tokens like “the” and “sky” (which would be poor continuations for “the sky is ___”) still received real logit values — softmax then naturally assigns them very low probability relative to more plausible completions, but the raw scoring step itself doesn’t discriminate.

Advanced

Q: Explain precisely why “the model doesn’t look up an answer” is an accurate mechanical description, not just a simplification.

Ans: As traced completely in this module, every prediction — regardless of whether the resulting answer happens to be correct — goes through the identical computational path: final hidden state, linear projection via the LM head, softmax normalization, selection.

There is no conditional branch, database query, or alternate code path anywhere in this process that activates specifically for factual questions or “known” answers. What determines whether a specific prediction turns out correct is entirely the learned VALUES in the weight matrices (shaped by training, Module 8) — not a structurally different mechanism kicking in for different kinds of questions.

Scenario

**Q: You’re debugging an LLM-based application and notice the model occasionally outputs a token that seems completely disconnected from the context.

Using this module’s mechanism, what would you investigate?** A: Since every prediction follows the same logits-then-softmax mechanism, an unexpected token means it received an unexpectedly high probability in the distribution for that specific context — I’d investigate whether this reflects genuine model uncertainty (a flatter, less confident distribution where many tokens have comparable probability, Module 15’s sampling settings could then produce a “surprising” pick) versus a systematic model weakness for this type of context.

I’d also check the sampling configuration itself (temperature, top-k, top-p, Module 15) — a higher-randomness sampling setting can select a lower-probability, less contextually appropriate token even when the underlying distribution correctly favors more sensible options.

AI Engineering

Q: Why does understanding the exact logits → softmax → selection mechanism matter for someone building production LLM applications, rather than treating the model as an opaque text generator? A: This precise understanding directly explains several practical, important behaviors: why sampling parameters (temperature, top-k, top-p — Module 15) meaningfully change output by operating on the SAME underlying probability distribution differently; why longer, higher-vocabulary-coverage generation is computationally more expensive (every step scores the full vocabulary); and why “hallucination” (Module 21) isn’t a special failure mode with its own mechanism — it’s the same prediction process producing a wrong answer with unwarranted apparent confidence, exactly as capable of happening as producing a correct one.

20. Next Step

Next: Module 6 — Language Modeling and Probability — the underlying mathematics: conditional probability, the chain rule, autoregressive modeling, cross-entropy, and perplexity, built on the mechanism just traced.

Author
TechByteByByte Editorial Team
Reviewed by
TechByteByByte Admin
Published
Last reviewed