Entropy-Regularized Probabilistic Gates for Sparse Model Discovery in Scarce-Data Federated Learning

技術要約：データ希薄な連合学習における疎なモデル発見のためのエントロピー正則化確率ゲート

問題提起
連合学習（FL）は、パラメータ数（$d$）が総サンプル数（$N$）を大幅に上回る（$d \gg N$）、データが乏しく高次元な領域で動作する場合、重大な課題に直面します。このような設定では、標準的な最適化は、未知のテストデータに対して汎化できないパラメータ構成をもたらすことがよくあります。疎なモデルは、FLにおける通信および計算効率の面で利点を提供しますが、これらの疎な構造を発見することは困難です。既存の手法、例えばマグニチュードベースのプルーニングや連邦反復ハードしきい値処理（Fed-IHT）は、パラメータ空間の探索が不十分な場合が多く、最適ではない疎なサポートへの早期のコミットメントを招く傾向があります。さらに、標準的な正則化手法（L1/L2）は、特定のパラメータ密度ではなくパラメータの大きさに依存するため、所望のスパース性レベル $\rho$ を持つモデルの発見を困難にしています。

手法
本論文では、不確実性の探索を維持しながら、$L_0$ 密度制約の下で疎なモデルを学習するために設計された分散アルゴリズムである E-FLoPS（Entropy-Regularized Federated Learning with Probabilistic Gates）を提案します。

1. 確率的ゲートと再パラメータ化: 本手法は、有効なパラメータが $\theta = \tilde{\theta} \odot z$ と定義される確率的ゲート $z \in \{0,1\}^d$ のフレームワークを採用しています。勾配ベースの最適化を可能にするため、離散的なゲートは、ベルヌーイ分布の連続緩和である Hard Concrete 分布 を用いて近似されます。これにより、$L_0$ 擬ノルム（非ゼロパラメータの数）を、活性化されたゲートの期待値 $\mathbb{E}_q[z_j]$ によって近似することが可能になります。
2. エントロピー正則化: 本手法の主要な革新は、最適化目的関数へのエントロピー正則化項の導入です。ラグランジュ関数には、ゲート分布の相対エントロピーを罰するカルバック・ライブラー（KL）ダイバージェンス項 $KL(q(z_j) || p(z_j))$ が含まれます。このメカニズムは、ゲート分布が学習中に決定論的な 0/1 の状態へと早期に崩壊することを防ぎ、競合する疎な構成の探索を促進します。温度パラメータ $T$ は、この罰則の強さを制御します。
3. 連合最適化: アルゴリズムは、中央サーバーのオーケストレーションの下で動作します。クライアントは、生のパラメータ $\tilde{\theta}$、ゲートパラメータ $\phi = \log \alpha$、およびラグランジュ乗数 $\lambda$（目標密度 $\rho$ を強制するもの）に対して、ローカルな確率的勾配降下法（SGD）の更新を実行します。
   • サーバーは更新を同期的に集約します。
   • 通信効率を確保するため、アルゴリズムは学習中にプルーニングメカニズムを採用しています。パラメータとその対応するゲート統計量のうち、上位 $m$ 個のインデックス（ここで $m = \lfloor \rho \cdot |\theta| \rfloor$）のみがシリアライズされ、通信されます。
   • 受信時、サーバーはグローバルモデルを再構成し、ノイズ成分を無視して集約されたゲート値からゲートパラメータを復元します。

主な貢献

• 定式化: 著者らは、連合学習の設定に特化した、微分可能でエントロピー正則化された $L_0$ 制約付きの目的関数を定式化しました。この定式化は、$L_0$ ノルムの非微分性と、$d \gg N$ の領域における最適化の不安定性に対処しています。
• アルゴリズム (E-FLoPS): 本論文は、確率的ゲートとエントロピー正則化を統合した Federated Averaging の変種を紹介しています。これにより、不確実性の探索を維持しながら、ユーザー定義の目標スパース性を実現できます。これは、劣悪な局所解への早期収束を回避します。
• 通信効率: 本手法は、必要なゲート統計量と疎なサポートのみを送信することで、大幅な通信圧縮を実現しており、これは反復ハードしきい値処理法と同等の通信量でありながら、学習後のプルーニングのオーバーヘッドを伴いません。

実験結果
著者らは、合成データおよび実世界のベンチマークを用いて、E-FLoPS を Fed-IHT および事後学習プルーニングを伴う高密度 Federated Averaging (FedAvg) と比較評価しました。

• 合成線形回帰: 高次元設定において、サンプル対パラメータの比率（$N/|\theta|$）が変化する場合でも、E-FLoPS は一貫して、Fed-IHT や FedAvg よりも高いテスト $R^2$ スコアと、より優れたスパース性回復精度（真の機能が復元された割合で測定）を達成しました。これは、特に最もデータが乏しい領域において顕著でした。
• MNIST 画像分類: 97.5% の目標スパース性を設定した CNN タスクにおいて、E-FLoPS は Fed-IHT と同等のテスト時の性能および通信効率を示しました。しかし、推論時の積和演算（MACs）の減少においてより優れた結果を示しており、より高い計算効率を示唆しています。
• 白血病分類: Golub 遺伝子発現データセット（72 サンプル、3571 特徴量、$N/|\theta| \approx 0.02$）を用いた実験において、E-FLoPS は 30 個の異なるランダムシードにわたって、Fed-IHT および FedAvg を有意に上回るテスト精度を示しました。決定的なことに、E-FLoPS によって選択された遺伝子（例：cystatin C や myeloperoxidase）は、がん転移に関する先行研究の知見と一致しており、他の手法が選択した特徴量よりも生物学的な妥当性が高いものでした。

意義と主張
本論文は、E-FLoPS が、FL における疎なモデル発見のための堅牢なメカニズムを提供することを主張しています。具体的には、小サンプル・高次元領域で見られる不安定性に対処しています。エントロピー正則化を利用することで、本手法は、最適ではない疎なサポートへの早期のコミットを防ぎ、より優れた汎化性能と、真の疎な構造のより正確な復元を実現します。著者らは、本アプローチが、既存の疎な手法である Fed-IHT と同等の通信コストを維持しながら、これらの統計的な利点を得られると断言しています。これは、プライバシー保護とリソース制約のある学習環境において、極めて有効なソリューションとなります。本研究は、マグニチュードベースのプルーニング技術ではしばしば見落とされがちな、疎な連合最適化における「不確実性の探索」の重要性を強調しています。