自壓縮神經(jīng)網(wǎng)絡(luò)

過去十年，人工智能研究主要集中在探索深度神經(jīng)網(wǎng)絡(luò)的潛力。我們近年來看到的進步至少可以部分歸因于網(wǎng)絡(luò)規(guī)模的不斷擴大。從使用GPT-3 [1] 的文本生成到使用 Imagen [2] 的圖像生成，研究人員付出了相當(dāng)大的努力來創(chuàng)建更大、更復(fù)雜的架構(gòu)，以實現(xiàn)越來越令人印象深刻的壯舉。此外，現(xiàn)代神經(jīng)網(wǎng)絡(luò)的成功使其在各種應(yīng)用中部署。就在我寫這篇文章的時候，一個神經(jīng)網(wǎng)絡(luò)正在施圖預(yù)測我即將寫的下一個單詞，盡管它不夠準確，不能很快取代我！

另一方面，性能優(yōu)化在該領(lǐng)域受到的關(guān)注相對較少，這是神經(jīng)網(wǎng)絡(luò)更廣泛部署的一個重大障礙。造成這種情況的一個可能原因是能夠同時在數(shù)千個GPU 或其他硬件上的數(shù)據(jù)中心中訓(xùn)練大型神經(jīng)網(wǎng)絡(luò)。這與計算機圖形領(lǐng)域形成鮮明對比，例如，必須在單臺計算機上實時運行的限制產(chǎn)生了在不犧牲質(zhì)量的情況下優(yōu)化算法的強大動力。

神經(jīng)網(wǎng)絡(luò)容量的研究表明，發(fā)現(xiàn)高精度解決方案所需的網(wǎng)絡(luò)容量大于表示這些解決方案所需的容量。Frankle和Carbin [3]在他們的論文《彩票假設(shè)：尋找稀疏、可訓(xùn)練的神經(jīng)網(wǎng)絡(luò)》 [3] 中發(fā)現(xiàn)，只需要網(wǎng)絡(luò)中權(quán)重的一小部分即可代表一個好的解決方案，但直接訓(xùn)練容量減少的網(wǎng)絡(luò)并不能達到相樣的精度。同樣，Hinton等人。[4] 發(fā)現(xiàn)，將“知識”從高精度網(wǎng)絡(luò)轉(zhuǎn)移到低容量網(wǎng)絡(luò)可以產(chǎn)生比使用、相同損失函數(shù)的高容量網(wǎng)絡(luò)更高精度的網(wǎng)絡(luò)。

在本篇博文中，我們查找是否可以在訓(xùn)練時動態(tài)減少網(wǎng)絡(luò)參數(shù)。雖然這樣做具有挑戰(zhàn)性，但由于實現(xiàn)的復(fù)雜性（ PyTorch不是為處理動態(tài)網(wǎng)絡(luò)架構(gòu)而設(shè)計的，例如，在訓(xùn)練期間移除整個通道），

我們希望實現(xiàn)以下優(yōu)點。

減少最終網(wǎng)絡(luò)中的權(quán)重數(shù)量。

減少剩余權(quán)重的位寬。

減少最終網(wǎng)絡(luò)的運行時間。

減少訓(xùn)練時間。

降低設(shè)計網(wǎng)絡(luò)架構(gòu)時選擇層寬度的復(fù)雜性。

• 不需要特殊的硬件來優(yōu)化（例如，不需要稀疏矩陣乘法）。

在這項工作中，我們通過引入一種新穎的量化感知訓(xùn)練（QAT）方案來實現(xiàn)這些目標(biāo)，該方案平衡了最大化網(wǎng)絡(luò)精度和最小化網(wǎng)絡(luò)規(guī)模的要求。我們同時最大限度地提高精度并最大限度地減少權(quán)重位深度，從而消除不太重要或不必要的通道，從而以現(xiàn)有硬件可以輕松利用的方式降低計算和帶寬需求。

可微量化這是通過可微量化實現(xiàn)的，正如我在之前的文章[5]中介紹的那樣。簡而言之，可微量化允許您同時學(xué)習(xí)數(shù)字格式的參數(shù)和權(quán)重。這允許以與網(wǎng)絡(luò)中的權(quán)重完全相同的方式學(xué)習(xí)量化，并啟用諸如自壓縮網(wǎng)絡(luò)之類的新技術(shù)——本文的主題。量化函數(shù)量化為可變比特率有符號定點格式：

