Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking

要約

最近のマスク拡散モデル（MDMS）は、言語モデリングの自己回帰モデル（ARMS）と比較して競争力のあるパフォーマンスを示しています。
ほとんどの文献は、パフォーマンスのサンプリング手順の向上に焦点を当てていますが、MDMSからの効率的なサンプリングはほとんど調査されていません。
しばしば、特定のマスクされたトークンの特定のシーケンスが決定論的に複数の未知のトークンの値を決定するという観察を行います。つまり、マスクされたモデルの単一の予測は、標準サンプリング手順で未使用の追加情報を保持します。
この観察結果に基づいて、既存のサンプラーの単純なドロップイン交換であるEB-Samplerを導入し、事前に定義された近似誤差耐性を備えた1つの関数評価で複数のトークンを動的にマスクするエントロピー境界のないアンマスキング手順を利用します。
私たちは、アルゴリズムの選択を動機付けるエラー分析を提供する幅広い適応サンプラーの一部としてeB-samplerを策定します。
EB-Samplerは、パフォーマンスを損なうことなく、標準コーディングおよび数学推論ベンチマークで、現在の最新MDMSからのサンプリングを約2〜3倍に加速します。
また、同じ手順を検証し、迷路のナビゲーションやSudokuなどの小さな推論タスクでうまく機能します。

要約(オリジナル)

Recent masked diffusion models (MDMs) have shown competitive performance compared to autoregressive models (ARMs) for language modeling. While most literature has focused on performance enhancing sampling procedures, efficient sampling from MDMs has been scarcely explored. We make the observation that often a given sequence of partially masked tokens determines the values of multiple unknown tokens deterministically, meaning that a single prediction of a masked model holds additional information unused by standard sampling procedures. Based on this observation, we introduce EB-Sampler, a simple drop-in replacement for existing samplers, utilizing an Entropy Bounded unmasking procedure that dynamically unmasks multiple tokens in one function evaluation with predefined approximate error tolerance. We formulate the EB-Sampler as part of a broad family of adaptive samplers for which we provide an error analysis that motivates our algorithmic choices. EB-Sampler accelerates sampling from current state of the art MDMs by roughly 2-3x on standard coding and math reasoning benchmarks without loss in performance. We also validate the same procedure works well on smaller reasoning tasks including maze navigation and Sudoku, tasks ARMs often struggle with.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google