技術要約：Mixture-of-Expertsを用いたドメイン汎化のための部分集合共有不変性の学習

1. 問題設定

ドメイン汎化（Domain Generalization, DG）は、訓練中にターゲットデータにアクセスすることなく、一つまたは複数のソースドメインで学習したモデルを、未知のターゲットドメインへ汎化させることを目的としている。DGにおける主流のアプローチは、ドメイン不変な表現を学習することに依存しており、通常、すべてのソースドメイン間で条件付き分布 $P(Z|Y)$ が一貫するように強制する。特徴量アライメント、Invariant Risk Minimization (IRM)、勾配マッチングなどの手法は、単一の、グローバルに共有された予測構造が存在するという仮定の下で動作している。

しかし、著者らはこの仮定がしばしば制限的すぎると主張している。現実世界のシナリオでは、ドメイン間において予測構造が部分的にしか重なっていない場合がある。すなわち、特定の機能は一部のドメインでは安定しているが、他のドメインでは変化する場合がある。ヘテロジニアス（異質）なドメインの数が増えるにつれて、グローバルな不変性を強制することは、実現可能な表現空間を漸進的に制限してしまう。理論的分析および経験的観察（図1）は、ソースドメインの数が増加するにつれて、学習された表現とラベルの間の相互情報量（$I(Z; Y)$）が単調に減少することを示している。その結果、特定のドメインの部分集合内でのみ安定している貴重な予測因子が破棄され、ドメインのヘテロジニティが高い場合に汎化性能の低下を招く。

2. 手法：MESSI

グローバルな不変性の限界に対処するため、本論文では MESSI (Mixture-of-Experts for Subset-Shared Invariances) を提案する。核心となる仮説は、不変な構造はグローバルに共有されるのではなく、**部分集合共有（subset-shared）**としてモデル化されるべきであるという点である。この手法は、ドメイン関係の構造化された分解を学習するために、Mixture-of-Experts (MoE) アーキテクチャを採用している。

2.1 アーキテクチャ

MESSIは、共有バックボーンエンコーダ（例：DeiT）に続いて、ルーターと複数のエキスパートネットワークを利用する：

1. 共有エンコーダ: 特徴表現 $u = b_\phi(x)$ を生成する。
2. ルーター: エキスパートに対するソフトな分布 $\pi(x) = \text{softmax}(g_\theta(u))$ を計算し、潜在的な部分集合変数 $S$ を近似する。
3. エキスパート: 各エキスパート $m$ は、特定の表現 $z^{(m)} = h_m(u)$ を生成する。
4. 集約: 最終的な表現は加重和 $z(x) = \sum \pi_m(x) z^{(m)}$ となり、部分集合不変な成分の適応的な合成を可能にする。

2.2 学習目的関数

モデルは、複合損失関数（式5）を最適化するように訓練される：
$$\mathcal{L} = \mathcal{L}_{cls} + \lambda_{ssi}\mathcal{L}_{OT}^{ssi} + \lambda_{sp}\mathcal{L}_{sp} + \lambda_{bal}\mathcal{L}_{bal} + \lambda_{div}\mathcal{L}_{div}$$

• ルーティング条件付き部分集合アライメント ($\mathcal{L}_{OT}^{ssi}$): すべてのドメインペアをアライメントする代わりに、不変性を選択的に強制する。各エキスパート $m$ について、そのエキスパートが特定のクラス $c$ に対して高いルーティング責任を持つドメインペア $(i, j)$ の間でのみ、クラス条件付き特徴分布をアライメントする。アライメントの重み $a_{ijc}^{(m)}$ はルーティング確率から導出される。不一致は、これらのルーティング質量によって重み付けされた Optimal Transport (OT)（またはMMDのバリアント）を用いて測定される。これにより、「ルーティング誘発型不変性」が作成され、モデルは「どこで」不変性が保持されるべきかを学習する。
• エキスパートの専門化 ($\mathcal{L}_{div}$): エキスパートが冗長な表現に崩壊するのを防ぐため、異なるエキスパート間の特徴方向のコサイン類似度を最小化する多様性正則化項を用いる。
• 確信度とバランスの取れたルーティング ($\mathcal{L}_{sp}, \mathcal{L}_{bal}$):
  • $\mathcal{L}_{sp}$ は、インスタンスレベルでのスパースなルーティング（低エントロピー）を促進し、各サンプルが特定のエキスパートの部分集合を活性化するようにする。
  • $\mathcal{L}_{bal}$ は、データセットレベルでのバランスの取れたエキスパート利用を保証し、少数のエキスパートがトラフィックを支配することを防ぐ。