這可以描述為以下步驟順序：

1. 使用指數(shù)縮放輸入值：
2. 使用位深度鉗位值：
3. 四舍五入到最接近的整數(shù)：

反轉(zhuǎn)步驟 1 中引入的縮放：

其中 b 是位深度，e 是指數(shù)，x 是被量化的值（或一組值） 。為了確保連續(xù)可微性，我們在訓(xùn)練期間使用實值位深度參數(shù)。

上述函數(shù)使用舍入運算。通過它傳播可用梯度的常用方法是將四舍五入操作的梯度定義為1 而不是 0。這類似于“直通估計器” [6] 。要了解其工作原理，請考慮下圖：

當(dāng)我們從函數(shù)中“縮小”時，您可以看到它是如何實現(xiàn)的；舍入函數(shù)似乎接近y=x 線。我們將取整函數(shù)的后向傳遞（梯度）替換為函數(shù) y=x 的梯度，即常數(shù)1。

可微量化進行自壓縮

在這項工作中，我們使用可微量化（1）來減少訓(xùn)練期間網(wǎng)絡(luò)參數(shù)的位寬（即壓縮），以及（2）發(fā)現(xiàn)哪些參數(shù)可以用 0 位表示。當(dāng)神經(jīng)網(wǎng)絡(luò)中的參數(shù)可以用 0 位表示而不影響網(wǎng)絡(luò)的精度時，就沒必要使用該參數(shù)。當(dāng)發(fā)現(xiàn)權(quán)重張量中的通道可以用0 位表示時，在訓(xùn)練期間將其從網(wǎng)絡(luò)中刪除。這樣做的一個附加好處是訓(xùn)練會隨著時間的推移而加速（見圖2）。

該過程可以描述如下：

1. 將網(wǎng)絡(luò)的參數(shù)拆分為通道。
2. 用位寬和指數(shù)的單個量化參數(shù)對每個通道進行量化。
3. 為原始任務(wù)訓(xùn)練網(wǎng)絡(luò)，同時最小化所有位寬參數(shù)。

當(dāng)位寬參數(shù)達到 0 時，從網(wǎng)絡(luò)中移除該參數(shù)編碼的網(wǎng)絡(luò)權(quán)重通道。由于消除了整個輸出通道，這減少了相應(yīng)卷積的大小以及消耗輸出張量的任何后續(xù)操作，而不會更改網(wǎng)絡(luò)輸出。

通過在訓(xùn)練期間從網(wǎng)絡(luò)中移除空（即0 位）通道，我們可以顯著加速訓(xùn)練而不改變訓(xùn)練結(jié)果：訓(xùn)練結(jié)果與我們在最后只移除空通道時得到的網(wǎng)絡(luò)相同。

盡管本文中描述的方法學(xué)習(xí)壓縮和消除通道，但它可以推廣到其他硬件可利用的學(xué)習(xí)稀疏模式。

網(wǎng)絡(luò)架構(gòu)

選擇的網(wǎng)絡(luò)架構(gòu)是David Page 的CIFAR-10[7]的DAWNbench條目，這是一個可以快速訓(xùn)練的淺ResNet 。

使用快速訓(xùn)練網(wǎng)絡(luò)有幾個優(yōu)點，包括：

• 使算法設(shè)計迭代更快，
• 縮短調(diào)試周期，
• 使在合理的時間內(nèi)在單個 GPU 上執(zhí)行實驗變得容易，

幫助重現(xiàn)這項工作的結(jié)果。

該網(wǎng)絡(luò)由兩種主要類型的塊組成：卷積塊（卷積→批量歸一化→激活→池化）和殘差塊（殘差分支由兩個卷積塊組成）。

以下部分描述了如何對這些模塊應(yīng)用可微量化以使其可壓縮。

優(yōu)化目標(biāo)

