Building a Conditional and Unconditional Reverse-Diffusion Sampler

Hello! In the previous lesson, you built reverse_update: given candidate tokens, confidence scores, and a reveal count, it commits predictions only at eligible masked positions. Now you will call that update repeatedly. The result is a complete sampler that can start with nothing but masks or keep a supplied prompt prefix fixed while it fills in the response.

This is an irreversible-commitment sampler: once a token is revealed, later passes use it as context but do not change it. That makes the generation loop easy to inspect before adding more decoding options.

From one update to a finished sequence

Watch how the repeated model calls fit together in “Zed Inferred: Diffusion Language Models” from Zed Industries:

The prompt-conditioned panels in the following diagram show the same starting arrangement: visible prompt tokens beside masked response positions. Our sampler will protect those prompt tokens on every pass.

For a response of $N$ positions, choose a number of steps $S$. At step $i$, set the cumulative reveal target to

The reveal count passed to reverse_update is $R_i-R_{i-1}$, with $R_0=0$. This guarantees that the scheduled counts add up to exactly $N$: every response position is visible at the end. If you request more steps than response positions, some steps have a reveal count of zero; we will skip those model calls.

Build the sampler

Run the cell containing reverse_update from the previous lesson first. The function below handles one sequence at a time, so the model input has batch shape [1, length]. Its logits_fn argument is a small adapter: it receives the current IDs, a Boolean valid-position mask, and a noise level, and returns logits of shape [1, length, vocabulary_size].

import torch

@torch.no_grad()
def sample_text(
    logits_fn, num_new_tokens, steps, mask_id, pad_id,
    prompt_ids=None, device="cpu"
):
    """Generate a fixed number of tokens, optionally after a fixed prefix."""
    if num_new_tokens < 1 or steps < 1:
        raise ValueError("num_new_tokens and steps must be positive")

    prefix = torch.as_tensor(
        [] if prompt_ids is None else prompt_ids,
        dtype=torch.long,
        device=device,
    )
    if prefix.ndim != 1:
        raise ValueError("prompt_ids must be a one-dimensional sequence")
    if ((prefix == mask_id) | (prefix == pad_id)).any().item():
        raise ValueError("The fixed prompt cannot contain MASK or PAD")

    prefix_length = prefix.numel()
    length = prefix_length + num_new_tokens

    # With no prompt, every position begins masked.
    x = torch.full((1, length), mask_id, dtype=torch.long, device=device)
    x[0, :prefix_length] = prefix

    generation_slots = torch.zeros((1, length), dtype=torch.bool, device=device)
    generation_slots[:, prefix_length:] = True

    # This simple sampler has no padding; every position is valid.
    valid = torch.ones_like(generation_slots)
    history = [x[0].tolist()]

    for i in range(1, steps + 1):
        previous_target = ((i - 1) * num_new_tokens + steps - 1) // steps
        target = (i * num_new_tokens + steps - 1) // steps
        reveal_count = target - previous_target
        if reveal_count == 0:
            continue

        # Fraction of valid positions currently masked, shape [1].
        t = ((x == mask_id) & valid).sum(dim=1).float()
        t = t / valid.sum(dim=1)

        logits = logits_fn(x, valid, t)
        if logits.ndim != 3 or logits.shape[:2] != x.shape:
            raise ValueError("Expected logits with shape [1, length, vocab_size]")
        if max(mask_id, pad_id) >= logits.shape[-1]:
            raise ValueError("MASK and PAD IDs must be within the vocabulary")

        # Neither special token is a permissible generated character.
        scores = logits.clone()
        scores[..., mask_id] = float("-inf")
        scores[..., pad_id] = float("-inf")

        probabilities = scores.softmax(dim=-1)
        confidence, candidate_ids = probabilities.max(dim=-1)

        x = reverse_update(
            x, candidate_ids, confidence, generation_slots,
            reveal_counts=[reveal_count],
            mask_id=mask_id, pad_id=pad_id,
        )
        history.append(x[0].tolist())

    assert not (x[generation_slots] == mask_id).any().item()
    assert torch.equal(x[0, :prefix_length], prefix)
    return x[0].tolist(), history

Notice the order inside the loop. logits_fn sees the current x; reverse_update commits some of that pass’s proposals; the next pass sees the changed x and produces fresh proposals. Uncommitted predictions are never carried over.

The generation_slots mask is what protects the prefix. Even if the model strongly predicts a different character at a prompt position, that position is ineligible for an update. With prompt_ids=None, the prefix is empty and all positions are generation slots.

The code sets t to the masked fraction of all valid positions. That matches a training setup in which noise level means “fraction of valid tokens masked.” If your denoiser was trained with a different convention—for example, the masked fraction of response slots only—calculate t using that same convention here. If your model does not take an explicit noise level, its adapter can simply ignore t.

Check the wiring, then attach your denoiser

A tiny fake logits function lets you check the sampler without relying on a trained model. It proposes ordinary character IDs, so its output is only a plumbing test, not meaningful generated language.

PAD, MASK = 0, 1

def toy_logits(x, valid, t):
    logits = torch.zeros((*x.shape, 5), device=x.device)
    for position in range(x.shape[1]):
        proposed_id = 2 + position % 3  # one of three ordinary tokens
        logits[:, position, proposed_id] = 5.0 + position / 10
    return logits

unconditional_ids, unconditional_history = sample_text(
    toy_logits, num_new_tokens=4, steps=2,
    mask_id=MASK, pad_id=PAD,
)

prompt_ids = [2, 3]
conditioned_ids, conditioned_history = sample_text(
    toy_logits, num_new_tokens=4, steps=2,
    mask_id=MASK, pad_id=PAD,
    prompt_ids=prompt_ids,
)

print("Without prompt:", unconditional_history)
print("With prompt:   ", conditioned_history)
assert conditioned_ids[:len(prompt_ids)] == prompt_ids

Look at the histories: each begins with four response masks, reveals two positions per step, and ends with none. In the second history, the first two IDs never change.

To connect your trained denoiser, replace toy_logits with an adapter to your model’s actual forward signature. For example, if your module accepts token IDs, a noise-level tensor, and a padding mask, the adapter could look like this:

model.eval()
device = next(model.parameters()).device

def trained_logits(x, valid, t):
    return model(
        x,
        t,
        src_key_padding_mask=~valid,
    )

generated_ids, history = sample_text(
    trained_logits, num_new_tokens=12, steps=4,
    mask_id=MASK, pad_id=PAD, device=device,
)

Change the arguments inside trained_logits if your module uses different names. In this unpadded example, ~valid is all False; it is included to show where your model’s padding mask belongs. Decode generated_ids with the character vocabulary you built earlier. For conditional generation, encode a prefix and pass its IDs as prompt_ids.

Takeaway

A full sampler needs only a few additions around your one-step update: initialize masked response slots, set cumulative reveal targets, obtain new logits at each active step, and stop with no response masks remaining. An empty prefix gives unconditional generation; a visible, ineligible prefix gives fixed-prefix generation. You still choose the output length in advance.

Next, you will keep this loop but change how candidate tokens are chosen, adding temperature-controlled sampling and reproducible random seeds.