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

news2024/10/23 23:20:21

论文地址：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}$ ，则二值量化过程为
$\widetilde{W}=\text{Sign}(W-\alpha) \tag{1} \\$

$\text{Sign}(W_{ij})=\begin{cases} +1,&&\text{if}\;W_{ij}>0 \\ -1,&&\text{if}\;W_{ij}\leq 0 \\ \end{cases} \tag{2} \\$

$\alpha=\frac{1}{nm}\sum_{ij}W_{ij} \tag{3} \\$

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

$\text{Clip}(x,a,b)=\max(a,\min(b,x)),\quad\gamma=\parallel x\parallel_\infty \tag{5} \\$

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

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

BitLinear的完成计算过程为
$y=\widetilde{W}\tilde{x}=\widetilde{W}\text{Quant}(\text{LN}(x))\times\frac{\beta\gamma}{Q_b}\tag{7} \\$

$\text{LN}(x)=\frac{x-E(x)}{\sqrt{\text{Var}(x)+\epsilon}},\quad\beta=\frac{1}{nm}\parallel W\parallel_1 \tag{8} \\$

二、BitNet b1.58

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

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