Part 2 · Module 1: What a Large Language Model Actually Is

Part B: Next-Token Prediction, the Only Thing the Model Does

14 min read·22 Sept 2026

The whole job in one sentence

A large language model does one thing: given a sequence of tokens (pieces of text: whole words, parts of words, punctuation, or single bytes), it outputs a probability distribution over its vocabulary (the fixed list of every token it knows) for what the next token should be. That is all. Chat, translation, code, and "reasoning" are all produced by calling that one function over and over.

A few terms you will see constantly:

Logits: the raw scores the model produces, one per vocabulary entry. They can be any number.
Softmax: the formula that turns logits into probabilities that are all positive and sum to 1.
Autoregressive generation: pick a token from the distribution, append it to the input, and ask again. The model's own output becomes its next input.
Parameters (or weights): the numbers inside the model that were adjusted during training. The model's "knowledge" is nothing more than their values.

TinyLM: a model small enough to open up

Hosted models do not let you look inside. So the course ships TinyLM, a real GPT-style model with 1,071,872 parameters, pretrained on this project's 0.7 MB support-domain corpus (data/corpus.txt). Production models have roughly 10,000 to 1,000,000 times more parameters and were trained on tens to hundreds of millions of times more text. TinyLM is useless as an assistant. It is excellent for seeing mechanisms with real numbers: probabilities, generation, training curves, memorization.

The corpus is synthetic: the 12 help-center articles (each appears once) plus 3,000 short templated conversations such as:

text

Customer (Ben): Good morning, how much does the Team plan cost?
Agent (Dara): Team costs 12 USD per user per month, or 10 USD billed annually. Business costs 24 USD per user per month, or 20 USD annually.
Customer (Ben): Thank you.

Code explained

In simple words: this is the kind of document TinyLM read, thousands of times over, during pretraining.
What happens: scripts/build_corpus.py fills 12 question templates with random names, greetings, and plans, and pairs each with a fixed agent answer taken from the help center. So a small number of facts appear hundreds of times each, always in the same words.
Comes out: a copy of the first conversation in data/corpus.txt. Remember this structure; it explains almost every behavior you will measure below.

Here is the model file. You do not need every line. This module teaches the architecture only to the depth needed for decisions; Module 3 returns to the sampling code. The explanation after the file says what each part is for.

supportdesk/tinylm.py

python

"""TinyLM: a very small GPT-style language model you can run and train on a laptop CPU.

It exists to make mechanisms visible with real numbers: next-token
probabilities, sampling parameters, prefill and decode timing, the KV cache,
fine-tuning with LoRA, and quantization. At about 1 million parameters trained
on a 0.7 MB synthetic corpus, it is roughly 100,000 times smaller than a small
production model. Never read its output quality as evidence about real LLMs;
read the mechanisms.
"""
from __future__ import annotations

import json
import math
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path

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

MODEL_DIR = Path(__file__).resolve().parents[1] / "models" / "tinylm-base"


@dataclass
class TinyConfig:
    vocab_size: int = 2048
    context: int = 128
    d_model: int = 128
    n_layers: int = 4
    n_heads: int = 4
    dropout: float = 0.0


