一文告訴你AI模型QAT量化遇到震蕩問題應(yīng)該如何解決呢？

1、簡介

量化是優(yōu)化神經(jīng)網(wǎng)絡(luò)以實現(xiàn)高效推理和設(shè)備執(zhí)行同時保持高精度的最成功方法之一。通過將常規(guī)32位浮點格式的權(quán)重和激活壓縮為更高效的低位定點表示，如INT8，這樣可以在邊緣設(shè)備上部署神經(jīng)網(wǎng)絡(luò)時降低功耗并加速推理。

盡管量化具有明顯的功率和延遲優(yōu)勢，但由于精度降低，量化是以增加噪聲為代價的。然而，近年來的研究人員已經(jīng)表明，神經(jīng)網(wǎng)絡(luò)對這種噪聲是魯棒的，并且可以使用訓(xùn)練后量化技術(shù)（PTQ）以最小的精度下降量化到8比特。

PTQ可能非常有效，通常只需要訪問一個小的校準數(shù)據(jù)集，但當應(yīng)用于神經(jīng)網(wǎng)絡(luò)的低位量化（≤4位）時會受到影響。同時，量化感知訓(xùn)練（QAT）已成為事實上的標準方法，用于實現(xiàn)低比特量化，同時保持接近全精度的精度。通過模擬訓(xùn)練或微調(diào)期間的量化操作，網(wǎng)絡(luò)可以適應(yīng)量化噪聲，并達到比PTQ更好的解決方案。

本文重點關(guān)注量化權(quán)重在量化感知訓(xùn)練過程中發(fā)生的振蕩。在量化神經(jīng)網(wǎng)絡(luò)的優(yōu)化中，這是一個鮮為人知且研究不足的現(xiàn)象，在訓(xùn)練期間和訓(xùn)練后對網(wǎng)絡(luò)產(chǎn)生了重大影響。當使用流行的直通估計器（STE）進行QAT時，權(quán)重似乎在相鄰量化級別之間隨機振蕩，從而在優(yōu)化過程中導(dǎo)致有害噪聲。有了這一見解，作者調(diào)研了QAT的最新進展，這些進展聲稱性能得到了改善，并評估了它們在解決這種振蕩行為方面的有效性。

權(quán)重振蕩的一個不利癥狀是，它們會破壞在訓(xùn)練期間收集的批量歸一化層的估計推斷統(tǒng)計數(shù)據(jù)，導(dǎo)致驗證準確性差。作者發(fā)現(xiàn)這種效應(yīng)在具有深度可分離層的高效網(wǎng)絡(luò)（如MobileNets或EfficientNets）的低位量化中尤為明顯，但可以通過在訓(xùn)練后重新估計批量歸一化統(tǒng)計數(shù)據(jù)來有效解決。

雖然批量歸一化重新估計克服了振蕩的一個顯著癥狀，但它并沒有解決其根本原因。為此，我們提出了兩種有效減少振蕩的新算法：振蕩抑制和迭代權(quán)重凍結(jié)。通過在振蕩源處處理振蕩，本文的方法提高了精度，超出了批量歸一化重新估計的水平。本文表明，這兩種方法在高效網(wǎng)絡(luò)的4-bit和3-bit量化方面都取得了最先進的結(jié)果，如MobileNetV2、MobileNetV3和ImageNet上的EfficientNet lite。

2、QAT中的振蕩

首先研究了為什么權(quán)值在量化感知訓(xùn)練中會振蕩，以及這種現(xiàn)象如何在實踐中影響神經(jīng)網(wǎng)絡(luò)訓(xùn)練。

2.1、量化感知訓(xùn)練

2.2、振蕩問題

在附錄A.3中表明降低學(xué)習(xí)率會降低振蕩的振幅，但不會影響其頻率。

值得注意的是，這種行為與隨機舍入有相似之處，其中潛在權(quán)重與量化水平的接近程度與舍入到該水平的概率有關(guān)。然而，在STE中，隨機性的來源于梯度的離散性質(zhì)，而不是采樣?？梢宰⒁獾剑袷幉⒉皇菢藴蔛TE獨有的，而是存在于文獻中提出的STE的幾種變體中，在圖1中給出了其中的一小部分。

2.3、實踐中的振蕩

