Alistarh, Dan, et al. "QSGD: Communication-efficient SGD via gradient quantization and encoding." Advances in Neural Information Processing Systems. 2017.
Abstract
SGD의 병렬 구현은 뛰어난 확장성 덕분에 상당한 관심을 받고 있다. SGD 병렬화의 bottleneck은 gradient update를 하는데 필요한 통신비용이다. 이를 해결하기 위해, 각 노드가 quantized gradient를 전달하는 손실 압축 휴리스틱이 제안되었다. 이는 통신비용을 줄이는데있어 효과적이지만 수렴성이 보장되지는 않는다.
이 논문에서 수렴을 보장하면서 성능이 우수한 압축 스키마인 QSGD (Quantized SGD)를 제안한다. QSGD를 사용하면 사용자가 communication bandwidth와 convergence time간의 상충관계를 조절할 수 있다. QSGD는 비동시성하에서 convex, non-convex 문제에 대한 수렴을 보장한다.
이 기술은 image classification, speech recognition에 대해 end-to-end training time을 상당히 줄인다. 이 논문의 실험에서는 ResNet-152 네트워크를 accuracy 손실 없이 ImageNet을 1.8배 빠른 속도로 학습시켰다.
Introduction
방대한 데이터의 급증은 분산 학습 알고리즘에 대해 상당한 관심을 불러 일으켰다. 그 중 SGD에 대해 많은 연구를 진행하였으며, SGD는 다음 절차를 반복하여 수렴한다.
$x_{t+1} = x_t - \eta_t \tilde{g}(x_t)$
이는 신경망 훈련과 같은 많은 fundamental task를 캡처한다
이 논문에서는 높은 확장성으로 인해 관심받고있는 병렬 SGD 방법에 초점을 맞춘다. Loss function을 최소화하는 K개 프로세서 사이에서 큰 dataset이 분할되는 설정을 고려한다. 각 프로세서는 local copy parameter를 가지고 있는다. 그리고 프로세서는 gradient 업데이트를 주변에게 broadcast하고 aggregation을 진행하여 새로운 $x_{t+1}$을 계산한다.
위의 구현에서는 gradient update를 다른 모든 프로세서에게 전달해야한다. Gradient vector가 dense한 경우 n개의 floating-point number를 주고 받아야한다. 실제로 gradient를 통신하는 것은 상당한 bottleneck이 된다.
QSGD는 두가지 아이디어를 기반으로한다. 첫번째는 stochastic quantization으로 확률적으로 반올림하여 gradient를 반올림함으로써, 통계적 속성을 보전한다. 두번째는 gradient quantization을 위한 효율적인 무손실 코드이다. 이는 precision-variance trade-off에 대한 tight bound를 제공한다.
중요한 점은 QSGD가 수렴을 위한 반복 오버헤드를 상쇄할 정도로 통신 비용을 줄일 수 있는지의 여부이다. ImageNet에서 Image classification task, LSTM의 Speech recognition에서 실험을 진행하였고, 정확도 손실없이 표준 매개변수 하에서 다중 GPU training을 수행할 때 통신 감소의 이점을 크게 누릴 수 있는 것으로 확인되었다. 예를들어, 16개의 다중 GPU 학습에서 AlexNet을 훈련할 때 통신 시간이 4배 정도 감소함을 보인다. 일부 설정에서는 오히려 정확도를 약간 향상시킬 수 있다.
Preliminaries
$\mathcal{X}$는 알려진 convex set, $f : \mathcal{X} \rightarrow \mathbb{R}$은 differentiable, convex, smooth, unknown이라 하자. $f$에 대해서 SGD method로 접근한다고 가정하자. $f$에 대한 stochastic gradient $\tilde{g}(x)$는 다음을 만족한다.
$\mathbb{E}[\tilde{g}(x)] = \nabla f(x)$
$\mathbb{E}[||\tilde{g}(x) - \nabla f(x)||_2^2] \le \sigma^2$
Minibatched SGD는 위에서 보였던 SGD를 수정하여 나타낼 수 있다.
$x_{t+1} = \prod_{\mathcal{X}}(x_t - \eta_t \tilde{G_t}(x_t))$
$\tilde{G_t}(x_t) = {1 \over m} \sum_{i=1}^{m}{\tilde{g_{t,i}}}$
Data-Parallel SGD
여기서는 synchronous data-parallel SGD를 고려한다. 전체적인 학습 프로세스는 아래 알고리즘으로 보일 수 있다.
한 iteration에서, 여러개의 프로세서로부터 random gradient update를 얻은다음 다른 노드들에게 update를 전달한다. 모든 노드들에 수신된 업데이트를 집계한 뒤 이를 local에 적용한다. 다른 노드와 통신하기 전에 한가지 step이 추가되는데, 이는 encoding/decoding step으로 gradient를 보내거나 받을 때 실행된다.
Quantized Stochastic Gradient Descent (QSGD)
Stochastic quantization
quantization function은 다음과 같이 정의된다.
$Q_s(v_i) = ||v||_2 \cdot sgn(v_i) \cdot \eta_i(v, s)$
$s \ge 1$는 tuning parameter으로, quantization level을 조절하는 변수이다. Quantization 된 변수들은 uniformly distribution을 가진다고 가정한다. $\eta_i(v, s)$는 independent random variable이고 아래의 식으로 정의한다.
$\eta_i(v, s) = \begin{cases} l/s, & \mbox{with probability} \,\, 1- p({{|v_i|} \over {||v||_2}} , s) \\ (l+1)/s & \mbox{otherwise} \end{cases}$
$p(a, s) = as - l$이다. 이 quantization은 아래 그림을 통해 쉽게 이해할 수 있다.
위의 quantization 식은 $\{ 0, 1/s, ... , 1 \}$ 분포에 대해 최소 분산을 갖는다. 또한, expectation $\mathbb{E} = |v_i| / ||v||_2$를 만족한다.
Gradient quantization에 대한 encoding/decoding을 효과적으로 실행하기 위해 Elias 방식을 채택할 수 있다. 자세한 내용은 논문에 언급되어있다.
QSGD Guarantees
QSGD는 smooth, convex 문제에 대해 communication과 variance를 보장할 수 있다. 예를들어, $s = \sqrt{n}$인 경우 communication cost는 $2.8n + 32$로 감소한다. 이는 smooth non-convex optimization에 대해서도 동일하게 증명할 수 있다.
QSGD Variants
딥러닝에 QSGD를 활용하기 위해, 즉 non-convex 문제에 대해 사용하기 위해 구현을 확장한다. 첫 번째로 고정 크기 d의 버킷으로 quantize하여 quantization의 분산을 제어할 수 있다. 두 번째로 2-norm을 사용하지 않고 vector의 maximum 값으로 scaling을 진행한다. 이는 더 조밀한 분포를 갖게 하여 조금 높은 accuracy에 도달한다.
Experiments
이 실험에서는 32bit precision에 최적화된 하이퍼 파라미터를 사용했다. 그리고 계산과 동시에 communication과 quantization을 실행하기 위해 double buffering 방법을 사용한다. 작은 gradient array의 경우 통신 비용 감소에 영향이 크지 않으므로 quantize하지 않는다. 이를 적용하여도 모든 parameter의 99%이상이 quantize 된 형태로 통신된다.
병렬 처리 하에, 학습중 computation과 communication cost를 비교한다. 대부분의 모델은 계산 집약적 모델 (Inception, ResNet), 통신 집약적 모델 (AlexNet, VGG, LSTM)으로 분류할 수 있다. 두 분류 모두 GPU수가 많을 수록 통신의 상대적인 비용이 증가한다. 그리고 QSGD (4bit)를 사용했을 때 communication cost가 크게 감소함을 볼 수 있다. (반투명한 box) 결과적으로, 1 epoch당 training 시간이 크게 감소함을 관찰할 수 있다.
그 다음 실험은 QSGD가 accuracy에 얼마나 큰 영향을 미치는지를 확인한다. 위 그림 (a), (b)는 학습 시간에 따라 accuracy를 측정한 결과이다. 최종 accuracy는 크게 손실되지 않았으며, AlexNet의 경우 8bit QSGD가 최종 accuracy에 도달할 때까지의 시간은 2배로 감소하였다. (c)는 epoch에 따른 accuracy 변화를 보여주는데, epoch에 대해서 크게 accuracy 손실이 없는 것을 확인할 수 있다. 이는 최종 accuracy의 손실 없이 학습 시간 대비 accuracy 수렴 속도는 QSGD가 뛰어나다는 것을 보인다.
