Confidence-Based Remasking vs. Irreversible Token Commitment

Hello! In the previous lesson, your sampler could draw candidate characters at a chosen temperature, but once it committed a character, that position was frozen. Today you will let it reconsider low-confidence commitments. The goal is a small change to the sampler—not a new model or training objective—and a clear comparison of the two generation trajectories.

What remasking changes

Imagine four character slots. An early pass commits c x _ _. Once c is visible, the denoiser might prefer a in the second slot, but the irreversible sampler never asks about that slot again. A remasking sampler can temporarily return it to MASK, giving the denoiser another opportunity to fill it using the richer context.

Watch iterative infilling in Zed Inferred: Diffusion Language Models by Zed Industries. Notice the initial distinction between discarding fresh proposals and committing a subset. Then watch self-correction for the idea of revisiting an earlier choice. Our code below does the latter: it masks a token that had already survived a previous step.

The Kuleshov Group’s How to Build a Diffusion Language Model makes an important qualification: irreversible commitment follows naturally from the standard absorbing-mask process, whereas revisiting a committed token requires a modified sampling procedure.

Our implementation uses a simple policy. Before each pass after the first, it remasks the committed generation position with the lowest confidence recorded when that token was selected. It then makes fresh proposals for all currently masked positions and reveals enough to meet the usual schedule. The number of revealed positions still grows, even though the identity of those positions can change.

Add a revisable sampler

Run the cell from the previous lesson that defines pick_candidates. The function below reuses it, including its temperature, special-token filtering, and seeded draws. remask_per_step=0 gives the irreversible comparison; remask_per_step=1 revisits one commitment per later pass.

import math
import torch

@torch.no_grad()
def sample_with_remasking(
    logits_fn, num_new_tokens, steps, mask_id, pad_id,
    prompt_ids=None, device="cpu", temperature=0.0,
    seed=None, remask_per_step=1,
):
    if num_new_tokens < 1 or steps < 1 or remask_per_step < 0:
        raise ValueError("Counts must be positive; remask_per_step may be zero")
    if not math.isfinite(temperature) or temperature < 0:
        raise ValueError("temperature must be finite and nonnegative")

    prompt_ids = list(prompt_ids or [])
    if any(token in (mask_id, pad_id) for token in prompt_ids):
        raise ValueError("The fixed prompt cannot contain PAD or MASK")

    dev = torch.device(device)
    x = torch.tensor(
        [prompt_ids + [mask_id] * num_new_tokens],
        dtype=torch.long, device=dev,
    )
    valid = torch.ones_like(x, dtype=torch.bool)
    generation_slots = torch.zeros_like(x, dtype=torch.bool)
    generation_slots[:, len(prompt_ids):] = True

    # Scores belong to committed generation slots, not to the fixed prompt.
    stored_confidence = torch.zeros(x.shape, device=dev)
    generator = None
    if seed is not None:
        generator = torch.Generator(device=dev).manual_seed(seed)

    history = [x[0].tolist()]

    for step in range(steps):
        committed = (x != mask_id) & generation_slots

        if step > 0 and remask_per_step > 0:
            positions = torch.where(committed[0])[0].tolist()
            # Position breaks confidence ties predictably.
            positions.sort(
                key=lambda pos: (stored_confidence[0, pos].item(), pos)
            )
            for pos in positions[:remask_per_step]:
                x[0, pos] = mask_id
                stored_confidence[0, pos] = 0.0

        eligible = (x == mask_id) & generation_slots
        if not eligible.any().item():
            break

        # Supply an approximate current noise level to the denoiser.
        t = eligible.sum().float().reshape(1) / valid.sum()
        logits = logits_fn(x, valid, t)
        if logits.shape[:2] != x.shape:
            raise ValueError("Expected logits of shape [1, length, vocabulary]")

        candidates, confidence = pick_candidates(
            logits, x, eligible, temperature, generator, mask_id, pad_id
        )

        target_revealed = math.ceil(
            num_new_tokens * (step + 1) / steps
        )
        already_revealed = ((x != mask_id) & generation_slots).sum().item()
        to_reveal = target_revealed - already_revealed

        positions = torch.where(eligible[0])[0].tolist()
        positions.sort(key=lambda pos: (-confidence[0, pos].item(), pos))
        for pos in positions[:to_reveal]:
            x[0, pos] = candidates[0, pos]
            stored_confidence[0, pos] = confidence[0, pos]

        history.append(x[0].tolist())

    return x[0].tolist(), history