這些振蕩不僅僅是這個toy例子的副作用。它們確實存在于更大的神經(jīng)網(wǎng)絡(luò)中，對它們的優(yōu)化具有重要意義。圖2顯示了使用ImageNet上的LSQ訓(xùn)練的MobileNetV2的深度可分離層中接近收斂的3位量化權(quán)重的進展。可以觀察到，許多權(quán)重似乎在2個相鄰的量化級別之間隨機振蕩。

在圖3中還可以看到，在假定的網(wǎng)絡(luò)收斂之后，很大一部分潛在權(quán)重正好位于網(wǎng)格點之間的決策邊界。這進一步強化了一個觀察結(jié)果，即很大一部分權(quán)重會振蕩而不收斂。

作者確定了與神經(jīng)網(wǎng)絡(luò)訓(xùn)練中的振蕩相關(guān)的2個主要問題：

批量歸一化推理統(tǒng)計的錯誤估計

對網(wǎng)絡(luò)優(yōu)化的不利影響

1、批量標準化的影響

在訓(xùn)練批量歸一化過程中，層跟蹤每層輸出的平均值和方差的指數(shù)移動平均值（EMA），以便在推理過程中用作真實樣本統(tǒng)計的近似值。當訓(xùn)練全精度神經(jīng)網(wǎng)絡(luò)時，可以預(yù)期權(quán)重在接近收斂時會非常緩慢地變化。因此，預(yù)計每一層的輸出統(tǒng)計數(shù)據(jù)在迭代過程中都相當穩(wěn)定，因此EMA是統(tǒng)計數(shù)據(jù)的一個很好的估計。

然而，QAT中的振蕩會導(dǎo)致整數(shù)權(quán)重的快速變化（見圖2），從而導(dǎo)致迭代之間的顯著分布偏移，甚至接近收斂。振蕩引起的輸出分布的突然而大的變化可能會破壞EMA統(tǒng)計數(shù)據(jù)，導(dǎo)致準確性的顯著下降。

事實上，有兩個因素放大了這種影響：

權(quán)重位寬

每個輸出通道的權(quán)重數(shù)量

比特寬度越低，量化級別之間的距離就越大，因為它與成比例。當振蕩權(quán)重從一個量化級別移動到另一個量化級時，它們在輸出分布中造成成比例的較大偏移。第二個重要因素是每個輸出通道的權(quán)重數(shù)量。權(quán)重的數(shù)量越小，單個權(quán)重對最終累積的貢獻就越大。當累積數(shù)量增加時，由于大數(shù)定律，振蕩的影響趨于平均。

在表1中，使用KL散度來量化population和 estimated statistics之間的差異，確實觀察到，與MobileNetV2中的逐點卷積和ResNet18中的全卷積相比，深度可分離層的KL發(fā)散要大得多。

這個問題的一個簡單直接的解決方案是在訓(xùn)練后用一小部分數(shù)據(jù)重新估計批量歸一化統(tǒng)計。這種方法作者稱之為批處理歸一化（BN）重新估計，有時用于隨機量化公式。然而，作者認為由于振蕩權(quán)重，它在確定性QAT公式中也是必不可少的。

在表2中可以看到BN重新估計不僅提高了MobileNetV2的最終量化精度，而且降低了不同種子之間的方差。作者進一步觀察到，對于MobileNetV2，隨著比特寬度的減小，精度上的差距變大，而對于ResNet18，情況并非如此。

2、對訓(xùn)練的影響

除了損害BN統(tǒng)計之外，振蕩還可能對訓(xùn)練過程本身產(chǎn)生負面影響。為了說明這一點，首先表明，如果在兩種振蕩狀態(tài)之間隨機采樣振蕩權(quán)重，則具有3位權(quán)重的收斂MobileNetV2可以實現(xiàn)較低的訓(xùn)練損失（和較高的驗證精度）。

在表3中給出了這個實驗的結(jié)果?？梢杂^察到，采樣網(wǎng)絡(luò)上的平均訓(xùn)練損失與最終收斂模型的平均訓(xùn)練損失相似。然而，許多樣本實現(xiàn)了較低的訓(xùn)練損失，并且最佳隨機采樣網(wǎng)絡(luò)顯示出顯著較低的訓(xùn)練損失。

