怎么使用蒸餾法對(duì)低精度推理出浮點(diǎn)網(wǎng)絡(luò)進(jìn)行高保真的轉(zhuǎn)換？

神經(jīng)網(wǎng)絡(luò)加速器快速、低功耗推理的一個(gè)主要挑戰(zhàn)是模型的大小。近年來(lái)，隨著模型尺寸的增加，推理時(shí)間和每次推理能耗的相應(yīng)增加，神經(jīng)網(wǎng)絡(luò)向更深的神經(jīng)網(wǎng)絡(luò)發(fā)展，激活和系數(shù)也在增加。這在資源受限的移動(dòng)和汽車(chē)應(yīng)用中尤為重要。

低精度推理有助于通過(guò)降低 DRAM 帶寬（這是影響設(shè)備能耗的一個(gè)重要因素）、計(jì)算邏輯成本和功耗來(lái)降低推理成本。在這種情況下，下面的問(wèn)題自然而然地出現(xiàn)了：編碼神經(jīng)網(wǎng)絡(luò)權(quán)重和激活的最佳位深度是多少？有幾個(gè)建議的數(shù)字格式可以減少位深度，包括Nvidia的TensorFloat，谷歌的8位非對(duì)稱(chēng)定點(diǎn)（Q8A）和bfloat16。

但是，雖然這些格式是朝著正確方向邁出的一步，但不能說(shuō)它們是最佳的：例如，大多數(shù)格式都為要表示的每個(gè)值存儲(chǔ)一個(gè)指數(shù)，當(dāng)多個(gè)值在同一區(qū)間時(shí)，這些指數(shù)可能是多余的。

更重要的是，他們沒(méi)有考慮神經(jīng)網(wǎng)絡(luò)的不同部分通常有不同的位深度要求這一事實(shí)。有些圖層可以用較低的位深度編碼，而其他圖層（如輸入和輸出層）需要更高的位深度。MobileNet v3就是一個(gè)例子，它可以從 32 位浮點(diǎn)轉(zhuǎn)換為大多數(shù)在 5-12 位區(qū)間內(nèi)的位深度（見(jiàn)圖 1）。

我們稱(chēng)使用不同的位深度來(lái)編碼原始浮點(diǎn)網(wǎng)絡(luò)的不同部分為可變位深度 （VBD） 壓縮。當(dāng)然，權(quán)重的位深度編碼和網(wǎng)絡(luò)準(zhǔn)確度也是需要權(quán)衡的。較低的位深度會(huì)導(dǎo)致更有效的推斷，但刪除過(guò)多的信息會(huì)損害準(zhǔn)確性，這意味著需要找到一個(gè)最佳的折中方案。VBD 壓縮的目標(biāo)是在壓縮和精度之間進(jìn)行平衡。

原則上，可以將其視為優(yōu)化問(wèn)題：我們希望用盡可能少的比特?cái)?shù)達(dá)到最佳精度的網(wǎng)絡(luò)。這是通過(guò)在損失函數(shù)中添加新的一項(xiàng)來(lái)實(shí)現(xiàn)的，該項(xiàng)表示網(wǎng)絡(luò)的大小，以便可以沿原始損失函數(shù)最小化，該函數(shù)可以大致表示為：總損失=網(wǎng)絡(luò)錯(cuò)誤+γ（網(wǎng)絡(luò)大?。┢渲笑檬且粋€(gè)權(quán)重因子，用于控制網(wǎng)絡(luò)大小和誤差之間的目標(biāo)權(quán)衡。

因此，要以這種方式壓縮網(wǎng)絡(luò)，我們需要兩件事：精度的可微度量（錯(cuò)誤）和網(wǎng)絡(luò)大小的可微度量（壓縮位深度）。網(wǎng)絡(luò)大小項(xiàng)和網(wǎng)絡(luò)誤差項(xiàng)相對(duì)于比特?cái)?shù)的可微性非常重要，因?yàn)樗刮覀兡軌騼?yōu)化（學(xué)習(xí)）比特深度。

可微網(wǎng)絡(luò)規(guī)模

對(duì)于可微量化，我們可以使用任何具有可微位深度參數(shù)的函數(shù)，該參數(shù)將一個(gè)或多個(gè)浮點(diǎn)值映射為硬件可表示的壓縮數(shù)字格式。例如，我們可以使用 Google Q8A 量化的可變位深度版本（其中可表示范圍不是以零為中心）：

其中：

?x?是 x 四舍五入到最近的整數(shù)（使用目標(biāo)硬件平臺(tái)的舍入模式）

B 是位深度。

E 是浮點(diǎn)表示的指數(shù)。

α是不對(duì)稱(chēng)參數(shù)。

