[논문 리뷰] PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization

Thijs Vogels, Sai Praneeth Karimireddy외 1명|arXiv (Cornell University)|2019. 05. 31.

Advanced Data Compression Techniques참고 문헌 35인용 수 96

한 줄 요약

PowerSGD는 오차 피드백과 all-reduce 집계를 기반으로 한 고유 차수의 그래디언트 압축기를 도입하여 분산 학습에서 SGD와 유사한 정확도를 달성하면서 통신 속도를 크게 향상시킵니다.

ABSTRACT

We study gradient compression methods to alleviate the communication bottleneck in data-parallel distributed optimization. Despite the significant attention received, current compression schemes either do not scale well or fail to achieve the target test accuracy. We propose a new low-rank gradient compressor based on power iteration that can i) compress gradients rapidly, ii) efficiently aggregate the compressed gradients using all-reduce, and iii) achieve test performance on par with SGD. The proposed algorithm is the only method evaluated that achieves consistent wall-clock speedups when benchmarked against regular SGD with an optimized communication backend. We demonstrate reduced training times for convolutional networks as well as LSTMs on common datasets. Our code is available at https://github.com/epfml/powersgd.

연구 동기 및 목표

• 깊은 학습을 위한 데이터 병렬 분산 최적화의 통신 병목 현상을 동기화하고 해결합니다.
• all-reduce 집계를 지원하는 확장 가능한 선형 그래디언트 압축기를 개발합니다.
• 오차 피드백을 통해 비편향 압축에서 수렴과 테스트 정확도를 유지합니다.
• CNN과 LSTM에 대한 전체 정밀도 SGD 대비 벽시계 시간 단축을 모든 데이터 세트 및 하드웨어에서 입증합니다.
• 오픈 소스 구현과 함께 실용적이고 적용 가능한 방법을 제공합니다.

제안 방법

• 각 그래디언트 행렬 M을 PQ^T로 근사하는 랭크-r 그래디언트 압축기를 도입합니다. 이전 단계에서의 워름 스타팅을 사용한 부분공간/전력 반복의 한 단계로 근사합니다.
• 선형성을 활용해 W개의 워커 간 압축된 그래디언트의 all-reduce 기반 집계를 가능하게 합니다.
• 후압축 모멘텀을 통한 오차 피드백을 적용해 편향 압축에서도 수렴을 가능하게 합니다.
• 레이어별로 그래디언트를 행렬로 분해하고 각 행렬을 독립적으로 압축합니다. r이 작은 경우 P∈R^{n×r} 및 Q∈R^{m×r}를 사용합니다(1–4).
• PQ^T로 디컴프레싱하고 모멘텀 SGD(EF-SGD with Momentum)로 분산 업데이트를 수행합니다.
• CIFAR-10/ResNet18 및 Wikitext-2/LSTM에서 SGD 및 다른 압축기(Signum, Spectral Atomo)와 비교하여 벽시계 시간, 전송 데이터, 정확도 등을 보고합니다.

실험 결과

연구 질문

• RQ1 전력 반복을 기반으로 한 저랭크 그래디언트 압축기가 정확도 손실 없이 통신을 크게 줄일 수 있을까요?
• RQ2 오차 피드백이 비편향 압축 방식의 수렴성과 비선형 딥러닝 설정에서의 로버스트함을 가능하게 하나요?
• RQ3 Power iteration의 워름 스타트가 근사 품질과 훈련 정확도에 시간에 따라 어떤 영향을 미치나요?
• RQ4 CNN과 RNN/LSTM 모델에서 랭크, 통신 오버헤드, 엔드투 엔드 학습 시간 사이의 트레이드오프는 어떻게 되나요?
• RQ5 PowerSGD가 워커 수 및 서로 다른 통신 백엔드(NCCL vs GLOO)에서 어떻게 확장되나요?

주요 결과

• PowerSGD는 최적화된 NCCL 백엔드를 사용한 16-GPU 설정에서 일반 SGD 대비 벽시계 속도 향상을 달성하는 동시에 그래디언트를 120배 이상 압축합니다.
• CIFAR-10의 ResNet-18에서 Rank-2 PowerSGD는 에폭당 8 MB를 전송하면서 테스트 정확도 94.4%를 달성했으며, SGD는 1023 MB로 94.3%를 달성했습니다.
• Wikitext-2의 LSTM에서 언어 모델링의 경우 Rank-4 PowerSGD가 SGD와 비슷한 perplexity를 달성하면서 에폭당 전송 데이터가 약 99% 감소하고 엔드투엔드 학습 시간이 약 55% 감소합니다.
• 오차 피드백이 있는 PowerSGD는 편향 압축에서의 테스트 정확도에서 비편향 랭크 기반 압축을 능가합니다(예: 편향 Rank-2는 94.4%에 도달하는 반면 비편향은 비슷한 설정에서 약 75–76% 수준).
• Power iteration의 워름 스타트는 최고 랭크-r 근사치까지의 간극을 줄이고 테스트 정확도를 유지하거나 향상시키면서 SVD 비용을 피합니다.
• 이 방법은 전체-reduce를 통해 워커 수가 증가할수록 확장성이 좋으며, 최적화된 백엔드에서 대안(Signum)보다 속도와 확장성 면에서 우수합니다.