データ多様体の幾何学を明示的にモデル化する拡散画像生成

画像生成モデルの本質的な目標は、データ多様体から高品質なデータ点をサンプリングすることです。しかし、そのためには高密度で低次元かつコンパクトなパラメータ化空間の学習とデコードが必要であり、従来の手法では構造表現と生成効率を両立することが困難でした。今回、研究チームはこれを解決する新しいフレームワーク「データ多様体認識画像拡散モデル（MIND）」を提案しました。

MINDの核心は、連続拡散モデルのスコア関数に離散パッチトークン化を統合することにあります。これにより、離散トークンの構造的定量化能力と連続拡散の並列生成柔軟性を同時に活用し、データ多様体の幾何構造を明示的にモデル化することが可能になりました。このアプローチにより、生成過程でデータの本来の構造をより正確に捉えることができます。

エンドツーエンドの微分可能トレーニングを実現するために、研究チームは新しいソフトtop-k集約メカニズムを導入しました。このメカニズムは、離散トークンの選択を効率的に処理しつつ、勾配がスムーズに伝播することを保証します。また、Transformerバックボーンが低次元入力に対して示すスペクトルバイアスを軽減するため、二重ブランチの高周波特徴埋め込み層を設計しました。この設計により、高周波信号の感度が向上し、スペクトルバイアスによる性能低下が緩和されます。

推論時には、タイムステップに基づいてサンプリング方式を動的に調整するマルチステージ遷移サンプリング方式を採用しています。これにより、生成品質と効率の間で最適なバランスを実現します。

実験はImageNet 256×256データセットで行われました。80エポックのトレーニング後、ベースモデルはガイダンスなしでFID 22.73を達成し、これはバニラDiT-B/2ベースラインの43.47からほぼ半分に削減された値です。DiTおよびSiTベースラインと比較して、平均FIDをそれぞれ15.95および9.06削減しました。ガイダンス付きの画像生成では、わずか1.3億パラメータのMIND-BがFID 2.06を達成し、31億パラメータのLlamaGen-3Bを上回りました。さらに、7.15億パラメータのMIND-XLはFIDを1.95まで低減しました。

MINDは拡散ベースの画像生成に新たな視点をもたらし、その明示的な幾何モデリングのアプローチは、今後の研究と革新の基盤となると期待されます。コードは公開される予定です。