イトウ密度推定器を用いた拡散モデルの重ね合わせ

The Superposition of Diffusion Models Using the Itô Density Estimator

December 23, 2024

著者: Marta Skreta, Lazar Atanackovic, Avishek Joey Bose, Alexander Tong, Kirill Neklyudov

cs.AI

要旨

容易にアクセス可能な事前トレーニング済み拡散モデルのカンブリア爆発は、複数の異なる事前トレーニング済み拡散モデルを組み合わせる方法への需要を示唆していますが、より大きな組み合わせモデルを再トレーニングすることなく、著しい計算負担を負うことなくこれらを組み合わせる方法に対する需要があります。本論文では、提案された新しいフレームワークである「superposition」の下で、生成段階で複数の事前トレーニング済み拡散モデルを組み合わせる問題を取り扱います。理論的には、我々はsuperpositionを厳密な第一原理から導出し、連続方程式から派生したものとし、SuperDiff内で拡散モデルを組み合わせるために特別に設計された2つの新しいアルゴリズムを開発します。SuperDiffは、拡散SDEの対数尤度のための新しいスケーラブルなIt\^o密度推定器を活用し、発散計算に必要なHutchinsonの推定器と比較して追加のオーバーヘッドが発生しません。SuperDiffは、推論中にのみ合成されるため、大規模な事前トレーニング済み拡散モデルにスケーラブルであり、異なる事前トレーニング済みベクトルフィールドを自動的な再重み付けスキームを介して組み合わせることで痛みを伴わず実装できます。特筆すべきは、SuperDiffが推論時に効率的であり、論理ORや論理ANDなどの従来の合成演算子を模倣していることです。我々は、SuperDiffを使用してCIFAR-10上でより多様な画像を生成する有用性、Stable Diffusionを使用したより忠実なプロンプト条件付き画像編集、およびタンパク質の改善された無条件のde novo構造設計のための実用性を経験的に示します。https://github.com/necludov/super-diffusion

English

The Cambrian explosion of easily accessible pre-trained diffusion models suggests a demand for methods that combine multiple different pre-trained diffusion models without incurring the significant computational burden of re-training a larger combined model. In this paper, we cast the problem of combining multiple pre-trained diffusion models at the generation stage under a novel proposed framework termed superposition. Theoretically, we derive superposition from rigorous first principles stemming from the celebrated continuity equation and design two novel algorithms tailor-made for combining diffusion models in SuperDiff. SuperDiff leverages a new scalable It\^o density estimator for the log likelihood of the diffusion SDE which incurs no additional overhead compared to the well-known Hutchinson's estimator needed for divergence calculations. We demonstrate that SuperDiff is scalable to large pre-trained diffusion models as superposition is performed solely through composition during inference, and also enjoys painless implementation as it combines different pre-trained vector fields through an automated re-weighting scheme. Notably, we show that SuperDiff is efficient during inference time, and mimics traditional composition operators such as the logical OR and the logical AND. We empirically demonstrate the utility of using SuperDiff for generating more diverse images on CIFAR-10, more faithful prompt conditioned image editing using Stable Diffusion, and improved unconditional de novo structure design of proteins. https://github.com/necludov/super-diffusion