Parallel Sampling from Masked Diffusion Models via Conditional Independence Testing

TL;DR

PUNT is a novel parallel sampling method for masked diffusion models that balances independence and confidence, enabling faster text generation with improved accuracy and hierarchical structure emergence.

Contribution

It introduces a model-agnostic approach that identifies token dependencies and removes conflicting tokens to enable effective parallel sampling in masked diffusion models.

Findings

Up to 16% higher accuracy on IFEval benchmark.

Better trade-off between accuracy and compute compared to baselines.

Induces hierarchical, planning-like generation strategy.

Abstract

Masked diffusion models (MDMs) offer a compelling alternative to autoregressive models (ARMs) for discrete text generation because they enable parallel token sampling, rather than sequential, left-to-right generation. This means potentially much faster inference. However, effective parallel sampling faces two competing requirements: (i) simultaneously updated tokens must be conditionally independent, and (ii) updates should prioritise high-confidence predictions. These goals conflict because high-confidence predictions often cluster and depend on each other, opportunities for parallel updates. We present PUNT, a model-agnostic sampler that reconciles this trade-off. Our method identifies token dependencies and removes lower-confidence tokens from conflicting groups. This produces sets of indices for unmasking that satisfy both independence and confidence criteria. Our approach ensures improved parallel unmasking through approximate conditional independence testing. Our experiments show that PUNT delivers a superior trade-off between accuracy and compute when compared to other strong training-free baselines, especially for generation of longer sequences. On the IFEval benchmark, it achieves up to 16\% higher accuracy over baseline methods, including sequential generation (one-by-one). These gains hold across different values of hyperparameters, mitigating the need for brittle hyperparameter tuning. Moreover, we observe that PUNT induces an emergent hierarchical generation strategy, where the model first establishes high-level paragraph structure before local refinement, suggesting a planning-like generation process that contributes to strong alignment performance.