由於良好的可擴展性，隨機梯度下降（SGD）的並行實現是最近研究的熱點。實現並行化SGD的關鍵障礙就是節點間梯度更新時的高帶寬開銷。因此，研究者們提出了一些啟發式的梯度壓縮方法，使得節點間只傳輸壓縮后的梯度。盡管這些啟發式方法在實踐中很有效，但它們有時並不會收斂。
本文提出了量化SGD（Quantization SGD，QSGD），它是一類具有收斂保證且在實踐中性能良好的壓縮模式。QSGD允許用戶平滑得權衡通信帶寬和收斂時間：節點可以在每輪迭代時調整發送的比特數，代價可能是更高的方差。這種權衡是固有的，因為將其提高到某個閾值會違反信息理論的下限。QSGD在異步情況下保證了凸與非凸目標函數的收斂性，並且可以使用隨機方差削減技術擴展。
當在圖像識別與語音識別任務上訓練深度神經網絡時，QSGD顯著地降低了端到端的訓練時間。
1. Introduction
目前，許多研究者正在研究如何擴展大規模SGD算法。SGD算法的定義如下。令\(f:\mathbb{R}^n\rightarrow\mathbb{R}\)是我們要最小化的目標函數。我們可以得到隨機梯度\(\widetilde{g}\)，即\(\mathbb{E}[\widetilde{g}(x)]=\triangledown f(x)\)。通過多次迭代，SGD可以收斂到最小值。
其中\(\boldsymbol{x}_t\)是當前，\(\eta_t\)是變量的步長。給出由未知分布\(D\)生成的獨立同分布的數據點\(X_1,X_2,\cdots,X_m\)，和一個衡量模型\(\theta\)在數據點\(X\)的損失函數\(l(X,\theta)\)。我們希望找到一個最優模型\(\theta^\ast\)，使得\(f(\theta)=\mathbb{E}_{X\sim D}[l(X,\theta)]\)。該框架可以解決神經網絡訓練等基礎性任務。
本文主要關注並行SGD方法，特別地，我們考慮將大型數據集分割到\(K\)個處理器上，它們共同優化目標函數\(f\)。每個處理器維護一個參數向量\(\boldsymbol{x}_t\)的本地副本，在每次迭代中，該處理器會獲得對應參數向量的梯度更新。隨后每個處理器向其他處理器廣播它們的梯度更新，並聚合其他處理器的梯度更新以計算得到下一次迭代的參數向量\(\boldsymbol{x}_{t+1}\)
在許多並行化的SGD算法的實現中，每次迭代時每個處理器都需要將所有的梯度更新發送給所有其他的處理器。如果梯度向量比較稠密，每次迭代時每個處理器需要發送或接收\(n\)個浮點數並維護參數向量\(\boldsymbol{x}\)。實踐中發現，梯度傳輸所產生的通信開銷是一個重要的性能瓶頸。一種通用的降低通信開銷的方法就是在略微損失訓練精度的情況下對梯度進行壓縮。一個簡單的實現是降低表示的精度，該方法在凸性和稀疏性假設下會收斂。更為激進的量化技術是1BitSGD，它將梯度的每個分量減少到1比特，通過對\(\widetilde{g}\)的坐標平均縮放，可以在本地累積誤差。在某些條件下，實驗觀察到1BitSGD可以保證收斂。然而，目前尚不清楚1BitSGD是否在強假設下也是如此，並且不清楚是否可以實現更高的壓縮率。
2. Preliminaries
令\(\mathcal{X} \subseteq \mathbb{R}^n\)為一個已知的凸集，\(f:\mathcal{X} \rightarrow \mathbb{R}\)是一個未知凸函數，它具有可微和平滑的性質。假設我們要重復訪問\(f\)在（可能是隨機的）輸入\(x\)處的隨機梯度，並期望在正確的方向上前進。
Definition 2.1固定\(f:\mathcal{X} \leftarrow \mathbb{R}\)，\(f\)的隨機梯度是一個隨機函數\(\widetilde{g}(x)\)且\(\mathbb{E}[\widetilde{g}(x)]=\triangledown f(x)\)。對於所有的\(x\in \mathcal{X}\)，若滿足\(\mathbb{E}[\left\|\widetilde{g} \right\|_2^2]\leq B\)則說明隨機梯度的二階矩最多為\(B\)；若滿足\(\mathbb{E}[\left\|\widetilde{g}(x)-\triangledown f(x) \right\|_2^2]\leq \sigma^2\)則說明方差最大為\(\sigma^2\)。
Theorem 2.1令\(\mathcal{X} \subseteq \mathbb{R}^n\)為一個已知的凸集，\(f:\mathcal{X} \leftarrow \mathbb{R}\)是一個未知凸函數，它是拉普拉斯平滑的。對於給定的\(x\in \mathcal{X}\)，令\(R^2 = \sup_{x\in \mathcal{X}}\left\|x-x_0\right\|^2\)，固定迭代次數\(T>0\)，且\(f\)的隨機梯度的方差上界為\(\sigma^2\)，每次迭代的步長為\(\eta_t = \frac{1}{L+1/\gamma}\)，其中\(\gamma = \frac{R}{\sigma}\sqrt{\frac{2}{T}}\)，那么就有：
Minibatch SGD小批量SGD（Minibatch SGD）是傳統SGD的一種變體。在小批量SGD中，更新公式為\(x_{t+1} = \prod_{\mathcal{X}}(x_t-\eta_t \widetilde{G}_t(x_t))\)，其中\(\widetilde{G}_t(x_t) = \frac{1}{m}\sum_{m}^{i=1}\widetilde{g}_{t,i}\)。顯然，如果隨機梯度\(\widetilde{g}_{t,i}\)的方差上界為\(\sigma^2\)，那么隨機梯度\(\widetilde{G}_t\)的方差上界就是\(\sigma^2/m\)。由Theorem2.1可知，只要等式\(\ref{2}\)中的第一項占主導地位，小批量SGD的收斂次數是傳統SGD的\(1/m\)。
Data-parallel SGD我們主要考慮多GPU環境下的數據並行SGD算法的通信開銷。我們有\(K\)個同步的處理器\(p_1,p_2,\cdots,p_K\)，使用點對點方式進行通信。其中每個處理器維護一個\(n\)維向量\(\boldsymbol{x}\)。
如算法1所示，在每次同步迭代中，每個處理器聚合\(\boldsymbol{x}\)的值，然后得到\(\boldsymbol{x}\)不同部分的隨機梯度更新，接着將這些更新發送給其他所有的處理器，最終在本地聚合並應用這些更新。重要的是，我們在算法第3行（廣播之前）與第7行（接收之后）添加了編碼與解碼的步驟。在不同的SGD算法的變體中，只有編碼與解碼函數是不同的。注意這里解碼函數無須覆蓋原始梯度\(\widetilde{g}^l\)，因為我們通常使用近似值。
對於原始的並行SGD（即沒有編碼與解碼函數）,在每個處理器中，如果\(\boldsymbol{x}_t\)是處理器在迭代\(t\)之前維護的\(\boldsymbol{x}\)值，那么在此迭代結束時\(\boldsymbol{x}\)的更新值是\(\boldsymbol{x}_{t+1}=\boldsymbol{x}_t -(\eta_t /K)\sum_{l=1}^K\widetilde{g}^l(\boldsymbol{x}_t)\)，其中每個\(\widetilde{g}^l\)都是隨機梯度。特別地，此更新是大小為\(K\)的小批量更新。因此，綜合上述討論與Theorem2.1，我們有如下推論：
Corollary 2.2令\(\mathcal{X},f,L,x_0\text{和}R\)的定義與Theorem2.1中相同。固定\(\epsilon > 0\)。假設我們在\(K\)個處理器上運行並行SGD，每個處理器訪問獨立的隨機梯度且二階矩的界為\(B\)，步長為\(\eta_t=1/(L+\sqrt{K}/\gamma)\)，其中\(\gamma\)定義與Theorem2.1中相同。如果
大多數情況下，式\(\ref{3}\)中最大的第一項將主導必要的迭代次數。具體來說，迭代次數與二階矩\(B\)線性相關。
3. Quantized Stochastic Gradient Descent(QSGD)
本節主要介紹QSGD，符號約定如下。\(\log\)是以2為底的對數函數，浮點數約定為32位。對於任意向量\(v\in \mathbb{R}^n\)，令\(\left\|v\right\|_0\)表示\(v\)的非零數。對於任意字符串\(\omega \in \{0,1\}^*\)，我們令\(|\omega|\)表示該字符串的長度。對於任意標量\(x\in \mathbb{R}\)，我們用\(\text{sng}(x)\in \{-1,+1\}\)表示它的符號，其中\(\text{sgn}(0)=1\)。
3.1. Generalized Stochastic Quantization and Coding
Stochastic Quantization首先考慮一個通用的、參數化的有損壓縮模式。量化函數定義為\(Q_s(v)\)，其中\(s \geq 1\)是一個需要調整的超參數，代表量化的程度。直觀地，我們定義\(s\)為0和1之間的均勻分布的級別，每個值以保持期望值不變的方式進行量化，並引入最小方差。如圖1所示。
對於任意\(v\in \mathbb{R}^n\)且\(v \neq 0\)，\(Q_s(v)\)的定義為：
其中\(\xi(v,s)\)是按照下式定義的隨機變量。令\(0\leq l < s\)是一個滿足\(|v_i|/\|v\|_2 \in [l/s,(l+1)/s]\)條件的整數。也就是說，\([l/s,(l+1)/s]\)是\(|v_i|/\|v\|_2\)對應的量化區間。那么
這里，\(p(a,s)=as-l\)對於所有的\(a\in[0,1]\)都成立。若\(v=0\)，則定義\(Q_s(v)=0\)。\(\xi_i(v,s)\)分布具有最小方差且它的期望滿足\(\mathbb{E}[\xi_i]=|v_i|/\|v\|_2\)。
Lemma 3.1對於所有的向量\(v\in\mathbb{R}^n\)，我們有（1）\(\mathbb{E}[Q_s(v)]=v\)（2）\(\mathbb{E}[\|Q_s(v)-v\|_2^2]\leq \min(n/s^2,\sqrt{n}/s)\)（3）\(\mathbb{E}[\|Q_s(v)\|_0]\leq s(s+\sqrt{n})\)
Efficient Coding of Gradients
對於任意向量\(v\)，\(Q_s(v)\)的輸出是一個三元組\((\|v\|_2,\sigma,\zeta)\)，其中\(\sigma\)是\(v_i\)的符號，而\(\zeta\)是由整數值\(s·\xi_i(v,s)\)構成的向量。編碼方案背后的關鍵思想是並非所有整數值\(s·\xi_i(v,s)\)都以相同的概率出現，具體來說，較大的整數出現的頻率會較低。我們將通過專門的Elias整數編碼來實現它。對於任何正整數k，其編碼（表示為Elias(k)）從k的二進制表示開始，其前綴為該表示的長度，然后它遞歸編碼這個前綴。對於任意正整數k，最終編碼長度為\(|Elias(k)|=\log k+\log\log k + \cdots + 1 \leq (1+o(1))\log k +1\)，Elias編碼與解碼都可以很高效地實現。
對於給定三元組\((\|v\|_2,\sigma,\zeta)\)表示的梯度向量，使用\(s\)表示量化程度，我們的編碼算法的輸出是如下定義的字符串\(S\)。首先，使用32位來編碼\(\|v\|_2\)。接着，它使用Elias遞歸編碼\(\zeta\)的第一個非零項的位置。然后用一位來表示符號\(\sigma_i\)，再后面是\(Elias(s·\xi_i(v,s))\)。迭代地，它繼續編碼從當前位置到下一個非零項的距離，並以相同的方式編碼\(\sigma_i\)和\(\zeta_i\)。解碼過程很直觀：我們首先讀取前32位來構建\(\|v\|_2\)，然后迭代地使用Elias遞歸解碼\(\zeta\)中非零項的位置、非零項的符號和值。
Theorem 3.2固定\(f:\mathbb{R}^n \rightarrow \mathbb{R}\)，然后取任意\(\boldsymbol{x} \in\mathbb{R}^n\)，固定量化級別\(s \geq 2\)。若\(\widetilde{g}(\boldsymbol{x})\)是\(f\)在\(\boldsymbol{x}\)處的隨機梯度且二階矩上限為\(B\)，則\(Q_s(\widetilde{g}(\boldsymbol{x}))\)是\(f\)在\(\boldsymbol{x}\)處的隨機梯度且方差上界為\(\min\left(\frac{n}{s^2},\frac{\sqrt{n}}{s}\right)B\)。此外，因為編碼函數的存在，所以\(Q_s(\widetilde{g}(\boldsymbol{x}))\)需要通信的比特數上界為：
3.2. QSGD Guarantees
將上面給出的通信和方差的界與在平滑凸函數上的SGD算法的收斂保證放在一起,產生以下結果：
Theorem 3.4(Smooth Convex QSGD)令\(\mathcal{X},f,L,x_0\text{和}R\)與Theorem2.1中定義相同。固定\(\epsilon > 0\)，假設我們在\(K\)個處理器上以\(s\)量化級別來運行並行QSGD且隨機梯度的二階矩的界為\(B\)，步長為\(\eta_t=1/(L+\sqrt{K}/\gamma)\)，其中\(\gamma\)在Theorem 2.1中定義；再令\(\sigma=B'\)，其中\(B'=\min\left(\frac{n}{s^2},\frac{\sqrt{n}}{s}\right)B\)。那么若\(T=O\left(R^2·\max\left(\frac{2B'}{K\epsilon^2},\frac{L}{\epsilon}\right)\right)\)，則\(\mathbb{E}\left[f\left(\frac{1}{T}\sum_{t=0}^T\boldsymbol{x}_t\right)\right] - \min_{\boldsymbol{x}\in\mathcal{X}}f(\boldsymbol{x})\leq\epsilon\)。此外，QSGD在每一輪需要\(\left(3+\left(\frac{3}{2}+o(1)\right)\log\left(\frac{2(s^2+n)}{s^2+\sqrt{n}}\right)\right)\)Bit來進行通信。當\(s=\sqrt{n}\)時，通信所需的比特數降至\(2.8n+32\)。
Theorem 3.5(QSGD for smooth non-convex optimization)令\(f:\mathbb{R}^n \rightarrow \mathbb{R}\)是一個拉普拉斯平滑（可能非凸）函數，令\(\boldsymbol{x}_1\)是任意的初始點。令\(T>0\)，\(s>0\)。然后在\(\{1,2,\cdots,N\}\)上有一個隨機停止時間\(R\)，因此擁有常數步長\(\eta=O(1/L)\)且\(f\)的隨機梯度的二階矩的界為\(B\)的QSGD滿足\(\frac{1}{L}\mathbb{E}[\|\triangledown f(\boldsymbol{x})\|^2_2]\leq O\left(\frac{\sqrt{L(f(\boldsymbol{x}_1)-f^*)}}{N}+\frac{\min(n/s^2,\sqrt{n}/s)B}{L}\right)\)。另外，通信開銷與Theorem3.4相同
3.3. Quantized Variance-Reduced SGD
假設我們有\(K\)個處理器和一個大於0的參數\(m\)，每個處理器\(i\)都訪問函數\(\{f_{im/K},\cdots,f_{(i+1)m/K-1}\}\)。目標是近似最小化\(f=\frac{1}{m}\sum_{i=1}^mf_i\)。對於處理器\(i\)，令\(h_i=\frac{1}{m}\sum_{j=im/K}^{(i+1)m/K-1}f_i\)是該處理器所知道的f的一部分，因此有\(f=\sum_{i=1}^K h_i\)。
那么應該如何在並行化SVRG算法上應用隨機量化方法以減少通信開銷呢？經過檢查，我們注意到最終的更新將破壞標准SVRG。我們最終證明可以使用隨機量化技術來量化SVRG更新，並獲得相同的收斂界限。
Algorithm Description令\(\widetilde{Q}(\boldsymbol{v})=Q(\boldsymbol{v},\sqrt{n}\)，其中\(Q(\boldsymbol{v},s)\)在3.1節中定義。給定任意初始點x0，我們令\(\boldsymbol{y}^{(1)}=\boldsymbol{x}_0\)。在第p輪次開始時，每個處理器廣播一個未量化的梯度\(\triangledown h_i(\boldsymbol{y}^{(p)})\)，並對其他處理器發來的梯度進行聚合\(\triangledown f(\boldsymbol{y}^{(p)})=\sum_{i=1}^m\triangledown h_i(\boldsymbol{y}^{(p)}\)。在每一輪次，對於每次迭代\(t=1,2,\cdots,T\)和每個處理器\(i=1,2,\cdots,K\)，我們令\(j_{i,t}^{(p)}\)是一個隨機整數，那么在每次迭代中，處理器\(i\)廣播更新向量\(\boldsymbol{u}_{t,i}^{(p)}=\widetilde{Q}\left(\triangledown f_{j_{i,t}^{(p)}}(\boldsymbol{x}_t^{(p)})-\triangledown f_{j_{i,t}^{(p)}}(\boldsymbol{y}_t^{(p)})+\triangledown f(\boldsymbol{y}_t^{(p)})\right)\)。然后，每個處理器計算所有的更新\(\boldsymbol{u}_t^{(p)}=\frac{1}{K}\sum_{i=1}^K \boldsymbol{u}_{t,i}\)，並令\(\boldsymbol{x}_{t+1}^{(p)}=\boldsymbol{x}_t^{(p)}-\eta \boldsymbol{u}_t^{(p)}\)。在第p輪次結束時，每個處理器會執行\(\boldsymbol{y}^{(p+1)}=\frac{1}{T}\sum_{t=1}^T\boldsymbol{x}_t^{(p)}\)。
Theorem 3.6令\(f(x)=\frac{1}{m}\sum_{i=1}^mf_i(\boldsymbol{x})\)，其中\(f\)是一個強凸函數，對於所有的\(i\)來說，\(f_i\)是一個強凸且拉普拉斯平滑的函數。令\(x^\ast\)是\(f\)在\(\mathbb{R}^n\)上的最小值點。若\(\eta=O(1/L)\)且\(T=O(L/l)\)，那么初始點為\(\boldsymbol{y}^{(1)}\)的QSVRG對任意\(p\geq1\)的輪次保證了\(\mathbb{E}[f(\boldsymbol{y}^{(p+1)})]-f(\boldsymbol{x}^{\ast})\leq0.9^p (f(\boldsymbol{y}^{(1)})-f(\boldsymbol{x}^{\ast}))\)。此外，每個輪次擁有\(T\)次迭代的QSVRG至多需要\((F+2.8n)(T+1)+Fn\)比特進行通信。
4. QSGD Variants
首先，我們可以通過量化到大小固定為d的buckets中來控制量化方差。如果我們把每個梯度都看成一維向量v，桶的定義是向量中連續的d個值，比如v[(i-1)d+1:i·d]。我們可以使用qsgd單獨量化每個桶中的值。d=1的時候相當於沒進行量化，也就是普通的sgd。d=n相當於前面小節提到的全量化。顯然，使用桶量化后，由lemma3.1得到的收斂保證將取決於d而不是完整的維度n。我們可以通過桶來控制方差，代價就是需要額外存儲縮放因子。舉例來說，如果我們令桶的大小為512，量化4位，那么方差的上界就變為\(\sqrt{512} / 2^4=1.41\)。
其次，與理論不同的是，我們會按照向量的最大值而不是2范數進行縮放。直觀地講，通過最大值進行歸一化可以保留更多的值，並且對於相同的迭代次數具有略高的精度。由於較低的比特寬度（例如，每個維度32Bit到2Bit），兩種方法在baseline上降低了相同的帶寬，但是通過最大值進行歸一化理論上不提供稀疏性保證。在實踐中，最大值歸一化會產生非平凡的稀疏性。
5. Experiments
Setup我們在amazon ec2p2.16xlarge實例上做的實驗，每個實例包括16塊tesla k80顯卡。實例間使用GPUdirect點對點通信，不支持NCCL。我們在CNTK上實現了QSGD，CNTK提供一個高效地基於MPI-Bcast的GPU到GPU的通信方式，並且實現了1BitSGD。代碼開源在GitLab上。
我們在圖像分類和語音識別任務上做了實驗。其中圖像分類使用ilsvrs2015（imagenet）,cifar10和mnist數據集，語音識別使用cmu an4數據集。對於圖像分類任務，我們實現了alexnet，vgg，resnet和帶有bn的inception-net。對於語音識別任務，我們訓練了lstm。實驗詳情見表1。
Protocol在實驗中我們使用的都是標准的網絡，除非特別說明，否則實驗中用到的網絡模型的超參數都是CNTK2.0中默認最優的。隨着GPU數量的增多，我們增大batchsize以平衡計算和通信開銷。我們采用雙緩沖技術在計算的同時進行通信和量化。量化技術在較低的學習率下效果很好，因此我們使用32Bit的學習率，並且減少桶大小以降低方差。我們並不量化較小的（元素數量小於10K）梯度矩陣，因為這樣會得不償失。然而，在實驗中99%的梯度都被量化
Communication vs. Compution圖2給出了不同網絡模型在圖像分類任務的結果。首先，基於通信與計算的比率，我們可以粗略地將神經網絡模型分為通信密集型（AlexNet，VGG，LSTM）和計算密集型（Inception，ResNet）。對於這兩種網絡類型，隨着GPU數量的增加，通信時間的占比會顯着增加。對於32位版本的傳統SGD，所有的網絡模型都可以從降低通信中獲益。舉例來說，在16塊GPU上訓練的AlexNet，當batchsize為1024時，超過80%的訓練時間都花在通信上。而對於在2塊GPU上訓練的LSTM模型，當batchsize為256時，通信占總訓練時間的71%。
接下來，我們看一下QSGD的通信時間與總訓練時間的關系（這里通信時間包括壓縮和解壓梯度的時間）。我們測試了使用2Bit量化與桶大小為128的QSGD以及使用4Bit量化、8Bit量化與桶大小為512的QSGD。這兩種變體的實驗結果非常相似，因為不同的桶大小意味着4Bit量化僅比2Bit量化發送的數據多77％（但是比32Bit少了約8倍）。本實驗中的桶大小僅保證不錯的收斂性，沒有進行仔細地調參。對於在16塊GPU上以batchsize為1024訓練的AlexNet，4BitQSGD減少了4倍的通信時間，總訓練時間減少了2.5倍。對於LSTM模型，通信量減少了6.8倍，總訓練時間減少了2.7倍。
Accuracy我們在imagenet數據集上訓練了alexnet和resent，在an4數據集上訓練lstm，在cifar10上訓練reset110，在mnist訓練了一個2層的感知機。實驗結果由圖3給出，完整的數據在表1中。qsgd在使用8塊gpu時性能最好，一般來說，4Bit或8Bit梯度量化足以達到甚至略微提高32Bit的傳統SGD的精度，同時確保非平凡的加速比。注意到在訓練深度網絡時向梯度添加噪聲的好處，而量化可以被視為一種零均值噪聲，這恰好使得通信更有效。與此同時，我們注意到更高程度的量化可能會損害准確性。比如，使用桶大小為8192的4BitQSGD訓練的AlexNet與使用32BitSDG相比，損失了0.57%的top5准確率和0.68%的top1准確率。
需要詳細地研究的一個問題是哪些層對量化更敏感。如果過度地量化卷積層（例如2BitQSDG）會導致精度損失，但使用4BitQSGD或8BitQSGD幾乎不損失精度。這表明，視覺任務的現代架構（如ResNet或Inception，幾乎完全是卷積層）可能比LSTM等網絡更少地受益於量化技術。
Additional Experiments論文的完整版本包含其他實驗，包括與1BitSGD的完整比較。簡言之，QSGD在本文考慮的網絡和參數方面優於或等於1BitSGD的性能和最終精度。
6. Conclusions and Future Work
QSGD是一類允許用戶在每輪的通信量和訓練時間進行平滑權衡的算法。實驗表明，QSGD在各種任務都表現良好。但是還有很多工作沒有完成，最重要的是利用QSGD創造的稀疏性。MPI的當前實現不提供對稀疏類型的支持，但將來的工作可以對這部分進行探索。此外，還可以研究一下QSGD在大規模應用（如超級計算機）中的潛力。