class Block(nn.Module):
    """One transformer layer: self-attention then a small feed-forward network."""

    def __init__(self, cfg: TinyConfig) -> None:
        super().__init__()
        self.n_heads = cfg.n_heads
        self.ln1 = nn.LayerNorm(cfg.d_model)
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
        self.proj = nn.Linear(cfg.d_model, cfg.d_model)
        self.ln2 = nn.LayerNorm(cfg.d_model)
        self.mlp = nn.Sequential(nn.Linear(cfg.d_model, 4 * cfg.d_model), nn.GELU(), nn.Linear(4 * cfg.d_model, cfg.d_model))

    def forward(self, x: torch.Tensor, cache: dict | None = None) -> torch.Tensor:
        b, t, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=2)
        q, k, v = (z.view(b, t, self.n_heads, d // self.n_heads).transpose(1, 2) for z in (q, k, v))
        if cache is not None:
            if "k" in cache:
                k = torch.cat([cache["k"], k], dim=2)
                v = torch.cat([cache["v"], v], dim=2)
            cache["k"], cache["v"] = k, v
        # Causal mask only matters when several new positions are processed at once.
        y = F.scaled_dot_product_attention(q, k, v, is_causal=(t > 1 and (cache is None or k.shape[2] == t)))
        x = x + self.proj(y.transpose(1, 2).contiguous().view(b, t, d))
        return x + self.mlp(self.ln2(x))


class TinyGPT(nn.Module):
    def __init__(self, cfg: TinyConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.context, cfg.d_model)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight  # weight tying saves parameters
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        """Small random starting weights, so the untrained model predicts roughly uniformly."""
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)

    def forward(self, idx: torch.Tensor, caches: list[dict] | None = None, start_pos: int = 0) -> torch.Tensor:
        """Return logits (unnormalized scores) for the next token at every position."""
        t = idx.shape[1]
        pos = torch.arange(start_pos, start_pos + t, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        for i, block in enumerate(self.blocks):
            x = block(x, None if caches is None else caches[i])
        return self.head(self.ln_f(x))

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())


# Loading and saving ----------------------------------------------------------

def save(model: TinyGPT, tokenizer: Tokenizer, directory: Path | str, meta: dict | None = None) -> None:
    directory = Path(directory)
    directory.mkdir(parents=True, exist_ok=True)
    torch.save(model.state_dict(), directory / "model.pt")
    tokenizer.save(str(directory / "tokenizer.json"))
    (directory / "config.json").write_text(json.dumps({"config": asdict(model.cfg), "meta": meta or {}}, indent=2) + "\n")


def load(directory: Path | str = MODEL_DIR) -> tuple[TinyGPT, Tokenizer]:
    """Load a saved TinyLM (the pretrained base model by default) in eval mode."""
    directory = Path(directory)
    info = json.loads((directory / "config.json").read_text())
    model = TinyGPT(TinyConfig(**info["config"]))
    model.load_state_dict(torch.load(directory / "model.pt", map_location="cpu", weights_only=True))
    model.eval()
    return model, Tokenizer.from_file(str(directory / "tokenizer.json"))


# Inspecting predictions -------------------------------------------------------

@torch.no_grad()
def next_token_distribution(model: TinyGPT, tokenizer: Tokenizer, text: str, top: int = 10) -> list[tuple[str, float]]:
    """The model's probability for each candidate next token, highest first."""
    ids = tokenizer.encode(text).ids[-model.cfg.context:]
    logits = model(torch.tensor([ids]))[0, -1]
    probs = torch.softmax(logits, dim=-1)
    values, indices = probs.topk(top)
    return [(tokenizer.decode([i]) if tokenizer.decode([i]) else tokenizer.id_to_token(i), round(v, 4))
            for v, i in zip(values.tolist(), indices.tolist())]


# Generation ---------------------------------------------------------------------

@dataclass
class SamplingParams:
    max_new_tokens: int = 40
    temperature: float = 1.0
    top_k: int | None = None
    top_p: float | None = None
    min_p: float | None = None
    frequency_penalty: float = 0.0
    presence_penalty: float = 0.0
    stop: list[str] = field(default_factory=list)
    logit_bias: dict[int, float] = field(default_factory=dict)
    seed: int | None = None


@dataclass
class Generation:
    text: str
    token_ids: list[int]
    prefill_ms: float
    decode_ms_per_token: float
    stop_reason: str  # "stop", "length", or "constraint"


def apply_sampling(logits: torch.Tensor, params: SamplingParams, generated: list[int]) -> torch.Tensor:
    """Turn raw logits into a probability distribution using the sampling parameters."""
    logits = logits.clone()
    for token_id, bias in params.logit_bias.items():
        logits[token_id] += bias
    if generated and (params.frequency_penalty or params.presence_penalty):
        counts = torch.bincount(torch.tensor(generated), minlength=logits.shape[0]).float()
        logits -= params.frequency_penalty * counts + params.presence_penalty * (counts > 0).float()
    if params.temperature <= 0:
        probs = torch.zeros_like(logits)
        probs[logits.argmax()] = 1.0
        return probs
    probs = torch.softmax(logits / params.temperature, dim=-1)
    if params.top_k is not None:
        kth = probs.topk(params.top_k).values[-1]
        probs = torch.where(probs >= kth, probs, torch.zeros_like(probs))
    if params.min_p is not None:
        probs = torch.where(probs >= params.min_p * probs.max(), probs, torch.zeros_like(probs))
    if params.top_p is not None:
        sorted_probs, order = probs.sort(descending=True)
        keep = sorted_probs.cumsum(0) - sorted_probs < params.top_p
        mask = torch.zeros_like(probs, dtype=torch.bool)
        mask[order[keep]] = True
        probs = torch.where(mask, probs, torch.zeros_like(probs))
    return probs / probs.sum()


@torch.no_grad()
def generate(model: TinyGPT, tokenizer: Tokenizer, prompt: str, params: SamplingParams | None = None,
             use_cache: bool = True, allowed=None) -> Generation:
    """Generate a continuation of `prompt`.

    `allowed`, if given, is a function (generated_ids) -> set of token ids that
    may come next; everything else is masked out (constrained decoding).
    """
    params = params or SamplingParams()
    generator = torch.Generator().manual_seed(params.seed) if params.seed is not None else None
    ids = tokenizer.encode(prompt).ids[-(model.cfg.context - params.max_new_tokens):]
    caches = [dict() for _ in model.blocks] if use_cache else None
    generated: list[int] = []
    started = time.perf_counter()
    logits = model(torch.tensor([ids]), caches)[0, -1]
    prefill_ms = (time.perf_counter() - started) * 1000
    decode_started = time.perf_counter()
    stop_reason = "length"
    for _ in range(params.max_new_tokens):
        probs = apply_sampling(logits, params, generated)
        if allowed is not None:
            mask = torch.zeros_like(probs)
            permitted = list(allowed(generated))
            if not permitted:
                stop_reason = "constraint"
                break
            mask[permitted] = 1.0
            probs = probs * mask
            probs = probs / probs.sum() if probs.sum() > 0 else mask / mask.sum()
        next_id = int(torch.multinomial(probs, 1, generator=generator)) if params.temperature > 0 else int(probs.argmax())
        generated.append(next_id)
        text = tokenizer.decode(generated)
        if any(s in text for s in params.stop):
            stop_reason = "stop"
            break
        position = len(ids) + len(generated) - 1
        if position >= model.cfg.context:
            break
        if use_cache:
            logits = model(torch.tensor([[next_id]]), caches, start_pos=position)[0, -1]
        else:
            logits = model(torch.tensor([ids + generated]))[0, -1]
    decode_ms = (time.perf_counter() - decode_started) * 1000
    text = tokenizer.decode(generated)
    for s in params.stop:
        if s in text:
            text = text[: text.index(s)]
    return Generation(text, generated, round(prefill_ms, 2), round(decode_ms / max(len(generated), 1), 3), stop_reason)


# Training ------------------------------------------------------------------------

def batches(ids: torch.Tensor, context: int, batch_size: int, generator: torch.Generator):
    """Yield random (input, target) windows; the target is the input shifted by one token."""
    while True:
        starts = torch.randint(0, len(ids) - context - 1, (batch_size,), generator=generator)
        x = torch.stack([ids[s: s + context] for s in starts])
        y = torch.stack([ids[s + 1: s + context + 1] for s in starts])
        yield x, y


def loss_on(model: TinyGPT, x: torch.Tensor, y: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
    """Average cross-entropy of predicting each next token (optionally only where mask is 1)."""
    logits = model(x)
    losses = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1), reduction="none")
    if mask is None:
        return losses.mean()
    mask = mask.reshape(-1).float()
    return (losses * mask).sum() / mask.sum().clamp(min=1)


