論文:Extremely Low Bit Neural Network: Squeeze the Last Bit Out with ADMM
論文地址:https://arxiv.org/pdf/1707.09870.pdf
研究背景
近年來,深度學習在人工智慧領域取得了重大的突破。在計算機視覺、語音識別等諸多領域,深度神經網路 (DNN, Deep Neural Network) 均被證明是一種極具成效的問題解決方式。如卷積神經網路 (CNN, Convolutional neural network) 在計算機視覺諸多傳統問題(分類、檢測、分割)都超越了傳統方法,迴圈神經網路 (RNN, Recurrent Neural Networks) 則在時序訊號處理,如機器翻譯,語音識別等超過傳統方法。
在利用深度網路解決問題的時候人們常常傾向於設計更為複雜的網路收集更多的資料以期獲得更高的效能。但是,隨之而來的是模型的複雜度急劇提升,直觀的表現是模型的層數越來越深,引數越來越多。這會給深度學習帶來兩個嚴重的問題:
(1) 隨著模型引數的增多,模型的大小越來越大,給嵌入式端模型的儲存帶來了很大的挑戰。
(2) 隨著模型的增大,模型 inference 的時間越來越長,latency 越來越大。
以上兩個問題給深度學習在終端智慧裝置上的推廣帶來了很大的挑戰。比如,經典的深度卷積網路 VGG-16 的模型大小達到 528M,使用者很難接受下載一個如此大的模型到手機或者其他終端裝置上。同時,在一般的智慧手機上,VGG-16 識別一張影像的時間高達 3000+ms,這個 latency 對於大多數使用者來說也是難以接受的。此外,由於深度網路的計算量很大,執行深度網路的能耗很高,這對於手機等終端裝置也是一個巨大的挑戰。
所提出的演算法
在這個工作中,我們提出一種基於低位元表示技術的神經網路壓縮和加速演算法。我們將神經網路的權重表示成離散值,並且離散值的形式為 2 的冪次方的形式,比如 {-4,-2,-1,0,1,2,4}。這樣原始 32 位元的浮點型權重可以被壓縮成 1-3 位元的整形權重,同時,原始的浮點數乘法操作可以被定點數的移位操作所替代。在現代處理器中,定點移位操作的速度和能耗是遠遠優於浮點數乘法操作的。
首先,我們將離散值權重的神經網路訓練定義成一個離散約束最佳化問題。以三值網路為例,其目標函式可以表示為:
更進一步,我們在約束條件中引入一個 scale 引數。對於三值網路,我們將約束條件寫成 {-a, 0, a}, a>0. 這樣做並不會增加計算代價,因為在卷積或者全連線層的計算過程中可以先和三值權重 {-1, 0, 1} 進行矩陣操作,然後對結果進行一個標量 scale。從最佳化的角度看,增加這個 scale 引數可以大大增加約束空間的大小,這有利於演算法的收斂。如下圖所示:
對於三值網路而言,scale 引數可以將約束空間從離散的 9 個點擴增到 4 條直線。
為了求解上述約束最佳化問題,我們引入 ADMM 演算法。在此之前,我們需要對目標函式的形式做一個等價變換。
其中 Ic 為指示函式,如果 G 符合約束條件,則 Ic(G)=0,否則 Ic(G) 為無窮大。該目標函式的增廣拉格朗日形式為:
ADMM 演算法將上述問題分成三個子問題進行求解,即
與其它演算法不同的是,我們在實數空間和離散空間分別求解,然後透過拉格朗日乘子的更新將兩組解聯絡起來。
第一個子問題需要找到一個網路權重最小化
在實驗中我們發現使用常規的梯度下降演算法求解這個問題收斂速度很慢。在這裡我們使用 Extra-gradient 演算法來對這個問題進行求解。Extra-gradient 演算法包含兩個基本步驟,分別是:
第二個子問題在離散空間中進行最佳化。透過簡單的數學變換第二個子問題可以寫成:
該問題可以透過迭代最佳化的方法進行求解。當 a 或 Q 固定時,很容易就可以獲得 Q 和 a 的解析解。
實驗結果
ImageNet 影像識別:我們分別在 Alexnet、VGG16、Resnet18、Resnet50、GoogleNet 等五個主流的 CNN 框架上驗證了所提出的演算法。實驗中我們分別嘗試了 Binary 網路、Ternary 網路、{-2, -1, 0, 1, 2}、{-4, -2, -1, 0, 1, 2, 4} 四種形式。在 Imagenet 上 Top-1 和 Top-5 準確度結果如下:
Alexnet 和 VGG16:
Resnet:
GoogleNet:
其中 BWN[1] 和 TWN[2] 為我們對比的兩種 Binary 網路和 Ternary 網路量化方法。從這些結果可以看出,在各個網路框架下,我們的演算法都顯著超過對比演算法。同時,當位元數達到 3 時,量化之後的網路精度相比於原始網路幾乎可以達到無損。在 Alexnet 和 VGG16 這兩個冗餘度比較高的網路上,量化之後的網路甚至可以取得超過原始網路的精度,這是因為量化操作可以起到一個正則的作用,從而提高這類網路的泛化效能。
Pascal VOC 目標檢測:我們在 SSD 檢測框架下對演算法進行驗證,分別採用了 VGG16+SSD 和 Darknet+SSD 兩種網路結構。對於檢測任務,嘗試了 Ternary 網路和 {-4, -2, -1, 0, 1, 2, 4} 兩種量化形式。實驗結果如下:
對於 Darknet 我們使用了兩種設定,第一種設定中所有的權重進行相同的量化;第二種設定中,1x1 的卷積核使用 INT8 量化,即括號中的結果。和識別中的結果類似,在 VGG+SSD 結構中,我們的演算法幾乎可以做到無失真壓縮。
參考文獻:
[1] Rastegari, M.; Ordonez, V.; Redmon, J.; and Farhadi, A. 2016. Xnor-net: Imagenet classification using binary convolutional neural networks. European Conference on Computer Vision.
[2] Li, F.; Zhang, B.; and Liu, B. 2016. Ternary weight networks. arXiv preprint arXiv:1605.04711.