【LLM】1-bit大模型，BitNet与BitNet b1.58

BitNet

BitNet 是一种用于大型语言模型的 1-bit Transformer 架构。主要目标是通过使用低精度二进制权重和量化激活来提高计算效率。

BitLinear

重点看下BitLinear，BitLinear 是 BitNet 的核心组件，它替代了传统 Transformer 中的矩阵乘法操作。BitLinear 使用二进制（1-bit）模型权重来进行计算。其计算步骤如下：

1. 权重二值化 ： 使用符号函数将权重二值化为 +1 或 -1：

其中， 是权重的均值，用于在二值化之前将权重集中到零均值：

二值化的具体实现为： 2. 激活量化 ： 激活值被量化到 b-bit 精度。这里使用了 absmax 量化方法，将激活值缩放到 范围内，其中 ：

其中， 是输入矩阵的最大绝对值， 是一个小的浮点数，用于防止裁剪时的溢出。 3. 非线性函数前的激活量化 ： 对于非线性函数（如 ReLU）之前的激活值，将其缩放到 范围内，使得所有值均为非负：

其中， 是输入矩阵的最小值。 4. 量化后矩阵乘法表示 ：

5. 方差估计 ： 假设权重 和输入 的元素相互独立且分布相同，且 和 也相互独立。输出 的方差可以估计为：

进一步展开：

由于 的方差可以通过 来调整： 6. LayerNorm ： 为了在量化后保持输出的方差，引入了 LayerNorm 函数。在激活量化之前应用 LayerNorm，使得输出 的方差估计为：

在 Transformer 的上下文中，LayerNorm 的具体实现为 SubLN 。 7. BitLinear 的最终计算过程 ：

其中，LayerNorm 的公式为：

的计算公式为：

小结：对于 BitNet，矩阵乘法的能耗主要由加法操作主导，因为权重是 1-bit 的；由于使用了低精度二进制权重和量化激活，BitNet 的内存占用显著减少。

BitNet b1.58

BitNet b1.58 是基于 BitNet 架构的一种改进版本，该架构将标准的 Transformer 中的 nn.Linear 层替换为 BitLinear 层。BitNet b1.58 的主要特点包括：

权重量化 ：BitNet b1.58 使用 1.58 位表示权重，并采用 absmean 量化函数来约束权重值为 -1、0 或 +1。

这个量化函数首先通过其平均绝对值缩放权重矩阵，然后将每个值四舍五入到最接近的整数 {-1, 0, +1}。

激活量化 ：与原始 BitNet 类似，但不同之处在于激活在非线性函数之前不缩放到[0, Qb]范围，激活值的量化范围为 [-Qb, Qb]，以消除零点量化的影响。

LLaMA-like 组件 为了更好地融入开源社区，BitNet b1.58 设计了类似于 LLaMA 的组件，如：

• RMSNorm ：用于层归一化。
• SwiGLU ：一种门控线性单元，用于增强模型的非线性能力。
• RoPE ：用于增强模型对序列数据的处理能力。
• 移除所有偏置项 ：简化模型结构，提高计算效率。

实验

困惑度与下游任务

显存优势与吞吐

参考文献

• BitNet: Scaling 1-bit Transformers for Large Language Models，

• The Era of 1-bit LLMs:All Large Language Models are in 1.58 Bits，https://arxiv.org/pdf/2402.17764