技術要約: ROMS-IMLE

問題提起

生成モデルは、変分オートエンコーダ（VAE）や敵対的生成ネットワーク（GAN）から、拡散モデルやフロー・マッチングのような確率的補間（Stochastic Interpolant; SI）モデルへと進化してきました。SIモデルは画像合成における最先端を定義していますが、それらは複雑な機構（例：Transformerなどの精緻なアーキテクチャ）、高度なサンプリング手法、および反復的な数値積分に依存しています。支配的な見解は、高品質な生成には、ノイズからデータへの変換を多くの小さな逐次的ステップへと分解することが必要であるというものです。しかし、この反復プロセスは計算コストが高く、単一のサンプルを生成するために複数のフォワードパス（高い関数評価回数、NFE）を必要とします。

著者らは、この反復的な分解が本当に必要なのかという疑問を投げかけています。彼らは、SIモデルの成功は反復的なサンプリングプロセスそのものからではなく、ミニマリストなシングルステップの枠組みに適応可能な特定の学習特性、すなわち**「ステージごとの監督（per-stage supervision）」と「スペクトル特化（spectral specialization）」**に由来していると仮定しています。目標は、反復的なデノイジングの計算オーバーヘッドなしに、高品質なサンプルを単一のフォワードパスで生成できる、競争力のある生成モデルを設計することです。

手法: ROMS-IMLE

提案手法である**ROMS-IMLE（Robust Multi-Stage IMLE）**は、**暗黙的最大尤度推定（Implicit Maximum Likelihood Estimation; IMLE）**に基づいています。IMLEは、生成されたサンプルの集合における実画像とその最近傍との間の距離を最小化することにより、事前分布からデータ分布へと写像するシングルステップの生成器です。著者らは、IMLEのシンプルさ（変分境界、敵対的学習、Transformerを排除すること）を維持しつつ、SIモデルの強みを模倣するための2つの決定的な修正を学習目的関数に導入しました。

1. スペクトルを考慮したステージごとの監督

著者らは、SIモデルに対する「テスト中心」の視点を採用しています。拡散を反復的なデノイジングプロセスとして捉えるのではなく、サンプラーを関数の合成 $h_\theta = g^1_\theta \circ g^2_\theta \circ \dots \circ g^T_\theta$ として捉えます。この視点において、SIモデルはステージごとの監督を提供します。すなわち、すべての中間層（特定のノイズレベルまたはタイムステップに対応）が直接的な学習信号を受け取ります。さらに、この監督はスペクトル構造を持っています。初期段階（高ノイズ）は粗い低周波構造を学習し、後半の段階（低ノイズ）は微細な高周波の詳細を学習します。

これをシングルステップのIMLE生成器 $f_\theta$ に適応させるため、著者らは以下を行います：

• アーキテクチャの足場（Architectural Scaffold）: 生成器をアップサンプリング・ブロックのスタック（$f_\theta = f_{\theta_L} \circ \dots \circ f_{\theta_1}$）として構築します。ここで、各ステージは段階的に高い解像度で動作します。
• ステージごとの損失（Per-Stage Loss）: すべてのステージ $l$ に出力ヘッドを導入します。学習目的関数は、ステージ $l$ の出力（$\tilde{x}_l$）を、そのステージの解像度にダウンサンプリングされた実画像（$K_l(x_i)$）と比較する損失の総和をとります。
• スペクトル特化（Spectral Specialization）: ダウンサンプリング・カーネル（ローパスフィルタ）をターゲットとして使用することで、モデルは初期段階で粗い構造を学習し、後半段階で高周波の詳細を学習することを暗黙的に強制され、拡散モデルのスペクトル的な進行を模倣します。

2. 不一致ペアに対するロバストな損失

シングルステップのIMLEでは、モデルは実画像を生成されたサンプルの最近傍と一致させます。サンプリングのランダム性により、実画像が実際には正しいモードを表していない生成サンプルと一致してしまう（「不一致ペア」となる）可能性があります。標準的な二乗 $\ell_2$ 損失は、これらの外れ値に対して敏感です。なぜなら、大きな残差が不釣り合いに大きな勾配を生み出し、学習を不安定にする可能性があるからです。

