一貫性軌道モデル:拡散過程の確率フローODE軌道の学習

一貫性モデル(Consistency Models, CM)(Songら, 2023)は、スコアベース拡散モデルのサンプリングを高速化する一方で、サンプル品質の低下を伴い、品質と速度のトレードオフを自然に制御する手段を欠いている。この制限を克服するために、本研究では一貫性軌道モデル(Consistency Trajectory Model, CTM)を提案する。CTMは、CMおよびスコアベースモデルを特別なケースとして含む汎用的な枠組みであり、単一のニューラルネットワークを学習させることで、1回の順伝播(forward pass)において、対象のサンプルの対数密度の勾配(すなわちスコア)を出力可能となる。これにより、拡散過程における確率フロー常微分方程式(Probability Flow ODE)の任意の初期時刻から終了時刻まで、制限のない軌道遷移が実現される。CTMは、敵対的学習とノイズ除去スコアマッチング損失を効率的に組み合わせることで性能を向上させ、CIFAR-10における1ステップ拡散モデルサンプリングで新記録のFID(1.73)、および64×64解像度のImageNetではFID(1.92)を達成した。また、CTMは、ODE解軌道に沿った長距離ジャンプを含む、決定論的および確率論的な新たなサンプリングスキームの家族を可能にした。計算リソースが増加するにつれて、サンプル品質が一貫して向上し、CMで見られる品質の低下を回避する。さらに、CMとは異なり、CTMはスコア関数へのアクセスを保有しているため、拡散モデルコミュニティで既に確立された制御可能・条件付き生成手法の導入を容易にする。また、このスコア関数の利用により、尤度(likelihood)の計算も可能となる。コードは、https://github.com/sony/ctm にて公開されている。