[論文レビュー] Practical and Asymptotically Exact Conditional Sampling in Diffusion Models

Luhuan Wu, Brian L. Trippe+3 more|UvA-DARE (University of Amsterdam)|Jun 30, 2023

Machine Learning in Materials Science被引用数 13

ひとこと要約

本稿は Twisted Diffusion Sampler (TDS) を紹介する。拡散モデルからの漸近的に正確な条件付きサンプリングのための SMC ベースの手法であり、twisting を用いてタスク固有の訓練を必要とせず条件付けを組み込む。さらに、inpainting、クラス条件付き MNIST、およびタンパク質モチーフのスキャフォールディングにおける経験的効果を示す。

ABSTRACT

Diffusion models have been successful on a range of conditional generation tasks including molecular design and text-to-image generation. However, these achievements have primarily depended on task-specific conditional training or error-prone heuristic approximations. Ideally, a conditional generation method should provide exact samples for a broad range of conditional distributions without requiring task-specific training. To this end, we introduce the Twisted Diffusion Sampler, or TDS. TDS is a sequential Monte Carlo (SMC) algorithm that targets the conditional distributions of diffusion models through simulating a set of weighted particles. The main idea is to use twisting, an SMC technique that enjoys good computational efficiency, to incorporate heuristic approximations without compromising asymptotic exactness. We first find in simulation and in conditional image generation tasks that TDS provides a computational statistical trade-off, yielding more accurate approximations with many particles but with empirical improvements over heuristics with as few as two particles. We then turn to motif-scaffolding, a core task in protein design, using a TDS extension to Riemannian diffusion models. On benchmark test cases, TDS allows flexible conditioning criteria and often outperforms the state of the art.

研究の動機と目的

• 条件付きサンプリングをタスク固有の訓練や重いヒューリスティクスなしに、無条件拡散モデルから実現可能にする動機づけ。
• twisting を用いて p_theta(x^0|y) を標的とする実用的な逐次モンテカルロ（SMC）アプローチを開発する。
• TDS をinpaintingおよび Riemannian diffusionモデルへ拡張し、より広範な条件付けタスクに対応する。
• MNIST のインペインティングおよびクラス条件付け、タンパク質モチーフスキャフォールディングにおける TDS の経験的優位性を実証する。

提案手法

• extended diffusion sequences and SMC を介して p_theta(x^0|y) からの条件付きサンプリングを定式化する。
• denoising 推定 hat{x}_theta(x^t) によって近似される p_theta(y|x^t) に基づく twisting 関数を導入する。
• 最適な SMC の選択を近似しつつ漸近的正確性を維持するように、twisted proposal と weights を定義する。
• インペインティングおよび Riemannian diffusion モデルへの拡張を提供し、自由度条件付けを含める。
• 正規性条件の下で粒子近似が p_theta(x^0|y) に収束することを（非公式に）証明する。

Figure 1: Errors of conditional mean estimations with 2 SEM error bars averaged over 25 replicates. TDS applies to all three tasks and provides increasing accuracy with more particles.

実験結果

リサーチクエスチョン

• RQ1 TDS はタスク固有の conditioning 訓練なしに p_theta(x^0|y) からの漸近的に正確なサンプルを提供できるか。
• RQ2 TDS は精度と粒子効率の観点で、ヒューリスティック法や条件付き訓練と比べてどうか。
• RQ3 TDS をインペインティング、自由度条件付け、および Riemannian diffusion モデルへどのように拡張できるか。

主な発見

• TDS は条件情報を伴うタスクに対して、拡散モデルから漸近的に正確な条件付きサンプルを提供する。
• 経験的に、少数の粒子を用いた TDS は MNIST のクラス条件付き生成およびインペインティング設定でヒューリスティック手法よりも性能が高いことが多い。
• TDS はタンパク質モチーフスキャフォールディングで柔軟な条件付けを実証し、これらの問題で最先端の条件付きモデルを凌駕できる。
• TDS は計算統計のトレードオフが好ましく、粒子数が増えるにつれて精度が向上する（試行で O(1/K) の速さ）。
• インペインティングおよびリーマニアン多様体への拡張は漸近的正確性を維持し、適用範囲を広げる。

(a) Approximate conditional samples of each method with 64 particles given class $y=7$ .