KLASS: KL-Gestuurde Snelle Inferentie in Gemaskeerde Diffusiemodellen

KLASS: KL-Guided Fast Inference in Masked Diffusion Models

November 7, 2025

Auteurs: Seo Hyun Kim, Sunwoo Hong, Hojung Jung, Youngrok Park, Se-Young Yun

cs.AI

Samenvatting

Gemaskerde diffusiemodellen hebben competitieve resultaten behaald bij diverse taken, waaronder taalgeneratie. Vanwege het iteratieve verfijningsproces wordt de inferentie echter vaak beperkt door een trage en statische samplesnelheid. Om dit probleem te ondervangen, introduceren we `KL-Adaptive Stability Sampling' (KLASS), een snelle maar effectieve samplingmethode die gebruikmaakt van token-level KL-divergentie om stabiele, hoogvertrouwensvoorspellingen te identificeren. Door meerdere tokens per iteratie te demaskeren zonder extra modeltraining, versnelt onze aanpak de generatie aanzienlijk met behoud van samplekwaliteit. Op redeneerbenchmarks behaalt KLASS tot 2,78x snelheidswinst in kloktijd en verbetert de prestaties ten opzichte van standaard gretige decodering, waarmee state-of-the-art resultaten worden bereikt onder op diffusie gebaseerde samplers. We valideren KLASS verder in diverse domeinen, waaronder tekst-, beeld- en moleculaire generatie, en tonen aan dat het een breed toepasbare sampler is voor verschillende modellen.

English

Masked diffusion models have demonstrated competitive results on various tasks including language generation. However, due to its iterative 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. On reasoning benchmarks, KLASS achieves up to 2.78times wall-clock speedups while improving performance over standard greedy decoding, attaining state-of-the-art results among diffusion-based samplers. We further validate KLASS across diverse domains, including text, image, and molecular generation, showing its effectiveness as a broadly applicable sampler across different models.