Denoising Steps: Balancing Sample Quality and Computation

Welcome back. Last time, you distinguished the denoiser’s masked-token training loss from the probability of a finished sentence. Now you’ll test a more visible property of your trained model: what changes when you give generation more denoising steps? This is a small, controlled experiment in your Colab notebook, not a retraining task.

What a step buys you

Your sampler starts with masked positions and calls the denoiser repeatedly. With an irreversible-commitment sampler, each call can reveal some positions; later calls predict the remaining blanks using the newly visible context. More steps therefore give the model more opportunities to condition on its earlier choices. They also require more model evaluations.

The Large Language Diffusion Models paper describes step count as an inference-time choice. Its reported results concern a much larger model, so use them to understand the trade-off—not as a quality target for your character-level notebook.

For your notebook, one model evaluation means one forward call to the denoiser. Wall-clock time is also worth measuring: two runs with the same evaluation count can take different amounts of time because of hardware and notebook overhead.

Sweep one variable: step count

Keep the trained weights, output length, prompt or lack of prompt, temperature, and sampler type unchanged. Try $1,2,4,8,16$ steps on a 32-position response. If your response length is shorter, choose step counts no larger than that length for an irreversible sampler: once every position is committed, extra calls cannot reveal anything.

The cell below assumes your existing generate function and decode function. Edit only sample_once to match their actual signatures. It should return decoded text; do not change the sampler’s reveal or remasking policy between rows.

import statistics
import time
import torch

STEP_COUNTS = [1, 2, 4, 8, 16]
SEEDS = [17, 29, 43]
LENGTH = 32

def sample_once(steps):
    # Adapt this call to your existing sampler.
    # Keep length, prompt, temperature, and sampler type fixed.
    ids = generate(
        model=denoiser,
        length=LENGTH,
        num_steps=steps,
        temperature=0.8,
    )
    return decode(ids)  # Adapt if ids has a batch dimension.

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

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

# Warm up the notebook before timing the sweep.
with torch.inference_mode():
    torch.manual_seed(SEEDS[0])
    sample_once(4)
sync_device()

call_count = [0]

def count_forward_calls(_module, _inputs, _output):
    call_count[0] += 1

# Attach the hook to the denoiser actually called by generate().
hook = denoiser.register_forward_hook(count_forward_calls)
results = []

try:
    with torch.inference_mode():
        for steps in STEP_COUNTS:
            times, evaluations, samples = [], [], []

            for seed in SEEDS:
                torch.manual_seed(seed)
                if device.type == "cuda":
                    torch.cuda.manual_seed_all(seed)

                call_count[0] = 0
                sync_device()
                start = time.perf_counter()

                text = sample_once(steps)

                sync_device()
                times.append(time.perf_counter() - start)
                evaluations.append(call_count[0])
                samples.append(text)

            results.append({
                "steps": steps,
                "evaluations": statistics.mean(evaluations),
                "median_seconds": statistics.median(times),
                "samples": samples,
            })
finally:
    hook.remove()

for row in results:
    print(
        f"{row['steps']:>2} steps | "
        f"{row['evaluations']:>4.1f} evaluations | "
        f"{row['median_seconds']:.3f} s | "
        f"{row['samples'][0]}"
    )

The hook should report approximately one evaluation per step if your sampler makes one denoiser call per step. If it does not, inspect the sampler: it may make an additional prediction call, stop early, or call a different model object. Report the measured count, rather than assuming it equals the requested steps.

Using the same seeds makes runs reproducible. It does not force different step counts to make identical intermediate random choices: their generation paths differ. If your sampler is entirely greedy, changing the seed may have no effect; in that case, compare several fixed prompt prefixes or simply report the single deterministic sample at each step count.

Read quality alongside cost

Put the generated strings side by side. For each, consider whether it contains recognizable words, resembles your training corpus’s short-text format, and avoids conspicuous repetition. You can mark each sample 0 (largely incoherent), 1 (partly plausible), or 2 (plausible for this tiny corpus), then compare the average rating with evaluations and median seconds. Keep the actual strings visible: a single score can hide how a sample failed.

You may see a large improvement from one to a few steps and little change afterward. You may also see a worse sample at a higher step count. Neither result is paradoxical. Additional calls cost computation, but your small model may not use their extra context well; with irreversible commitment, it also cannot repair an early committed mistake. Your result establishes a trade-off for this model, sampler, and corpus, not a universal best number of diffusion steps.

Finally, resist comparing wall-clock time directly with a left-to-right model from this experiment alone. A diffusion step predicts positions in parallel, but it still runs a model forward pass; architecture, sequence length, batching, and caching all matter to a fair speed comparison.

Takeaways

A controlled step sweep holds the generation setup fixed while measuring sample quality, actual denoiser evaluations, and elapsed time. More steps offer additional rounds of context-dependent prediction, at an additional compute cost; better quality is possible, not guaranteed.

Next, you’ll connect the clean text, masked state, noise level, denoiser output, and reverse update in your notebook to the notation in the linked paper.