KLASS: マスク拡散モデルにおけるKL誘導高速推論

要旨

マスク拡散モデルは、言語生成を含む様々なタスクで競争力のある結果を示している。しかし、その反復的な精緻化プロセスのため、推論は遅く静的なサンプリング速度によってしばしばボトルネックとなる。この問題を克服するため、我々は`KL適応型安定性サンプリング`（KLASS）を提案する。これはトークンレベルのKLダイバージェンスを利用して安定した高信頼度の予測を特定する、高速かつ効果的なサンプリング手法である。追加のモデル学習を必要とせず各反復で複数のトークンのマスクを解除することで、本手法はサンプル品質を維持しつつ生成速度を大幅に向上させる。推論ベンチマークでは、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