KLASS : Inférence Rapide Guidée par KL dans les Modèles de Diffusion à Masques

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

Les modèles de diffusion masquée ont démontré des résultats compétitifs sur diverses tâches, y compris la génération de langage. Cependant, en raison de leur processus de raffinement itératif, l'inférence est souvent limitée par une vitesse d'échantillonnage lente et statique. Pour résoudre ce problème, nous présentons `KL-Adaptive Stability Sampling` (KLASS), une méthode d'échantillonnage rapide et efficace qui exploite la divergence KL au niveau des tokens pour identifier les prédictions stables et de haute confiance. En démasquant plusieurs tokens à chaque itération sans entraînement supplémentaire du modèle, notre approche accélère significativement la génération tout en préservant la qualité des échantillons. Sur des benchmarks de raisonnement, KLASS atteint des accélérations en temps réel allant jusqu'à 2,78 fois tout en améliorant les performances par rapport au décodage glouton standard, obtenant des résultats state-of-the-art parmi les échantillonneurs basés sur la diffusion. Nous validons en outre KLASS dans divers domaines, incluant la génération de texte, d'images et de molécules, démontrant son efficacité en tant qu'échantillonneur largement applicable à travers différents modèles.

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.