@torch.no_grad()
def perplexity(model: TinyGPT, tokenizer: Tokenizer, text: str) -> float:
    """exp(average next-token loss) on `text`: lower means the text looks more familiar to the model."""
    ids = torch.tensor(tokenizer.encode(text).ids)
    ctx = model.cfg.context
    total, count = 0.0, 0
    for start in range(0, max(len(ids) - 1, 1), ctx):
        chunk = ids[start: start + ctx + 1]
        if len(chunk) < 2:
            break
        loss = loss_on(model, chunk[:-1].unsqueeze(0), chunk[1:].unsqueeze(0))
        total += loss.item() * (len(chunk) - 1)
        count += len(chunk) - 1
    return math.exp(total / max(count, 1))

Code explained

In simple words: a transformer is a stack of identical layers that repeatedly lets every token "look at" earlier tokens and then transforms what it gathered. At the end, a final layer scores every vocabulary entry as the possible next token.
What happens (what each part is for):
TinyConfig: the spec. vocab_size (rows in the token table), context (the most tokens the model can see at once), d_model (the length of the vector that represents each token inside the model), n_layers, n_heads. Part D reads these numbers as a spec sheet.
Block: one transformer layer. Self-attention (qkv, proj, scaled_dot_product_attention) lets each position gather information from earlier positions; the causal mask forbids looking at later ones, which is what makes it a next-token predictor. n_heads splits attention into several parallel lookups. The MLP (a small feed-forward network) then transforms each position on its own; roughly half of TinyLM's parameters live here. LayerNorm keeps numbers in a stable range. The cache argument stores past keys and values (the KV cache) so generation does not recompute the whole prompt for every new token; Module 3 measures this.
TinyGPT: token embeddings (a lookup table from token id to vector), position embeddings (so the model knows word order), the stack of blocks, and the output head that produces logits. The head shares its weights with the token table (weight tying), a common trick to save parameters. num_parameters() counts them.
save and load: write and read model.pt (weights), tokenizer.json, and config.json. load() defaults to models/tinylm-base and puts the model in eval mode.
next_token_distribution: runs the model once and returns the top candidates for the next token with their probabilities. This is the function to reach for when you want to see what a model "thinks".
SamplingParams, Generation, apply_sampling, generate: the generation loop from the diagram above, plus the knobs (temperature, top-k, top-p, min-p, penalties, stop strings, logit bias, seed) that decide how a token is picked from the distribution. generate also times prefill (reading the prompt) and decode (writing each new token). Module 3 teaches all of these; here we only use greedy decoding (temperature=0, always take the most likely token) and stop.
batches, loss_on, perplexity: training utilities. batches cuts random windows from the corpus where the target is the input shifted by one token: that shift is the entire training objective. loss_on computes cross-entropy (how surprised the model was by the real next token, averaged). perplexity is exp(loss), a friendlier scale: "on average the model was as unsure as if choosing among this many tokens".
Comes out: nothing when imported. The next examples call load, next_token_distribution, generate, and perplexity.

