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 greedy decoding would make it judge the decoding instead.

Greedy decoding loops

Greedy decoding takes the argmax at every step. It is 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. 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$ approaches 0 it becomes greedy, which is why octlm treats temperature 0 as argmax rather than 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. With thousands of unlikely tokens, their combined probability can be large enough that the model regularly picks nonsense. Top-k cuts that tail.

What happened?

• At temperature 0, every draw is " happy". That is 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 gained temperature, top_k and a torch.Generator. The default temperature is 0, so every earlier day's greedy output is unchanged.

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

def sampling_distribution(logits: Tensor, temperature: float, top_k: int) -> 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"))
    return torch.softmax(logits / temperature, dim=-1)

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 and cannot check the draws, because PyTorch's random number generator cannot be matched 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

Generation still reruns the whole context for every new token. The KV cache that avoids this is EXP-071 on Day 5.

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.

What happened?

• All twenty are grammatical, stay with the prompt's characters, and follow a story shape: a problem, then a turn. The errors are logic slips, not broken text.
• Greedy repeats itself in 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 odder 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 locks into one loop.

The note's reading, confirmed by the user on 2026-09-28: 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 did not 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 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.