Commit Selected Token Predictions in a Reverse-Denoising Update

Hello! Last time, you calculated a reveal budget: how many additional positions should become visible at each denoising step. You also have candidate tokens and confidence scores for the positions that are still masked. Today you will put those pieces together in one reverse-denoising update: select eligible positions and write their proposed tokens into the sequence.

This is one step of generation, not the entire sampler. The next lesson will call the update repeatedly, obtaining fresh predictions after each change.

What changes in one reverse step?

The training picture below separates the partially masked sequence from the denoiser’s estimate of the clean sequence. During generation, an estimate is not automatically the next sequence: we may commit only some of its predicted tokens and leave the other positions masked.

For a visual account of that selection, watch this short passage from “Zed Inferred: Diffusion Language Models” by Zed Industries:

Suppose the current sequence is [prefix, MASK, MASK, MASK]. The denoiser proposes three tokens, but your reveal budget is one. You keep one proposal, leaving two masks for a later pass. The newly visible token then becomes context when the denoiser runs again.

Define exactly which positions may change

For each row of a batch, the update needs four things:

• Current token IDs: the sequence before this step.
• Candidate token IDs and confidence scores: one proposal and score per position, obtained from the denoiser as in the earlier lesson.
• Generation slots: a Boolean mask marking positions that belong to the output, rather than a fixed prefix or padding.
• Reveal count: the number of new positions to commit in this step.

A position is eligible only if it is a generation slot and its current token is MASK. Among eligible positions, choose the highest-confidence proposals up to the reveal count. Copy their candidate IDs into a new sequence; leave everything else unchanged.

The resulting invariant is useful for debugging: a fixed prefix, padding, and previously committed tokens must retain their current IDs. This version makes commitments irreversible. Revisiting an earlier choice through remasking is a separate technique you will meet later.

Implement the update in PyTorch

The function below takes predictions as inputs. It does not call the denoiser: keeping those operations separate makes it easier to inspect one update before building a full generation loop.

import torch

@torch.no_grad()
def reverse_update(
    x, candidate_ids, confidence, generation_slots,
    reveal_counts, mask_id, pad_id
):
    """Commit the most confident eligible predictions in each batch row."""
    if not (
        x.shape == candidate_ids.shape
        == confidence.shape == generation_slots.shape
    ):
        raise ValueError("All position-wise tensors need shape [batch, length]")
    if generation_slots.dtype != torch.bool:
        raise ValueError("generation_slots must be Boolean")
    if len(reveal_counts) != x.shape[0]:
        raise ValueError("Provide one reveal count per batch row")

    next_x = x.clone()
    eligible = generation_slots & (x == mask_id) & (x != pad_id)

    for b in range(x.shape[0]):
        k = int(reveal_counts[b])
        remaining = int(eligible[b].sum().item())
        if not 0 <= k <= remaining:
            raise ValueError("Reveal count exceeds eligible masks")

        if k == 0:
            continue

        scores = confidence[b].masked_fill(
            ~eligible[b], float("-inf")
        )
        _, positions = torch.topk(scores, k=k)

        chosen_ids = candidate_ids[b, positions]
        if ((chosen_ids == mask_id) | (chosen_ids == pad_id)).any().item():
            raise ValueError("A committed proposal cannot be MASK or PAD")

        next_x[b, positions] = chosen_ids

    return next_x

The crucial line before topk replaces ineligible scores with negative infinity. A prefix token might have a very high confidence score, but it must not consume a reveal slot. topk then chooses positions separately for each batch row, since rows can have different numbers of masks and different reveal budgets. If two eligible scores tie, you should not rely on a particular tie-breaking order.

Candidate IDs and confidence scores should come from the same denoiser pass over x. Before that pass becomes an update, proposals intended for commitment must be ordinary vocabulary tokens—not MASK or PAD. The check above catches either special token if it is selected. Keep the tensors on the same device when you connect this function to your Colab model.

Follow one update by hand—and in Colab

Here is a small batch with two different reveal counts. IDs 0 and 1 mean PAD and MASK; the other numbers stand for ordinary character IDs.

PAD, MASK = 0, 1

x = torch.tensor([
    [7, 8, MASK, MASK, MASK, PAD],
    [9, MASK, MASK, MASK, PAD, PAD],
])

generation_slots = torch.tensor([
    [False, False, True, True, True, False],
    [False, True,  True, True, False, False],
])

candidate_ids = torch.tensor([
    [7, 8, 11, 12, 13, 0],
    [9, 21, 22, 23,  0, 0],
])

confidence = torch.tensor([
    [0.99, 0.97, 0.55, 0.90, 0.72, 0.80],
    [0.88, 0.60, 0.30, 0.95, 0.10, 0.10],
])

next_x = reverse_update(
    x, candidate_ids, confidence, generation_slots,
    reveal_counts=[2, 1], mask_id=MASK, pad_id=PAD
)
print(next_x)

The result is:

tensor([[ 7,  8,  1, 12, 13,  0],
        [ 9,  1,  1, 23,  0,  0]])

In the first row, the two most confident eligible proposals are 12 and 13. In the second, only 23 is committed. Notice that the prefix scores of 0.99 and 0.97 do not matter, and padding remains untouched.

For the following step, you would feed next_x back to the denoiser. You would not reuse the uncommitted proposals 11, 21, and 22: the visible context has changed, so the model gets another chance to predict those masked positions.

Takeaway

One reverse-denoising update combines the reveal budget with the denoiser’s proposals: identify masked generation slots, select the required number by confidence, and commit only those tokens. Every other position stays as it was. Next, you will compose these updates into a complete sampler, including an optional fixed prompt prefix.