Asking TinyLM what comes next

python

"""Module 1: next-token prediction, made visible with TinyLM."""
import torch

from supportdesk.tinylm import load, next_token_distribution

torch.set_num_threads(1)  # one thread is plenty for a 1M-parameter model
model, tokenizer = load()  # models/tinylm-base

prompts = [
    "Customer (Maya): Hi, I was charged twice for the Team plan this month. Can you refund the duplicate?\nAgent (Dara):",
    "Customer (Maya): Hello, my account is locked after too many attempts.\nAgent (Dara):",
    "Customer (Maya): Hi, does the Team plan include",
]
for prompt in prompts:
    print(repr(prompt[-60:]))
    for token, p in next_token_distribution(model, tokenizer, prompt, top=5):
        print(f"   {token!r:18} {p:.4f}")

# Generate a reply by hand: predict, pick the most likely token, append, repeat.
prompt = prompts[1]
ids = tokenizer.encode(prompt).ids
print("\nstep  chosen token          p(chosen)  runner-up")
with torch.no_grad():
    for step in range(12):
        probs = torch.softmax(model(torch.tensor([ids]))[0, -1], dim=-1)
        top = probs.topk(2)
        best, second = top.indices.tolist()
        ids.append(best)
        print(f"{step:>4}  {tokenizer.decode([best])!r:22} {top.values[0]:.4f}     "
              f"{tokenizer.decode([second])!r} {top.values[1]:.4f}")
