[論文レビュー] Efficient Wireless Federated Learning with Partial Model Aggregation

Zhixiong Chen, Wenqiang Yi+2 more|arXiv (Cornell University)|Apr 20, 2022

Privacy-Preserving Technologies in Data被引用数 5

ひとこと要約

本稿では、パラメータサーバーで下位（特徴抽出用）の層のみをアグリゲートし、上位（予測用）の層をローカルに保って個人化を実現する、部分的モデルアグリゲーションによるフェデレーテッドラーニング（PMA-FL）を提案する。Lyapunov最適化を用いてデバイススケジューリング、帯域幅、計算時間割り当てを共同最適化することで、通信負荷を低減し、データ非独立同分布（non-IID）に対するロバスト性を向上させ、MNISTでは3.13%、CIFAR-10では11.8%の高い精度を達成した。ベンチマークと比較して、エネルギー消費を最大29%、時間消費を最大20%削減した。

ABSTRACT

The data heterogeneity across devices and the limited communication resources, e.g., bandwidth and energy, are two of the main bottlenecks for wireless federated learning (FL). To tackle these challenges, we first devise a novel FL framework with partial model aggregation (PMA). This approach aggregates the lower layers of neural networks, responsible for feature extraction, at the parameter server while keeping the upper layers, responsible for complex pattern recognition, at devices for personalization. The proposed PMA-FL is able to address the data heterogeneity and reduce the transmitted information in wireless channels. Then, we derive a convergence bound of the framework under a non-convex loss function setting to reveal the role of unbalanced data size in the learning performance. On this basis, we maximize the scheduled data size to minimize the global loss function through jointly optimize the device scheduling, bandwidth allocation, computation and communication time division policies with the assistance of Lyapunov optimization. Our analysis reveals that the optimal time division is achieved when the communication and computation parts of PMA-FL have the same power. We also develop a bisection method to solve the optimal bandwidth allocation policy and use the set expansion algorithm to address the device scheduling policy. Compared with the benchmark schemes, the proposed PMA-FL improves 3.13\% and 11.8\% accuracy on two typical datasets with heterogeneous data distribution settings, i.e., MINIST and CIFAR-10, respectively. In addition, the proposed joint dynamic device scheduling and resource management approach achieve slightly higher accuracy than the considered benchmarks, but they provide a satisfactory energy and time reduction: 29\% energy or 20\% time reduction on the MNIST; and 25\% energy or 12.5\% time reduction on the CIFAR-10.

研究の動機と目的

• フェデレーテッドラーニングにおけるデータ非独立同分布（non-IID）と限られた無線通信リソースの課題に対処する。
• 制限された帯域幅とエネルギー条件下で、非独立同分布データにおけるモデル収束性と精度を向上させる。
• 無線FLシステムにおける通信オーバーヘッドとエネルギー消費を低減する。
• デバイススケジューリング、帯域幅割り当て、計算時間割り当ての共同最適化フレームワークを構築する。
• 計算と通信のバランスを取ることで、低いエネルギーと時間コストでより優れた学習性能を達成する。

提案手法

• パラメータサーバーでニューラルネットワークの下位層のみをアグリゲートし、上位層をローカルに保持して個人化を実現する、新しいPMA-FLフレームワークを提案する。
• 非凸損失関数下でのPMA-FLの収束バウンドを導出することで、データサイズの不均衡が与える影響を分析する。
• Lyapunov最適化を用いて通信と計算エネルギーのバランスを保ちつつ、スケジューリング可能なデータサイズを最大化する共同最適化問題を定式化する。
• リソース制約下での最適帯域幅割り当て方針を求めるために、二分探索法を設計する。
• 最適なデバイススケジューリング方針を決定するための集合拡張アルゴリズムを実装する。
• 通信フェーズと計算フェーズ間の動的時間割り当てを統合し、両フェーズのパワーが等しいときに最適な性能を示すことを示した。

実験結果

リサーチクエスチョン

• RQ1 データ非独立同分布下で、部分的モデルアグリゲーションは無線フェデレーテッドラーニングの収束性と性能にどのように影響を与えるか？
• RQ2 PMA-FLにおいて、通信と計算エネルギーの最適なバランスは、グローバル損失を最小化するためにどのように設定されるか？
• RQ3 デバイススケジューリング、帯域幅割り当て、時間割り当ての共同最適化は、学習精度を向上させるとともにエネルギーと時間コストを削減できるか？
• RQ4 PMA-FLは、既存のベンチマークと比較して、厳格な遅延およびエネルギー制約下でどのように性能を発揮するか？
• RQ5 不均衡なデータサイズは、PMA-FLの収束性と精度にどのような影響を与えるか？

主な発見

• PMA-FLは、非独立同分布データ分布下で、ベンチマーク手法と比較してMNISTで3.13%、CIFAR-10で11.8%の高いテスト精度を達成した。
• 提案された共同最適化により、MNISTではエネルギー消費を29%、時間消費を20%削減し、CIFAR-10ではエネルギー消費を25%、時間消費を12.5%削減したが、精度は維持または向上させた。
• 20%の時間予算削減（例：1.5秒 vs. 2.5秒）の下で、提案アルゴリズムはOCEANよりも高い精度を達成した。
• 高いエネルギー制約（V=0.001）下でも、OCEANと同等の性能を発揮したが、よりエネルギー効率がよかった。
• 通信と計算パワーがバランスしている場合に最適な時間割り当て方針が得られ、全体のシステムコストが最小化された。
• Lyapunov重みパrameter V の調整により、精度とエネルギー消費のトレードオフを制御でき、V を適切にチューニングすることで最適な性能が得られた。