在實(shí)踐中，量化參數(shù) B、E 和α用于壓縮多個(gè)權(quán)重，例如，單個(gè)參數(shù)用于壓縮整個(gè)神經(jīng)網(wǎng)絡(luò)層（激活或權(quán)重張量）或?qū)觾?nèi)的通道。我們還可以通過(guò)將不對(duì)稱(chēng)參數(shù)設(shè)置為 0 來(lái)使用對(duì)稱(chēng)量化（有效地將其轉(zhuǎn)換為縮放的 B位無(wú)符號(hào)整數(shù)格式）：為了實(shí)現(xiàn)對(duì)所有參數(shù)的反向傳播，我們使用直通估計(jì)器，將圓形函數(shù)的梯度作為 1。這使得公式中的所有操作都是可微的，使所有參數(shù)（包括位深度）都可學(xué)習(xí)！此時(shí)，我們可以選擇要訓(xùn)練的參數(shù)：1. 權(quán)重和最大壓縮位深度。2. 只有權(quán)重和指數(shù)（對(duì)于固定的位深度）。3. 只有量化參數(shù)，它有幾個(gè)好處（如下所述），成本可能更低的壓縮比。在生成本文的結(jié)果時(shí)選擇了選項(xiàng) 3。

另一個(gè)需要考慮的方面是優(yōu)化位深度參數(shù) B（對(duì)于某些格式的指數(shù) E）：任何硬件都需要 B 是整數(shù)。要找到整數(shù)解，我們有多種選擇：

將四舍五入與直通估計(jì)器一起應(yīng)用于 B 參數(shù)（例如，使用公式）。但是，這給優(yōu)化表面帶來(lái)了不連續(xù)性，雖然可以處理，但超出了本文的范圍。

這里選擇的替代方案是在訓(xùn)練的第一階段優(yōu)化浮點(diǎn)數(shù)B，“保守”地將其四舍五入到最近的整數(shù)?B?（ 否則激活和權(quán)重張量的重要部分可能會(huì)被鉗制），將其固定為常量并繼續(xù)訓(xùn)練。這樣平均會(huì)損失大約 0.5 位的潛在壓縮，但保證不會(huì)發(fā)生不適當(dāng)?shù)牟眉簟?/p>

可微精度測(cè)度

為了測(cè)量網(wǎng)絡(luò)的準(zhǔn)確性，我們可以簡(jiǎn)單地使用網(wǎng)絡(luò)最初訓(xùn)練的相同損失函數(shù)。然而，在許多應(yīng)用中，目標(biāo)是壓縮一個(gè)已經(jīng)用32 位浮點(diǎn)訓(xùn)練的網(wǎng)絡(luò)，這意味著我們可以用蒸餾損失代替。這意味著壓縮網(wǎng)絡(luò)的精度是根據(jù)原始網(wǎng)絡(luò)的輸出來(lái)衡量的。在這項(xiàng)工作中，選擇（輸出）logits之間的絕對(duì)差異作為蒸餾損失，但也可以采用其他措施。

使用的蒸餾損失定義為：

（等式1）

蒸餾損失有許多優(yōu)點(diǎn)：

無(wú)標(biāo)簽訓(xùn)練：我們不需要標(biāo)簽數(shù)據(jù)來(lái)壓縮網(wǎng)絡(luò)，因?yàn)槲覀冇迷季W(wǎng)絡(luò)的輸出取代了標(biāo)簽。這意味著我們可以壓縮網(wǎng)絡(luò)，而無(wú)需訪問(wèn)（可能是專(zhuān)有的）原始數(shù)據(jù)集：我們只需要具有代表性的輸入和原始網(wǎng)絡(luò)。

通用性：神經(jīng)網(wǎng)絡(luò)壓縮機(jī)具有通用性：它不需要特定于網(wǎng)絡(luò)。

更少的訓(xùn)練數(shù)據(jù)：由于我們只是訓(xùn)練量化參數(shù)，過(guò)度擬合的范圍大大縮小，所以我們可以使用更少的訓(xùn)練數(shù)據(jù)。有時(shí)候一張圖片就足夠了！

軟目標(biāo)：由于從輸入得到的蒸餾損失比標(biāo)簽包含更多的量化誤差信息，因此它允許更快、更準(zhǔn)確的收斂。

我們可以使用蒸餾損失將權(quán)重與量化參數(shù)一起訓(xùn)練。然而，在這種情況下，我們需要更多的訓(xùn)練數(shù)據(jù)來(lái)防止過(guò)度擬合。兩全其美的方法是以很小的學(xué)習(xí)率來(lái)訓(xùn)練權(quán)重，并提前停止。這樣，權(quán)重可以抵消量化誤差，而不會(huì)過(guò)度擬合數(shù)據(jù)集。

可微壓縮將通用的、可微的精度度量與可微量化相結(jié)合，得到可微壓縮的損失函數(shù)：