print("\nreply so far:", tokenizer.decode(ids[len(tokenizer.encode(prompt).ids):]))

Code explained

In simple words: show the model three support conversations that stop just before the agent speaks, and read off its top five guesses for the next token. Then write the agent's reply by hand, one token at a time.
What happens: load() reads the pretrained base model. next_token_distribution(..., top=5) returns (token_text, probability) pairs. Note the leading spaces: ' Sorry' is a single token that includes the space before the word, which is how most modern tokenizers work. The second half is the generation loop with nothing hidden: run the model on all tokens so far, take softmax of the last position's logits, pick the highest (topk(2) also shows the runner-up), append its id, repeat 12 times.
Comes out: real output from this build (probabilities are reproducible on the same library versions and may differ in the last digit elsewhere).

text

  'plan this month. Can you refund the duplicate?\nAgent (Dara):'
     ' Sorry'           0.9971
     ' Automations'     0.0005
     ' Please'          0.0002
     ' Export'          0.0001
     ' workspace'       0.0001
  ' my account is locked after too many attempts.\nAgent (Dara):'
     ' After'           0.9976
     ' Automations'     0.0002
     ' The'             0.0001
     ' Monthly'         0.0001
     ' they'            0.0001
  'Customer (Maya): Hi, does the Team plan include'
     ' SSO'             0.9814
     ' comments'        0.0013
     ' card'            0.0013
     ' workspace'       0.0010
     ' automation'      0.0007

  step  chosen token          p(chosen)  runner-up
     0  ' After'               0.9976     ' Automations' 0.0002
     1  ' 5'                   0.9990     ' 14' 0.0001
     2  ' failed'              0.9892     ',' 0.0016
     3  ' attempts'            0.9960     ' certificates' 0.0001
     4  ' an'                  0.9977     ' for' 0.0001
     5  ' account'             0.9984     ' existing' 0.0003
     6  ' is'                  0.9921     ' before' 0.0020
     7  ' locked'              0.9983     ' available' 0.0002
     8  ' for'                 0.9988     ' after' 0.0003
     9  ' 15'                  0.9969     ' security' 0.0004
    10  ' minutes'             0.9967     ' features' 0.0002
    11  '.'                    0.9980     ' and' 0.0007

  reply so far:  After 5 failed attempts an account is locked for 15 minutes.

Read the numbers, not just the words. For a question the corpus contains hundreds of times, TinyLM puts more than 99 percent of its probability on the right first word and keeps that confidence for every following token. The runner-up tokens (' Automations', ' certificates', ' security') are fragments of other memorized answers. This is what "knowing" looks like inside a base model: a very peaked distribution. You will see what "not knowing" looks like in Part C, and it is less obvious than you might hope.

Why one objective produces translation, code, and apparent reasoning

TinyLM only learned the support desk because that is all it read. The same objective applied to trillions of tokens of web pages, books, code, and math forces a far larger model to learn much more, because every one of those skills lowers next-token loss somewhere:

To predict the next line of a Python file, it helps to track variable names, types, and indentation.
To predict the second half of a bilingual web page, it helps to learn translation.
To predict the last line of a worked math solution, it helps to compute the steps before it.
To predict what a helpful agent writes next, it helps to model what the customer needs.

Nothing in the objective says "learn grammar" or "learn arithmetic". Those skills emerge because they are compressions of the training data: a model that has them predicts text better than a model that does not. That is also why the skills are uneven. Skills that are common in text and easy to verify locally (fluent prose, common code patterns, translation between widely written languages) are strong. Skills that text rarely spells out step by step, or that need exact symbol manipulation, are weaker. Part F measures some of these limits.

Keep one sentence from this part: the model outputs a plausible continuation, and plausible is not the same as true. Every capability and every failure in this course traces back to it.