這項工作的目標(biāo)是減少神經(jīng)網(wǎng)絡(luò)的推理和訓(xùn)練時間。為了實現(xiàn)這一點，應(yīng)該在損失函數(shù)中體現(xiàn)推理時間，以便將其最小化，從而產(chǎn)生更快的網(wǎng)絡(luò)。在這種情況下使用的指標(biāo)是網(wǎng)絡(luò)規(guī)模，定義為用于表示網(wǎng)絡(luò)中權(quán)重的總位數(shù)。作為網(wǎng)絡(luò)性能的體現(xiàn)，計算層輸出所需的激活張量大小或操作數(shù)也可以最小化。單個權(quán)重張量的大小可以用四個張量維度的乘積表示：輸出通道、輸入通道、濾波器高度和濾波器寬度（0、I、H、W）。由于我們使用一個單獨的數(shù)字格式量化每個輸出通道，并為層提供一個可學(xué)習(xí)的位數(shù)，因此用于表示張量的總位數(shù)由下式給出：

當(dāng) 為 0 時，ith通道變得不必要，減少了權(quán)重張量中的輸出通道總數(shù)，以及下一個卷積的權(quán)重張量中相應(yīng)的輸入通道數(shù)。因此最小化通過 最小化輸出通道的數(shù)量，可以最小化權(quán)重張量中的元素數(shù)量。這有效地最小化了權(quán)重張量的輸出維度。認識到一層的輸入通道數(shù)等于前一層的輸出通道數(shù)，可以使壓縮損失更好地反映網(wǎng)絡(luò)的大小。這樣一個權(quán)重張量的輸入維度也可以最小化：

一旦通道可以被壓縮到0 位，它就可能在訓(xùn)練期間被刪除。然而，需要克服的實際問題是，從卷積層中移除一個輸出通道并不一定意味著可以從下一層的輸入中安全地移除相應(yīng)的輸入通道，因為可以將偏差添加到層的輸出0中，在這種情況下刪除它可能會顯著改變網(wǎng)絡(luò)的輸出。為了處理這個問題，識別達到 0 位的加權(quán)通道（過濾器），并對其輸出應(yīng)用L1 損耗，以將其推至 0 位。只有當(dāng)偏差減少到0 時，這些過濾器才會被移除，因為此時移除這樣的通道不會改變網(wǎng)絡(luò)的輸出。

整個網(wǎng)絡(luò)的大小是所有層大小的總和：

為了平衡網(wǎng)絡(luò)的準確性和規(guī)模，我們簡單地使用兩項的線性組合：

其中L0是網(wǎng)絡(luò)的原始損失， 是壓縮因子。較大的 會生成較小但不太準確的網(wǎng)絡(luò)。

處理分支

壓縮網(wǎng)絡(luò)時出現(xiàn)的另一個問題是網(wǎng)絡(luò)分支的處理，例如，在殘差塊中。解決這個問題最簡單的方法是分別考慮這兩個分支。

更新優(yōu)化器

實現(xiàn)細節(jié)涉及使優(yōu)化器隨著網(wǎng)絡(luò)的變化而更新的問題。優(yōu)化器跟蹤網(wǎng)絡(luò)中每個參數(shù)的信息（元參數(shù)），當(dāng)網(wǎng)絡(luò)參數(shù)被動態(tài)刪除時，相應(yīng)的元參數(shù)也必須從優(yōu)化器中刪除。

結(jié)果

自壓縮網(wǎng)絡(luò)允許在規(guī)模和精度之間進行權(quán)衡，可以在規(guī)模準確度圖中可視化（參見圖1）。該圖中的每個點都表示一個神經(jīng)網(wǎng)絡(luò)的大小和精度，該神經(jīng)網(wǎng)絡(luò)經(jīng)過隨機壓縮率，從覆蓋范圍的對數(shù)均勻分布中采樣 。圖1 顯示了在使用隨機壓縮率訓(xùn)練網(wǎng)絡(luò)時，用于表示網(wǎng)絡(luò)權(quán)重的位數(shù)與32 位每權(quán)重基線（對應(yīng)于 32 位浮點）之間的關(guān)系。這是通過保留權(quán)重的百分比乘以剩余權(quán)重的平均位寬來計算的。網(wǎng)絡(luò)的基線精度（未壓縮精度）為95.69 ± 0.22。

圖 1：當(dāng)使用隨機壓縮率訓(xùn)練網(wǎng)絡(luò)時，用于表示網(wǎng)絡(luò)權(quán)重的位數(shù)與32位/權(quán)重基線之間的關(guān)系。

