Temperature-Controlled Sampling and Reproducible Random Seeds

Hello! Last time, you assembled a sampler that repeatedly asks the denoiser for predictions and commits some of them to the masked sequence. It always chose the highest-probability token at each position. Today you will give it another option: draw a token from the denoiser’s distribution, with a temperature controlling how adventurous that draw is and a seed making runs repeatable.

The reveal schedule and fixed-prefix protection stay exactly as they were. Only candidate selection changes.

From logits to a token draw

The denoiser produces vocabulary logits at every position. For one masked position, let $z_v$ be the logit for token $v$. Temperature $T$ changes the distribution from which we draw:

where $A$ is the set of allowed tokens. In our character sampler, PAD and MASK are not in $A$: neither should appear as a generated character.

At $T=1$, we sample from the model’s ordinary softmax distribution. A lower positive temperature concentrates probability on its favorites; a higher temperature spreads probability toward alternatives. Neither changes the ranking of the logits. We will define $T=0$ separately as greedy selection, rather than divide by zero.

There is a second choice to make. In the previous sampler, confidence meant the probability of the most likely token, which was also the token selected. Once we sample, those can differ. We will score a candidate by the temperature-scaled probability of the token actually drawn. reverse_update can then continue to reveal the highest-confidence proposals first. A surprising draw may be left masked and proposed again on the next pass.

One seed for the whole generation

Sampling consumes random numbers. A seed initializes their generator, so rerunning the same sampler with the same inputs, model, temperature, and seed can reproduce its draws.

Create the generator once, before the denoising loop, not once per step. Recreating it at every step would restart the same random stream repeatedly. With seed=None, our code uses PyTorch’s existing random state instead.

A seed controls the sampler’s draws, not every possible source of variation. For a meaningful comparison, keep the model in evaluation mode; exact results across different hardware or nondeterministic GPU operations are not guaranteed.

Add sampling to the existing Colab sampler

Run the previous lesson’s cells defining reverse_update and sample_text. First, add this helper in a new cell:

import math
import torch

def pick_candidates(
    logits, x, eligible, temperature, generator, mask_id, pad_id
):
    """Choose candidates for masked generation slots only."""
    candidate_ids = x.clone()
    confidence = torch.zeros(
        x.shape, dtype=logits.dtype, device=x.device
    )

    # Shape: [number of eligible positions, vocabulary size].
    allowed = logits[eligible].clone()
    if allowed.shape[0] == 0:
        return candidate_ids, confidence

    allowed[:, mask_id] = float("-inf")
    allowed[:, pad_id] = float("-inf")

    if (
        not torch.isfinite(allowed).any(dim=-1).all().item()
        or torch.isposinf(allowed).any().item()
    ):
        raise ValueError("Each eligible position needs a finite allowed logit")

    if temperature == 0:
        probabilities = allowed.softmax(dim=-1)
        drawn = allowed.argmax(dim=-1)
    else:
        probabilities = (allowed / temperature).softmax(dim=-1)
        drawn = torch.multinomial(
            probabilities, num_samples=1, generator=generator
        ).squeeze(-1)

    candidate_ids[eligible] = drawn
    confidence[eligible] = probabilities.gather(
        dim=1, index=drawn[:, None]
    ).squeeze(1)
    return candidate_ids, confidence

logits[eligible] gathers only positions still awaiting a token. That avoids drawing tokens for the protected prefix or positions committed on earlier passes. Setting the two special-token logits to negative infinity gives them zero sampling probability. Notice also that gather retrieves the probability of the drawn ID—not necessarily the largest probability.

Now make three small edits to your existing sample_text definition.

1. Extend its signature, preserving the previous arguments:

def sample_text(
    logits_fn, num_new_tokens, steps, mask_id, pad_id,
    prompt_ids=None, device="cpu", temperature=0.0, seed=None
):

The default temperature=0.0 keeps the old greedy behavior.

2. Just after the existing positive-count validation, validate temperature and initialize the run’s generator:

    if not math.isfinite(temperature) or temperature < 0:
        raise ValueError("temperature must be finite and nonnegative")

    generator = None
    if seed is not None:
        generator = torch.Generator(device=torch.device(device))
        generator.manual_seed(seed)

3. Inside the loop, leave the model call and its shape checks in place. Replace the old block beginning # Neither special token... and ending confidence, candidate_ids = probabilities.max(dim=-1) with:

        eligible = (x == mask_id) & generation_slots
        candidate_ids, confidence = pick_candidates(
            logits, x, eligible, temperature, generator, mask_id, pad_id
        )

Keep the existing call to reverse_update immediately after this block. It still determines how many positions to commit and which positions are eligible. With $T=0$, candidates and their confidence scores match the earlier greedy approach. With positive temperature, every still-masked response position gets a fresh draw on each active step, but only the scheduled number are committed.

Check the behavior before using the trained model

This fake denoiser deliberately assigns very high logits to PAD and MASK. Successful generation therefore also checks that the sampling helper excludes them.

PAD, MASK = 0, 1

def varied_logits(x, valid, t):
    logits = torch.zeros((*x.shape, 5), device=x.device)
    logits[..., PAD] = 20.0
    logits[..., MASK] = 20.0
    logits[..., 2] = 1.2
    logits[..., 3] = 1.0
    logits[..., 4] = 0.6
    return logits

def run_demo(temperature, seed):
    return sample_text(
        varied_logits,
        num_new_tokens=8,
        steps=4,
        mask_id=MASK,
        pad_id=PAD,
        prompt_ids=[2, 3],
        temperature=temperature,
        seed=seed,
    )

first_ids, first_history = run_demo(temperature=1.0, seed=7)
again_ids, again_history = run_demo(temperature=1.0, seed=7)

assert first_ids == again_ids
assert first_history == again_history
assert first_ids[:2] == [2, 3]
assert all(token not in (PAD, MASK) for token in first_ids[2:])

print("Repeatable history:", first_history)
for temperature in (0.0, 1.0, 2.0):
    ids, _ = run_demo(temperature=temperature, seed=7)
    print(f"T={temperature}: {ids}")

The two runs at $T=1$ should have identical histories, not merely identical final sequences. At $T=0$, every candidate is the highest-logit ordinary token; a seed has no effect on these greedy draws. Higher temperatures give the other ordinary IDs more opportunity to appear, though two short runs are not guaranteed to differ. Keep the seed fixed when comparing temperatures so that you are changing one setting at a time.

For your trained denoiser, use the trained_logits adapter from the previous lesson and set model.eval() before calling the updated sampler. Temperature changes token choices—and potentially which proposals are confident enough to commit—but it does not change your reveal counts or allow a fixed prompt prefix to be overwritten.

Takeaway

Temperature reshapes each masked position’s allowed-token distribution before categorical sampling: zero means greedy choice, one uses ordinary softmax, and higher positive values give alternatives more weight. A per-run PyTorch generator makes the sampling draws repeatable without resetting the rest of your notebook’s random state. The sampler still protects the prompt and reveals the scheduled number of positions.

Next, you will revisit a limitation of this irreversible sampler: a committed token cannot be corrected. Confidence-based remasking will let later denoising passes reconsider earlier choices.