[논문 리뷰] Parallel Restarted SPIDER -- Communication Efficient Distributed Nonconvex Optimization with Optimal Computation Complexity

Pranay Sharma, Kafle, Swatantra외 4명|arXiv (Cornell University)|2019. 12. 12.

Stochastic Gradient Optimization Techniques참고 문헌 43인용 수 13

한 줄 요약

이 논문은 분산 비볼록 최적화 문제에 대해 통신 효율적인 알고리즘인 PR-SPIDER를 제안한다. 이 알고리즘은 병렬 재시작 SGD 프레임워크에 SPIDER 기반의 분산 추정 기반 분산 감소 기법을 통합하여, 유한합 및 온라인 설정 모두에서 최적의 $O(\tilde{\rho}^{-1})$ 통신 복잡도와 최적의 $O(\tilde{\rho}^{-1})$ 계산 복잡도를 달성한다. 이는 이전의 $O(\tilde{\rho}^{-2})$ IFO 복잡도에 비해 크게 향상되었으며, 선형 속도 향상 기능을 유지하면서 비.i.i.d. 데이터 분포를 지원한다.

ABSTRACT

In this paper, we propose a distributed algorithm for stochastic smooth, non-convex optimization. We assume a worker-server architecture where $N$ nodes, each having $n$ (potentially infinite) number of samples, collaborate with the help of a central server to perform the optimization task. The global objective is to minimize the average of local cost functions available at individual nodes. The proposed approach is a non-trivial extension of the popular parallel-restarted SGD algorithm, incorporating the optimal variance-reduction based SPIDER gradient estimator into it. We prove convergence of our algorithm to a first-order stationary solution. The proposed approach achieves the best known communication complexity $O(ε^{-1})$ along with the optimal computation complexity. For finite-sum problems (finite $n$), we achieve the optimal computation (IFO) complexity $O(\sqrt{Nn}ε^{-1})$. For online problems ($n$ unknown or infinite), we achieve the optimal IFO complexity $O(ε^{-3/2})$. In both the cases, we maintain the linear speedup achieved by existing methods. This is a massive improvement over the $O(ε^{-2})$ IFO complexity of the existing approaches. Additionally, our algorithm is general enough to allow non-identical distributions of data across workers, as in the recently proposed federated learning paradigm.

연구 동기 및 목표

• 대규모 분산 데이터 환경에서 분산 비볼록 최적화의 높은 통신 및 계산 비용을 해결하기 위해.
• 선형 속도 향상을 유지하면서 최적 수렴 복잡도를 달성하는 확장 가능한 알고리즘을 개발하기 위해.
• 작업자 간 비동일한 데이터 분포를 지원하여 페더레이티드 러닝 환경에 적용 가능하게 하기 위해.
• 유한합 및 온라인 비볼록 문제에 대해 반복 복잡도(IFO 복잡도)를 최적 속도로 감소시키기 위해.

제안 방법

• 알고리즘은 SPIDER 기반의 분산 감소 기법을 통합한 병렬 재시작 SGD를 확장하여 분산 환경에서의 분산 감소를 구현한다.
• 각 워커는 미니배치를 사용하여 국소 SPIDER 기울기를 계산하고, 서버는 각 에포크 종료 시 반복값과 기울기를 집계한다.
• 알고리즘은 매 $m$ 반복 후 새로운 초기값으로 재시작하는 메커니즘을 사용하여 수렴 성능을 향상시킨다.
• 통신은 각 에포크당 한 번만 발생하며, 분산 감소를 통해 기울기 정확도를 유지하면서 통신 횟수를 최소화한다.
• 유한합 및 온라인 문제 모두에 대해 배치 크기와 기울기 추정기 조정을 통해 알고리즘을 적응시킨다.
• 각 워커가 자체 로컬 함수와 기울기 추정기를 유지할 수 있도록 하여 비.i.i.d. 데이터 분포를 지원한다.

실험 결과

연구 질문

• RQ1 분산 비볼록 최적화 알고리즘이 최적의 통신 복잡도 $O(\tilde{\rho}^{-1})$ 를 달성하면서도 최적의 계산 복잡도를 유지할 수 있는가?
• RQ2 알고리즘은 유한합 및 온라인 비볼록 문제 모두에서 $O(\tilde{\rho}^{-1})$ IFO 복잡도를 달성할 수 있는가?
• RQ3 비.i.i.d. 데이터 분포가 존재하더라도 알고리즘이 워커 간 선형 속도 향상을 유지할 수 있는가?
• RQ4 SPIDER 기반 분산 감소 기법이 비볼록 최적화를 위한 분산 및 재시작 프레임워크로 효과적으로 확장될 수 있는가?

주요 결과

• 제안된 PR-SPIDER 알고리즘은 유한합 및 온라인 문제 모두에 대해 최적의 통신 복잡도 $O(\tilde{\rho}^{-1})$ 를 달성한다.
• 유한합 문제의 경우 IFO 복잡도는 $O\left(\frac{\sqrt{Nn}}{\epsilon}\right)$ 로서 최적이며, 이는 이전의 $O(\tilde{\rho}^{-2})$ 복잡도에 비해 크게 향상된 것이다.
• 온라인 문제의 경우 IFO 복잡도는 $O\left(\frac{\sigma}{\epsilon^{3/2}} + \frac{\sigma^2}{\epsilon}\right)$ 로서 최적이며, 이는 이전의 $O(\tilde{\rho}^{-2})$ 속도에 비해 큰 향상이다.
• 알고리즘은 $N$ 명의 워커 간 선형 속도 향상을 유지하여, 워커 수에 비례하여 수렴 속도가 유리하게 스케일링된다.
• 비.i.i.d. 데이터 분포에 대해 강건하여 페더레이티드 러닝 및 기타 분산 학습 환경에 적합하다.
• 이론적 분석을 통해 표준 smoothness 및 유한 분산 조건 하에 $\epsilon$-1차 정류점으로 수렴하는 것이 확인되었다.