理論介紹

相比于訓練后量化方法，將量化過程插入到訓練中可以彌補量化產(chǎn)生的誤差，但是帶來的問題可能是增加了訓練的時間。在tansformer的量化實現(xiàn)中，我們采用了訓練中量化的方法，在網(wǎng)絡前向傳輸中，對權(quán)重等參數(shù)進行線性量化。反向傳播中，對scale和權(quán)重參數(shù)的求導采用Hinton的strait-through estimator的方式。在CPU上訓練花費了10天的時間，在這期間又review了最近的量化方法的文章。所以先總結(jié)一下，然后再分析一下transformer量化的結(jié)果。

1） PACT

這是一種實現(xiàn)對activation量化的方法，基本思想是通過訓練來獲得ReLU的一個clip參數(shù)a。a的動態(tài)調(diào)整能夠在減少量化誤差和保證反向傳播有效進行之間獲得平衡。PACT重新定義了ReLU過程如下：

參數(shù)a限定了activation的范圍為[0, a]。然后獲得的激活值y在進行線性映射到k bit的表示空間，如下：

在這里[0, a]是y值的一個限定，a>=y。所以其范圍比y值的實際范圍要大，這可以對y的量化誤差有一些彌補。采用strait-through estimator方法計算其相對于a的梯度為：

當a趨向于無窮大的時候，就接近于ReLU函數(shù)，所以訓練過程一定是往a增大方向移動。通過在loss中增加a的L2 規(guī)范化可以尋求一個合適的a值。

2） quantization-aware training

谷歌采用量化和訓練分離的方法，在前向計算使用量化數(shù)據(jù)，而在訓練的時候還是浮點訓練。量化方法為如下公式：

其中S為scale參數(shù)，z是零點偏移，q是量化后參數(shù)。Z值的存在會導致矩陣或者卷積運算中有交叉項。這會增加一部分加法和乘法項。這在CPU等通用處理器上容易實現(xiàn)，只是一個時間復雜度的問題，但是實際上不利于在FPGA等硬件上實現(xiàn)。所以FPGA等平臺的量化一般都讓z值為0。消除交叉項計算。對于一個矩陣乘法，量化導致了scale的組合，比如：

在這里M是浮點數(shù)據(jù)，在這里作者對其又做了一次量化，首先將M數(shù)據(jù)映射到[0.5, 1)空間，然后在使用32bit數(shù)據(jù)來表達為整數(shù)。

32bit的表達能夠降低量化精度。

在量化整個網(wǎng)絡的過程中，作者也提供了一些處理技巧。在進行線性量化的時候，采用了對稱的量化區(qū)間，比如8bit量化，正常取值范圍在[-128, 127]，作者取了對稱空間[-127, 127]。這樣做的目的和實現(xiàn)的平臺有關(guān)。在量化activation的時候，使用EMA來處理收集到的數(shù)值范圍，這樣做可以在初始訓練中，完全屏蔽掉對activation的量化，使得訓練進入到一個比較穩(wěn)定的狀態(tài)后在進行量化。BN是一個復雜的計算，但是可以將其折疊到之前的卷積層和FC層中，如下圖所示：

3） 訓練后量化，基于KL發(fā)散性。

基于訓練后的量化方法的優(yōu)勢就是量化花費時間短。在tensorRT中使用了KL發(fā)散性來描述量化后的數(shù)據(jù)和浮點數(shù)據(jù)之間的信息損失程度。通過最小化這個值來達到量化后數(shù)據(jù)包含的信息接近浮點數(shù)據(jù)的信息。這種方法的出發(fā)點是，為了保證量化后模型的精度損失較小，應該讓量化后的數(shù)據(jù)和原始浮點數(shù)據(jù)表達的信息最一致。具體的做法是：

對每層網(wǎng)絡，先收集activation的數(shù)值區(qū)間，這樣就生成一個activation值的分布；采用不同的量化區(qū)間[a,b]來對activation進行線性映射，這樣就形成了針對參數(shù)a和b的多種不同分布，然后找到和原始數(shù)據(jù)分布KL最小的分布，這個時候得到的a和b的值就是量化activation時所采用的threshold值。

Transformer量化結(jié)果

還是決定由簡入難，先進行16bit的量化，量化內(nèi)容包括transformer中的dense層，F(xiàn)C層。對權(quán)重和數(shù)據(jù)都進行16bit的量化，即將量化節(jié)點插入到計算圖中。梯度采用strait-through estimator來估計。對于embedding，softmax，layer normalization還是使用浮點值。因為擔心對這些的量化可能會導致精度降低。選擇batch size為256，epoch為20，數(shù)據(jù)集使用英語德語翻譯數(shù)據(jù)集。這個數(shù)據(jù)集有460萬個句子。在服務器上使用CPU跑了10天，以下是結(jié)果：

對比一下github上作者浮點模型的訓練結(jié)果：

發(fā)現(xiàn)存在以下問題：

1 loss下降很慢，浮點模型在訓練到達5k次的時候，loss已經(jīng)下降到4了，但是量化的訓練loss在5k次的時候才到5.4。經(jīng)歷了前幾次快速下降之后，后邊更加緩慢。

2 BLEU得分很低，訓練了10K次后得分才有0.11。得分低的原因也是loss值很低。

第一次做沒有什么經(jīng)驗，猜測可能有以下幾種原因：

1 對所有的scale我都使用了常數(shù)2作為初始值，為什么選擇2，并沒有什么原因，就是隨便選擇的?；蛟S初始值的不當導致了loss訓練很慢。設想通過以下方式來改進，先進行warmup，通過計算參數(shù)的范圍來計算出一個scale值。進行了幾輪warmup之后再進行量化訓練。

2 因為看到loss也一直是下降的趨勢，那么猜測可能是量化訓練是比正常訓練收斂慢。因為量化參數(shù)的梯度在參數(shù)超過閾值會為0，這個可能導致梯度更新較慢。

編輯：hfy