1-bit大模型還能再突破！新一代BitNet架構啟用4位激活值
量化到1 bit的LLM還能再突破？
這次，他們對激活值下手了！
近日，BitNet系列的原班人馬推出了新一代架構：BitNet a4.8，為1 bit大模型啟用了4位激活值：
圖片
論文地址：https://arxiv.org/pdf/2411.04965
眾所周知，激活值量化通常是比較難辦的。
本次的BitNet a4.8采用混合量化和稀疏化策略，來減輕異常通道引入的量化誤差。
簡單來說就是，對注意力層和FFN層的輸入采用4位量化，同時用8位整數稀疏化中間狀態。
大量實驗表明，BitNet a4.8在相同的訓練成本下，實現了與前代BitNet b1.58相當的性能，同時因為可以吃到4位（INT4/FP4）內核的計算紅利，實現了更快的推理速度。
BitNet a4.8僅激活55%的參數，并支持3 bit KV cache，進一步提升了大規模LLM部署和推理的效率。
BitNet a4.8
圖片
模型架構
模型的整體架構如圖1所示，BitNet a4.8采用了與BitNet b1.58相同的布局。
作者使用BitLinear替換注意力（MHA）和前饋網絡（FFN）中的線性投影，以從頭開始學習1.58 bit權重。對于激活值，采用混合量化和稀疏化策略來減輕異常值維度引入的誤差。
圖片
圖2說明了模型大小為7B的BitNet b1.58中，每個模塊輸入的分布。
注意力層和FFN層的輸入通常類似高斯分布，而在FFN下采樣之前的激活值和注意力中的輸出投影中，發現了很多異常值通道和大量接近零的條目（全精度LLM也有類似觀察結果）。
圖片
如圖3所示，直接將低位量化應用于這些中間狀態會引入很大的量化誤差。
因此，作者使用Q-Sparse的稀疏化方法，將這些中間狀態保持在8位（同時消除了計算瓶頸）。
對于自注意層的輸出投影，使用sparsify-then-quantize函數：
兩個Q分別表示權重W和激活X的量化函數，M是掩碼，根據激活X的絕對值取topK，⊙是元素乘法。
具體來說，權重量化和激活值量化函數可以表述為：
對于FFN，這里采用squared ReLU和門控線性單元（GLU）來進一步提高激活的稀疏性：
根據初步實驗的結果，使用squared ReLU時，下采樣輸入的稀疏性超過了80%，且對性能的影響最小。
此外，作者還觀察到gate + squared ReLU的輸出也表現出高激活稀疏性（7B模型為67.5%）。通過首先計算gate projection，然后僅在非零通道上執行up projection，可以進一步減少推理的計算量。
相比之下，attention和FFN的輸入中包含的異常值特征要少得多，可以使用absmean函數將激活值量化為4位整數：
模型訓練
初始化
BitNet a4.8使用BitNet b1.58的權重開始訓練，分為W1.58A8與W1.58A4兩階段。
第一階段使用8位激活和GLU + squared ReLU訓練模型；第二階段采用上面介紹過的混合量化和稀疏化。
圖片
BitNet a4.8只需少量訓練，即可快速適應4bit位寬和稀疏激活，同時性能損失可以忽略不計。
梯度近似
作者使用直通估計器（STE）對BitNet a4.8進行梯度逼近，使用混合精度訓練來更新參數。
圖片
這里直接繞過了不可微函數，包括反向傳播過程中的量化函數和topK稀疏函數。對于混合精度訓練，保持全精度latent weight來累積參數更新。
模型量化
浮點量化提供了比基于整數的量化更寬的動態范圍，這對于處理激活值的長尾分布至關重要。
研究人員將FFN下采樣層的輸入保留為8位整數，其他激活值使用MinMax量化器量化為FP4：
公式中E和M分別表示指數和尾數部分的位寬。這里采用E2M1格式，因為它的動態范圍更大。
實驗
本文將BitNet a4.8、BitNet b1.58，以及各種參數量大小的FP16精度LLaMA進行了比較。
其中的1.58 bit模型，遵循BitNet b1.58的訓練方案，采用了兩階段權重衰減和學習率調度。
圖片
所有模型都使用RedPajama數據集中的100B token進行訓練，以確保公平比較。
對于BitNet a4.8，作者首先使用95B token來訓練8位激活值的模型。然后重用優化器狀態，并使用5B token進行混合量化和稀疏化的訓練。實驗將topK設置為50%（attention的輸出投影位置）。
作者使用lm-evaluation-harness工具包，評估模型在一系列語言任務上的zero-shot準確性，包括ARC-Easy（ARCe）、ARCChallenge（ARCc）、Hellaswag（HS）、Winogrande（WGe）和PIQA（PQ）。另外還測試了在C4數據集（測試集）上的困惑度。
主要結果
圖片
表1總結了BitNet a4.8、BitNet b1.58和FP16 LLaMA的詳細測試結果。
全精度（FP16）LLaMA和BitNet b1.58之間的性能差距，隨著模型大小的增長而縮小。對于7B模型，BitNet b1.58在語言模型困惑度和任務的平均準確性方面與LLaMA相當。
此外，相比于BitNet b1.58，BitNet a4.8的平均精度幾乎沒有損失。
圖片
表2展示了各種大小的BitNet a4.8、BitNet b1.58 和 FP16 LLaMA中每個模塊的詳細稀疏性（使用C4驗證集上的非嵌入參數計算）。
值得注意的是，BitNet a4.8的稀疏性明顯高于BitNet b1.58和LLaMA。
比如在7B模型中，BitNet a4.8的整體稀疏性達到了44.5%，只有3.4B的活躍參數。down projection層的輸入顯示出特別高的稀疏性，且中間狀態分布以零為中心。
此外，gate projection的輸出非常稀疏，導致了up projection的高稀疏性（因為只需要在從Gate中選擇非零通道來執行投影）。
具體來說，對于7B BitNet a4.8，Gate和up projection的稀疏率分別為67.5%和12.0%。
圖片
表3顯示了BitNet a4.8在3B和7B模型大小下，low-bit attention的詳細情況。模型使用4位KV或QKV頭，精度損失可忽略不計，同時KV cache可以量化為3位整數。
low-bit attention對于高效的長序列建模至關重要，它減少了KV cache的內存占用和IO，并加速了注意力計算。
在本文的實驗中，作者采用RoPE后量化。使用absmax函數將QKV頭直接量化為無符號整數，無需任何校準數據集。
對于3 bit KV量化，研究人員將bos token的頭保留為4 bit，因為它包含更多的異常值特征。
消融實驗
圖片
圖4顯示了700M BitNet a4.8的訓練損耗曲線，比較了使用完整的INT4/FP4量化，以及本文的混合量化和稀疏化。
完整的INT4量化會導致發散，而混合架構在訓練困惑度方面明顯優于完整的FP4架構。
使用RedPajama數據集中25B token，來進行模型的第一階段訓練，采用absmean和MinMax量化器分別進行完整的INT4和FP4量化。
對于完整的INT4量化，由于其輸入具有更大的異常值，這里設置β = 2*mean（|X|）。
圖片
接下來為1.3B BitNet a4.8的down projection層輸入，設置不同的量化或激活函數。
所有模型都使用RedPajama數據集中的50B token進行第一階段訓練。為了確保公平比較，其他激活值都保留在8位。
圖5顯示了這些模型的訓練損失曲線。Squared ReLU的訓練困惑度比Swish略好，同時實現了更高的稀疏性。
此外，對down projection的輸入應用FP4量化會導致性能顯著下降，而將INT4激活與STE一起使用會導致發散。
參考資料：
