[논문 리뷰] A Communication-Efficient Distributed Gradient Clipping Algorithm for Training Deep Neural Networks

Mingrui Liu, Zhenxun Zhuang외 2명|arXiv (Cornell University)|2022. 05. 10.

Stochastic Gradient Optimization Techniques인용 수 6

한 줄 요약

이 논문은 기울기 폭주 문제로 인해 기울기 폭주가 발생하기 쉬운 RNN 및 LSTMs와 같은 모델을 위한 통신 효율적인 분산 기울기 클리핑 알고리즘을 제안한다. 각 머신에서 국소 기울기 클리핑을 적용하고 모델 가중치를 주기적으로만 집계함으로써, 이 방법은 $O(1/N\epsilon^4)$ 반복 복잡도와 $O(1/\epsilon^3)$ 통신 복잡도를 달성하며, 이는 느슨한 부드러움 조건 하에서 선형 속도 향상과 통신 라운드 수 감소를 가능하게 한다.

ABSTRACT

In distributed training of deep neural networks, people usually run Stochastic Gradient Descent (SGD) or its variants on each machine and communicate with other machines periodically. However, SGD might converge slowly in training some deep neural networks (e.g., RNN, LSTM) because of the exploding gradient issue. Gradient clipping is usually employed to address this issue in the single machine setting, but exploring this technique in the distributed setting is still in its infancy: it remains mysterious whether the gradient clipping scheme can take advantage of multiple machines to enjoy parallel speedup. The main technical difficulty lies in dealing with nonconvex loss function, non-Lipschitz continuous gradient, and skipping communication rounds simultaneously. In this paper, we explore a relaxed-smoothness assumption of the loss landscape which LSTM was shown to satisfy in previous works, and design a communication-efficient gradient clipping algorithm. This algorithm can be run on multiple machines, where each machine employs a gradient clipping scheme and communicate with other machines after multiple steps of gradient-based updates. Our algorithm is proved to have $O\left(\frac{1}{Nε^4} ight)$ iteration complexity and $O(\frac{1}{ε^3})$ communication complexity for finding an $ε$-stationary point in the homogeneous data setting, where $N$ is the number of machines. This indicates that our algorithm enjoys linear speedup and reduced communication rounds. Our proof relies on novel analysis techniques of estimating truncated random variables, which we believe are of independent interest. Our experiments on several benchmark datasets and various scenarios demonstrate that our algorithm indeed exhibits fast convergence speed in practice and thus validates our theory.

연구 동기 및 목표

• RNN 및 LSTMs와 같은 깊은 네트워크를 훈련할 때 기울기 폭주 문제로 인해 발생하는 분산 SGD의 느린 수렴 문제를 해결하기 위해.
• 매번 단계마다 전역 기울기 평균을 취할 필요 없이, 드물게 통신하는 분산 환경에서도 기울기 클리핑을 가능하게 하기 위해.
• 비-Lipschitz, 비볼록, 느슨한 부드러움 조건이 적용된 손실 곡면 하에서 선형 속도 향상과 통신 복잡도 감소를 보장하는 이론적 보장을 수립하기 위해.
• Wikitext-2, Penn Treebank, ImageNet과 같은 벤치마크 데이터셋을 통해 이론적 주장의 타당성을 실증적으로 검증하기 위해.

제안 방법

• 각 머신에서 통신 이전에 다수의 국소 SGD 단계와 함께 기울기 클리핑을 수행하는 분산 알고리즘을 제안한다.
• 전역 평균 기울기 대신 각 머신의 국소 기울기 노름 기반 기울기 클리핑을 적용하여 통신 효율성을 유지한다.
• 표준 부드러움 조건을 일반화하고 LSTMs의 행동을 더 잘 반영하는 $(L_0, L_1)$-부드러움 조건을 도입한다.
• 비-Lipschitz 및 비볼록 설정 하에서 수렴 보장을 가능하게 하기 위해, 잘라낸 랜덤 변수를 추정하는 새로운 분석 기법을 사용한다.
• 모델 가중치를 $I$개의 국소 업데이트 단계 이후에만 집계하는 통신 전략을 도입하여 통신 빈도를 감소시킨다.
• 이론적 분석을 통해 $\epsilon$-정류점(\epsilon$-stationary point)을 찾는 데 $O(1/N\epsilon^4)$ 반복 복잡도와 $O(1/\epsilon^3)$ 통신 복잡도를 증명한다.

실험 결과

연구 질문

• RQ1 드물게 통신하는 환경에서 기울기 클리핑을 효과적으로 확장할 수 있으며, 수렴 보장이 유지되는가?
• RQ2 느슨한 부드러움 조건 하에서 국소 기울기 클리핑과 주기적 전역 집계가 분산 훈련에서 선형 속도 향상을 달성하는가?
• RQ3 비볼록, 비-Lipschitz 설정에서 기계적 병렬 기울기 클리핑에 비해 통신 복잡도를 감소시킬 수 있는가?
• RQ4 이 알고리즘은 이질적, 부분 참여, 동질적 데이터 설정에서 어떻게 성능을 발휘하는가?
• RQ5 실제 벤치마크인 AWD-LSTM과 ResNet-50에서 빠른 수렴과 일반화 성능을 유지하는가?

주요 결과

• 제안된 알고리즘은 $O(1/N\epsilon^4)$ 반복 복잡도와 $O(1/\epsilon^3)$ 통신 복잡도를 달성하여 선형 속도 향상과 통신 횟수 감소를 나타낸다.
• Wikitext-2 데이터셋에서 $I=32$일 때 검증 퍼플렉서티가 76.72로, 단일 GPU 기준선인 75.63과 유사한 성능을 기록했으며, 분산 훈련을 사용함에도 불구하고 성능을 유지했다.
• Penn Treebank에서 $I=1024$일 때 검증 퍼플렉서티는 66.5976을 기록하여 동질적 설정에서 기계적 병렬 기준선인 66.0435를 초월했다.
• 부분 참여 설정에서 $I=16$일 때 검증 퍼플렉서티는 67.3863로 강력한 성능을 유지했으며, 기계적 병렬 기준선인 65.5083보다 열등하지 않았다.
• 이질적 설정에서 $I=32$일 때 검증 퍼플렉서티는 70.4912를 기록하여 기계적 병렬 기준선인 68.5727을 초월했다.
• 모멘타임을 함께 사용한 알고리즘은 $I=4$일 때 검증 퍼플렉서티가 64.1062를 기록하여 기계적 병렬 기준선인 68.2008보다 뚜렷이 뛰어난 성능을 보였다.