Sampling

A language model predicts a probability for every token in the vocabulary. Generation turns that into text one token at a time, it picks a token, appends it, and runs the model again. Up to Day 3, octlm always picked the most likely token. Day 4 added sampling, because the exit check judges the model's stories, and with greedy decoding it would be judging the decoding instead.

Greedy decoding loops

Greedy decoding takes the argmax at every step. It's deterministic, and on story data it repeats itself. Once a phrase like "and they were very happy" becomes the most likely continuation of itself, the model says it again, and again, and again. Sampling breaks the loop by sometimes picking a less likely token, and two settings control how often and how far.

Divide the scores before the softmax

Temperature $\tau$ divides every logit before the softmax:

At $\tau = 1$ the distribution is the model's own. Below 1 the gaps between scores grow, and the distribution sharpens toward the top token. Above 1 the gaps shrink and it flattens toward uniform. As $\tau$ goes to 0 it becomes greedy, which is why octlm treats temperature 0 as argmax instead of dividing by zero.

Keep only the k best tokens

Top-k sets every logit outside the $k$ largest to $-\infty$ before the softmax, so those tokens get exactly zero probability. Temperature alone never removes a token. And with thousands of unlikely tokens, their combined probability can be big enough that the model regularly picks nonsense. Top-k cuts that tail.

So what happened ?

• At temperature 0, every draw is " happy". That's greedy decoding.
• At 0.5 the top token takes most of the probability. At 1.5 the less likely tokens get real shares, and the entropy rises toward the 3 bits of a uniform choice among eight.
• Top-k 3 zeroes five tokens outright, whatever the temperature. The remaining three renormalize to sum to 1.
• The draw counts follow the probabilities, but not exactly. With 200 draws, a token at p = 0.1 lands anywhere from about 12 to 28 times.

Two functions and a generator

Decoder.generate got temperature, top_k and a torch.Generator. The default temperature is 0, so every earlier day's greedy output stays the same.

def sample_token(
    logits: Tensor,
    temperature: float,
    top_k: int,
    generator: torch.Generator | None,
    top_p: float = 1.0,
) -> Tensor:
    if temperature == 0:
        return logits.argmax(dim=-1, keepdim=True)
    probabilities = sampling_distribution(logits, temperature, top_k, top_p)
    return torch.multinomial(probabilities.cpu(), 1, generator=generator).to(logits.device)

def sampling_distribution(
    logits: Tensor, temperature: float, top_k: int, top_p: float = 1.0
) -> Tensor:
    if top_k:
        cutoff = logits.topk(min(top_k, logits.shape[-1]), dim=-1).values[:, -1:]
        logits = logits.masked_fill(logits < cutoff, float("-inf"))
    probabilities = torch.softmax(logits / temperature, dim=-1)
    if top_p >= 1:
        return probabilities
    ordered, order = probabilities.sort(dim=-1, descending=True)
    ordered = ordered.masked_fill(ordered.cumsum(-1) - ordered >= top_p, 0.0)
    kept = torch.zeros_like(probabilities).scatter(-1, order, ordered)
    return kept / kept.sum(-1, keepdim=True)

sampling_distribution is the deterministic half, it filters, scales and applies softmax. The draw itself uses torch.multinomial with the seeded generator, so a given seed reproduces the same story. This page can check the distribution against PyTorch but not the draws, because you can't match PyTorch's random number generator in the browser.

@torch.inference_mode()
def generate(
    self,
    token_ids: Tensor,
    max_new_tokens: int,
    temperature: float = 0.0,
    top_k: int = 0,
    generator: torch.Generator | None = None,
) -> Tensor:
    if token_ids.ndim != 2 or token_ids.shape[0] != 1:
        raise ValueError("generation accepts one sequence")
    if temperature < 0 or top_k < 0:
        raise ValueError("temperature and top_k must not be negative")
    for _ in range(max_new_tokens):
        context = token_ids[:, -self.config.context_length :]
        logits = self(context)[:, -1].float()
        next_token = sample_token(logits, temperature, top_k, generator)
        token_ids = torch.cat((token_ids, next_token), dim=1)
    return token_ids

These Day 4 samples used the generator that reruns the whole context. Day 5's KV-cached generator now reuses past keys and values, and the trained T4 run measured a 10.6 to 15.8 percent speed gain with exact greedy-token parity.

Twenty stories from the trained model

octlm.day4 samples loaded the step-24,000 checkpoint from the main run and wrote ten fixed prompts twice, once greedy, once at temperature 0.8 with top-k 40, seed 0. Each continuation is 200 new tokens, so most stories stop mid-sentence. All twenty are below, read from the run's samples.jsonl.

So what happened ?

• All twenty are grammatical, stay with the prompt's characters, and follow a story shape, a problem and then a turn. The errors are logic slips, not broken text.
• Greedy repeats itself exactly the way the problem section predicted. Prompt 9 lists "a necklace, a necklace, a ring, and a necklace". In prompt 3 the cat dances, then "was so happy that he started to dance" again. That story alone has 13 of greedy's 23 repeated four-word runs.
• Sampling removes the loops and adds weirder events instead. The dog eats the cat, "Lily thanked Lily", and a dog named Spot talks and drives.
• Counted over all ten prompts, 2.6% of greedy's four-word runs repeat an earlier run in the same story, against 1.5% for sampling. The gap is real but small, and most of it is one story. At 200 tokens, greedy decoding on this model repeats phrases but never gets stuck in one loop.

The note's reading, which I checked myself on 2026-09-28, is that both sets pass as coherent for a 26M TinyStories model. That closes the last Day 4 exit check.

Decision, keep temperature 0.8 and top-k 40 as the default for generation.

What we didn't build, and why

• Top-p, or nucleus, sampling. It keeps the smallest set of tokens whose probabilities sum to p, which adapts to how confident the model is. Top-k with temperature is enough to stop greedy loops, and adding top-p is one line if the samples ever need it.
• Repetition penalties. They hide a weak model's repetition instead of measuring it.
• Beam search. It searches for the single most likely continuation, which makes repetition worse on open-ended text.