There are two separate selections here. First, a low-scored old commitment is returned to MASK. After the model call, high-scored new proposals compete for the slots being revealed. A remasked position might be filled immediately with a different token, or it might remain masked until another pass. Fixed prompt positions are in neither selection.

stored_confidence is deliberately modest: it remembers how confident the model was at commitment time, not how likely it considers that token now. That makes this an inexpensive remasking heuristic, not a guarantee of error detection.

See two trajectories from the same denoiser

Use a tiny deterministic denoiser so that a correction is easy to spot. Its first pass strongly proposes c, but also proposes an erroneous x before it can see that committed c. On later passes it proposes a for the second slot if that slot is masked and c is visible.

PAD, MASK = 0, 1
C, A, X, T_CHAR, EXCL, OTHER = 2, 3, 4, 5, 6, 7

def toy_logits(x, valid, t):
    logits = torch.zeros((*x.shape, 8), device=x.device)
    logits[..., PAD] = float("-inf")
    logits[..., MASK] = float("-inf")

    first_is_c = (x[0, 0].item() == C)
    choices = [
        (C, 0.99),
        (A, 0.97) if first_is_c else (X, 0.80),
        (T_CHAR, 0.90) if first_is_c else (T_CHAR, 0.45),
        (EXCL, 0.25),
    ]
    # Among six ordinary IDs, give the chosen ID probability q.
    for pos, (token, q) in enumerate(choices):
        logits[0, pos, token] = math.log(5 * q / (1 - q))
    return logits

def show(history):
    symbols = {
        MASK: "_", C: "c", A: "a", X: "x",
        T_CHAR: "t", EXCL: "!", OTHER: "?",
    }
    for state in history:
        print("".join(symbols[token] for token in state))

common = dict(
    logits_fn=toy_logits, num_new_tokens=4, steps=3,
    mask_id=MASK, pad_id=PAD, temperature=0.0, seed=7,
)

fixed, fixed_history = sample_with_remasking(
    **common, remask_per_step=0
)
revised, revised_history = sample_with_remasking(
    **common, remask_per_step=1
)

print("Irreversible:")
show(fixed_history)
print("\nWith remasking:")
show(revised_history)

assert fixed == [C, X, T_CHAR, EXCL]
assert revised == [C, A, T_CHAR, EXCL]

Both histories start at ____ and initially reach cx__. The irreversible path must keep x; the revisable path can end at cat!. In both cases, the target number of revealed characters after successive passes is 2, 3, then 4. Remasking changes which characters occupy those slots, not the final length or the reveal budget.

With your trained model, replace toy_logits with your earlier trained_logits adapter and call model.eval() first. Keep the temperature, seed, prompt, and step count the same between runs. Inspect the histories, not just the final strings: a changed final result alone does not show when a commitment was revisited. A fixed prompt should remain untouched, and neither generated output should contain PAD or MASK.

Takeaway

Irreversible sampling only fills positions that are still masked. Confidence-based remasking also returns selected committed positions to masks, allowing the denoiser to propose replacements using later context. Our low-cost policy uses confidence recorded at commitment time; it can correct an early mistake, but confidence is imperfect and extra revisiting does not guarantee better text.

Next, you will run the complete miniature model on a small corpus and use its loss curve and generated samples to judge what it has actually learned.