[논문 리뷰] MARINA: Faster Non-Convex Distributed Learning with Compression

Eduard Gorbunov, Konstantin Burlachenko외 2명|arXiv (Cornell University)|2021. 02. 15.

Stochastic Gradient Optimization Techniques참고 문헌 45인용 수 15

한 줄 요약

MARINA는 비볼록 분산 학습에서 기울기 차이 압축을 사용하는 통신 효율적인 1차 방법을 제안하며, 뉘앙스 있는 기울기 추정기를 도입하여 이론적 및 실험적으로 뛰어난 성능을 달성한다. 통신 복잡도를 $\mathcal{O}\left(\frac{1 + \omega / \sqrt{n}}{\varepsilon^2}\right)$로 감소시켜, $\varepsilon$-정류점(ε-stationary point)을 찾는 데 있어 기존 방법보다 통신 및 오라클 복잡도에서 뛰어난 성능을 발휘한다.

ABSTRACT

We develop and analyze MARINA: a new communication efficient method for non-convex distributed learning over heterogeneous datasets. MARINA employs a novel communication compression strategy based on the compression of gradient differences that is reminiscent of but different from the strategy employed in the DIANA method of Mishchenko et al. (2019). Unlike virtually all competing distributed first-order methods, including DIANA, ours is based on a carefully designed biased gradient estimator, which is the key to its superior theoretical and practical performance. The communication complexity bounds we prove for MARINA are evidently better than those of all previous first-order methods. Further, we develop and analyze two variants of MARINA: VR-MARINA and PP-MARINA. The first method is designed for the case when the local loss functions owned by clients are either of a finite sum or of an expectation form, and the second method allows for a partial participation of clients -- a feature important in federated learning. All our methods are superior to previous state-of-the-art methods in terms of oracle/communication complexity. Finally, we provide a convergence analysis of all methods for problems satisfying the Polyak-Lojasiewicz condition.

연구 동기 및 목표

• 분산 비볼록 최적화에서의 통신 병목 현상을 해결하며, 특히 페더레이티드 및 이질적 학습 환경에서의 적용을 고려한다.
• 효율적인 압축을 통해 통신 비용을 최소화하면서도 높은 수렴 속도를 유지하는 방법을 개발한다.
• 기존 방법의 한계를 극복하기 위해 압축된 통신에 특화된 뉘앙스 있는 기울기 추정기를 도입한다.
• 기존 최신 기술 대비 개선된 이론적 통신 및 오라클 복잡도 한계를 달성한다.
• 부분적 클라이언트 참여와 같은 실용적 구현 시나리오, 예를 들어 유한 합 또는 기대 기반 국소 목표 함수를 지원한다.

제안 방법

• 압축된 기울기 차이를 바탕으로 한 뉘앙스 있는 기울기 추정기를 사용하는 1차 방법인 MARINA를 도입하며, DIANA의 무편향 접근과는 다릅니다.
• 기대 밀도 $\zeta_{\mathcal{Q}}$를 갖는 양자화 연산자 $\mathcal{Q}$를 활용하여 기울기를 압축하면서 편향을 제어합니다.
• 비볼록 문제에 대해 선형 수렴을 보장하는 폴리악-로자예프스키 조건(Polyak-Łojasiewicz, PL) 하에서 수렴 보장을 도출합니다.
• 두 가지 변형을 제안: 유한 합 및 기대 기반 국소 손실을 위한 VR-MARINA, 그리고 페더레이티드 학습에서 부분적 클라이언트 참여를 위한 PP-MARINA입니다.
• 수렴 분석과 오차 전파 제어를 위해 라플라스 함수 $\Phi_k = f(x^k) - f(x^*) + \frac{\gamma}{p}\|g^k - \nabla f(x^k)\|^2$를 사용합니다.
• 수렴 속도와 오차 누적 간 균형을 맞추기 위해 스텝 사이즈 $\gamma$를 최적화하여 $\mathbf{E}[\Phi_{k+1}] \leq (1 - \gamma\mu)\mathbf{E}[\Phi_k]$를 확보합니다.

실험 결과

연구 질문

• RQ1 비볼록 분산 학습에서 뉘앙스 있는 기울기 추정기가 무편향 방법보다 통신 효율성이 뛰어날 수 있는가?
• RQ2 비볼록 분산 최적화에서 $\varepsilon$-정류점(ε-stationary point)을 찾는 데 최적의 통신 복잡도는 무엇인가?
• RQ3 부분적 클라이언트 참여는 압축을 적용한 페더레이티드 학습에서 수렴과 통신 비용에 어떤 영향을 미치는가?
• RQ4 기울기 차이 압축이 표준 기울기 압축보다 더 나은 이론적 한계를 달성할 수 있는가?
• RQ5 PL 조건은 압축된 분산 방법에 대해 더 탴튼한 수렴 한계를 가능하게 하는가?

주요 결과

• MARINA는 통신 복잡도 $\mathcal{O}\left(\frac{1 + \omega / \sqrt{n}}{\varepsilon^2}\right)$를 달성하며, DIANA의 $\mathcal{O}\left(\frac{1 + (1+\omega)\sqrt{\omega/n}}{\varepsilon^2}\right)$보다 향상된 성능을 보인다.
• VR-MARINA는 압축된 업데이트에서 분산 감소를 활용하여, 유한 합 및 기대 기반 문제의 오라클 복잡도를 감소시킨다.
• PP-MARINA는 라운드당 $r$명의 클라이언트를 샘플링하여 부분 참여를 지원하며, 통신 비용으로 $\mathcal{O}\left(dn + \max\left\{dn, \frac{L}{\mu}\left(\zeta_{\mathcal{Q}}r + \sqrt{(1+\omega)\zeta_{\mathcal{Q}}(dn - \zeta_{\mathcal{Q}}r)}\right)\right\} \log \frac{\Delta_0}{\varepsilon}\right)$을 달성한다.
• PL 조건 하에서, MARINA는 $\mathbf{E}[f(x^K) - f(x^*)] \leq \varepsilon$를 만족하는 선형 수렴을 보장하며, $K = \mathcal{O}\left(\frac{dn}{\zeta_{\mathcal{Q}}r} \frac{L}{\mu} \left(1 + \sqrt{\frac{1+\omega}{r}\left(\frac{dn}{\zeta_{\mathcal{Q}}r} - 1\right)}\right)\right)$라운드 내에 이를 달성한다.
• 이 방법의 뉘앙스 있는 추정기는 압축 오차와 수렴 속도 사이의 더 나은 트레이드오���을 가능하게 하여, DIANA나 FedCOMGATE와 같은 무편향 방법보다 이론적·실제로 뛰어난 성능을 발휘한다.
• 실험 결과는 딥러닝 벤치마크에서 MARINA가 고압축 조건에서도 통신 효율성과 수렴 속도에서 뛰어난 성능을 보임을 확인한다.