隨著語(yǔ)言模型規(guī)模日益龐大，設(shè)備端推理變得越來(lái)越緩慢且耗能巨大。一個(gè)直接且效果出人意料的解決方案是剪除那些對(duì)任務(wù)貢獻(xiàn)甚微的完整通道（channel）。我們?cè)缙诘难芯刻岢隽艘环N訓(xùn)練階段的方法 —— 自壓縮（Self-Compression）[1, 4]，它通過(guò)反向傳播自動(dòng)決定每個(gè)通道的比特寬度，從而逐步“淡化”那些無(wú)用的通道。這種方法可同時(shí)減少模型參數(shù)、數(shù)值精度，甚至可調(diào)超參數(shù)的數(shù)量，而不會(huì)影響模型的預(yù)測(cè)質(zhì)量。

當(dāng)我們將這一想法擴(kuò)展應(yīng)用于 Transformer 架構(gòu) [5] 時(shí)，觀察到了一個(gè)耐人尋味的現(xiàn)象：當(dāng)某個(gè)通道的學(xué)習(xí)精度降低至零比特，所得模型的緊湊程度甚至超過(guò)了使用固定三值編碼（ternary code）的模型。由于該方法僅作用于標(biāo)準(zhǔn)的線性層，壓縮后的網(wǎng)絡(luò)無(wú)需修改運(yùn)行時(shí)堆棧，便可在 CPU、GPU、DSP 和 NPU 上直接獲得性能提升，從而實(shí)現(xiàn)一個(gè)輕量模型在多種硬件平臺(tái)上的通用部署。

在本項(xiàng)工作中，我們更進(jìn)一步，引入了基于塊的稀疏性模式（block-based sparsity pattern）。接下來(lái)的章節(jié)將介紹如何將自壓縮機(jī)制整合進(jìn)基礎(chǔ)模型、它所產(chǎn)生的權(quán)重分布模式，以及這一方法在資源受限部署環(huán)境中的潛在影響。

自壓縮大型語(yǔ)言模型（LLM）

我們的參考模型是 nanoGPT [2]，這是一個(gè)精簡(jiǎn)版的 GPT 變體，訓(xùn)練數(shù)據(jù)集為 shakespeare_char 語(yǔ)料庫(kù)。該模型擁有約 1100 萬(wàn)個(gè)可訓(xùn)練參數(shù)，規(guī)模足夠小以實(shí)現(xiàn)快速運(yùn)行，同時(shí)又足夠大以展現(xiàn)完整的 Transformer 計(jì)算模式。

該模型包含以下結(jié)構(gòu)：

詞嵌入層：將每個(gè) token 映射為一個(gè)多維向量；

6 個(gè)相同的 Transformer 塊，每個(gè)塊包括：因果型多頭注意力機(jī)制（包含輸入和輸出的線性層），層歸一化（Layer Normalisation），一個(gè)前饋模塊，內(nèi)部又包含兩個(gè)線性層；

輸出部分：最后的層歸一化，一個(gè)線性層，一個(gè) Softmax 層，用于輸出各個(gè)候選 token 的概率分布。

在 Transformer 網(wǎng)絡(luò)中，90% 以上的權(quán)重——也就是大部分的內(nèi)存帶寬、DRAM 占用以及功耗——集中在 Transformer 塊內(nèi)的四個(gè)大型線性層中。因此，在我們的實(shí)驗(yàn)中，自壓縮僅針對(duì)這幾層線性層進(jìn)行，其余的基準(zhǔn)模型部分保持不變。

量化后的 Transformer 架構(gòu)，基于文獻(xiàn) [3] 中的圖示，并由作者修改。

后續(xù)章節(jié)將分析在各個(gè)塊（blocks）和通道（channels）中出現(xiàn)的稀疏性特征。了解哪些層最先變得稀疏，可以為我們提供有關(guān)大型語(yǔ)言模型（LLM）中哪些層相對(duì)不那么重要的有價(jià)值見(jiàn)解。這一發(fā)現(xiàn)可能有助于未來(lái)優(yōu)化工作的定向開(kāi)展，特別是在那些冗余自然積累的部分。

自壓縮的工作原理

自壓縮方法 [1, 4] 使網(wǎng)絡(luò)在常規(guī)神經(jīng)網(wǎng)絡(luò)訓(xùn)練過(guò)程中自主學(xué)習(xí)其通道寬度和數(shù)值精度。每個(gè)輸出通道都通過(guò)一個(gè)可微分的函數(shù)進(jìn)行量化。

其中，比特深度 b≥0 和縮放指數(shù) e 是可學(xué)習(xí)參數(shù)，其地位與神經(jīng)網(wǎng)絡(luò)權(quán)重相同。我們使用直通估計(jì)器（straight-through estimator），將取整操作的導(dǎo)數(shù)視為 1，從而使 b 和 e 能夠接收正常的梯度。這種做法在如 PyTorch 等深度學(xué)習(xí)框架中實(shí)現(xiàn)起來(lái)非常簡(jiǎn)單。

訓(xùn)練過(guò)程旨在最小化原始任務(wù)損失 L(0)，但我們額外引入了一個(gè)模型規(guī)模懲罰項(xiàng) Q：

其中，Q 表示模型中每個(gè)通道平均使用的比特?cái)?shù)，

γ 是由用戶設(shè)定的懲罰系數(shù)。

在適當(dāng)選擇 γ 的情況下，該方法能夠在保留基線精度的同時(shí)顯著壓縮模型總比特?cái)?shù)，并且整個(gè)流程仍屬于標(biāo)準(zhǔn)的訓(xùn)練范式。

