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