Path Planning for Masked Diffusion Model Sampling

要約

このホワイトペーパーでは、トークンのマスキング順序がマスクされた拡散モデル（MDM）の生成品質にどのように影響するかを探ります。
拡張されたエビデンス下限（Elbo）を導き出し、各ステップでマスクするトークンを選択するプランナーを導入します。
私たちの分析により、代替のアンマスキング戦略が生成パフォーマンスを向上させることができることが明らかになりました。
これに基づいて、Path Planning（P2）を提案します。これは、事前に訓練されたBERTモデルまたは悪魔自体を使用して、決定を導くためにemoiser自体を使用するサンプリングフレームワークを提案します。
P2は、既知のすべてのMDMサンプリング戦略を一般化し、言語生成（コンテキスト学習、コード生成、ストーリーの浸透、数学的推論、逆呪い補正）および生物学的配列生成（タンパク質およびRNA配列）を含む多様なドメインのパフォーマンスを大幅に向上させます。

要約(オリジナル)

In this paper, we explore how token unmasking order influences generative quality in masked diffusion models (MDMs). We derive an expanded evidence lower bound (ELBO) that introduces a planner to select which tokens to unmask at each step. Our analysis reveals that alternative unmasking strategies can enhance generation performance. Building on this, we propose Path Planning (P2), a sampling framework that uses a pre-trained BERT model or the denoiser itself to guide unmasking decisions. P2 generalizes all known MDM sampling strategies and significantly improves performance across diverse domains, including language generation (in-context learning, code generation, story infilling, mathematical reasoning, reverse curse correction) and biological sequence generation (protein and RNA sequences).

arxiv情報

提供元, 利用サービス

arxiv.jp, Google