Generating Candidate Tokens from a Fully Masked Sequence

Hello! You have already seen how masking corrupts text and how a bidirectional denoiser learns to predict hidden tokens. This module turns that trained model into a generator. Today we will do only the first inference pass: create an entirely masked sequence, ask the denoiser for predictions, and inspect its candidate-token distributions. Later lessons will decide which predictions to reveal and when.

By the end, you should be able to identify the tensors involved in that first pass and run it in your Colab notebook without yet writing a sampling loop.

A blank canvas with a fixed number of slots

Suppose you want to generate 12 characters. Before predicting any characters, allocate 12 positions and put the mask token in every one:

This is a fully masked canvas, not an empty tensor. Its length already specifies how many character positions the denoiser must predict. There is no padding here: all 12 positions are intended to become output characters.

The diagram shows the entire journey, but pause at its top row. At this point, the model has no visible words to condition on. It can still make predictions because training taught it about patterns in the corpus; position embeddings also let it distinguish one slot from another. Its first guesses may be uncertain or rather generic. That is expected.

Watch the following short portions of “Zed Inferred: Diffusion Language Models” from Zed Industries to connect the canvas to the denoiser’s output.

Ask the trained denoiser for candidates

Your denoiser’s forward pass returns logits: unnormalized scores for every vocabulary token at every position. For batch size $B$, canvas length $L$, and vocabulary size $V$, their shape is

For example, $\text{logits}[0,3,:]$ contains all $V$ candidate scores for position 3 of the first sequence. Applying softmax across the vocabulary axis converts those scores into a distribution:

Thus, for each particular $b,i$, the probabilities across vocabulary entries sum to one. Softmax does not decide which position to reveal.

Here is a compact Colab cell. It assumes model is your trained denoiser and that mask_id and pad_id are the token IDs established earlier. It also assumes model(x) returns a logits tensor; use the noise-level call shown in the comment if that is how you trained your model.

import torch

length = 12
batch_size = 1
device = next(model.parameters()).device

# Every slot is a valid generation position, initially hidden.
x = torch.full(
    (batch_size, length),
    fill_value=mask_id,
    dtype=torch.long,
    device=device,
)

model.eval()
with torch.no_grad():
    logits = model(x)

    # If your trained forward method requires a noise level instead:
    # t = torch.ones(batch_size, device=device)
    # logits = model(x, t)  # only if t=1 meant full masking in training

    # PAD and MASK are input/control symbols, not output characters.
    candidate_logits = logits.clone()
    candidate_logits[..., pad_id] = -torch.inf
    candidate_logits[..., mask_id] = -torch.inf

    probabilities = torch.softmax(candidate_logits, dim=-1)

print("canvas:", x.shape)
print("logits:", logits.shape)
print("distributions:", probabilities.shape)
print("one distribution sums to:", probabilities[0, 0].sum().item())

For a single 12-position canvas, expect x.shape to be [1, 12] and both output shapes to be [1, 12, V]. The printed sum should be approximately 1.0. Excluding PAD and MASK before softmax ensures their probability is zero and the remaining candidate probabilities are normalized. If your vocabulary contains other control symbols that should never be generated as characters, exclude those too.

The noise-level comment matters only if your trained model expects that input. Use the same noise-level convention you used during training: t=1 is appropriate only when it denotes full masking in your implementation. A model never trained on fully or nearly fully masked inputs may produce poor first-pass predictions, even if this code runs correctly.

Predictions are not yet generated text

You can inspect the strongest candidate at each position without changing the canvas:

top_probability, top_token_id = probabilities.max(dim=-1)

print("candidate IDs:", top_token_id[0].tolist())
print("candidate probabilities:", top_probability[0].tolist())
print("canvas is still masked:", x[0].tolist())

Notice that x remains masked. top_token_id is a collection of proposals, not the generated sequence. For a toy three-character vocabulary, the first two distributions might look like this:

The leading candidates are a and b, but committing both immediately could produce a poor result. Once some positions become visible, the denoiser can make new predictions using that context. Obtaining candidate distributions and choosing which candidates enter the canvas are separate operations.

Takeaway

Generation begins with a fixed-length tensor filled with mask_id. One denoiser call produces logits of shape $[B,L,V]$; softmax over $V$ gives a candidate-token distribution at every position. Nothing has been revealed yet—the canvas is still masked. Next, you will use these distributions to calculate token confidence scores, which can help decide which predictions to trust first.