3. 主な貢献

本論文は、主に3つの貢献を述べている：

1. 根本的な限界の特定: 著者らは、不変性を強制することが予測情報を減少させることを実証した。彼らは 部分集合共有不変性（subset-shared invariance） を導入し、DGを、単一のグローバルな不変表現ではなく、構造化された部分集合依存の不変性を学習する問題へと再定義した。
2. ルーティングベースのMoEフレームワーク: ドメイン関係のデータ駆動的な分解を誘導するために、ルーティングを使用する手法を提案した。エキスパートが担当するドメインの部分集合に対してのみアライメントを行うことで、明示的な部分集合ラベルを必要とせずに、適応的な不変成分の合成を可能にする。
3. 実証的検証: このアプローチは、標準的なベンチマーク、特にヘテロジニアスなドメインシフトが存在し、ソースドメインの数が増加するシナリオにおいて、優れたアウトオブドメイン汎化性能を示す。

4. 実験結果

実験は、標準的なDGベンチマーク（DomainBed: PACS, OfficeHome, TerraIncognita, DomainNet）および、ドメイン成長をテストするために設計された制御されたRotated-Colored MNISTデータセットを用いて行われた。

• 標準ベンチマーク: MESSI（特にOTを用いたDeiT-S/16のバリアント）は、DomainBedにおいて平均精度66.2%を達成し、ERM、CORAL、Fishr、SAGM、およびGMoEやOMoEといった他のMoEベースの手法を含む強力なベースラインを上回った。
• ドメイン成長に対する堅牢性: 固定予算ドメイン成長プロトコルにおけるRotated-Colored MNISTにおいて、MESSIは優れた堅牢性を示した。グローバルなアライメント手法（CORALなど）は、ソースドメインが増加するにつれて（過剰な制約により）性能低下に見舞われたが、MESSIは高い精度を維持した（最小の低下は1.86%、最高のCORALバリアントは2.58%の低下）。
• アブレーション研究:
  • 部分集合条件付きアライメント ($\mathcal{L}_{ssi}$) を削除すると最大の性能低下が発生し、これが決定的な要素であることを裏付けた。
  • 多様性またはルーティングの正則化項を削除すると、エキスパートの冗長性とルーティングの拡散が生じ、性能が低下した。
  • コントロール群（Global-MoE, Random-Subset-MoE）との比較により、MESSIが単にグローバルな目的関数をスパース化しているのではなく、明確に異なるアライメント構造を学習していることが示された。
• 効率性: OTベースのバリアントは訓練オーバーヘッド（ERMの約12.39倍）を導入するが、MMDバリアントはこれを約4.24倍に軽減する。推論コストは他のMoE手法と同等である。

5. 意義と主張

本論文は、ドメイン汎化は単一のグローバルな不変性を強制するというパラダイムを超えて進むべきであると主張している。代わりに、予測構造は本質的に 部分集合依存的（subset-dependent） であると論じている。

• 理論的洞察: 本研究は、相互情報量分析を通じて、なぜグローバルな不変性がヘテロジニティが増すにつれて失敗するのかについて理論的な正当化を提供している。それは、グローバルな不変性が、普遍的に共有されていない予測因子を破棄することを強いるためである。
• 実用的な意味合い: 部分集合にわたる部分的に共有された構造を通じて不変性をモデル化することで、MESSIはより高い堅牢性を達成する。ルーティングメカニズムにより、明示的な部分集合ラベルや定義済みのドメイングループを必要とせずに、データから直接これらの部分集合を発見できる。
• 限界: 著者らは、この手法がルーターの専門化能力に依存していることを謙虚に認めている。ソースドメインが少ない場合やクラスのカバー範囲が不均衡な場合、ルーティング信号がノイズになる可能性がある。さらに、特にOptimal Transportを用いる場合のペアワイズアライメントの計算コストは、標準的なグローバルアライメント手法よりも高い。

結論として、MESSIは部分集合レベルの不変性を保持し構成するためのメカニズムを提供しており、DGの未来は、一様なグローバルアライメントではなく、構造化された適応的な分解にあることを示唆している。