技術要約：長尾分布における拡散モデルのためのコントラスト型条件付き・無条件整合（CCUA）

問題提起

クラス条件付き画像合成の学習データは、多くの場合、「テイル（裾）」クラスの画像が「ヘッド（頭部）」クラスに比べて極端に少ない、ロングテール分布を示します。この不均衡により、拡散モデルはテイルクラスの生成において**モード崩壊（mode collapse）**や多様性の低下を引き起こします。標準的なデノイジング拡散確率モデル（DDPM）やスコアベース拡散モデル（SBDM）を含む既存の手法は、ヘッドクラスの生成品質を損なうことなく、これら過小代表なクラスに対して高忠実度かつ多様な画像を生成することに失敗することがよくあります。コントラスト学習は、ロングテール認識やGANベースの生成において効果的であることが証明されていますが、拡散モデルへの適用には、デノイジングプロセスの特定のメカニズムに対処するための新たな適応が必要です。

手法：コントラスト型条件付き・無条件整合（CCUA）

著者らは、テイルクラスの画像の多様性と忠実度を向上させつつ、ヘッドクラスの品質を維持するために設計された、2つの相乗的な損失関数からなるフレームワークであるCCUAを提案しています。この手法は、U-NetベースのアーキテクチャとDiffusion Transformer（DiT/SiT）の両方と互換性があります。

1. 無条件生成のための非教師ありコントラスト損失（UCL）

モード崩壊に対処するため、著者らは無条件生成の潜在変数に適用される非教師ありコントラスト損失を導入しています。

• メカニズム: ノイズ推定ネットワークを、潜在エンコーダ $e(\cdot)$ とデコーダ $d(\cdot)$ に分割します。エンコーダは、ノイズを含んだ画像 $x_t$ を低次元の潜在変数 $h$ にマッピングします。
• 損失の定式化: InfoNCEスタイルの定式化を用い、負例（negative samples）のみを使用します。ミニバッチ内では、各画像がアンカーとして機能し、バッチ内の他のすべての画像が負例として機能します。データ拡張による正例ペアは作成されず、アンカー自身が自身の正のベクトルとして扱われます。
• 目的: これにより、潜在空間に「反発力」を生み出し、クラスに関係なく合成画像の間の距離を最大化します。これにより、モデルが無条件の潜在変数を多様化させ、一定のベクトルへの崩壊や限られた学習サンプルへの過学習を防ぐように強制します。
• 理論的正当化: コントラスト損失を直接条件付き潜在変数に適用すると、モデルがクラス間の多様性（クラスの分離）を最大化する一方で、クラス内の多様性を無視するという自明な解に陥る可能性があります。潜在変数に（無条件の）損失を適用することで、モデルはすべてのサンプルを多様化させることを強制され、これが第2の損失を通じて条件付き生成へと蒸留されます。

2. 条件付き・無条件整合のための整合損失（AL）

第2のコンポーネントは、多様化された無条件生成と条件付き生成を整合させます。

• メカニズム: この損失は、条件付き分布 $p_\theta(x_{t-1}|x_t, c)$ と無条件分布 $p_\theta(x_{t-1}|x_t)$ の間のKLダイバージェンスを最小化します。実際には、これは推定されたノイズ $\epsilon_\theta(x_t, t, c)$ と $\epsilon_\theta(x_t, t, \emptyset)$ の間のMSE損失へと簡略化されます。
• 時間重み付け: 決定的なのは、この整合が大きなタイムステップ（逆プロセスの初期段階）においてより重く設定されることです。
• 動機: 著者らは、デノイジングの初期段階において、異なるクラス間の画像（および条件付き生成と無条件生成の間）の低周波成分が非常に類似していることを観察しています。これらの初期段階を整合させることで、モデルは**知識共有（knowledge sharing）**を促進します。これにより、テイルクラスは、生成の初期段階においてヘッドクラス（豊富なデータ）が学習した豊かな多様性の恩恵を受けることができます。

3. 相乗効果

