Comparing Model Evaluations and Generation Time in Full-Sequence and Blockwise Decoding

Welcome back. Last time, you made the sampler finish one block before moving to the next. The denoiser still saw the entire canvas on every pass, but only positions in the active block could change.

Today you’ll measure two things in your Colab notebook: model evaluations (forward calls to the denoiser) and wall-clock generation time. They answer different questions. A sampler can make fewer forward calls yet take longer if its calls or surrounding Python work are more expensive.

Count the work before timing it

A model evaluation is one call to model(tokens, fraction). Predicting several characters in that call still counts as one evaluation.

Consider 12 characters with no prefix. A full-sequence sampler using four reveal steps can make four model calls. With blocks of three characters and two steps per block, the blockwise sampler has four blocks and can make eight calls. Both produce 12 characters, but the second sampler makes more sequential calls.

For the reveal schedules used here, a step with no new positions to fill is skipped. Thus, if the requested step count exceeds the number of available positions, it does not necessarily add a model call. The count worth reporting is the count your code actually executes.

Before benchmarking, read this short note on PyTorch timing:

Give the full-sequence sampler the same canvas

Keep the generate_blockwise function from the previous lesson. Add this compact full-sequence counterpart after training. It uses the same model, vocabulary, confidence-based selection, and masked-fraction input. Unlike the blockwise function, any still-masked position can be selected on a step.

@torch.no_grad()
def generate_full_countable(total_length, steps=4):
    if not (0 < total_length <= max_len):
        raise ValueError("total_length must fit the model")
    if steps < 1:
        raise ValueError("steps must be positive")

    model.eval()
    tokens = torch.full(
        (1, total_length), MASK, dtype=torch.long, device=device
    )

    for step in range(steps):
        remaining = tokens[0].eq(MASK)
        target_filled = (step + 1) * total_length // steps
        already_filled = (~remaining).sum().item()
        new_count = target_filled - already_filled

        if new_count == 0:
            continue

        fraction = tokens.eq(MASK).float().mean(dim=1)
        logits = model(tokens, fraction)[:, :, 2:]
        probabilities = logits.softmax(dim=-1)
        confidence, prediction = probabilities.max(dim=-1)

        confidence = confidence.masked_fill(
            ~remaining[None, :], -1.0
        )
        positions = confidence[0].topk(new_count).indices
        tokens[0, positions] = prediction[0, positions] + 2

    assert tokens.ne(MASK).all()
    return "".join(itos[int(token)] for token in tokens[0])

This is an unconditional, greedy comparison: no prefix or random token sampling changes between runs. Both functions pass a canvas of the same length to the same bidirectional model. The difference is which positions each sampler permits itself to reveal.

Measure actual forward calls and elapsed time

A PyTorch forward hook can count calls without editing either sampler. The helper below also warms up each function, synchronizes if you are using CUDA, and reports the median of five complete generations. The timer includes Python selection and decoding work, not just the neural network.

import statistics
import time

L = min(max_len, 12)
assert L >= 4, "Use a model with room for at least 4 positions"

def synchronize_if_cuda():
    if torch.device(device).type == "cuda":
        torch.cuda.synchronize()

def measure_generation(label, generate_once, expected_calls, repeats=5):
    model.eval()

    # Warm up model execution before collecting measurements.
    for _ in range(2):
        generate_once()
    synchronize_if_cuda()

    times_ms = []
    counts = []

    for _ in range(repeats):
        seen = [0]

        def count_forward(module, inputs, output):
            seen[0] += 1

        handle = model.register_forward_hook(count_forward)
        try:
            synchronize_if_cuda()
            start = time.perf_counter()
            text = generate_once()
            synchronize_if_cuda()
            elapsed_ms = (time.perf_counter() - start) * 1000
        finally:
            handle.remove()

        assert len(text) == L
        assert seen[0] == expected_calls, (label, seen[0], expected_calls)
        counts.append(seen[0])
        times_ms.append(elapsed_ms)

    print(
        f"{label:12s} calls={counts[0]:2d}  "
        f"median={statistics.median(times_ms):.2f} ms  "
        f"range={min(times_ms):.2f}–{max(times_ms):.2f} ms"
    )

full_calls = min(4, L)

# Each complete block uses two calls; a final short block may use fewer.
block_calls = (L // 3) * 2 + min(2, L % 3)

measure_generation(
    "full 4 steps",
    lambda: generate_full_countable(L, steps=4),
    full_calls,
)
measure_generation(
    "blockwise",
    lambda: generate_blockwise(
        L, block_size=3, steps_per_block=2, prefix=""
    ),
    block_calls,
)

If L is 12, expect four full-sequence calls and eight blockwise calls. Your times will depend on your Colab runtime; there is no correct millisecond value. If an assertion fails, inspect the sampler’s skipped steps or check whether the notebook is using the previous lesson’s version of generate_blockwise.

The first comparison answers, “What happens with these two chosen schedules?” It does not isolate block structure: four calls versus eight is a substantial difference. To compare them at approximately the same evaluation budget, add one more run:

measure_generation(
    "full matched",
    lambda: generate_full_countable(L, steps=block_calls),
    block_calls,
)

Because block_calls cannot exceed L with this schedule, that full-sequence run makes the same number of model calls as the blockwise run. Its intermediate text states will still differ, so matching call counts is not the same as making the algorithms identical.

What the numbers can—and cannot—tell you

Report both measurements when comparing samplers. Forward-call count is a useful measure of sequential denoising work; elapsed time tells you what this particular implementation actually costs. A single call here processes the whole canvas even during blockwise decoding. Completed blocks are frozen as outputs, but this notebook’s Transformer does not cache their attention computations. Consequently, blockwise decoding has no built-in speed advantage in this implementation.

At a length of roughly 12 characters, Python overhead, GPU synchronization, and run-to-run variation can also be large relative to model computation. Treat a small timing difference cautiously, especially if the printed ranges overlap. The experiment establishes the cost of your two samplers on your runtime, not a general speed claim about diffusion language models.

Takeaways

You now have a reproducible way to count actual denoiser calls and time complete generations. For the 12-character example, four full-sequence steps mean four calls, while four blocks with two steps each mean eight. Warm-up and CUDA synchronization make the wall-clock comparison more meaningful, and the matched-call run helps separate evaluation count from decoding structure.

Next, you’ll examine a different distinction: the denoiser was trained on independently corrupted sequences, whereas generation feeds it states created by its own earlier predictions.