技術的サマリー：半教師あり学習における拡散モデルのための多目的学習

1. 問題定義

拡散モデルは、異質なプロンプトドメインにわたるテキストから画像への合成や、複数の環境にわたるロボティクス制御など、多様なタスクにおける条件付き生成器として、ますます展開されるようになっている。これらのシナリオは本質的に、単一のモデルが複数の目標分布にわたってパレート最適な性能を達成しなければならない「多目的学習（MOL）」の問題を構成する。

この設定における中心的な課題は、「汎用モデル」を訓練する際の「統計的コスト」である。複数の目的を処理するために、汎用モデルクラス（$\mathcal{F}$）は、通常、個々のタスクに必要な「専門モデル」クラス（$\mathcal{H}_k$）よりもはるかに大きな容量を必要とする。完全教師あり設定では、汎用モデルの訓練に必要なサンプル複雑性は、その自身の複雑性（$C_\mathcal{F}$）に比例して増大し、これはしばしば許容不可能なほど高くなる。

さらに、実世界での展開では、しばしば「非対称なデータ環境」に直面する。

• ペアデータ $(x, y)$（例えば、キュレーションされたテキスト - 画像ペアや専門家によるロボット実演）は希少であり、取得に高額を要する。
• 非ペアの条件データ $y$（例えば、生テキストプロンプトや状態観測）は豊富に存在する。

本論文は、以下の問いに答えるものである：弱教師あり（ラベルなしの条件）データを活用することで、サンプル効率的な方法でパレート最適な拡散モデルを学習できるか？

2. 手法

著者らは、高容量の汎用モデルの必要性とペアデータの希少性を両立させる「2段階の半教師あり訓練フレームワーク」を提案する。

ステージ 1：専門モデルの訓練

$K$個の目的（目標分布 $P_k$）のそれぞれについて、利用可能な限られたペアデータ $\{(x_i^k, y_i^k)\}_{i=1}^n$ のみを用いて、軽量な「専門モデル」スコア関数 $h_k \in \mathcal{H}_k$ が訓練される。

• 専門モデルクラス $\mathcal{H}_k$ は、汎用モデルクラスに比べて著しく低い複雑性（$C_{\mathcal{H}_k} \ll C_\mathcal{F}$）を持つと仮定される。
• 訓練は、ペアデータ上のノイズ除去スコアマッチング目的関数を最小化する。

ステージ 2：疑似サンプリングによる汎用モデルの訓練

すべての目的にわたるパレート最適集合を近似するために、大規模な「汎用モデル」$f \in \mathcal{F}$ が訓練される。

• 疑似サンプル生成： 訓練済みの専門モデル $\{h_k\}$ を用いて疑似ペアデータを生成する。各豊富なラベルなし条件 $y_i^k \sim P_Y^k$ に対して、専門モデルが疑似ターゲット $\tilde{x}_i^k \sim P_{h_k}(\cdot|y_i^k)$ を生成する。
• 目的： 汎用モデルは、元のラベル付きデータと生成された疑似ペアデータの組み合わせに基づいて訓練される。訓練目的関数は、すべての $K$ 個の目的にわたるスコアマッチング損失のスカラー化 $S$ を最小化する：
  $$\hat{f}_S = \arg\min_{f \in \mathcal{F}} S(\tilde{L}_1(f), \dots, \tilde{L}_K(f))$$
  ここで、$\tilde{L}_k$ は疑似サンプル上で計算された経験損失であり、専門モデルのバイアスを考慮するための補正項を含む。

逐次意思決定への拡張

このフレームワークは、模倣学習における「拡散方策」に拡張される。ここでは、「条件」は状態 $s$ であり、「ターゲット」は行動 $a$ である。著者らは、オンポリシーロールアウトに内在する「分布シフト」に対処する。すなわち、専門家状態分布は直接サンプリングすることが困難であるため、第 2 段階では、環境ダイナミクス下で学習された専門モデル方策によって生成された軌跡を用いて、汎用方策のための疑似データを収集する。

3. 主要な理論的貢献