作者還使用AdaRound的自適應(yīng)來執(zhí)行振蕩權(quán)重的二進制優(yōu)化。在最終任務(wù)丟失時同時優(yōu)化所有層的舍入，類似于文獻中使用模擬退火來解決二進制優(yōu)化問題?？梢钥吹?，這種二進制優(yōu)化在最佳隨機樣本和原始收斂網(wǎng)絡(luò)的基礎(chǔ)上顯著改進。這表明，權(quán)重振蕩會阻止網(wǎng)絡(luò)在訓(xùn)練過程中收斂到最佳局部最小值，并可能對優(yōu)化過程不利。

最后，在訓(xùn)練的早期使用振蕩凍結(jié)技術(shù)防止振蕩會導(dǎo)致比振蕩權(quán)重的二元優(yōu)化更高的驗證精度。這表明，振蕩不僅會阻止QAT在訓(xùn)練結(jié)束時收斂到最佳局部最小值，而且還會導(dǎo)致優(yōu)化器在訓(xùn)練早期朝著次優(yōu)方向發(fā)展。

3、克服QAT的振蕩問題

既然已經(jīng)確定了振蕩在優(yōu)化過程中可能會產(chǎn)生負面影響，特別是對于低位量化，所以就把重點放在如何克服它們上。首先，引入了一種用于量化振蕩的度量，然后提出了兩種新的技術(shù)，旨在防止在量化感知訓(xùn)練過程中出現(xiàn)振蕩。

3.1、量化振蕩

在解決振蕩之前，需要一種在訓(xùn)練中檢測和測量振蕩的方法。建議使用指數(shù)移動平均（EMA）來計算隨時間變化的振蕩頻率。然后，可以將最小頻率定義為振蕩權(quán)重的閾值。對于在迭代t中發(fā)生的振蕩，需要滿足兩個條件：

然后，使用指數(shù)移動平均（EMA）來跟蹤隨時間變化的振蕩頻率：

3.2、振蕩阻尼

當權(quán)重振蕩時，它們總是在兩個量化區(qū)間之間的決策閾值附近移動。這意味著振蕩權(quán)重總是接近量化bin的邊緣。為了抑制振蕩行為，使用了一個正則化項，該項鼓勵潛在權(quán)重靠近bin的中心而不是邊緣。將類似于權(quán)重衰減的阻尼損失定義為：

獨立于標度s，因此間接獨立于位寬。進一步將潛在權(quán)重剪裁到量化網(wǎng)格的范圍，使得只有在量化期間沒有被剪裁的權(quán)重才會獲得正則化效果。這對于避免在基于LSQ的范圍學(xué)習(xí)中與量化尺度梯度的任何有害相互作用很重要。這種正則化的缺點是，它不僅影響振蕩的權(quán)重，而且還會阻礙不處于振蕩狀態(tài)的權(quán)重的移動。

3.3、振蕩權(quán)值的迭代凍結(jié)

作者提出了另一種更有針對性的方法，通過在訓(xùn)練中凍結(jié)權(quán)重來防止權(quán)重振蕩。在這種方法中，跟蹤訓(xùn)練期間每個權(quán)重的振蕩頻率，如方程（4）所述。如果任何權(quán)重的振蕩頻率超過閾值，該權(quán)重將被凍結(jié)，直到訓(xùn)練結(jié)束。在整數(shù)域中應(yīng)用凍結(jié)，以便在優(yōu)化過程中scale 的潛在變化不會導(dǎo)致不同的舍入。

當一個權(quán)重振蕩時，它不一定在兩種振蕩狀態(tài)下花費相等的時間。正如在第2.2節(jié)的示例中所示，權(quán)重處于每個狀態(tài)的可能性線性地取決于該量化狀態(tài)與最優(yōu)值的距離。結(jié)果，隨著時間的推移，所有量化值的期望將對應(yīng)于最優(yōu)值。一旦權(quán)重的頻率超過閾值，它就可能處于兩種量化狀態(tài)中的任何一種。為了將權(quán)重凍結(jié)到更頻繁的狀態(tài)，使用指數(shù)移動平均值（EMA）記錄以前的整數(shù)值。然后，通過四舍五入EMA將最頻繁的整數(shù)狀態(tài)分配給凍結(jié)權(quán)重。

