清華、哈工大把大模型壓縮到了1bit,把大模型放在手機里跑的愿望就快要實現(xiàn)了!
自從大模型火爆出圈以后,人們對壓縮大模型的愿望從未消減。這是因為,雖然大模型在很多方面表現(xiàn)出優(yōu)秀的能力,但高昂的的部署代價極大提升了它的使用門檻。這種代價主要來自于空間占用和計算量?!改P土炕?nbsp;通過把大模型的參數(shù)轉(zhuǎn)化為低位寬的表示,進(jìn)而節(jié)省空間占用。目前,主流方法可以在幾乎不損失模型性能的情況下把已有模型壓縮至 4bit。然而,低于 3bit 的量化像一堵不可逾越的高墻,讓研究人員望而生畏。
圖 1 : 量化模型的困惑度在 2bit 時迅速上升
近期,一篇由清華大學(xué)、哈爾濱工業(yè)大學(xué)合作發(fā)表在 arXiv 上的論文為突破這一阻礙帶來了希望,在國內(nèi)外學(xué)術(shù)圈引起了不小的關(guān)注。這篇論文也在一周前登上 huggingface 的熱點論文,并被著名論文推薦師 AK 推薦。研究團(tuán)隊直接越過 2bit 這一量化級別,大膽地進(jìn)行了 1bit 量化的嘗試,這在模型量化的研究中尚屬首次。
論文標(biāo)題:OneBit: Towards Extremely Low-bit Large Language Models
論文地址:https://arxiv.org/pdf/2402.11295.pdf
作者提出的方法稱作 「OneBit」,非常貼切地形容了這一工作的本質(zhì):把預(yù)訓(xùn)練大模型壓縮到真正的 1bit。該論文提出了模型參數(shù) 1bit 表示的新方法,以及量化模型參數(shù)的初始化方法,并通過量化感知訓(xùn)練(QAT)把高精度預(yù)訓(xùn)練模型的能力遷移至 1bit 量化模型。實驗表明,這一方法能夠在極大幅度壓縮模型參數(shù)的同時,保證 LLaMA 模型至少 83% 的性能。
作者指出,當(dāng)模型參數(shù)壓縮至 1bit 后,矩陣乘法中的 「元素乘」將不復(fù)存在,取而代之的是更快速的 「位賦值」操作,這將大大提升計算效率。這一研究的重要意義在于,它不但跨越了 2bit 量化的鴻溝,也使在 PC 和智能手機上部署大模型成為可能。
已有工作的局限性
模型量化主要通過把模型的 nn.Linear 層(Embedding 層和 Lm_head 層除外)轉(zhuǎn)化為低精度表示實現(xiàn)空間壓縮。此前工作 [1,2] 的基礎(chǔ)是利用 Round-To-Nearest(RTN)方法把高精度浮點數(shù)近似映射到附近的整數(shù)網(wǎng)格。這可以被表示成。
然而基于 RTN 的方法在極低位寬時(3bit 以下)存在嚴(yán)重的精度損失問題,量化后的模型能力損失十分嚴(yán)重。特別是,量化后參數(shù)以 1bit 表示時,RTN 中的縮放系數(shù) s 和零點 z 會失去實際意義。這導(dǎo)致基于 RTN 的量化方法在 1bit 量化時幾乎失效,難以有效地保留原模型的性能。
此外,此前的研究中也曾對 1bit 模型可能采用什么結(jié)構(gòu)進(jìn)行過探索。幾個月前的工作 BitNet [3] 通過讓模型參數(shù)通過 Sign (?) 函數(shù)并轉(zhuǎn)為 + 1/-1 來實現(xiàn) 1bit 表示。但這一方法存在性能損失嚴(yán)重、訓(xùn)練過程不穩(wěn)定的問題,限制了其實際應(yīng)用。
OneBit 框架
OneBit 的方法框架包括全新的 1bit 層結(jié)構(gòu)、基于 SVID 的參數(shù)初始化方法和基于量化感知知識蒸餾的知識遷移。
1. 新的 1bit 結(jié)構(gòu)
OneBit 的終極目標(biāo)是將 LLMs 的權(quán)重矩陣壓縮到 1bit。真正的 1bit 要求每個權(quán)重值只能用 1bit 表示,即只有兩種可能的狀態(tài)。作者認(rèn)為,在大模型的參數(shù)中,有兩個重要因素都必須被考慮進(jìn)來,那就是浮點數(shù)的高精度和參數(shù)矩陣的高秩。
因此,作者引入兩個 FP16 格式的值向量以補償由于量化導(dǎo)致的精度損失。這種設(shè)計不僅保持了原始權(quán)重矩陣的高秩,而且通過值向量提供了必要的浮點精度,有助于模型的訓(xùn)練和知識遷移。1bit 線性層的結(jié)構(gòu)與 FP16 高精度線性層的結(jié)構(gòu)對比如下圖:
圖 3 : FP16 線性層與 OneBit 線性層的對比
左側(cè)的 (a) 是 FP16 精度模型結(jié)構(gòu),右側(cè)的 (b) 是 OneBit 框架的線性層??梢姡?OneBit 框架中,只有值向量 g 和 h 保持 FP16 格式,而權(quán)重矩陣則全部由 ±1 組成。這樣的結(jié)構(gòu)兼顧了精度和秩,對保證穩(wěn)定且高質(zhì)量的學(xué)習(xí)過程很有意義。
OneBit 對模型的壓縮幅度究竟如何?作者在論文中給了一個計算。假設(shè)對一個 4096*4096 的線性層進(jìn)行壓縮,那么 OneBit 需要一個 4096*4096 的 1bit 矩陣,和兩個 4096*1 的 16bit 值向量。這里面總的位數(shù)為 16,908,288,總的參數(shù)個數(shù)為 16,785,408,平均每個參數(shù)占用僅僅約 1.0073 個 bit。這樣的壓縮幅度是空前的,可以說是真正的 1bit 大模型。
2. 基于 SVID 初始化量化模型
為了使用充分訓(xùn)練好的原模型更好地初始化量化后的模型,進(jìn)而促進(jìn)更好的知識遷移效果,作者提出一種新的參數(shù)矩陣分解方法,稱為 「值 - 符號獨立的矩陣分解(SVID)」。這一矩陣分解方法把符號和絕對值分開,并把絕對值進(jìn)行秩 - 1 近似,其逼近原矩陣參數(shù)的方式可以表示成:
這里的秩 - 1 近似可以通過常見的矩陣分解方式實現(xiàn),例如奇異值分解(SVD)和非負(fù)矩陣分解(NMF)。而后,作者在數(shù)學(xué)上給出這種 SVID 方法可以通過交換運算次序來和 1bit 模型框架相匹配,進(jìn)而實現(xiàn)參數(shù)初始化。并且,論文還證明了符號矩陣在分解過程中確實起到了近似原矩陣的作用。
3. 通過知識蒸餾遷移原模型能力
作者指出,解決大模型超低位寬量化的有效途徑可能是量化感知訓(xùn)練 QAT。在 OneBit 模型結(jié)構(gòu)下,通過知識蒸餾從未量化模型中學(xué)習(xí),實現(xiàn)能力向量化模型的遷移。具體地,學(xué)生模型主要接受教師模型 logits 和 hidden state 的指導(dǎo)。
訓(xùn)練時,值向量和矩陣的值會被更新。模型量化完成后,直接把 Sign (?) 后的參數(shù)保存下來,在推理和部署時直接使用。
實驗及結(jié)果
OneBit 與 FP16 Transformer、經(jīng)典的訓(xùn)練后量化強基線 GPTQ、量化感知訓(xùn)練強基線 LLM-QAT 和最新的 2bit 權(quán)重量化強基線 OmniQuant 進(jìn)行了比較。此外,由于目前還沒有 1bit 權(quán)重量化的研究,作者只對自己的 OneBit 框架使用了 1bit 權(quán)重量化,而對其他方法采取 2bit 量化設(shè)置,屬于典型的 「以弱勝強」。
在模型選擇上,作者也選擇了從 1.3B 到 13B 不同大小、OPT 和 LLaMA-1/2 不同系列的模型來證明 OneBit 的有效性。在評價指標(biāo)上,作者沿用了以往模型量化的兩大評價維度:驗證集的困惑度和常識推理的 Zero-shot 準(zhǔn)確度。
表 1 : OneBit 與基線方法的效果比較(OPT 模型與 LLaMA-1 模型)
表 2 : OneBit 與基線方法的效果比較(LLaMA-2 模型)
表 1 和表 2 展示出了 OneBit 相比于其他方法在 1bit 量化時的優(yōu)勢。就量化模型在驗證集的困惑度而言,OneBit 與 FP16 模型最為接近。就 Zero-shot 準(zhǔn)確度而言,除 OPT 模型的個別數(shù)據(jù)集外,OneBit 量化模型幾乎取得了最佳的性能。其余的 2bit 量化方法在兩種評價指標(biāo)上呈現(xiàn)較大的損失。
值得注意的是,OneBit 在模型越大時,效果往往越好。也就是說,隨著模型規(guī)模增大,F(xiàn)P16 精度模型在困惑度降低上收效甚微,但 OneBit 卻表現(xiàn)出更多的困惑度下降。此外,作者還指出量化感知訓(xùn)練對于超低位寬量化或許十分有必要。
圖 4 : 常識推理任務(wù)對比
圖 5 : 世界知識對比
圖 6 : 幾種模型的空間占用和平均位寬
圖 4 - 圖 6 還對比了幾類小模型的空間占用和性能損失,它們是通過不同的途徑獲得的:包括兩個充分訓(xùn)練的模型 Pythia-1.0B 和 TinyLLaMA-1.1B,以及通過低秩分解獲得的 LowRank Llama 和 OneBit-7B??梢钥闯?,盡管 OneBit-7B 有最小的平均位寬、占用最小的空間,它在常識推理能力上仍然優(yōu)于不遜于其他模型。作者同時指出,模型在社會科學(xué)領(lǐng)域面臨較嚴(yán)重的知識遺忘??偟膩碚f,OneBit-7B 展示出了其實際應(yīng)用價值。正如圖 7 所展示的,OneBit 量化后的 LLaMA-7B 模型經(jīng)過指令微調(diào)后,展示出了流暢的文本生成能力。
圖 7 : OneBit 框架量化后的 LLaMA-7B 模型的能力
討論與分析
1. 效率
表 3 : OneBit 在不同 LLaMA 模型的壓縮比
表 3 給出的是 OneBit 對不同規(guī)模 LLaMA 模型的壓縮比??梢钥闯觯琌neBit 對模型的壓縮比均超過 90%,這一壓縮能力是史無前例的。其中值得注意的是,隨著模型增大,OneBit 的壓縮比越高,這是由于 Embedding 層這種不參與量化的參數(shù)占比越來越小。前文提到,模型越大,OneBit 帶來的性能增益越大,這顯示出 OneBit 在更大模型上的優(yōu)勢。
圖 8 : 模型大小與性能的權(quán)衡
雖然超低比特量化可能會導(dǎo)致一定的性能損失,但如圖 8 所示,它在大小和性能之間達(dá)到了良好的平衡。作者認(rèn)為,壓縮模型的大小十分重要,特別是在移動設(shè)備上部署模型時。
此外,作者還指出了 1bit 量化模型在計算上的優(yōu)勢。由于參數(shù)是純二進(jìn)制的,可以用 0/1 在 1bit 內(nèi)表示,這毫無疑問地節(jié)省大量的空間。高精度模型中矩陣乘法的元素相乘可以被變成高效的位運算,只需位賦值和加法就可以完成矩陣乘積,非常有應(yīng)用前景。
2. 魯棒性
二值網(wǎng)絡(luò)普遍面臨訓(xùn)練不穩(wěn)定、收斂困難的問題。得益于作者引入的高精度值向量,模型訓(xùn)練的前向計算和后向計算均表現(xiàn)的十分穩(wěn)定。BitNet 更早地提出 1bit 模型結(jié)構(gòu),但該結(jié)構(gòu)很難從充分訓(xùn)練的高精度模型中遷移能力。如圖 9 所示,作者嘗試了多種不同的學(xué)習(xí)率來測試 BitNet 的遷移學(xué)習(xí)能力,發(fā)現(xiàn)在教師指導(dǎo)下其收斂難度較大,也在側(cè)面證明了 OneBit 的穩(wěn)定訓(xùn)練價值。
圖 9 : BitNet 在多種不同學(xué)習(xí)率下的訓(xùn)練后量化能力
論文的最后,作者還建議了超低位寬未來可能得研究方向。例如,尋找更優(yōu)的參數(shù)初始化方法、更少的訓(xùn)練代價,或進(jìn)一步考慮激活值的量化。
更多技術(shù)細(xì)節(jié)請查看原論文。