KLASS: KL-geführte schnelle Inferenz in maskierten Diffusionsmodellen

KLASS: KL-Guided Fast Inference in Masked Diffusion Models

November 7, 2025

papers.authors: Seo Hyun Kim, Sunwoo Hong, Hojung Jung, Youngrok Park, Se-Young Yun

cs.AI

papers.abstract

Masked Diffusion Models haben wettbewerbsfähige Ergebnisse bei verschiedenen Aufgaben, einschließlich der Sprachgenerierung, gezeigt. Allerdings wird der Inferenzprozess aufgrund des iterativen Optimierungsverfahrens oft durch eine langsame und statische Abtastgeschwindigkeit limitiert. Um dieses Problem zu überwinden, führen wir `KL-Adaptive Stability Sampling' (KLASS) ein, eine schnelle und dennoch effektive Abtastmethode, die die tokenweise KL-Divergenz nutzt, um stabile Vorhersagen mit hoher Konfidenz zu identifizieren. Indem in jeder Iteration mehrere Token entschleiert werden, ohne dass ein zusätzliches Modelltraining erforderlich ist, beschleunigt unser Ansatz die Generierung erheblich, bei gleichbleibender Probenqualität. Auf Reasoning-Benchmarks erzielt KLASS bis zu 2,78-fache Beschleunigungen in der Echtzeit und übertrifft dabei die Leistung der standardmäßigen greedy-Decodierung, was zu state-of-the-art Ergebnissen unter diffusionsbasierten Samplern führt. Wir validieren KLASS weiterhin in verschiedenen Domänen, einschließlich Text-, Bild- und Molekülgenerierung, und zeigen dessen Wirksamkeit als ein breit anwendbarer Sampler für verschiedene Modelle.

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.