Variational Consensus Monte Carlo for Bayesian Mixture

技術要約：ベイズ混合モデルのための変分コンセンサス・モンテカルロ法

問題提起
本論文は、フェデレーテッドラーニング（FL）の設定、特にデータが自然に複数のノード（例：医療センター）に分割されており、プライバシーや規制上の理由から統合できない「クロスサイロ」シナリオにおける、混合モデルのベイズ推論という課題に取り組んでいる。コンセンサス・モンテカルロ（CMC）は、局所的な事後分布のサンプルを集計することによる分散型ベイズ推論の枠組みを提供しているが、既存の手法を非教師ありクラスタリングに適用する場合、重大な制限に直面する。具体的には、Rabinovich, Angelino, and Jordan (2015) が提案した変分コンセンサス・モンテカルロ（VCMC）のアプローチは、クラスター数と主要な混合パラメータが既知であり、かつすべてのクラスターがすべてのデータシャードに存在することを前提としている。これらの仮定は、ヘルスケアデータのアプリケーションにおいて、クラスターの出現頻度がサイトごとに異なり、真の潜在的サブグループ数が未知である場合など、実務上ほとんど成立しない。さらに、標準的なCMC集計（例：加重平均）は、グローバルな事後分布の最適な近似を保証するものではなく、既存のフェデレーテッド・クラスタリング手法は、共役事前分布を必要とする平均場近似に依存しているか、全データにアクセスせずに非共役の設定を扱うことができない。

手法
著者らは、VCMCをカテゴリカル分布のオーバーフィットされたベイズ混合モデルへと拡張する包括的なパイプラインを提案している。このフレームワークは、主に2つのフェーズで動作する：

1. 適用ステップ（局所推論）：

   • 各データシャード内で独立してMCMCアルゴリズム（具体的にはギブスサンプラー）を実行する。
   • クラスター数の未知性を扱うため、モデルを「オーバーフィット」させ、混合成分の数（$C'$）を想定される真のクラスター数（$C$）よりも大きく設定する。
   • フラクショネイテッド・プライア（分画事前分布）を使用し、グローバルな事前分布 $p(\theta)$ をシャード間で $p(\theta)^{1/S}$ として分割する。これにより、局所的な事後分布の積がグローバルな事後分布を近似することを保証する。
   • 局所推論を加速させるために、並列化されたC++17ギブスサンプラーを実装している。

2. 集計ステップ（グローバル推論）：

   • クラスター・マッチング： 混合モデルはラベルの置換に対して不変であるため、クラスター間のマッチングは極めて重要なステップとなる。本論文では、以下の3つのアルゴリズムを導入し、評価している：
     • ハンガリアン・マッチング（Hungarian Matching）： 一対一の対応を強制するが、すべてのクラスターがすべてのシャードに存在しない場合には失敗する。
     • 最小ダイバージェンス・マッチング（Minimum Divergence Matching）： 局所的な周辺事後分布間のカルバック・ライブラー（KL）ダイバージェンスを最小化する。
     • ボール・マッチング（Ball Matching）： パラメータ空間において、オーバーフィットされた混合モデルの収束率に基づき、クラスターが「近い」（半径 $r$ 以内）場合にそれらをグループ化する新しいアルゴリズム。この手法は、識別不可能なクラスターをシャード内でマージすることを可能にし、局所的なクラスター数の過大評価を軽減する。
   • 変分集計： 集計は、近似された事後分布と真のグローバル事後分布との間のKLダイバージェンスを最小化する変分推論問題として定式化される。集計関数は、局所パラメータの加重和（$\sum \upsilon_s \theta_s$）に制限される。
   • 最適化： 最適な集計ウェイトは、射影確率的勾配降下法を用いて求められる。本論文では、異なるFLの制約に対応するために、以下の4つの最適化戦略を詳述している：
     1. 生データをグローバルノードに共有する（最も制約が少ない）。
     2. フェデレーテッド勾配降下法（勾配のみを共有）。
     3. 限定的な通信（少数のイテレーションと複数の開始点を用いた並列実行）。
     4. コンティンジェンシー行列（要約統計量）をグローバルノードに共有する。

主な貢献
本論文は、主に4つの方法論的貢献を行っている：

