Parallel Sampling from Masked Diffusion Models via Conditional Independence Testing
Iskander Azangulov, Teodora Pandeva, Niranjani Prasad, Javier Zazo, Sushrut Karmalkar
TL;DR
This work tackles the inefficiency of autoregressive text generation by enabling parallel sampling in discrete diffusion models through contextual independence testing. It introduces PUNT, a training-free sampler that divides masked tokens into contextually independent groups, pruning dependent, low-confidence tokens, and unmasking high-confidence ones in parallel with $O( obreak ext{ } obreak ext{ } ext{log}|M|)$ forward evaluations per step. Empirically, PUNT delivers superior accuracy–compute trade-offs on long-form text and instruction-following benchmarks (e.g., MTBench and IFEval) and demonstrates robustness across hyperparameters, with emergent hierarchical generation patterns that resemble planning. The approach is grounded in transformer-attention properties (independence-stability) and validated through KL-divergence analyses and attention-based metrics, offering a practical, training-free path to faster and more reliable MDM-based generation. Future work may include adaptive thresholds, distillation to single-pass predictors, and combining PUNT with complementary efficiency techniques like KV-caching.
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.