[논문 리뷰] Communication-efficient distributed SGD with Sketching

Nikita Ivkin, Daniel Rothchild외 4명|arXiv (Cornell University)|2019. 03. 12.

Stochastic Gradient Optimization Techniques참고 문헌 37인용 수 16

한 줄 요약

이 논문은 전체 기울기 대신 크기가 𝒪(log d)인 기울기 스케치를 전송함으로써 통신 효율성을 높인 분산 확률적 기울기 하강 알고리즘인 Sketched-SGD를 제안한다. 이는 하위선형 통신 복잡도를 달성하며, 트랜스포머, LSTM, 잔차 신경망에서 최대 40배의 총 통신 비용 절감을 이끌어내며 모델 정확도에 손실가지 않고 256명의 워커까지 효과적으로 확장된다.

ABSTRACT

Large-scale distributed training of neural networks is often limited by network bandwidth, wherein the communication time overwhelms the local computation time. Motivated by the success of sketching methods in sub-linear/streaming algorithms, we introduce Sketched SGD, an algorithm for carrying out distributed SGD by communicating sketches instead of full gradients. We show that Sketched SGD has favorable convergence rates on several classes of functions. When considering all communication -- both of gradients and of updated model weights -- Sketched SGD reduces the amount of communication required compared to other gradient compression methods from $\mathcal{O}(d)$ or $\mathcal{O}(W)$ to $\mathcal{O}(\log d)$, where $d$ is the number of model parameters and $W$ is the number of workers participating in training. We run experiments on a transformer model, an LSTM, and a residual network, demonstrating up to a 40x reduction in total communication cost with no loss in final model performance. We also show experimentally that Sketched SGD scales to at least 256 workers without increasing communication cost or degrading model performance.

연구 동기 및 목표

• 네트워크 대역폭이 학습 속도를 제한하는 대규모 분산 딥 러닝에서의 통신 병목 현상 해결.
• 각 워커의 통신 비용을 𝒪(d) 및 𝒪(W) 이하로 줄여 𝒪(log d)로 낮추어 많은 수의 워커에서 효율적인 확장성 확보.
• 높은 압축률을 달성하면서도 수렴 보장과 모델 성능을 유지하는 방법 개발.
• 워커 수(W) 증가에 따라 성능이 급격히 떨어지는 기존 기울기 압축 기법의 한계 극복.
• 스트리밍 알고리즘에서 유래한 스케칭을 활용해 통신 오버헤드를 최소화하고 실용적이고 확장 가능한 분산 학습 구현.

제안 방법

• Count Sketch와 희소 복구 기법을 사용해 각 워커의 기울기를 크기가 𝒪(log d)인 스케치로 압축.
• 전체 기울기 대신 로컬 기울기의 스케치만 전송함으로써 통신량을 𝒪(d)에서 𝒪(log d)로 감소.
• 오차가 제한된 스케칭 복구 알고리즘을 사용해 파rameter 서버에서 근사적인 전체 기울기를 복원.
• 스케칭이 효과적인 최적화를 위해 충분한 정보를 유지하도록 하여 수렴 성질 유지를.
• 표준 동기 데이터 병렬 SGD에 스케칭 기반 메커니즘을 통합하여 표준 학습 워크플로우 유지.
• 편향이 없는 및 편향이 있는 기울기 추정 모두 지원하며, 적절한 가정 하에 이론적 수렴 보장 제공.

실험 결과

연구 질문

• RQ1 스케칭을 통해 분산 SGD에서 통신 비용을 𝒪(d) 이하로 낮출 수 있으며, 수렴 성질을 유지할 수 있는가?
• RQ2 제안된 방법이 워커 수가 많아질수록 통신 비용 증가 없이 또는 성능 저하 없이 효과적으로 확장되는가?
• RQ3 스케칭을 통해 실제 모델에서 높은 압축 비율(예: 40배)을 달성하면서도 최종 모델 정확도가 떨어지지 않는가?
• RQ4 기존 기울기 압축 기법들(예: top-k, 양자화)과 비교해 Sketched-SGD는 통신 효율성과 확장성 측면에서 어떻게 다른가?
• RQ5 스케칭에 의해 유도된 기울기 근사화 하에 Sketched-SGD의 이론적 수렴 거동은 어떠한가?

주요 결과

• Sketched-SGD는 트랜스포머, LSTM, 잔차 신경망 모델에서 최대 40배의 총 통신 비용 절감을 이끌어내며 최종 모델 성능에 손실가지 않는다.
• 256명의 워커까지 효과적으로 확장되며 통신 비용 증가 없이 또는 모델 정확도 저하 없이 작동한다.
• 워커당 통신 비용은 𝒪(log d)로 줄어들며, 모델 크기 d에 대해 하위선형이고 워커 수 W에 대해 일정하다.
• MNIST에 대한 실험에서 Sketched-SGD는 일반 SGD와 동일한 학습 및 테스트 오차율을 달성하였으며, 이론적 수렴 속도와 일치한다.
• 소규모 실험에서 스케치 크기 280(40열, 7행)이 k=10 및 P=10 조건에서 약 4배의 압축 비율을 달성하였다.
• 이론적 분석을 통해 Sketched-SGD가 부드럽고 볼록 함수에 대한 표준 가정 하에 유리한 수렴 속도를 유지함을 확인하였다.