當前位置: 妍妍網 > 辦公

讓大模型「瘦身」90%!清華&哈工大提出極限壓縮方案

2024-03-28辦公

來源 | 量子位

對大模型進行量化、剪枝等 壓縮 操作,是部署時最常見不過的一環了。

不過,這個 極限 究竟有多大?

清華大學和哈工大的一項聯合研究給出的答案是:

90%。

他們提出了大模型1bit極限壓縮框架 OneBit ,首次實作大模型權重壓縮超越90%並保留大部份 (83%) 能力。

可以說,玩兒的就是「既要也要」~

一起來看看。

大模型1bit量化方法來了

從剪枝、量化,到知識蒸餾、權重低秩分解,大模型已經可以實作壓縮四分之一權重而幾乎無失真。

權重量化通常是指把大模型的參數轉化為低位寬的表示,可以透過對充分訓練後的模型進行轉換(PTQ)或在訓練中引入量化步驟(QAT)來實作。

然而,現有量化方法在低於3bit時面臨嚴重的效能損失,這主要是由於:

  1. 現有的參數低位寬表示方法在1bit時存在嚴重的精度損失。基於Round-To-Nearest方法的參數以1bit表示時,其轉換的縮放系數s和零點z會失去實際意義。

  2. 現有的1bit模型結構沒有充分考慮到浮點精度的重要性。浮點參數的缺失可能影響模型計算過程的穩定性,嚴重降低其本身的學習能力。

為了克服1bit超低位寬量化的阻礙,作者提出一種全新的1bit模型框架:OneBit,它包括全新的1bit線性層結構、基於SVID的參數初始化方法和基於量化感知知識蒸餾的深度遷移學習。

這種新的1bit模型量化方法能夠以極大的壓縮振幅、超低的空間占用和有限的計算成本,保留原模型絕大部份的能力。這對於實作大模型在PC端甚至智慧型手機上的部署意義非凡。

整體框架

OneBit框架總體上可以包括:全新設計的1bit模型結構、基於原模型初始化量化模型參數的方法以及基於知識蒸餾的深度能力遷移。

這種全新設計的1bit模型結構能夠有效克服以往量化工作在1bit量化時嚴重的精度損失問題,並且在訓練、遷移過程中表現出出色的穩定性。

量化模型的初始化方法能為知識蒸餾設定更好的起點,加速收斂的同時獲得更加的能力遷移效果。

1、1bit模型結構

1bit要求每個權重值只能用1bit表示,所以最多只有兩種可能的狀態。

作者選用±1作為這兩種狀態,好處就是,它代表了數位系統中的兩種符號、功能更加完備,同時可以透過Sign(·)函式方便地獲得。

作者的1bit模型結構是透過把FP16模型的所有線性層(嵌入層和lm_head除外)替換為1bit線性層實作的。

這裏的1bit線性層除透過Sign(·)函式獲得的1bit權重之外,還包括另外兩個關鍵元件—FP16精度的值向量。

FP16線性層與OneBit線性層的對比

這種設計不僅保持了原始權重矩陣的高秩,而且透過值向量提供了必要的浮點精度,對保證穩定且高品質的學習過程很有意義。

從上圖可以看出,只有值向量g和h保持FP16格式,而權重矩陣則全部由±1組成。

作者透過一個例子可以一觀OneBit的壓縮能力。

假設壓縮一個4096 4096的FP16線性層,OneBit需要一個4096 4096的1bit矩陣和兩個4096*1的FP16值向量。

這裏面總的位數為16,908,288,總的參數個數為16,785,408,平均每個參數占用僅僅約1.0073 bit。

這樣的壓縮振幅是空前的,可以說是真正的1bit LLM。

2、參數初始化和遷移學習

為了利用充分訓練好的原模型更好地初始化量化後的模型,作者提出一種新的參數矩陣分解方法,稱為「值-符號獨立的矩陣分解(SVID)」。

這一矩陣分解方法把符號和絕對值分開,並把絕對值進行秩-1近似,其逼近原矩陣參數的方式可以表示成:

秩-1近似可以透過常用矩陣分解方法實作,例如奇異值分解(SVD)和非負矩陣分解(NMF)。

