KLASS: 마스크 확산 모델에서 KL-지도 기반 고속 추론

초록

마스크 확산 모델은 언어 생성 등 다양한 과제에서 경쟁력 있는 성능을 입증해왔습니다. 그러나 반복적 정제 과정으로 인해 추론 속도가 느리고 고정적인 샘플링 속도에 병목 현상이 자주 발생합니다. 이러한 문제를 해결하기 위해 본 논문에서는 토큰 수준 KL 발산을 활용해 안정적이고 높은 신뢰도를 가진 예측을 식별하는 빠르면서도 효과적인 샘플링 기법인 `KL-Adaptive Stability Sampling`(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.

KLASS: 마스크 확산 모델에서 KL-지도 기반 고속 추론

KLASS: KL-Guided Fast Inference in Masked Diffusion Models

초록

Support