[논문 리뷰] ProxSkip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!

Konstantin Mishchenko, Grigory Malinovsky외 2명|arXiv (Cornell University)|2022. 02. 18.

Stochastic Gradient Optimization Techniques인용 수 10

한 줄 요약

ProxSkip는 대부분의 반복에서 비용이 많이 드는 프락시얼 연산 평가를 건너뛰어 분산 학습에서 통신을 증가시키는 새로운 최적화 방법이다. 이는 오직 $\mathcal{O}(\sqrt{\kappa}\log(1/\varepsilon))$회의 프락시얼 평가로 $\varepsilon$-정확도 해를 달성하며, 이는 이질성 가정 없이 피어드 및 분산 환경에서 통신 비용을 크게 감소시킨다.

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}$ 설정에서 다른 방법들을 능가한다.
• 스토하스틱 환경에서 클라이언트 수가 증가함에 따라 선형 속도 향상을 달성하며, 이론적 기대를 초월한다.