在算法1中總結(jié)了提出的迭代權(quán)重凍結(jié)。注意，該算法可以與任何基于梯度的優(yōu)化器組合使用，并且不限于特定的量化公式或梯度估計器。在迭代級別上凍結(jié)權(quán)重的想法與迭代修剪密切相關(guān)，在迭代修剪中，小權(quán)重被迭代修剪（凍結(jié)為零）。

4、實驗

4.1、消融實驗

1、振蕩阻尼

在表4中，作者研究了阻尼損失的強度如何影響網(wǎng)絡(luò)的最終精度以及訓(xùn)練結(jié)束時振蕩權(quán)重的比例。在前3行中可以觀察到，隨著系數(shù)λ的增加，振蕩權(quán)重的比例降低，BN重新估計前后的精度差距縮小。然而，過多的阻尼會損害最終的精度，這表明過度的正則化會抑制權(quán)重在量化級別之間的有益移動。

解決這個問題的方法是在訓(xùn)練過程中逐漸增加正規(guī)化權(quán)重。這允許潛在權(quán)重在訓(xùn)練的第一階段更自由地移動，同時通過應(yīng)用更強的正則化來減少接近收斂的有害振蕩。

作者發(fā)現(xiàn)λ的余弦退火計劃在實踐中效果良好。Han等人也注意到，這種規(guī)則化在訓(xùn)練的早期階段是有害的，但實際上采用了兩階段優(yōu)化過程。這樣的策略可以顯著抑制振蕩，同時不會損害準確性。最佳阻尼配置比BN后重新估計基線提高了近1%，比BN前重新估計基線改善了5%以上。

在圖4（左）中還看到了阻尼對圖3中相同深度可分離層的潛在重量分布的影響。正如預(yù)期的那樣，潛在權(quán)重現(xiàn)在聚集在量化bin中心周圍，在決策邊界幾乎沒有任何權(quán)重。

2、迭代權(quán)重凍結(jié)

在表5中展示了迭代權(quán)重凍結(jié)算法對各種凍結(jié)閾值的有效性。在整個訓(xùn)練過程中使用恒定的閾值，可以看到殘差振蕩的數(shù)量隨著閾值的降低而顯著減少，并且網(wǎng)絡(luò)中只保留一些低頻振蕩。還可以看到，前BN重新估計精度更接近后BN重新估計準確性，正如人們在訓(xùn)練結(jié)束時振蕩較少時所預(yù)期的那樣。

然而，如果振蕩閾值變得太低，那么在訓(xùn)練的早期階段，太多的權(quán)重會被凍結(jié)，從而降低最終的準確性。為了解決這個問題，對凍結(jié)閾值應(yīng)用了一個類似于阻尼中使用的退火計劃。這能夠使用更強的凍結(jié)閾值，并在訓(xùn)練結(jié)束時凍結(jié)幾乎所有的振蕩，此時它們最具破壞性。

最佳凍結(jié)閾值比BN后重新估計基線提高了近1%，比BN前重新估計基線改善了5%以上。它的精度與振蕩阻尼相當，同時殘差振蕩顯著減少（0.04%對1.11%）。

在圖4（右）中可以看到迭代權(quán)值凍結(jié)如何改變MobileNetV2的層conv.3.1的潛在權(quán)值分布。大部分潛在權(quán)值現(xiàn)在被凍結(jié)在bin中心，去除在圖3中決策邊界觀察到的峰值。

4.2、與其他QAT方法的比較

將克服振蕩的方法與其他QAT替代方案進行了比較，并證明了它們在流行的高效神經(jīng)網(wǎng)絡(luò)的低比特量化中的有效性。為了與文獻中現(xiàn)有的方法進行比較，作者對權(quán)重和激活進行了量化。在表6中展示了MobileNetV2的結(jié)果，并證明兩種算法在3-bits和4-bits量化方面都優(yōu)于文獻中所有競爭的QAT技術(shù)。

還在表7和表8中分別獲得了MobileNetV3 Small和EfficientNet lite的最新結(jié)果。在所有情況下，本文的振蕩預(yù)防方法都比常用的LSQ基線顯著提高（>1%），表明本文的方法對其他高效網(wǎng)絡(luò)的普遍適用性。

可以注意到，與LSQ基線相比，振蕩抑制導(dǎo)致訓(xùn)練時間增加了約33%。另一方面，迭代權(quán)重凍結(jié)在實現(xiàn)類似性能的同時，計算開銷可以忽略不計。