KLASS: KL-Adaptive Stability Sampling for Fast Inference in Masked Diffusion Models
Seo Hyun Kim, Sunwoo Hong, Youngrok Park, Hojung Jung, Se-Young Yun
Masked diffusion models have demonstrated competitive results on various tasks including language generation. However, due to its multiple refinement process, the inference is often bottlenecked by slow and static sampling speed. To overcome this problem, we introduce `KL-Adaptive Stability Sampling' (KLASS), a fast yet effective sampling method that exploits token-level KL divergence to identify stable, high-confidence predictions. By unmasking multiple tokens in each iteration without any additional model training, our approach speeds up generation significantly while maintaining sample quality. We test our sampler on diverse domains including $2\times$ faster in wall-clock time while improving performance in reasoning benchmarks (GSM8K and MATH) compared to standard greedy decoding, achieving state-of-the-art result among diffusion sampler.