[論文レビュー] FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning

Hong-You Chen, Wei‐Lun Chao|arXiv (Cornell University)|Sep 4, 2020

Privacy-Preserving Technologies in Data参考文献 55被引用数 28

ひとこと要約

FedBEは、非i.i.i.d.データおよび深層ニューラルネットワーク下でのロバスト性を向上させるためのベイジアンモデルアンサンブルに基づくアグリゲーション手法を提案する。局所モデルに適合した分布から高品質なグローバルモデルをサンプリングし、確率的重み平均化を用いた知識蒸留により、特に深層ネットワークおよび非均質な設定下でFedAvgを上回る優れた精度を達成する。

ABSTRACT

Federated learning aims to collaboratively train a strong global model by accessing users' locally trained models but not their own data. A crucial step is therefore to aggregate local models into a global model, which has been shown challenging when users have non-i.i.d. data. In this paper, we propose a novel aggregation algorithm named FedBE, which takes a Bayesian inference perspective by sampling higher-quality global models and combining them via Bayesian model Ensemble, leading to much robust aggregation. We show that an effective model distribution can be constructed by simply fitting a Gaussian or Dirichlet distribution to the local models. Our empirical studies validate FedBE's superior performance, especially when users' data are not i.i.d. and when the neural networks go deeper. Moreover, FedBE is compatible with recent efforts in regularizing users' model training, making it an easily applicable module: you only need to replace the aggregation method but leave other parts of your federated learning algorithm intact. Our code is publicly available at https://github.com/hongyouc/FedBE.

研究の動機と目的

• 非i.i.i.d.データ設定下でモデルドリフトや一般化性能の低下が生じるFedAvgの性能劣化を解消する。
• 単純なモデル平均化の限界を克服し、不確実性を捉えるためにベイジアンモデルアンサンブルを活用して予測のロバスト性を向上させる。
• 複数ラウンドのフェデレーテッドラーニングを効果的に行うために、アンサンブルの予測を1つのグローバルモデルに蒸留し、クライアントが再利用できるようにする。
• クライントレーニング手順を変更せずに、アグリゲーションステップのみを置き換えることで、既存のフェデレーテッドラーニングフレームワークとの互換性を確保する。
• サーバのラベルなしデータとテストデータの間で分布シフトが生じても、そのロバスト性を保証し、実用的応用性を高める。

提案手法

• 局所クライントレーニングモデルに適合したガウス分布またはディリクレ分布を用いてモデル分布を構築し、多様なグローバルモデル候補をサンプリング可能にする。
• 複数回にわたるサンプリングで得られたグローバルモデルの予測を集約することで、ベイジアンモデルアンサンブルを実現し、ロバスト性と精度を向上させる。
• アンサンブル予測に対してラベルなしサーバデータを用いて偽ラベルを生成し、1つのグローバルモデルへの知識蒸留を可能にする。
• 蒸留過程で確率的重み平均化（SWA）を適用し、ノイズの多いアンサンブル予測への過剰適合を防ぎ、一般化性能を向上させる。
• アンサンブル予測を教師信号として用い、学生グローバルモデルを蒸留損失関数に基づいて学習させる。
• クライアント側のトレーニングとモーメンタム機構を維持したまま、アグリゲーションステップのみを置き換えることで、既存のFLフレームワークにFedBEを統合する。

実験結果

リサーチクエスチョン

• RQ1 非i.i.i.d.データ条件下において、FedAvgと比較してベイジアンモデルアンサンブルがフェデレーテッドラーニングのモデルアグリゲーションをどのように改善できるか？
• RQ2 ResNetのようなより深いニューラルネットワークを用いた場合、FedAvgが通常劣化する状況でFedBEはどのように性能を発揮するか？
• RQ3 サーバのラベルなしデータを蒸留に使用した場合、特にテストデータと分布が異なる場合にどのような影響が生じるか？
• RQ4 部分的参加（1ラウンドあたり100人中10人のクライアント）やシステムの非均質性といった実用的制約下でも、FedBEは性能を維持できるか？
• RQ5 既存の正則化手法（例：FedProx）と組み合わせた場合、全体の性能にどのような影響があるか？

主な発見

• 非i.i.i.d. Tiny-ImageNetでResNet20を用いた場合、FedAvg（32.4%）に対してFedBEは5.9%高い精度（35.4%）を達成。MobileNetV2では3.4%高い（25.5% vs. 27.8%）。
• 非i.i.i.d. CIFAR-10では、ResNet20を用いたFedBEが77.1%の精度を達成し、FedAvg（69.9%）およびFedProx（69.4%）を大きく上回った。
• サーバのラベルなしデータがテストデータと異なる分布に属しても、FedBEは高い精度を維持し、ドメインシフトに対してロバストであることが示された。
• 部分的参加（100人中10人）のシステム下でも、Tiny-ImageNet（ResNet20）で35.4%の精度を達成し、FedAvg（32.4%）を上回った。
• システムの非均質性（異なるローカルトレーニングエポック数）下でも、CIFAR-10（ResNet20）で77.1%の精度を達成し、FedAvg（69.9%）およびFedProx（69.4%）を上回った。
• FedProxと組み合わせた場合、さらに性能が向上し、CIFAR-10（ResNet20）で77.5%の精度を達成した。既存の正則化手法との互換性および相乗効果が示された。