[論文レビュー] FedBN: Federated Learning on Non-IID Features via Local Batch Normalization

Xiaoxiao Li, Meirui Jiang+3 more|arXiv (Cornell University)|Feb 15, 2021

Privacy-Preserving Technologies in Data参考文献 40被引用数 77

ひとこと要約

FedBN keeps Batch Norm parameters local in federated learning to address feature shift non-IID data, improving convergence and accuracy over FedAvg and FedProx across benchmarks and real-world datasets.

ABSTRACT

The emerging paradigm of federated learning (FL) strives to enable collaborative training of deep models on the network edge without centrally aggregating raw data and hence improving data privacy. In most cases, the assumption of independent and identically distributed samples across local clients does not hold for federated learning setups. Under this setting, neural network training performance may vary significantly according to the data distribution and even hurt training convergence. Most of the previous work has focused on a difference in the distribution of labels or client shifts. Unlike those settings, we address an important problem of FL, e.g., different scanners/sensors in medical imaging, different scenery distribution in autonomous driving (highway vs. city), where local clients store examples with different distributions compared to other clients, which we denote as feature shift non-iid. In this work, we propose an effective method that uses local batch normalization to alleviate the feature shift before averaging models. The resulting scheme, called FedBN, outperforms both classical FedAvg, as well as the state-of-the-art for non-iid data (FedProx) on our extensive experiments. These empirical results are supported by a convergence analysis that shows in a simplified setting that FedBN has a faster convergence rate than FedAvg. Code is available at https://github.com/med-air/FedBN.

研究の動機と目的

• 連携クライアント間で局所的な特徴分布が異なる federated learning における特徴シフト非 IID を動機づけ、対処する。
• BN パラメータを局所的に保持（平均化しない）ようにしつつ、非 BN 層には FedAvg を用いる軽量な手法を提案する。
• 特徴シフトの下で収束を速める理論的収束解析を提供する。
• ベンチマークおよび実データセットで FedBN を実証し、FedAvg および FedProx より明確な改善を示す。

提案手法

• FedBN を導入し、局所更新を行い非 BN 層を集約する一方、BN パラメータは厳密に局所的で伝達されない。
• アーキテクチャの一部として二層ネットワークと Batch Normalization をモデル化し、BN パラメータの平均化を省略する FedBN の目的関数を定義する。
• ニューラルタンジェントカーネル（NTK）フレームワークを用いた収束解析を行い、FedBN の収束は補助的な Gram 行列の最小固有値によって支配され、特徴シフトの設定において FedBN は FedAvg より高速な収束率を達成する。
• ベンチマークおよび実データセットにおける実験プロトコルでは、FedBN を FedAvg、FedProx、SingleSet と比較し、データレジームおよびヘテロジニティの変化に対する収束挙動と精度を評価する。

実験結果

リサーチクエスチョン

• RQ1 フェデレーション学習における特徴シフトは FedAvg を劣化させるのか、局所 BN によってこの劣化を緩和できるのか。
• RQ2 特徴シフト非 IID データの下で FedBN は FedAvg および FedProx と比べて収束速度と最終精度を向上させることができるか。
• RQ3 特徴シフトの度合いが異なるベンチマーク領域および実データセットで、FedBN はどのように性能を発揮するか。
• RQ4 局所更新頻度やクライアントデータサイズの違いに対して FedBN はロバストか。

主な発見

• FedBN は、特徴シフト非 IID 設定の下でベンチマークおよび実データセットにおいて一貫して FedAvg および FedProx を上回る。
• FedBN は実験で FedAvg よりも収束が速く、損失曲線が安定している。
• FedBN は特徴シフトが顕著なデータセット（例: SVHN）でより大きな精度向上を示し、実行間の分散を低く保つ。
• 実世界タスク（Office-Caltech10、DomainNet、ABIDE）では、複数のクライアント/データセットに渡って他の手法より高い平均精度を達成している。
• 理論解析は、特徴シフト下の G-dominated レジームにおいて FedBN が FedAvg より速い収束率を持つことを示唆している。