論文通過DBTD方法計算過濾閾值，再結合隨機剪枝演算法對特徵值梯度進行裁剪，稀疏化特徵值梯度，能夠降低迴傳階段的計算量，在CPU和ARM上的訓練分別有3.99倍和5.92倍的加速效果
來源：曉飛的演算法工程筆記 公眾號
論文: Accelerating CNN Training by Pruning
Activation Gradients
Introduction
在訓練過程中，特徵值梯度的回傳和權值梯度的計算佔了大部分的計算消耗。由於這兩個操作都是以特徵值梯度作為輸入，而且零梯度不會佔用計算資源，所以稀疏化特徵值梯度可以降低迴傳階段的計算消耗以及記憶體消耗。論文的目標在於高效地降低訓練負載，從而在資源有限的平臺進行大規模資料集的訓練。
論文假設特徵值梯度服從正態分佈，基於此計算閾值$\tau$，隨後使用隨機剪枝演算法(stochastic pruning)將小於閾值的特徵值梯度隨機置為零或$\pm \tau$。經理論推理和實驗證明，這種方法不僅能夠有效地稀疏化特徵值梯度，還能在加速訓練的同時，不影響訓練的收斂性。
General Dataflow
卷積層通常包含4個階段：推理、特徵值梯度回傳、權值梯度計算和權值更新。為了表示這些階段的計算，論文定義了一些符號：
卷積層的四個訓練階段的總結為：
論文通過視覺化發現，回傳階段的特徵值梯度幾乎全是非常小的、接近於零的值，自然而然地想到將這些值去掉不會對權值更新階段造成很大的影響，所以論文認為剪枝特徵值梯度能夠加速卷積層在訓練時的計算。
Sparsification Algorithms
Distribution Based Threshold Determination (DBTD)
剪枝操作最關鍵的步驟是決定選擇哪些元素進行消除，先前有研究使用最小堆進行元素選擇，但這會帶來較大的額外計算開銷。為此，論文采用簡單的閾值過濾進行元素選擇。
論文首先分析了兩種經典的卷積網路結構的特徵值梯度分佈：Conv-ReLU結構和Conv-BN-ReLU結構：
- 對於Conv-ReLU結構，輸出的特徵值梯度$dO$是稀疏的，但其分佈是無規律的，而結構的輸入特徵值梯度$dI$幾乎全是非零值。通過統計發現，$dI(\cdot)$的分佈以零值對稱分佈，且密度隨著梯度值的增加而下降。
- 對於Conv-BN-ReLU結構，BN層設定在卷積層與ReLU層中間，改變了梯度的分佈，且$dO$的分佈與$dI$類似，。
所以，上述的兩種結構的梯度都可認為服從零均值、方差為$\sigma^2$的正態分佈。對於Conv-ReLu結構，由於ReLU不會降低稀疏性，$dO$能夠繼承$dI$的稀疏性，將$dI$是作為Conv-ReLU結構中的剪枝目標梯度$g$。而對於Conv-BN-ReLU結構，則將$dO$作為剪枝目標$g$。這樣，兩種結構的剪枝目標都可統一為正態分佈。假設$g$的數量為$n$，可以計算梯度的絕對值的均值，並得到該均值的期望為：
這裡的期望為從分佈中取樣$n$個點的期望，而非分佈的整體期望，再定義以下公式
將公式2代入公式1中，可以得到：
從公式3可以看出$\tilde{\sigma}$為引數$\sigma$的無偏估計，接近於真實的均值，且$\tilde{\sigma}$的整體計算消耗是可以接受的。基於上面的分析，論文結合正態分佈的累積函式$\Phi$、剪枝率$p$和$\tilde{\sigma}$計算閾值$\tau$：
Stochastic Pruning
剪枝少量值較小的梯度幾乎對權值的更新沒有影響，但如果將這些值較小的梯度全部設為零，則會對特徵值梯度的分佈影響很大，進而影響梯度更新，造成嚴重的精度損失。參考Stochastic Rounding演算法，論文采用隨機剪枝來解決這個問題。
隨機剪枝邏輯如演算法1所示，對於小於閾值$\tau$的梯度值，隨機取樣一個縮放權重來計算新閾值，再根據新閾值將梯度值置為零或$\pm \tau$。
隨機剪枝的效果如圖2所示，能夠在保持梯度分佈的數學期望的情況下進行剪枝，與當前的方法相比，論文提出的方法的優點如下：
- Lower runtime cost：DBTD的計算複雜度$O(n)$小於top-k演算法$O(nlogk)$，且DBTD對硬體更友好，能夠在異構平臺實現。
- Lower memory footprint：隨機裁剪能保持收斂性，且不需要儲存而外的記憶體。
至此，Sparsification Algorithms在梯度回傳時的特徵值梯度計算為：
Experimental Results
在CIFAR-10、CIFAR-100以及ImageNet上進行準確率驗證。
在CIFAR-10和ImageNet上進行收斂性驗證。
在不同的裝置上進行加速效果驗證。
Conclustion
論文通過DBTD方法計算過濾閾值，再結合隨機剪枝演算法對特徵值梯度進行裁剪，稀疏化特徵值梯度，能夠降低迴傳階段的計算量，在CPU和ARM上的訓練分別有3.99倍和5.92倍的加速效果。論文提出的特徵值稀疏化演算法看似很簡單，其實進行了充分的理論推導以及實驗驗證，才得到最終合理的過濾方法，唯一可惜的是沒在GPU裝置上進行實驗驗證。論文對演算法的收斂性以及期望有詳細的理論驗證，不過這裡沒有列出來，有興趣的可以去看看原文。
如果本文對你有幫助，麻煩點個贊或在看唄～
更多內容請關注 微信公眾號【曉飛的演算法工程筆記】
