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

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

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

神經(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)重的一小部分即可代表一個(gè)好的解決方案，但直接訓(xùn)練容量減少的網(wǎng)絡(luò)并不能達(dá)到相樣的精度。同樣，Hinton等人。[4] 發(fā)現(xiàn)，將“知識”從高精度網(wǎng)絡(luò)轉(zhuǎn)移到低容量網(wǎng)絡(luò)可以產(chǎn)生比使用、相同損失函數(shù)的高容量網(wǎng)絡(luò)更高精度的網(wǎng)絡(luò)。

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

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

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

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

減少最終網(wǎng)絡(luò)的運(yùn)行時(shí)間。

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

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

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

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

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

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

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

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

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

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

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

可微量化進(jìn)行自壓縮

在這項(xiàng)工作中，我們使用可微量化（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í)，就沒必要使用該參數(shù)。當(dāng)發(fā)現(xiàn)權(quán)重張量中的通道可以用0 位表示時(shí)，在訓(xùn)練期間將其從網(wǎng)絡(luò)中刪除。這樣做的一個(gè)附加好處是訓(xùn)練會(huì)隨著時(shí)間的推移而加速（見圖2）。

該過程可以描述如下：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

處理分支

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

更新優(yōu)化器

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

結(jié)果

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

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

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

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

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

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

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

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

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

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