離散拡散における償却逐次モンテカルロのための対照的分布マッチング

要旨

離散拡散モデルは、構造化されたカテゴリデータを生成するための強力な枠組みとして登場した。しかし、報酬に傾斜した分布からの効率的なサンプリングは依然として重要な課題である。Twisted Sequential Monte Carlo (SMC) はこの課題に対して漸近的な正確性を提供するが、離散状態空間における最適なねじれ関数の推定には高コストなモンテカルロ近似が必要であり、推論時に深刻な計算ボトルネックを引き起こす。この制約を克服するため、我々は Contrastive Distribution Matching (CDM) を導入する。これは、正例と負例のサンプルを通じてパラメータ化されたねじれ関数を学習することにより、SMC推論のコストを償却する新しい枠組みである。効率的な学習のために、勾配推定器を再構成し、離散拡散モデルの閉形式の前方カーネルを活用する。実際、学習したねじれ関数の評価にかかる追加計算オーバーヘッドは、ベースモデルの1回の順伝搬と比較して5%未満である。広範な実験評価を通じて、CDMが同一実経過時間条件下で既存のベースラインを一貫して上回ることを示す。我々は、有害テキスト生成、制御DNA配列設計、タンパク質設計可能性、拡散大規模言語モデルのアライメントなど、多様な応用課題において本手法の有効性と汎用性を検証する。

English

Discrete diffusion models have emerged as powerful frameworks for generating structured categorical data. However, efficiently sampling from reward-tilted distributions remains a fundamental challenge. While Twisted Sequential Monte Carlo (SMC) offers asymptotic exactness for this task, estimating the optimal twist function in discrete state spaces necessitates costly Monte Carlo approximations, resulting a severe computational bottleneck at inference. To overcome this limitation, we introduce Contrastive Distribution Matching (CDM), a novel framework that amortizes the cost of SMC inference by learning a parameterized twist function via positive and negative samples. For efficient training, we reformulate the gradient estimator to leverage the closed-form forward kernels of discrete diffusion models. In practice, evaluating our learned twist function incurs less than 5% additional computational overhead compared to a single forward pass of the base model. Through extensive empirical evaluations, we demonstrate that CDM consistently outperforms existing baselines under matched wall-clock time. We validate the effectiveness and versatility of our approach across a diverse range of applications, including toxic text generation, regulatory DNA sequence design, protein designability, and diffusion large language model alignment.