【自然语言处理】BitNet b1.58：1bit LLM时代

论文地址：https://arxiv.org/pdf/2402.17764.pdf

相关博客
【自然语言处理】【大模型】BitNet：用1-bit Transformer训练LLM
【自然语言处理】BitNet b1.58：1bit LLM时代
【自然语言处理】【长文本处理】RMT：能处理长度超过一百万token的Transformer
【自然语言处理】【大模型】MPT模型结构源码解析(单机版)
【自然语言处理】【大模型】ChatGLM-6B模型结构代码解析(单机版)
【自然语言处理】【大模型】BLOOM模型结构源码解析(单机版)

一、BitNet

​ BitNet采用了与Transformer基本一致的模型架构，仅将标准矩阵乘法层换成了BitLinear，其他组件仍然是高精度的。BitLinear主要是包含的操纵：权重量化、激活量化以及LayerNorm。

​ 权重量化。通过减均值实现0中心化，然后用sign实现二值化。假设全精度权重为$W\in {\mathcal{R}}^{n\times m}$，则二值量化过程为
$$\begin{gathered} \tilde{W}=\text{Sign}(W-\alpha ) \\ \text{\hspace{1em}(1)} \end{gathered}$$

$$\begin{gathered} \text{Sign}(W_{ij})=\{\begin{array}{llc}+1,& & \text{if}{W}_{ij}>0\\ -1,& & \text{if}{W}_{ij}\le 0\end{array} \\ \text{\hspace{1em}(2)} \end{gathered}$$

$$\begin{gathered} \alpha =\frac{1}{nm}\sum _{ij}{W}_{ij} \\ \text{\hspace{1em}(3)} \end{gathered}$$

​ 激活量化。使用absmax的方式将激活量化至b-bit。具体的实现方式是乘以$Q_b$再除以输入矩阵的最大绝对值，从而将激活缩放至$[-Q_b,Q_b](Q_b=2^{b-1})$，即
$$\begin{gathered} \tilde{x}=\text{Quant}(x)=\text{Clip}(x\times \frac{Q_b}{\gamma },-Q_b+ϵ,Q_b-ϵ) \\ \text{\hspace{1em}(4)} \end{gathered}$$

$$\begin{gathered} \text{Clip}(x,a,b)=\max (a,\min (b,x)),\gamma =\parallel x{\parallel }_{\infty } \\ \text{\hspace{1em}(5)} \end{gathered}$$

其中$ϵ$是防止裁剪时溢出的小浮点数。

​ 对于非线性函数之前的激活值则采用不同的量化方式，通过减轻最小值的方式将其缩放至$[0,Q_b]$，从而保证所有值均为非负：
$$\begin{gathered} \tilde{x}=\text{Quant}(x)=\text{Clip}((x-\eta )\times \frac{Q_b}{\gamma },ϵ,Q_b-ϵ),\eta =\underset{i,j}{\min }{x}_{ij} \\ \text{\hspace{1em}(6)} \end{gathered}$$
​ LayerNorm。在对激活值量化前，为了保证量化后的方差稳定，采用了SubLN。

​ BitLinear的完成计算过程为
$$\begin{gathered} y=\tilde{W}\tilde{x}=\tilde{W}\text{Quant}(\text{LN}(x))\times \frac{\beta \gamma }{Q_b} \\ \text{\hspace{1em}(7)} \end{gathered}$$

$$\begin{gathered} \text{LN}(x)=\frac{x-E(x)}{\sqrt{\text{Var}(x)+ϵ}},\beta =\frac{1}{nm}\parallel W{\parallel }_1 \\ \text{\hspace{1em}(8)} \end{gathered}$$

二、BitNet b1.58

​ BitNet b1.58在BitNet的基础上做了一些修改。

​ 权重量化。采用absmean的方式将权重约束在 { − 1 , 0 , 1 } \{-1,0,1\} {−1,0,1}中，而BitNet则将权重约束为二值 { − 1 , 1 } \{-1,1\} {−1,1}。具体来说，先使用平均绝对值来缩放权重，然后通过舍入的方式转换为 { − 1 , 0 , 1 } \{-1,0,1\} {−1,0,1}：
$$\begin{gathered} \tilde{W}=\text{RoundClip}(\frac{W}{\gamma +ϵ},-1,1) \\ \text{\hspace{1em}(9)} \end{gathered}$$

$$\begin{gathered} \text{RoundClip}(x,a,b)=\max (a,\min (b,\text{round}(x))) \\ \text{\hspace{1em}(10)} \end{gathered}$$

$$\begin{gathered} \gamma =\frac{1}{nm}\sum _{ij}|W_{ij}| \\ \text{\hspace{1em}(11)} \end{gathered}$$

​ 激活量化。同BitNet一样，但是对于非线性函数前的激活不再量化至$[0,Q_b]$，而是都量化至$[-Q_b,Q_b]$。

​ 此外，为了能够方便于开源软件兼容，整体结构采用类似LLaMA的结构。具体来说，使用RMSNorm、SwiGLU、RoPE并移除所有偏置。

三、实验

1. 困惑度

​ BitNet b1.58在3B大小时，困惑度与LLaMA相匹配，但是速度快2.71倍且显存使用减少3.55倍。当BitNet b1.58大小为3.9B时，速度快2.4倍且显存减少3.32倍，并且效果显著优于LLaMA 3B。

2. 下游任务

​ 随着模型尺寸的增加，BitNet b1.58和LLaMA在下游任务上的差距逐步缩小。在尺寸达到3B时，BitNet b.158能够与全精度相匹配。

3. 显存和延时

​ 随着模型尺寸的增加，BitNet b1.58的速度优势和显存优势会更加明显。

4. 能耗

​ 矩阵乘法是LLM中能耗最高的部分。BitNet b1.58主要是INT8的加法计算，而LLaMA则是由FP16加法和乘法组成。在7nm芯片上，BitNet b1.58能够节约71.4倍的计算能耗。随着模型尺寸的增加，BitNet b1.58在能耗方面会越来越高效。

5. 吞吐

​ 相同机器下，BitNet b1.58的batch size是LLaMA LLM的11倍，吞吐则是8.9倍。