サンプル複雑性の境界

本論文は、提案された半教師ありアプローチに対する非漸近的な汎化境界を確立する。

• 主要結果（定理 3.2）： 目標誤差 $\epsilon$ を達成するために必要な「ペア（ラベル付き）サンプル」の数は、「専門モデルクラス」の複雑性（$C_{\mathcal{H}_k}$）のみに依存し、「汎用モデルクラス」の複雑性（$C_\mathcal{F}$）には依存しない。
  • 必要なラベル付きサンプル数：$\tilde{O}\left(\frac{K C_\mathcal{H}}{\epsilon^2}\right)$。
  • 必要なラベルなし条件数：$\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$。
• 比較： 対照的に、汎用モデルの従来の完全教師あり訓練には、$\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$ のラベル付きサンプルが必要となる。これは、$C_\mathcal{F} \gg C_\mathcal{H}$ の場合、サンプル効率において著しい改善を示している。
• 線形スカラー化（定理 3.5）： 線形（重み付き和）スカラー化および凸関数クラスに対して、著者らはより鋭い収束率を導出し、近似誤差の境界をさらに精緻化している。

分布推定と方策の保証

• 分布推定： 本論文は、スコアマッチング誤差と分布推定（全変動距離）の間のギャップを埋める。学習された汎用モデルが、目標条件付き分布を近似する際にパレート最適な性能を達成することを証明する。
• 拡散方策（定理 3.6）： 著者らは、半教師あり設定における模倣学習の拡散方策の最適性ギャップに関する「最初の理論的保証」を提供する。彼らは、このフレームワークがオンポリシーロールアウトによって引き起こされる分布シフトを成功裡に処理しつつ、専門モデルから汎用モデルへのパイプラインが持つサンプル効率の利点を維持することを示す。

4. 実験結果

著者らは、2 つの異なるドメインで理論を検証する。

1. ロボティクス操作（StackCube & PickCube）：

   • 設定： ドメインランダム化（シーン、照明、カメラポーズ）が変化するタスクに対する拡散方策の訓練。
   • 結果： 半教師あり MOL 法は、ラベル付きデータのみを基にしたベースラインを大幅に上回った。特に「分布外（OOD）」設定において顕著であった（例：StackCube Level 2 で成功率が +16%、Level 3 で +14% 向上）。これは、疑似ロールアウトが分布シフトの処理に不可欠な監督情報を提供することを示唆している。

2. 画像復元（CelebA-HQ 画像修復）：

   • 設定： 3 種類の異なるマスクタイプ（交互の線、顔領域、広幅のストローク）を用いた顔の画像修復。
   • 結果： 半教師ありアプローチは、ラベル付きデータのみを基にしたベースラインと比較して、すべてのマスクタイプにおいて PSNR、SSIM、LPIPS を改善し、生成画像タスクにおける手法の有効性を示した。

5. 意義と主張

本論文は、半教師あり環境下における拡散モデルの多目的学習のための「原理的な統計的フレームワーク」を提供すると主張する。その主な意義は以下の点にある。

• 理論と実践の架橋： 「専門モデルから汎用モデル」への蒸留パイプラインを形式化し、ラベル付きデータが希少な場合、このアプローチが直接のマルチタスク学習よりもサンプル効率的である理由に対する理論的根拠を提供する。
• サンプル効率： 複雑な汎用モデルを学習する統計的コストは、その自身の容量から切り離され、代わりに限られたデータで訓練された軽量な専門モデルの複雑性に依存することを証明する。
• 分布シフトへの対処： これらの保証を逐次意思決定（模倣学習）に拡張し、ロボティクスにおける重要な実用的課題に対処する。

著者らは謙虚な立場を維持し、彼らの保証が「実現可能性の仮定」（真のスコア関数が専門モデルクラス内に存在するとの仮定）に依存しており、現在、テストされた特定のドメインランダム化レベルを超えた分布外（OOD）汎化には対処していないことを認めている。彼らは、この研究を生成モデルに対する半教師あり MOL の統計的性質を理解するための基礎的なステップと見なしている。