本フレームワークは、標準的なDDPM損失と提案されたCCUA損失を組み合わせたものです：
$$\mathcal{L}_{total} = \mathcal{L}_{ddpm} + \alpha \cdot \mathcal{L}_{ucl} + \gamma \cdot \mathcal{L}_{al}$$
相乗効果は、UCLが無条件の潜在空間を多様化させ、ALが条件付き潜在空間をその多様化された無条件空間へと整合させるという点にあります。この間接的な条件付き潜在変数の多様化は、条件付き潜在変数に直接コントラスト損失を適用する場合よりも効果的であることが経験的に示されています。

主な貢献

1. CCUAフレームワーク: 無条件潜在変数に対する非教師ありコントラスト学習と、条件付き・無条件整合損失を組み合わせた、クラス不均衡な拡散モデルのための新しい手法。
2. 整合戦略: 低周波成分の類似性を利用し、知識転送を可能にするために、条件付き生成と無条件生成を特に大きなタイムステップ（初期デノイジングステップ）で整合させる提案。
3. 非教師ありコントラスト正則化: コントラストペアに対してクラスラベルを必要とせずに、クラス内およびクラス間の多様性を最大化するために、負例のみを用いた無条件潜在変数への非教師ありコントラスト損失の使用。
4. アーキテクチャ非依存性: 本手法がU-NetベースのアーキテクチャとDiffusion Transformer (SiT) の両方において有効であることが示されている。

実験結果

本手法は、ImageNet-LT (256×256)、TinyImageNet-LT (64×64)、Places-LT (64×64)、および CIFAR-LT データセットで評価されました。

• 定量的パフォーマンス:
  • ImageNet-LT において、CCUAはSiT、CBDM、REPA、およびDispersive Lossを含むベースラインを大幅に上回りました。
  • 900kステップの学習において、CCUAは FID 15.1（SiTの19.9に対し）および FID$_{tail}$ 22.5（SiTの33.9に対し）を達成し、テイルクラスの忠実度において約11.4ポイントの改善を示しました。
  • また、本手法はRecall（多様性）とPrecision（忠実度）の指標を全体的に向上させました。
  • TinyImageNet-LT において、CCUAは最高のFID (12.0) および FID$_{tail}$ (20.8) を達成し、DDPM、CBDM、OCLTを上回りました。
• 定性的結果: 可視化により、CCUAは、標準的なSiT（テイルクラスに対して反復的または低品質な画像を生成することが多い）と比較して、テイルクラス（例：'red wine', 'espresso', 'window shade'）に対して高い多様性と忠実度を持つ画像を生成することが示されました。
• アブレーション研究:
  • 両方の損失を組み合わせることが必要であり、一方のみを使用すると最適ではない結果となります。
  • 整合損失を大きなタイムステップで適用すること（線形重み付け）は、定数または二次重み付けよりも優れています。
  • トランスフォーマーのN/4番目のブロックからの潜在表現を使用することが、全体のFIDと空間的FID/Recallの間の最適なトレードオフを提供します。

重要性と主張

本論文は、CCUAが、ロングテールGANの原理（具体的には、低解像度または初期段階での無条件学習）を拡散パラダイムに適応させることで、ロングテール生成の問題を効果的に解決することを主張しています。多様化のタスク（無条件コントラスト学習によって処理）と条件付けのタスク（整合によって処理）を切り離すことで、本手法は直接的な条件付きコントラスト正則化の落とし穴を回避しています。

著者らは、彼らのアプローチが以下の点において重要であることを強調しています：

• ヘッドクラスの品質を低下させることなく、テイルクラスの生成を改善する。
• 既存のDDPMおよびSBDMパイプラインに容易に実装できる。
• 極端なクラス不均衡（例：不均衡係数が0.001のImageNet-LT）のシナリオにおいて、Dispersive LossやDiffROPなどの同時期の研究よりも優れた性能を示す。
• 逆プロセスの初期ステップで条件付き生成と無条件生成を整合させることが、ロングテール設定における知識共有にどのように有益であるかについて、理論的および経験的な正当性を提供する。