Blockwise Semi-Autoregressive Decoding with Parallel Denoising

Welcome back. In the previous lesson, you trained a character-level denoiser and generated text by gradually filling masks across an entire sequence. Today you’ll change where that sampler is allowed to work: it will finish one short block before starting the next, while keeping completed blocks fixed.

This is semi-autoregressive generation. Blocks are generated in order, but several positions within the active block can be revealed in one denoising pass. You can implement it using the trained model already in your Colab notebook.

One canvas, one active block

Watch this short explanation of block diffusion before changing the code.

Imagine generating 13 characters with the fixed prefix "the " and a block size of three. The sampler works on positions 4–6 first, then 7–9, then 10–12. During the first block, positions 7–12 are still masks. During the second, the first block contains real characters and the third remains masked.

That gives us two rules to enforce in code:

1. Only active-block positions may change. Earlier characters are frozen, whether they came from the prompt or a completed block; future positions stay masked.
2. Predict the active block together. On each pass, the denoiser outputs candidate distributions for its positions. We reveal the most confident remaining positions according to a small reveal schedule.

The denoiser still receives the whole canvas. Its bidirectional attention can use the fixed prefix and previously filled blocks as context.

Complete the sampler scaffold

Continue after training in the Colab notebook from the previous lesson. Keep its model, device, MASK, stoi, itos, and max_len; the code below reuses all of them.

The scaffold is the familiar generate loop with one new outer loop: each start:stop slice defines the only block that can be updated. The completed version also checks that earlier and future positions obey that rule.

@torch.no_grad()
def generate_blockwise(total_length, block_size=3,
                       steps_per_block=2, prefix=""):
    if not (0 < total_length <= max_len):
        raise ValueError("total_length must fit the model's position embeddings")
    if block_size < 1 or steps_per_block < 1:
        raise ValueError("block_size and steps_per_block must be positive")
    if len(prefix) > total_length:
        raise ValueError("prefix is longer than the requested output")
    if any(char not in stoi for char in prefix):
        raise ValueError("prefix contains a character outside the vocabulary")

    model.eval()
    tokens = torch.full(
        (1, total_length), MASK, dtype=torch.long, device=device
    )
    if prefix:
        tokens[0, :len(prefix)] = torch.tensor(
            [stoi[char] for char in prefix],
            dtype=torch.long,
            device=device,
        )

    # The prompt is fixed; divide only the remaining positions into blocks.
    for start in range(len(prefix), total_length, block_size):
        stop = min(start + block_size, total_length)
        width = stop - start
        earlier = tokens[0, :start].clone()

        for step in range(steps_per_block):
            remaining = tokens[0, start:stop].eq(MASK)

            # Cumulative number of positions to have filled in this block.
            target_filled = (step + 1) * width // steps_per_block
            already_filled = (~remaining).sum().item()
            new_count = target_filled - already_filled
            if new_count == 0:
                continue

            # There is no padding on this generation canvas:
            # every position is either a character or MASK.
            fraction = tokens.eq(MASK).float().mean(dim=1)
            logits = model(tokens, fraction)[:, start:stop, 2:]
            probabilities = logits.softmax(dim=-1)
            confidence, prediction = probabilities.max(dim=-1)

            # Completed positions cannot be selected again.
            confidence = confidence.masked_fill(
                ~remaining[None, :], -1.0
            )
            local_positions = confidence[0].topk(new_count).indices
            tokens[0, start + local_positions] = (
                prediction[0, local_positions] + 2
            )

        # Check the blockwise rules before advancing.
        assert tokens[0, start:stop].ne(MASK).all()
        assert torch.equal(tokens[0, :start], earlier)
        assert tokens[0, stop:].eq(MASK).all()

    return "".join(itos[int(token)] for token in tokens[0])

Look closely at the two slices around the model call. model(tokens, fraction) predicts at every position, but [:, start:stop, 2:] keeps logits only for the active block and excludes the PAD and MASK vocabulary entries. local_positions indexes that short block; adding start converts those indices back into positions on the full canvas.

The schedule uses integer division. For a three-character block and two steps, its cumulative filled counts are one and then three. Thus the second pass can use the first revealed character as context before committing the other two. A short final block also works: if a step has nothing new to reveal, it is skipped.

Run a prefix-conditioned sample and an unconditional one:

output = generate_blockwise(
    13, block_size=3, steps_per_block=2, prefix="the "
)
print("With prefix:", repr(output))
assert output.startswith("the ")
assert len(output) == 13

print(
    "Without prefix:",
    repr(generate_blockwise(13, block_size=4, steps_per_block=2))
)

The exact characters depend on your trained model. The assertions inside the sampler test something more specific than whether the text sounds good: each block finishes before the next begins, and neither the prompt nor an earlier block is overwritten.

What this implementation does—and does not—change

This is a decoding change, not a newly trained block-diffusion model. The tiny denoiser was trained on independently masked character sequences, not specifically on blockwise generation trajectories. It may therefore produce imperfect text even when the block boundaries and freezing logic are correct.

Nor does freezing characters automatically provide a reusable attention cache. Our bidirectional Transformer recomputes representations from the whole canvas on every model call; it has no block-specific caching machinery. What you have built is the essential sampling behavior: ordered blocks, parallel predictions within the active block, and irreversible commitment of completed blocks.

Takeaways

Blockwise decoding adds an outer block loop to the iterative sampler you already know. At each step, the model sees the current canvas, but only masked positions in the active block are eligible for commitment. Once that block is full, its characters become fixed context for the next one.

In the next lesson, you’ll count model evaluations and time full-sequence versus blockwise decoding. That will separate the structure of this sampler from any claim that it is faster in this small Colab implementation.