圖 2 僅顯示了網(wǎng)絡(luò)中使用的權(quán)重數(shù)量的減少。在不影響精度的情況下，可以移除大約 75% 的權(quán)重。

圖 2 顯示了使用隨機壓縮率訓(xùn)練網(wǎng)絡(luò)時，網(wǎng)絡(luò)中保留的權(quán)重百分比與精度之間的關(guān)系。

圖 3 顯示了通過在訓(xùn)練期間移除權(quán)重對訓(xùn)練時間的影響。一個世代的訓(xùn)練時間不僅取決于網(wǎng)絡(luò)的大小，還取決于系統(tǒng)的其他部分，例如輸入數(shù)據(jù)通道。為了確定基線訓(xùn)練開銷，對于同一網(wǎng)絡(luò)進行訓(xùn)練，每個層僅使用一個通道。每個訓(xùn)練世代大約需要7.5 秒。

圖 3：隨著參數(shù)從網(wǎng)絡(luò)中移除，神經(jīng)網(wǎng)絡(luò)訓(xùn)練時間加快。訓(xùn)練結(jié)束時移除了 86% 的權(quán)重。圖 4 顯示了 時使用壓縮率訓(xùn)練的網(wǎng)絡(luò)架構(gòu)。訓(xùn)練將移除除殘差層中的快捷分支。其余九個通道在訓(xùn)練結(jié)束時已經(jīng)達到 0 位，并且正在消除它們的偏差。預(yù)計它們會隨著更長的訓(xùn)練而消失。第二個殘差層中的快捷分支與它相關(guān)的損失非常低（由于它對網(wǎng)絡(luò)規(guī)模的貢獻最?。?，因此它的減少速度太慢，無法在訓(xùn)練結(jié)束時消失。圖 4：訓(xùn)練前后的層大小和每層平均位寬的示例。這里刪除了 86% 的權(quán)重和 97.6% 的位。每個方塊代表一個卷積。方塊中的值表示卷積的輸出或輸入（“in”）通道的總數(shù)，其中需要此類信息（在分支處）。

圖 5 顯示了整個訓(xùn)練過程中的網(wǎng)絡(luò)規(guī)格。它在早期迅速收縮，然后逐漸減少。

圖 5：網(wǎng)絡(luò)規(guī)模在訓(xùn)練早期快速縮小，之后逐漸減小。

優(yōu)化您的網(wǎng)絡(luò)

在本篇博文中，我們分享了一個通用框架，用于優(yōu)化神經(jīng)網(wǎng)絡(luò)的典型固定特征——通道數(shù)和位寬——以使網(wǎng)絡(luò)在訓(xùn)練過程中學(xué)會自我壓縮。這樣做的主要優(yōu)點是更快的執(zhí)行時間和更快的生成網(wǎng)絡(luò)訓(xùn)練。以前的許多工作都集中在通過創(chuàng)建稀疏層來減少網(wǎng)絡(luò)規(guī)模，這需要軟件和/或硬件的特殊支持才能更有效地運行。簡單地減少層的寬度不需要專門支持。通過減少 DRAM 帶寬，支持可變位寬可以提高多種架構(gòu)的性能。

參考

[1] T. B. Brown and al, “Language Models are Few-Shot Learners,” 2020.

[2] C. Saharia and al, “Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding,” 2022.

[3] J. Frankle and M. Carbin, “The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks,” 2018.

[4] G. Hinton, O. Vinyals and J. Dean, “Distilling the Knowledge in a Neural Network,” 2015.

[5] Cséfalvay, S, “High-Fidelity Conversion of Floating-Point Networks for Low-Precision Inference using Distillation,” 25 May 2021. [Online]. Available: https://blog.imaginationtech.com/low-precision-inference-using-distillation/.

[6] G. Hinton, “Lecture 9.3 — Using noise as a regularizer [Neural Networks for Machine Learning],” 2012. [Online]. Available: https://www.youtube.com/watch?v=LN0xtUuJsEI&list=PLoRl3Ht4JOcdU872GhiYWf6jwrk_SNhz9.

[7] Page, D, “How to Train Your ResNet 8: Bag of Tricks,” 19 Aug 2019. [Online]. Available: https://myrtle.ai/how-to-train-your-resnet-8-bag-of-tricks/.

本文作者：Szabolcs Cséfalvay