西風(fēng) 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI
微軟又有“1 bit LLM”新成果了——
發(fā)布BitNet v2框架，為1 bit LLM實(shí)現(xiàn)了原生4 bit激活值量化，由此可充分利用新一代GPU（如GB200）對4 bit計(jì)算的原生支持能力。
同時(shí)減少內(nèi)存帶寬&提升計(jì)算效率
之前，微軟持續(xù)研究BitNet b1.58，把LLM的權(quán)重量化到1.58-bit，顯著降低延遲、內(nèi)存占用等推理成本。
然鵝BitNet b1.58激活值還是8-bit，這就導(dǎo)致沒辦法充分利用新一代硬件的4 bit計(jì)算能力，計(jì)算環(huán)節(jié)出現(xiàn)效率瓶頸。
還有個(gè)問題也很關(guān)鍵：
研究發(fā)現(xiàn)注意力層和前饋網(wǎng)絡(luò)層的輸入激活值分布還比較接近高斯分布，量化起來相對容易。
中間狀態(tài)的激活值有很多異常值，這就給低bit量化帶來了很大阻礙。
此前的BitNet a4.8相關(guān)研究，嘗試過用4 bit量化輸入，8 bit稀疏化處理中間狀態(tài)。
這種方法雖然性能損失不大，可稀疏化在批量推理場景里不太適合用來提高吞吐量，因?yàn)橛布m合密集計(jì)算。
這次，團(tuán)隊(duì)最新推出了BitNet v2，通過引入H-BitLinear模塊，該模塊在激活量化前應(yīng)用Hadamard變換。
研究中能有效將注意力層和前饋網(wǎng)絡(luò)中尖銳的激活分布重塑為更接近高斯分布的形式，顯著減少異常通道，使4 bit激活量化成為可能。
下面來看看究竟是怎么做的。
引入H-BitLinear模塊
對于權(quán)重量化，根據(jù)團(tuán)隊(duì)此前研究，使用per-tensor absmean函數(shù)將權(quán)重量化為三元值{-1, 0, 1}
對于低bit激活，團(tuán)隊(duì)引入了H-BitLinear。
H-BitLinear被用于注意力層的權(quán)重矩陣Wo和前饋網(wǎng)絡(luò)（FFN）層的Wdown中，這兩處正是激活值異常值最為集中的位置
該模塊在激活值量化前應(yīng)用Hadamard變換，滿足以下條件：
通過Hadamard變換，中間狀態(tài)分布更接近高斯分布，顯著減少了異常值數(shù)量，使其更適合INT4量化。
對于8 bit和4 bit激活，分別采用per-token absmax和absmean函數(shù)，激活量化可以表示為：
最終，H-BitLinear的矩陣乘法運(yùn)算可形式化表示為：
其中LN表示層歸一化（layer normalization）操作。
另外，訓(xùn)練策略方面，研究人員用了STE來近似梯度，還采用混合精度訓(xùn)練更新參數(shù)。在反向傳播的時(shí)候，會(huì)繞過量化里那些不可微函數(shù)，并且利用Hadamard變換矩陣的正交性，對梯度也進(jìn)行變換。
團(tuán)隊(duì)還表示，4 bit激活的BitNet v2可以從8 bit激活的版本接著訓(xùn)練，用少量數(shù)據(jù)微調(diào)就行，性能損失基本可以忽略不計(jì)，優(yōu)化器的狀態(tài)還能接著用。
4bit激活版本相比8bit激活性能幾乎不降
實(shí)驗(yàn)階段，研究者將BitNet v2與BitNet b1.58、BitNet a4.8在不同模型規(guī)模（400M、1.3B、3B和7B）上進(jìn)行了對比，所有模型都使用1.58bit權(quán)重訓(xùn)練
主要實(shí)驗(yàn)結(jié)果顯示，引入Hadamard變換的BitNet v2（8 bit激活）相比BitNet b1.58在各規(guī)模模型上都有所提升，在7B規(guī)模上，平均準(zhǔn)確率提高了0.61%。
當(dāng)降至4 bit激活時(shí)，BitNet v2的困惑度與BitNet a4.8相當(dāng)，下游任務(wù)表現(xiàn)甚至更優(yōu)。
研究者還對BitNet v2進(jìn)行了低bit注意力狀態(tài)的詳細(xì)實(shí)驗(yàn)，采用后RoPE量化處理QKV狀態(tài)。采用3 bit KV緩存的BitNet v2在3B、7B模型上達(dá)到了與全精度KV緩存版本相當(dāng)?shù)臏?zhǔn)確率：
與后訓(xùn)練量化方法SpinQuant、QuaRot相比較，BitNet v2表現(xiàn)更優(yōu)：
消融實(shí)驗(yàn)則進(jìn)一步驗(yàn)證了Hadamard變換對低bit激活的關(guān)鍵作用，沒有旋轉(zhuǎn)變換則模型會(huì)發(fā)散。
更多研究細(xì)節(jié)，感興趣的童鞋可以查看原論文。
論文鏈接：https://arxiv.org/pdf/2504.18415
熱門跟貼