CVPR 2020丨8比特數值也能訓練模型？商湯提出訓練加速新演算法
本文介紹的是CVPR2020入選論文《Towards Unified INT8 Training for Convolutional Neural Network》，作者為商湯研究院鏈接與編譯團隊、高性能計算團隊和北航劉祥龍老師團隊。
作者 | 商 湯
論文地址：//arxiv.org/pdf/1912.12607.pdf
在CVPR 2020上，商湯研究院鏈接與編譯團隊、高性能計算團隊和北航劉祥龍老師團隊合作提出了用於加速卷積神經網路訓練過程的INT8訓練技術。該工作通過將網路的輸入、權重和梯度量化到8比特來加速網路的前向傳播和反向傳播過程，縮短卷積神經網路訓練時間。論文觀察到梯度的獨特分布給量化訓練帶來了極大挑戰，為了解決梯度量化帶來的精度損失和不穩定問題，該論文進行了量化訓練收斂穩定性的理論分析並基於此提出了誤差敏感的學習率調節和基於方向自適應的梯度截斷方法。同時為了保證更高的加速比，該論文還提出使用周期更新、量化卷積融合等技術來減少量化操作帶來的時間開銷。應用了上述方法之後，INT8訓練在影像分類任務和檢測任務上都僅僅損失微小的精度，且訓練過程相比浮點訓練加速了22%。卷積神經網路被廣泛應用在多種電腦視覺任務中並且取得了優異的精度。由於擁有龐大的參數量，訓練和部署卷積神經網路需要耗費大量計算資源和漫長的訓練時間，如何用更少資源訓練卷積神經網路一直是一個學術研究熱點，也是工業界關心的話題。神經網路量化技術是一種使用定點計算代替浮點的加速技術，目前被廣泛地應用在神經網路部署中，可以極大地提升部署速度，並降低記憶體資源佔用。現有很多工作均表明將網路前向過程的浮點計算替換成INT8計算，不會帶來明顯的精度下降[1][2]。下圖展示了現代神經網路加速晶片對於不同精度計算的理論計算峰值對比，可以看到，INT8算力相比於FP32和FP/INT16均能有超過2倍峰值性能提升。
當考慮將神經網路量化技術應用在卷積神經網路訓練中時，為了加速卷積的反向梯度傳播過程，不得不對梯度進行量化操作。在將浮點的梯度量化到INT8數值範圍內之後，訓練過程變得極其不穩定，並且收斂到非常差的精度。如何解決量化梯度給訓練帶來的收斂穩定性問題，是十分重要的問題。與此同時，在提升訓練精度的同時，也不應當進入過多額外的計算，否則加速效果將會大打折扣。
一方面是高效的計算峰值保障，一方面是困難重重的演算法設計，這是INT8訓練技術的機遇與挑戰。標準的線性量化操作指的是，將一個浮點張量（tensor）進行線性映射，變換到整數空間中[3]。這個整數空間的大小由於量化比特數來決定，比如常見的8bit量化數，就有256個取值，本文中使用的是對稱量化，因此量化數的取值是從-128到127。具體公式如下，其中x是被量化的數據，q是量化後的數據，s是量化係數，clip是截斷函數：
在8bit的場景里，截斷函數和量化係數的計算公式如下：
為了降低量化帶來的誤差，一個常見做法是對取整過程進行隨機化，使得取整函數從期望上更接近原始的數，具體隨機取整的公式如下：
相反的，將8bit量化數變換回浮點的過程稱之為反量化。反量化公式如下所示，其中q為量化計算結果，s為量化係數，
為反量化後的結果。
上圖的上半部分展示了標準的卷積神經網路量化計算前向過程，該過程被廣泛應用在INT8部署加速中。在卷積計算之前，量化器會對輸入和權重進行量化操作，將浮點數量化到8bit數值上，通過INT8卷積計算核心，即可完成一次INT8前向計算，最終將求和得到的32bit數進行反量化操作回算到浮點數域中，以供給下一層計算使用。INT8訓練的一個核心的加速點在於卷積計算的反向過程，上圖展示了INT8訓練中卷積計算在反向傳播過程中的計算細節。在卷積的反向梯度傳播過程，同樣的將梯度進行浮點量化操作，不過為了降低量化的誤差，針對梯度的量化採用了隨機取整操作。通過INT8的反向卷積計算核心，可以得到下一層所需的回傳梯度，以及當前層的權重所需的梯度。由於INT8反向卷積輸出的是32bit數，與前傳類似，需要引入一次反量化操作，將32bit數反算回到浮點數域中。為什麼對梯度進行量化會給網路訓練帶來如此大的影響？我們可以觀察訓練過程中的梯度分布情況來進一步的分析。
通過圖（a）中對比梯度和輸入、權重的分布，可以發現：梯度分布相比輸入和權重分布更加尖銳，同時範圍更大。相比於輸入和權重，梯度有更多的值集中在0附近，但同時梯度還有許多較大值，讓梯度的分布範圍變得相當廣，這些特徵都會導致梯度量化的量化誤差比輸入和權重更大。
圖（b）展示的是layers16隨著訓練，其梯度從epoch 0到epoch 300的變化情況。從中可以看出，隨著訓練的進行，梯度分布越變得更加尖銳，同時仍然保持著較廣的分布範圍，這意味著梯度量化的誤差會隨著訓練的進行變得越來越大。
梯度的分布隨網路深度變化情況從圖（c）中可以看出。很容易發現，卷積層的深度越淺，梯度分布越尖銳，這也會導致梯度量化的誤差更大。
從圖（d）中可以看出卷積的結構也會影響梯度分布，對於MobileNetV2來說，conv2為depthwise卷積其相比conv1和conv3具有更加尖銳的分布。
由於卷積神經網路的梯度具有如上四個特點，所以當我們直接在訓練中對梯度進行量化時，訓練精度非常容易出現突發的崩潰情況。下圖展示了在CIFAR-10數據集上進行實驗的精度和損失函數變化曲線，以MobileNetv2在CIFAR-10數據集上訓練為例，其訓練的精度曲線和loss曲線如下圖，從圖中可以發現INT8訓練的loss在訓練初期正常下降，但隨後迅速上升，對應的精度也不斷下降。
根據以上的觀察和初步啟發，我們希望通過理論的分析和推導，對量化訓練的收斂穩定性進行建模。根據Adam等相關論文的經驗和優化理論中的Regret analysis，不失一般性地定義R(T)為：
其中f是損失函數，t是訓練輪數，T是訓練總輪數，
為t輪的權重，
是最優權重。
其中
為t輪的學習率，d 為權重的維度，
為t輪的量化誤差，
是輪的量化後梯度。為了確保網路能夠穩定收斂，
在T變大時需要能夠達到足夠小。通過上式可以發現，在T趨於無窮大時，第（1）項可以忽略不計，主要考慮減小第（2）項和第（3）項。我們發現，第（2）項與量化誤差正相關，第（3）項與學習率以及量化後的梯度大小有關。因此我們不難得到兩個直觀的提升訓練收斂穩定性的策略：通過調節量化函數中的截斷減小量化誤差
通過適當調低學習率來提高量化訓練精度
依據以上分析，我們針對量化誤差和學習率提出了基於方向自適應的梯度截斷和誤差敏感的學習率調節兩個方法來解決量化訓練帶來的精度損失問題。同時，為了減少量化操作帶來的額外開銷，本文還提出了周期更新和量化卷積融合的方法。
1、基於方向自適應的梯度截斷：調整截斷值，讓梯度方向保持正確。為了最小化量化誤差，之前有很多研究提出優化截斷值的方法，其中就有研究提出通過假設數據分布直接求解最優截斷值。但是已有的研究都針對於權重量化的截斷值進行優化。就如本文觀察所顯示，梯度的分布特徵與權重區別較大，無法直接使用。本文通過KS檢驗發現梯度的分布並不符合常見的高斯分布、拉普拉斯分布和學生t分布，因此很難通過假設梯度分布來直接求解最優的截斷值。
基於以上的分析，本文採用梯度下降的方法來自適應地學習最優截斷值，常見的目標函數有均方誤差函數，但是由於梯度的分布特徵，均方誤差的大小會受到梯度的影響，影響優化過程；同時對於梯度來說，均方誤差並不能很好地體現梯度的量化誤差對於優化過程的影響，因此本文提出使用能夠體現梯度方向的餘弦距離來衡量梯度的量化誤差，並以餘弦距離為目標函數來優化求解最優截斷值。餘弦距離定義如下：
其中，
是梯度，
是量化後的梯度。
2、誤差敏感的學習率調節：在錯誤的方向上盡量少更新。根據上述的理論分析，降低學習率能夠有助於模型量化訓練的收斂。針對學習率的調整，本文提出誤差敏感的學習率調節方法，使用學習率係數對原學習率進行調整，學習率係數與餘弦距離負相關，學習率係數定義如下：
其中
和
是超參數，用於控制衰減程度和調節下界。
由於量化操作需要的統計數據範圍和計算截斷值等操作十分耗時，為了減少這些操作的時間開銷，本文採用周期更新的方式，周期性地統計數據範圍和計算截斷值。通過周期更新的方法能夠有效地提高減少因量化引入的額外時間開銷。下表為ResNet50在ImageNet數據集上不同周期的單次訓練時間統計表。
4、量化卷積融合：減少訪存次數、節省cuda kernel launch次數通過將量化和反量化操作融合入卷積計算的CUDA核函數里，可以減少一次數據的訪存，有效地減少量化和反量化操作的時間開銷。
影像分類任務：本文在CIFAR10和ImageNet等影像分類數據集進行INT8訓練實驗。從下表結果中可以看出，在大多數網路結構中均取得了比現有最好方法更優的精度，並且首次在MobileNet、Inception等網路上進行量化訓練實驗，精度損失也在1.5%以內。
目標檢測任務：同時，本文也首次嘗試在PASCAL和COCO等目標檢測數據集上進行INT8訓練實驗，精度損失也在2%以內。
已有的少量探究梯度量化的論文[4]均未報告演算法在實際訓練任務中的真實加速性能，為了最大限度將方法實用化，本文在GeForce GTX1080TI顯示卡上編寫並優化了用於支援INT8訓練的卷積前向和後向計算核心。實測結果表明，使用INT8卷積計算的前向和後向過程相比於浮點計算有明顯的加速，其中前向過程平均加速1.63倍，後向過程平均加速1.94倍。如下圖所示：
同時，本文在實際訓練過程中進行了完整的端到端測試，可以看到，INT8訓練可以將ResNet50的一輪訓練過程從0.360秒降低到0.293秒，整體訓練過程提速了22%。
該論文目前已進行直播解讀，回放地址如下：
//www.mooc.ai/open/course/786
References
[1]. Ruihao Gong,Xianglong Liu, Shenghu Jiang, Tianxiang Li,Peng Hu, Jiazhen Lin, Fengwei Yu,and Junjie Yan. Differen-tiable soft quantization: Bridging full-precision and low-bitneuralnetworks. In ICCV, October 2019.
[2]. Rundong Li, Yan Wang, Feng Liang,Hongwei Qin, Junjie Yan, and Rui Fan. Fully quantized network for objectdetection. In The IEEE Conference on Computer Vision and Pattern Recognition(CVPR), June 2019.
[3]. Benoit Jacob, Skirmantas Kligys, Bo Chen, MenglongZhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko.Quantization and training of neural networks for efficientinteger-arithmetic-only inference. 2018 IEEE Conference on Computer Vision andPattern Recognition(CVPR), June 2018.
[4]. Yukuan Yang, Shuang Wu, Lei Deng, Tianyi Yan, YuanXie, and Guoqi Li. Training high-performance and large-scale deep neuralnetworks with full 8-bit integers, 2019.
為反量化後的結果。
為t輪的權重，
是最優權重。
為t輪的學習率，d 為權重的維度，
為t輪的量化誤差，
是輪的量化後梯度。為了確保網路能夠穩定收斂，
在T變大時需要能夠達到足夠小。通過上式可以發現，在T趨於無窮大時，第（1）項可以忽略不計，主要考慮減小第（2）項和第（3）項。
是梯度，
是量化後的梯度。
和
是超參數，用於控制衰減程度和調節下界。