作者在數學上給出,這種SVID方法可以透過交換運算次序來和1bit模型框架相匹配,進而實作參數初始化。

此外,符號矩陣在分解過程中對近似原矩陣的貢獻也被證明,詳情見論文。

作者認為,解決大模型超低位寬量化的有效途徑可能是量化感知訓練QAT。

因此,在SVID給出量化模型的參數起點後,作者把原模型作為教師模型並透過知識蒸餾從中學習。

具體而言,學生模型主要接受教師模型的logits和hidden state的指導。

訓練時,值向量和參數矩陣的值會被更新,而在部署時,則可以直接使用量化後的1bit參數矩陣進行計算。

模型越大,效果越好

作者選擇的基線是FP16 Transformer、GPTQ、LLM-QAT和OmniQuant。

後三個都屬於量化領域中經典的強基線,特別是OmniQuant是自作者之前最強的2bit量化方法。

由於目前還沒有1bit權重量化的研究,作者只對OneBit框架使用1bit權重量化,而對其他方法采取2bit量化設定。

對於蒸餾數據,作者仿照LLM-QAT利用教師模型自采樣的方式產生數據。

作者從1.3B到13B不同大小、OPT和LLaMA-1/2不同系列的模型來證明OneBit的有效性。在評價指標上,使用驗證集的困惑度和常識推理的Zero-shot準確度。詳情見論文。


上表展示了OneBit相比於其他方法在1bit量化時的優勢。 值得註意的是,模型越大時,OneBit效果往往越好。

隨著模型規模增大,OneBit量化模型降低的困惑度比FP16模型降低的困惑度要多。

以下是幾種不同小模型的常識推理、世界知識和空間占用情況:



作者還比較了幾種不同型別小模型的大小和實際能力。

作者發現,盡管OneBit-7B平均位寬最小、占用的空間最小、訓練的步數也相對少,但它在常識推理能力上不遜於其他模型。

同時作者也發現,OneBit-7B模型在社會科學領域出現較嚴重的知識遺忘。

FP16線性層與OneBit線性層的對比一個OneBit-7B指令微調後的文本生成例子

上圖還展示了一個OneBit-7B指令微調後的文本生成例子。可見,OneBit-7B有效地受到了SFT階段的能力增益,可以比較流暢地生成文本,盡管總參數只有1.3GB(與FP16的0.6B模型相當)。總的來說,OneBit-7B展示出了其實際套用價值。

分析與討論

作者展示了OneBit對不同規模LLaMA模型的壓縮比,可以看出,OneBit對模型的壓縮比均超過驚人的90%。

特別是,隨著模型增大,OneBit的壓縮比越高。

這顯示出作者方法在更大模型上的優勢:以更高的壓縮比獲得更大的邊際收益(困惑度)。此外,作者的方法在大小和效能之間做到了很好的權衡。

1bit量化模型在計算上具有優勢,意義十分重大。參數的純二進制表示,不但可以節省大量的空間,還能降低矩陣乘法對硬體的要求。

高精度模型中矩陣乘法的元素相乘可以被變成高效的位運算,只需位賦值和加法就可以完成矩陣乘積,非常有套用前景。

此外,作者的方法在訓練過程中保持了出色的穩定學習能力。

事實上,二值網路訓練的不穩定問題、對超參數的敏感性和收斂困難一直受到研究人員關註。

作者分析了高精度值向量在促進模型穩定收斂過程中的重要意義。

有前人工作提出過1bit模型架構並用於從頭訓練模型(如BitNet[1]),但它對超參數敏感並且難以從充分訓練的高精度模型中遷移學習。作者也嘗試了BitNet在知識蒸餾中的表現,發現其訓練還不夠穩定。

總結

作者提出了一種用於1bit權重量化的模型結構和相應的參數初始化方法。

在各種大小和系列的模型上進行的廣泛實驗表明,OneBit在代表性的強基線上具有明顯的優勢,並實作了模型大小與效能之間的良好折中。

此外,作者進一步分析了這種極低位元量化模型的能力和前景,並為未來的研究提供了指導。

論文地址:
https://arxiv.org/pdf/2402.11295.pdf

-END-