技術要約：最適輸送によるOne-Shot Federated Learningのための分布アライメント

1. 問題提起

本論文は、クライアントが中央サーバーと正確に一度だけ通信を行う通信レジームである**One-Shot Federated Learning (OSFL)**における課題に対処している。標準的なFederated Learning (FL)は、反復的なクライアント・サーバー間の最適化を通じてデータの不均一性を緩和するが、OSFLにはこの洗練（refinement）能力が欠けている。

特定された核心的な問題は、以下の2種類の統計的不均一性がもたらす累積効果である：

1. ドメインシフト (Domain Shift): 取得条件や環境の違いにより、クライアントは異なる入力分布 $P_k(x)$ を観測する。
2. ラベルシフト (Label Shift): クライアント間で異なるラベル周辺分布 $P_k(y)$ が存在し、クラス比率の差異を生じさせる。

反復的な修正が行われない場合、これらのシフトは**特徴表現の不整合（misaligned feature representations）**と、クライアント間での一貫性のないクラス条件付き決定関数を誘発する。蒸留、データ合成、またはアンサンブルリングに依存する既存のOSFL手法は、多くの場合、アライメントされた表現を前提としているか、あるいはこれらのシフトを個別に扱っている。その結果、特徴量の幾何学的構造やクラス条件付きバイアスの不一致が修正されず、集約されたグローバルモデルの性能を低下させてしまう。

2. 手法：SLOT-Align

著者らは、ダウンストリームのOSFLトレーニングの前にクライアントの特徴分布を調和させるために設計された、幾何学に配慮したフレームワークであるSLOT-Align (Single-round, Learning-free Optimal Transport Alignment) を提案する。この手法は前処理ステップとして機能し、共通の潜在空間にすべてのローカル入力をマッピングするために、（外部データセットで事前学習された）共有の凍結エンコーダ $f_\theta$ に依存する。

ワークフローは、以下の4つの主要な段階で構成される：

A. ローカル特徴統計量の抽出

各クライアント $k$ は、ローカルデータから抽出された特徴 $z = f_\theta(x)$ のコンパクトな1次および2次の統計量を計算する：

• 平均 (Mean): $\mu_k = \mathbb{E}[f_\theta(x)]$
• 共分散 (Covariance): $\Sigma_k$ （正定値性を保証するためにLedoit–Wolf収縮法を用いて正規化される）。
  これらの統計量は、ドメインシフトおよびラベルシフトによって引き起こされる位置（location）と広がり（spread）の主要な差異を捉える、ガウス近似 $N(\mu_k, \Sigma_k)$ として機能する。

B. グローバル参照の構築

サーバーは、これらのローカル統計量を集約してグローバルな参照分布を構築する：

• グローバル平均: クライアント平均の加重平均。
• グローバル共分散: クライアントの共分散の**Bures–Wasserstein重心（Barycenter）**として計算される。これは、Bures距離の二乗和の最小化因子を見つけるための固定点反復によって解かれ、クライアント固有の変動性とグローバルな一貫性を効果的にバランスさせる。

C. 測地線最適輸送マッピング (Geodesic Optimal Transport Mapping)

各クライアントは、自身のローカルなガウス近似をグローバルな参照へと整列させるための輸送マップを構築する。

• 最適マップ: 二次コストの下では、2つのガウス分布間の最適輸送マップはアフィンであり、共分散行列の平方根を含む閉形式の解を持つ。
• 補間 (Interpolation): 真の特徴分布は厳密にはガウス分布ではないため、過剰な修正を防ぐ目的で、補間パラメータ $\tau \in [0, 1]$ を導入する。最終的な変換 $T^{(\tau)}_k$ は、恒等写像と完全な最適輸送マップの凸結合となる。これにより、**Wasserstein変位補間（Wasserstein displacement interpolation）**が生成され、クライアントの分布がグローバルな参照に向かって一定速度の測地線に沿って移動する。

