[論文レビュー] STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning

Prashant Khanduri, Pranay Sharma+5 more|arXiv (Cornell University)|Jun 19, 2021

Privacy-Preserving Technologies in Data参考文献 35被引用数 20

ひとこと要約

本稿では、非凸な設定下でε-停留解を求めるために、近似的に最適なサンプル複雑度$ ilde{\mathcal{O}}(\epsilon^{-3/2})$と通信複雑度$ ilde{\mathcal{O}}(\epsilon^{-1})$を達成する、分散学習のための確率的二面的モーメンタムアルゴリズムであるSTEMを提案する。作業者とサーバーの両方のモーメンタムに基づく更新を、適応的ミニバッチサイズと局所更新頻度と併せて共同最適化することにより、STEMはこれらの複雑度を維持するトレードオフ曲線を確立し、FedAvgや先行手法を凌駕する性能を示す。

ABSTRACT

Federated Learning (FL) refers to the paradigm where multiple worker nodes (WNs) build a joint model by using local data. Despite extensive research, for a generic non-convex FL problem, it is not clear, how to choose the WNs' and the server's update directions, the minibatch sizes, and the local update frequency, so that the WNs use the minimum number of samples and communication rounds to achieve the desired solution. This work addresses the above question and considers a class of stochastic algorithms where the WNs perform a few local updates before communication. We show that when both the WN's and the server's directions are chosen based on a stochastic momentum estimator, the algorithm requires $ ilde{\mathcal{O}}(ε^{-3/2})$ samples and $ ilde{\mathcal{O}}(ε^{-1})$ communication rounds to compute an $ε$-stationary solution. To the best of our knowledge, this is the first FL algorithm that achieves such {\it near-optimal} sample and communication complexities simultaneously. Further, we show that there is a trade-off curve between local update frequencies and local minibatch sizes, on which the above sample and communication complexities can be maintained. Finally, we show that for the classical FedAvg (a.k.a. Local SGD, which is a momentum-less special case of the STEM), a similar trade-off curve exists, albeit with worse sample and communication complexities. Our insights on this trade-off provides guidelines for choosing the four important design elements for FL algorithms, the update frequency, directions, and minibatch sizes to achieve the best performance.

研究の動機と目的

• 非凸な分散学習における、作業者とサーバーの両方の更新方向、ミニバッチサイズ、局所更新頻度を同時に最適化するという未解決問題に取り組む。
• 非均質なデータを伴う確率的で分散型の設定において、サンプル複雑度と通信複雑度の両方を同時に最小化する。
• 局所更新頻度とミニバッチサイズの間の理論的トレードオフを確立し、近似的に最適な収束レートを維持する。
• FedAvg（Local SGD）に対する新たな洞察を提供し、その内在的なトレードオフと非最適な複雑度境界を示す。
• 標準的な仮定の下で、作業者およびサーバー両方のモーメンタムに基づく更新が、最適な収束を達成できることを示す。

提案手法

• 両方の作業者ノードとサーバーが、勾配のモーメンタム推定値を用いてモデル更新を行う、確率的二面的モーメンタムアルゴリズムであるSTEMを提案する。
• 局所更新頻度$I$、ミニバッチサイズ$b$、モーメンタムに基づく更新方向を同時に最適化する統一フレームワークを導入する。
• 勾配のノイズを低減し、学習を安定化させるために、モーメンタムによる分散低減を伴う確率的勾配推定器を採用する。
• 非凸な滑らかさおよび勾配の有界性の仮定の下で、リャプノフ解析と強凸性に類似した不等式を用いて収束境界を導出する。
• $ ilde{\mathcal{O}}(\epsilon^{-3/2})$のサンプルと$ ilde{\mathcal{O}}(\epsilon^{-1})$の通信ラウンドを達成するように、$I$と$b$の間のトレードオフ曲線を確立する。
• STEMの特殊ケース（モーメンタムなし）としてのFedAvgを分析し、比較のためのその非最適な複雑度境界を導出する。

実験結果

リサーチクエスチョン

• RQ1 非凸な設定下で、分散学習アルゴリズムが同時に近似的に最適なサンプル複雑度と通信複雑度を達成できるか？
• RQ2 確率的勾配を伴う分散学習において、局所更新頻度とミニバッチサイズの最適なトレードオフは何か？
• RQ3 作業者およびサーバー両方の側にモーメンタムを組み込むと、非凸なFLにおける収束複雑度にどのような影響を与えるか？
• RQ4 同じ仮定のもとで、FedAvg（Local SGD）の理論的サンプルおよび通信複雑度は何か？
• RQ5 最小限のリソース使用を実現するため、更新方向、ミニバッチサイズ、局所更新頻度を同時に最適化する統一フレームワークを開発できるか？

主な発見

• STEMは、非凸な分散学習において、ε-停留解を求めるための$\tilde{\mathcal{O}}(\epsilon^{-3/2})$のサンプル複雑度と$\tilde{\mathcal{O}}(\epsilon^{-1})$の通信複雑度を達成する。
• ミニバッチサイズ$b$と局所更新頻度$I$の間には、STEMがこの複雑度を近似的に最適に維持するためのトレードオフ曲線が存在する。
• FedAvg（STEMのモーメンタムなしバージョン）では、サンプル複雑度が$\mathcal{O}(\epsilon^{-2})$、通信複雑度が$\mathcal{O}(\epsilon^{-3/2})$となり、STEMと比較して非最適である。
• $b = \mathcal{O}(1)$かつ$I = (T/b^3K^3)^{1/4}$のとき、STEMは$\tilde{\mathcal{O}}(\epsilon^{-3/2})$のサンプル複雑度と$\tilde{\mathcal{O}}(\epsilon^{-1})$の通信複雑度を維持する。
• 解析により、確率的分散学習における最適な収束レートを達成するには、作業者およびサーバー両方の側でモーメンタムが不可欠であることが明らかになった。
• 導出された境界は標準的な仮定のもとでタイトであり、STEMはこの設定で同時に近似的に最適なサンプルおよび通信複雑度を達成する最初のアルゴリズムである。