[論文レビュー] ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!

Konstantin Mishchenko, Grigory Malinovsky+2 more|arXiv (Cornell University)|Feb 18, 2022

Stochastic Gradient Optimization Techniques被引用数 10

ひとこと要約

ProxSkip は、大多数の反復で高コストな近位作用素の評価を飛ばすことで、分散学習における通信を確実に高速化する新しい最適化手法である。非同期的でない条件下でも、$\epsilon$-精度解を $\epsilon(\kappa\log(1/\epsilon))$ 反復で達成でき、近位作用素の評価回数は $\epsilon(\sqrt{\kappa}\log(1/\epsilon))$ 回に抑えられる。これにより、非同一データ分布の仮定がない状況下でも、フェデレーテッドおよび分散型環境における通信コストを大幅に削減できる。

ABSTRACT

We introduce ProxSkip -- a surprisingly simple and provably efficient method for minimizing the sum of a smooth ($f$) and an expensive nonsmooth proximable ($ψ$) function. The canonical approach to solving such problems is via the proximal gradient descent (ProxGD) algorithm, which is based on the evaluation of the gradient of $f$ and the prox operator of $ψ$ in each iteration. In this work we are specifically interested in the regime in which the evaluation of prox is costly relative to the evaluation of the gradient, which is the case in many applications. ProxSkip allows for the expensive prox operator to be skipped in most iterations: while its iteration complexity is $\mathcal{O}\left(κ\log \frac{1}{\varepsilon} ight)$, where $κ$ is the condition number of $f$, the number of prox evaluations is $\mathcal{O}\left(\sqrtκ \log \frac{1}{\varepsilon} ight)$ only. Our main motivation comes from federated learning, where evaluation of the gradient operator corresponds to taking a local GD step independently on all devices, and evaluation of prox corresponds to (expensive) communication in the form of gradient averaging. In this context, ProxSkip offers an effective acceleration of communication complexity. Unlike other local gradient-type methods, such as FedAvg, SCAFFOLD, S-Local-GD and FedLin, whose theoretical communication complexity is worse than, or at best matching, that of vanilla GD in the heterogeneous data regime, we obtain a provable and large improvement without any heterogeneity-bounding assumptions.

研究の動機と目的

• 分散最適化における近位作用素の高コストな通信を軽減すること、特に勾配の平均化が高コストなフェデレーテッドラーニング環境を想定する。
• 収束保証を維持しつつ、近位作用素の評価回数を著しく削減する手法を開発すること。
• 特に非同一データ環境下で、既存の局所勾配法（例：FedAvg や SCAFFOLD）に比べ、通信複雑度の理論的改善を達成すること。
• 収束速度や精度を損なわずに、近位ステップの飛ばし方を理論的に裏付けたフレームワークを提供すること。

提案手法

• 大多数の反復で近位作用素の評価を飛ばすことができる近位勾配降下法の変種である ProxSkip を導入する。
• 確率的メカニズムを用い、近位ステップは確率 $p$ で実行され、それ以外は飛ばされるが、収束を維持するための補正項が導入される。
• 前回の反復に基づく制御変数を用いた確率的近似フレームワークを採用し、更新方向を補正する。
• 反復列の収束レートを $\mathcal{O}(\kappa \log(1/\varepsilon))$ と導出する一方で、近位作用素の評価回数は $\mathcal{O}(\sqrt{\kappa} \log(1/\varepsilon))$ に抑える。
• 共通化に基づく分散問題に適用し、$f(x) + \psi(\mathbf{L}x)$ を最小化する問題に再定式化する。ここで $\psi$ は共通化のための指示関数である。
• 混合行列 $\mathbf{W}$ を用いた分散型バージョンを導出し、通信が行列-ベクトル乗算 $\mathbf{L}x$ に対応することを示し、収束レート $\tilde{\mathcal{O}}(\kappa + 1/(p^2\delta))$ を理論的に証明する。

実験結果

リサーチクエスチョン

• RQ1 局所勾配ステップを用いることで、収束保証を損なわず通信コストを削減できるか？
• RQ2 高コストな近位作用素の評価回数を理論的に削減しつつ、同じ収束速度を維持できるか？
• RQ3 非同一データ環境下で、ProxSkip は FedAvg や SCAFFOLD よりも優れた通信複雑度を達成できるか？
• RQ4 収束速度の観点から、局所ステップ数と通信頻度の最適なトレードオフは何か？
• RQ5 一般の混合行列を用いた分散環境にも適用可能であり、依然として最適な収束を達成できるか？

主な発見

• ProxSkip は反復複雑度 $\mathcal{O}(\kappa \log(1/\varepsilon))$ を達成し、近位作用素の評価回数は $\mathcal{O}(\sqrt{\kappa} \log(1/\varepsilon))$ に抑えられ、通信コストの大幅な削減が可能である。
• 従来の ProxGD や FedAvg や SCAFFOLD などの局所法に比べ、理論的に通信複雑度の改善が保証され、特に非同一データ環境下で顕著である。
• 分散環境下では収束レートが $\tilde{\mathcal{O}}(\kappa + 1/(p^2\delta))$ となり、最適な $p = 1/\sqrt{\delta\kappa}$ が局所ステップと通信のバランスを取る。
• ネットワークが接続が悪い場合（$\delta \leq 1/\kappa$）、$p=1$（全反復で通信）が最適であり、既知の下界と一致する。
• w8a データセットを用いたロジスティック回帰の実験により、Scaffnew（ProxSkip の変種）が他の手法を上回り、理論的ハイパーパrameterと最適な $p \approx 1/\sqrt{\kappa}$ を用いた場合に顕著に優位であることが確認された。
• 確率的環境下では、クライアント数の増加に伴い線形加速を達成し、理論的期待を超える性能を示した。