D. 統合

クライアントは、変換 $z' = T^{(\tau)}_k(z)$ をローカルな特徴に適用する。これらの整列された特徴は、ダウンストリームの最適化手順を変更することなく、既存のOSFLアルゴリズム（例：FedAvg, FedCGS, FedPFT）に供給される。

3. 主な貢献

1. 定式化: 本論文は、ドメインおよびラベルの共同シフト下における分布アライメント問題としてOSFLを定式化し、これらのシフトがいかにワンショット手法の有効性を制限するかを明らかにしている。
2. アルゴリズムの提案: 統計量の単一の交換によってクライアントの特徴分布を調和させる、**学習フリー（learning-free）**で最適輸送に基づいたフレームワークであるSLOT-Alignを導入する。これには補助的なデータ、学習可能な生成器、または反復的な最適化は必要ない。
3. 実証的検証: 広範な実験により、SLOT-Alignが、様々なプリトレーニング済みバックボーン（ViT, ResNet）および最先端のOSFLベースラインに対し、複数のベンチマーク（Office-Home, Digits, DomainNet）において、一貫して精度と堅牢性を向上させることが示された。

4. 実験結果

著者らは、様々な程度のラベルシフト（ディリクレ集中パラメータ $\alpha$ によって制御）およびドメイン多様性の条件下でSLOT-Alignを評価した。

• 性能向上: ドメインおよびラベルの共同シフト（$\alpha=0.1$）において、SLOT-Alignはテストされたすべてのベースライン（O-FedAvg, FedCGS, FedPFT）に対して、マクロ平均精度を一貫して向上させた。例えば、Office-Homeデータセットにおいて、O-FedAvgを**+7.86%、FedPFTを+6.55%**向上させた。
• 堅牢性: 強力なラベルスキュー（$\alpha=0.05$ および $\alpha=0.01$）の下でも効果を維持し、深刻な不均一性レジームにおいても利得を示した。
• アライメント強度に関するアブレーション解析: 補間パラメータ $\tau$ が極めて重要であることが判明した。中間的な値（例：$\tau=0.4$）が最良のトレードオフをもたらし、$\tau=1.0$（完全な輸送）では性能が低下する場合があり、部分的なアライメントが完全な修正よりも好ましいことが示唆された。
• ドメインのみのシフト: ラベルが均衡している設定（ドメインシフトのみ）では、SLOT-Alignはわずかな利得または中立的な性能を示し、本手法がドメインとラベルのシフトが複雑に相互作用する状況を特に対処していることを裏付けた。
• 効率性: 通信コストは最小限であり（クライアントあたり平均と共分散に対して $O(m + m^2)$）、計算オーバーヘッドは特徴抽出およびダウンストリームのトレーニングと比較して無視できる程度である。

5. 意義と主張

本論文は、SLOT-Alignが、OSFLにおける決定的なボトルネック、すなわち単一の通信ラウンド内で分布の不整合を修正できないという問題に対する、モジュール型で計算効率の高い解決策を提供すると主張している。

• モジュール性: 既存のOSFLパイプラインを置き換えるものではなく、互換性のある前処理レイヤーとして機能するため、凍結エンコーダに依存する幅広い現在および将来の手法に適用可能である。
• 幾何学的認識: 2-Wasserstein空間の幾何学とBures–Wasserstein重心を活用することで、特徴分布の位置（平均）と形状（共分散）の両方の不一致を明示的に修正する。
• 実用性: データ合成や反復的な蒸留に伴うプライバシーリスクや計算負荷を回避し、コンパクトな統計的要約のみに依存する。

著者らは、OSFLは厳格な通信予算によって制約を受けるものの、SLOT-Alignは、特にドメインシフトとラベルシフトが相互作用して特徴統計量を歪ませる場合に、不均一なデータ分布によって引き起こされる性能低下を効果的に軽減できると結論付けている。