自壓縮后的權(quán)重表現(xiàn)

隨著訓(xùn)練的進(jìn)行，網(wǎng)絡(luò)中的平均比特?cái)?shù)逐步下降，同時(shí)驗(yàn)證損失也不斷降低。訓(xùn)練初期，我們?yōu)槊總€(gè)權(quán)重分配了 4 比特，但在數(shù)百個(gè) epoch 內(nèi)，這一數(shù)值便下降了一半，最終逐漸穩(wěn)定在約 每個(gè)權(quán)重 0.55 比特左右。

圖示：訓(xùn)練過(guò)程中平均比特寬度的變化（藍(lán)色，左軸）與驗(yàn)證損失（紅色，右軸），在施加壓縮懲罰項(xiàng) γ 的條件下。

更有趣的現(xiàn)象出現(xiàn)在稀疏率及其隨模型深度的變化上。我們觀察到：稀疏性在模型更深層逐漸增強(qiáng)，這表明后期層的“信息密度”較低。在注意力模塊中，深層 block 的線性層變得極為稀疏，第 4 和第 5 個(gè) block 中有超過(guò) 95% 的權(quán)重被移除；前饋網(wǎng)絡(luò)（feed-forward）中的線性層也變得非常稀疏，約有 85% 的權(quán)重被剪除；相比之下，第 0 層（最前層）保留了超過(guò)一半的權(quán)重，這可能是因?yàn)闇\層在捕捉數(shù)據(jù)中的基本模式時(shí)至關(guān)重要。

所有 Transformer 塊中各線性層的稀疏率（即被清零權(quán)重所占百分比）。

如果這一現(xiàn)象能夠推廣到其他語(yǔ)言模型和任務(wù)中，那么即使不使用自壓縮機(jī)制，也有可能通過(guò)在網(wǎng)絡(luò)后段減少特征維度來(lái)獲得性能收益。

第0、1和5號(hào)塊中，第一個(gè)前饋線性層的量化權(quán)重直方圖。

當(dāng)我們觀察這三個(gè)塊中權(quán)重的直方圖時(shí)，會(huì)發(fā)現(xiàn)非零權(quán)重主要集中在靠近零的小數(shù)值附近，尤其是在較深的塊中。這表明即使權(quán)重未被置零，模型也傾向于保持權(quán)重較小。塊越深，剩余的大權(quán)重越少。這說(shuō)明模型并不會(huì)通過(guò)增大剩余權(quán)重的幅度來(lái)彌補(bǔ)被剪除通道的損失。

第0、1和5號(hào)塊中，第一個(gè)前饋線性層的二值掩碼圖，黑色表示被剪除的權(quán)重，白色表示保留的權(quán)重。

稀疏掩碼展示了被剪除權(quán)重的分布情況。在第0塊，剪枝較為分散，呈現(xiàn)小間隙和細(xì)線條狀，反映出個(gè)別通道被移除；在第1塊，較大范圍的權(quán)重被同時(shí)剪除，形成了橫向帶狀區(qū)域，顯示整個(gè)輸出通道被刪除；到第5塊時(shí)，該層大部分權(quán)重已被剪除，只剩下少數(shù)幾個(gè)通道保留。

結(jié)論

自壓縮（Self-Compression）[1, 4] 同時(shí)降低了權(quán)重的比特寬度和活躍權(quán)重?cái)?shù)量，同時(shí)形成了易于理解且在硬件上高效利用的通道稀疏模式。淺層大多保持稠密，以保留重要信息，而深層則變得高度稀疏。剩余的少數(shù)權(quán)重保持較小且接近零。這些結(jié)果表明，自壓縮有助于構(gòu)建更小、更快的模型，使其適合在資源受限的環(huán)境中運(yùn)行，如邊緣設(shè)備。

本文所展示的實(shí)驗(yàn)驗(yàn)證了自壓縮方法能夠成功縮減 Transformer 模型（此處以在字符級(jí)莎士比亞數(shù)據(jù)集上訓(xùn)練的 nanoGPT [3] 為例），且不會(huì)損害其預(yù)測(cè)質(zhì)量。通過(guò)讓模型自主決定保留哪些通道和權(quán)重，該方法避免了繁瑣的手動(dòng)調(diào)優(yōu)，同時(shí)生成了結(jié)構(gòu)清晰的塊稀疏（block-sparse）模型，便于在 CPU、GPU、NPU 及其他硬件上高效部署。這意味著同一個(gè)緊湊模型可以無(wú)須額外修改，即可在整個(gè)邊緣計(jì)算棧中通用。

未來(lái)的工作可以探索將該方法應(yīng)用于更大型的語(yǔ)言模型、多模態(tài) Transformer，或針對(duì)特定任務(wù)微調(diào)的模型，同時(shí)也可以嘗試將自壓縮與知識(shí)蒸餾等其他技術(shù)結(jié)合，以進(jìn)一步提升效率。

關(guān)于作者

Jakub Przybyl是 Imagination Technologies 的中期實(shí)習(xí)生。他在弗羅茨瓦夫理工大學(xué)（Wroc?aw University of Science and Technology，簡(jiǎn)稱(chēng) WUST）完成了機(jī)電一體化（Mechatronics）學(xué)士學(xué)位，同時(shí)攻讀 IT 自動(dòng)化系統(tǒng)碩士學(xué)位，專(zhuān)攻人工智能方向。他的研究興趣包括機(jī)器學(xué)習(xí)、語(yǔ)言建模以及先進(jìn)的網(wǎng)絡(luò)壓縮技術(shù)。