（等式2）第一個(gè)項(xiàng)是誤差，第二項(xiàng)是網(wǎng)絡(luò)大小的成本。B 是網(wǎng)絡(luò)的平均位深度，可根據(jù)整個(gè)網(wǎng)絡(luò)的深度參數(shù)來(lái)計(jì)算：

（等式3）其中c_i是使用位深度參數(shù)B_i量化的網(wǎng)絡(luò)參數(shù)（權(quán)重或激活）的數(shù)量。請(qǐng)注意，此測(cè)量方式取決于批次大小。例如，如果對(duì)32個(gè)批次網(wǎng)絡(luò)進(jìn)行評(píng)估，則激活張量的大小實(shí)際上比使用1 批次大小的高出 32 倍。如果目標(biāo)是將網(wǎng)絡(luò)權(quán)重存儲(chǔ)在盡可能小的空間中，則此指標(biāo)中也可以忽略激活。

選擇量化粒度一些神經(jīng)網(wǎng)絡(luò)表示，如谷歌的Q8A格式，允許將不同的比例系數(shù)（與上面的指數(shù)E的2次冪相關(guān)）應(yīng)用于權(quán)重張量（過(guò)濾器）的不同通道。這種更精細(xì)的粒度可提高給定壓縮級(jí)別的網(wǎng)絡(luò)精度。

通過(guò)對(duì)每個(gè)通道應(yīng)用單獨(dú)的 E 和α參數(shù)，同時(shí)對(duì)整個(gè)張量使用相同的 B 參數(shù)，可以通過(guò)可變位深度壓縮實(shí)現(xiàn)相同的目標(biāo)。然而，每個(gè)通道的量化會(huì)導(dǎo)致更慢的收斂，因此根據(jù)我們的經(jīng)驗(yàn)，使用學(xué)習(xí)前張量參數(shù)的訓(xùn)練階段更快，然后將這些參數(shù)分解為每通道參數(shù)，并讓它們?cè)诹硪粋€(gè)訓(xùn)練階段收斂。

這最終導(dǎo)致三個(gè)階段的訓(xùn)練計(jì)劃：

根據(jù)張量訓(xùn)練所有量化參數(shù)

切換到每通道指數(shù)和移位參數(shù)

將位深度舍入到整數(shù)，并將其固定到常量，然后訓(xùn)練指數(shù)和移位參數(shù)α。此外，權(quán)重也以較小的學(xué)習(xí)率進(jìn)行訓(xùn)練。

結(jié)果對(duì)象分類(lèi)

壓縮分類(lèi)網(wǎng)絡(luò)的精確度。

圖像分割

應(yīng)用于分割網(wǎng)絡(luò)的不同壓縮方法之間的比較。第二個(gè)列基于啟發(fā)式算法，該算法試圖在不使用反向傳播的情況下確定固定位深度的最佳指數(shù)。

風(fēng)格轉(zhuǎn)換

最后一列使用上述啟發(fā)式給出10 比特時(shí)的完全空白輸出。

結(jié)論用于編碼神經(jīng)網(wǎng)絡(luò)權(quán)重和激活的位深度對(duì)推理性能有顯著影響。將大小精度權(quán)衡作為損失函數(shù)的一部分，可以在神經(jīng)網(wǎng)絡(luò)訓(xùn)練過(guò)程中學(xué)習(xí)任意粒度的最佳比特深度。此外，當(dāng)優(yōu)化將 0 位分配給網(wǎng)絡(luò)的一部分時(shí)，它會(huì)有效地從架構(gòu)中刪除該部分，作為一種架構(gòu)搜索，從而降低計(jì)算成本和帶寬成本。今后的工作將探索這方面的可微網(wǎng)絡(luò)壓縮。我們提出了一種基于微分量化和蒸餾的通用而靈活的方法，允許在不影響精度的情況下為各種任務(wù)優(yōu)化位數(shù)。我們的方法有幾個(gè)優(yōu)點(diǎn)，包括訓(xùn)練時(shí)間短，重復(fù)使用訓(xùn)練過(guò)的網(wǎng)絡(luò)，不需要標(biāo)簽，可調(diào)整的大小精度權(quán)衡和問(wèn)題無(wú)關(guān)的損失功能。通過(guò)這種方式，我們可以將網(wǎng)絡(luò)壓縮為有效的可變位深度表示，而不犧牲對(duì)原始浮點(diǎn)網(wǎng)絡(luò)的保真度。

［i］ https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-精確格式/

［ii］https://www.tensorflow.org/lite/performance/quantization_spec

［iii］https://arxiv.org/abs/1905.12322

［iv］ https://arxiv.org/abs/1308.3432

［v］https://arxiv.org/abs/1503.02531

作者：Szabolcs Cséfalvay

原文鏈接：https://www.imaginationtech.com/blog/low-precision-inference-using-distillation/

編輯：jq