1. オーバーフィット・モデルへの拡張： クラスター数の既知性を必要とせず、クラスター数とすべてのモデルパラメータを同時に推論するために、VCMCを拡張した。
2. 新規クラスター・マッチング・アルゴリズム： クラスターが一部のシャードに存在しない可能性があるクロスサイロの設定に適したマッチング戦略（特に最小ダイバージェンスおよびボール・マッチング）を導入した。これは標準的なハンガリアン・マッチングが失敗するシナリオである。
3. 柔軟な推論戦略： さまざまなフェデレーテッド制約（通信帯域の制限や厳格なプライバシー要件など）に合わせて調整された複数の集計戦略を提案した。
4. 実践的なガイドライン： データ特性（クラスターの分離性やシャードの不均一性など）に基づいた、マッチングアルゴリズムと推論戦略の選択に関する経験的なガイドラインを提供した。

結果
本フレームワークは、包括的なシミュレーション研究と、実世界の電子健康記録（EHR）データへの適用を通じて評価された。

• シミュレーション研究：

  • マッチング性能： すべてのクラスターがすべてのシャードに現れない不均一な設定では、ハンガリアン・アルゴリズムは正しいマッチングの回復に失敗した（非入れ子シナリオにおいて成功率0%）。ボール・マッチングと最小ダイバージェンス・マッチングは、大幅に優れた性能を示した。特にボール・マッチングは、局所的なMCMCノイズによって引き起こされるクラスター数の過大評価を軽減するのに効果的であった。
  • クラスタリングの質： VCMCは、最先端の代替手法（FedMerDelおよびFederated K-modes）と同等、あるいはそれ以上に優れた調整ランド指数（ARI）スコアを達成し、「ゴールドスタンダード」であるフルデータMCMCに匹敵した。特筆すべきは、シャード間のデータ分布が不均一な場合、VCMCが小さなクラスターの回復においてフルデータMCMCを上回ったことである。
  • パラメータ推定： VCMCは、クラスターの重みに対して低い相対バイアスを提供した。クラスターの分離性が低い、あるいはシャードの分布が不均一なシナリオでは、VCMCはフルデータMCMCよりも小さなクラスターの重みをより正確に推定した。これは、これらの小さなクラスターが特定のシャード内では、より低い不確実性を持って識別可能なほど十分に大きかったためと考えられる。
  • 代替手法との比較： FedMerDelは高速であったが、平均場近似に依存しており、事後分布の分散を過小評価していた。Federated K-modesは高いバイアスを示し、カテゴリカルデータにおけるクラスター構造の推定において劣悪な性能を示した。

• 実世界への適用（THINコホート）：

  • 本手法は、多疾患併存パターンを特定するために、英国のTHINデータベースから取得した289,821人の高齢患者に適用された。
  • アルゴリズムは27のクラスターを特定した。脳卒中、認知症、または特定の癌種が支配的なクラスターなど、臨床的に意味のあるグループを正常に特徴づけた。
  • 決定的なことに、本手法は、小さながらも明確なクラスタ（例：膵炎、関節リウマチ、勃起不全によって特徴づけられるわずか31人からなるグループ）を特定した。これは文書化された医学的相関関係と一致しており、データのプール分析において希釈される可能性のある希少なサブグループを検出できる能力を実証している。

意義と主張
著者らは、本フレームワークがプライバシーに配慮した環境における非教師ありクラスタリングの堅牢なソリューションを提供すると主張している。その意義は、以下の能力にある：

• データの集約なしに、クラスター数とモデルパラメータを推論すること。
• ローカルのデータシャードに特定のサブグループに対する十分なシグナルが存在する場合、標準的なプールされたMCMCよりも正確に、小さく希少なクラスターを回復すること。
• 共役事前分布やフルデータへのアクセスを必要とせず、厳格なプライバシー制約（クロスサイロFL）の下で動作すること。

本論文は、本フレームワークが他の混合分布にも理論的に拡張可能であることを控えめに述べているが、現在の検証はカテゴリカル分布に特化したものである。また、クラスター・マッチング・アルゴリズムの選択は特定の目的に依存することも強調している。すなわち、クラスター数の推定が目的であればボール・マッチングが好まれ、明確なクラスターの結合を避けることが優先される場合は最小ダイバージェンス・マッチングが適している。著者らは、VCMCは変分法（FedMerDelなど）と比較して計算負荷が高い可能性があるものの、複雑で不均一なフェデレーテッド設定におけるパラメータ推定とクラスター回復において、より優れた精度を提供すると結論づけている。