著者らは、標準的な距離指標を**ロバストな損失関数（具体的には Geman-McClure）**に置き換えます。この関数は残差に対して劣二次的に成長するため、不一致ペアの影響を軽減し、適切に一致したペアが勾配更新を支配することを可能にします。これにより、最近傍の照合が不完全であっても、安定した学習が保証されます。

3. ラウンドトリップ拒絶（潜在空間）

高解像度生成（ImageNet 256）のために、モデルは凍結されたオートエンコーダの潜在空間で動作します。生成器がデコードに失敗する可能性のある潜在量を生み出すという分布の不一致に対処するため、著者らは**ラウンドトリップ拒絶（round-trip rejection）**を導入しています。推論時、生成された画像は再度エンコードおよびデコードされます。元の生成画像と再デコードされた画像の間の距離（ラウンドトリップ・コスト）が高い場合、そのサンプルは破棄されます。これは、高忠実度のサンプルのみが出力されるようにするためのフィルタとして機能します。

主な貢献

1. ミニマリストなシングルステップ設計: 本論文は、単純な ConvNeXt ベースのアーキテクチャとシングルステップのフォワードパスを用いることで、競争力のある生成モデルを構築できることを示しており、Transformer や反復的なサンプリングの複雑さを回避しています。
2. IMLE におけるステージごとの監督: 著者らは、反復的な SI モデルにおける中間的な監督の概念をシングルステップの IMLE フレームワークに正常に変換し、複数の解像度での直接的な監督が品質にとって極めて重要であることを示しました。
3. ロバストな学習目的関数: ロバストな損失関数（Geman-McClure）の統合は、IMLE における最近傍の不一致による不安定性を効果的に緩和します。
4. 効率性: 結果として得られるモデルは、拡散ベースのベースラインと比較して、大幅に少ないパラメータ数と関数評価回数で、最先端の性能を達成しています。

実験結果

モデルは CIFAR-10、CelebA-HQ、および ImageNet 256 の3つのデータセットで評価されました。

• CIFAR-10 (無条件): ROMS-IMLE は、シングル NFE で FID 4.01 を達成しました。特筆すべきは、比較されたすべての手法の中で最高の Precision (0.91) と Recall (0.80) を達成し、100 ステップの DDIM ベースラインをカバー範囲と忠実度の両面で上回ったことです。
• CelebA-HQ (無条件): モデルは FID 6.70 を達成し、シングルステップ手法の中で最高の Precision (0.96) と Recall (0.61) を記録し、Shortcut (128 ステップ) のようなマルチステップのベースラインを両方の指標で上回りました。
• ImageNet 256 (クラス条件付き): 310M パラメータ（DiT-XL/2 の 675M よりも少ない）を用いて潜在空間で動作し、モデルはシングルフォワードパスかつ分類器フリーガイダンスなしで FID 2.56 に到達しました。これはマルチステップの拡散モデル（例：250 ステップの SiT-XL/2）に匹敵する性能であり、大幅に効率的です（250倍少ない関数評価回数）。

意義と主張

本論文は、現代の高品質な生成モデルの性質は、本質的に反復的なサンプリングに結びついているのではなく、むしろスペクトルを考慮した、ステージごとの監督に由来すると主張しています。これらの特定の学習信号をミニマリストな IMLE フレームワークに注入することで、著者らは、シングルステップの生成器が複雑な反復モデルのサンプル品質と多様性に匹敵できることを証明しています。

その意義は、画像生成をより安価でシンプルにできる可能性にあります。ROMS-IMLE は、反復的な推論の計算コストなしに、競争力のある FID、Precision、および Recall を達成しており、生成を段階的な変換へと分解する必要性に疑問を投げかけています。著者らは、適切な目的関数を用いて訓練された単純なシングルステップモデルは、人気のある SI モデルの強みを保持できると結論付けています。