KLASS: Быстрый вывод с использованием направляющей KL в маскированных диффузионных моделях

KLASS: KL-Guided Fast Inference in Masked Diffusion Models

November 7, 2025

Авторы: Seo Hyun Kim, Sunwoo Hong, Hojung Jung, Youngrok Park, Se-Young Yun

cs.AI

Аннотация

Маскированные диффузионные модели продемонстрировали конкурентоспособные результаты в различных задачах, включая генерацию языка. Однако из-за итеративного процесса уточнения вывод часто ограничивается медленной и статичной скоростью семплирования. Для решения этой проблемы мы представляем `KL-адаптивный стабильный семплинг` (KLASS) — быстрый, но эффективный метод семплирования, который использует расхождение Кульбака-Лейблера на уровне токенов для идентификации стабильных прогнозов с высокой уверенностью. Путем одновременного раскрытия нескольких токенов на каждой итерации без дополнительного обучения модели наш подход значительно ускоряет генерацию, сохраняя качество образцов. На тестах рассуждений KLASS достигает ускорения в реальном времени до 2.78 раз при одновременном улучшении производительности по сравнению со стандартным жадным декодированием, устанавливая современные результаты среди диффузионных сэмплеров. Мы дополнительно проверяем эффективность KLASS в различных областях, включая генерацию текста, изображений и молекул, демонстрируя его применимость в качестве универсального сэмплера для разных моделей.

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.