BitNet: Scaling 1-bit Transformers for Large Language Models

Abstract

大型语言模型的规模不断增加，给部署带来了挑战，并因高能耗而引发了对环境影响的担忧。在这项工作中，我们介绍了BitNet，一种为大型语言模型设计的可扩展且稳定的1位Transformer架构。具体来说，我们引入了BitLinear，作为nn.Linear层的替代品，以从头开始训练1位权重。语言建模的实验结果表明，BitNet在显著减少内存占用和能耗的同时，实现了与最先进的8位量化方法和FP16 Transformer基线相竞争的性能。此外，BitNet表现出类似于全精度Transformers的缩放规律，这表明它在保持效率和性能优势的同时，有潜力有效扩展到更大的语言模型。

1 Introduction

大型语言模型的快速增长 [BMR+20, Ope23, CND+22, ADF+23, TLI+23, TMS+23] 在各种任务中取得了显著的改进。然而，由于高昂的推理成本和能耗，托管大型语言模型的费用非常高。随着这些模型规模的扩大，访问和处理模型参数所需的内存带宽成为主要瓶颈，限制了整体推理性能。此外，当在分布式系统或多设备平台上部署这些模型时，设备间通信的开销会显著影响推理延迟和能量消耗。模型量化 [FAHA23, CCKS23, XLS+23] 已经成为一种有前景的解决方案，因为它可以显著减少大规模模型的内存占用和计算成本，同时保持竞争力的性能。
大多数现有的大型语言模型的量化方法都是训练后量化。它们简单易用，因为不需要对训练流程进行任何修改或重新训练模型。然而，当精度降低时，该方法会导致更显著的准确性损失，因为模型在训练期间未针对量化表示进行优化。
另一种量化深度神经网络的方法是量化感知训练。与训练后量化相比，它通常会带来更好的准确性，因为模型从一开始就考虑了精度降低。此外，它允许模型继续训练或进行微调，这对大型语言模型尤为重要。量化感知训练的挑战主要在于优化，即随着精度的降低，模型变得更难收敛。此外，尚不清楚量化感知训练是否遵循神经语言模型的缩放规律。
在这项工作中，我们专注于二值化（即1位），这是量化的极端情况，应用于大型语言模型。以前关于二值化神经网络的研究 [RORF16, BT19] 主要集中在卷积神经网络上。最近，有一些关于二值化Transformers的研究。然而，这些研究主要集中在机器翻译或BERT预训练上，与大型语言模型有很大不同。例如，机器翻译采用编码器-解码器架构，BERT预训练使用双向编码器，而大型语言模型使用单向解码器。此外，大型语言模型通常扩展到更大的模型规模，而BERT和机器翻译模型并没有经历如此广泛的扩展。
据我们所知，这项工作是首次研究1位大型语言模型的量化感知训练。我们提出了BitNet，一种用于大型语言模型的1位Transformer架构，旨在在内存和计算方面高效扩展。BitNet在训练期间使用低精度的二进制权重和量化激活，同时保持优化器状态和梯度的高精度。我们的方法设计为可扩展且稳定，能够高效处理大型语言模型。BitNet架构的实现非常简单，只需替换Transformer中的线性投影（即PyTorch中的nn.Linear）。此外，它还可以与其他大型语言模型的加速方法互补，如PagedAttention [KLZ+23]、FlashAttention [DFE+22, Dao23] 和推测解码 [LKM23]。
我们在一系列语言建模基准上评估了BitNet，与最先进的量化方法和FP16 Transformers进行比较。实验结果表明，BitNet在困惑度和下游任务准确性方面表现出竞争力。更重要的是，BitNet显著减少了内存占用和能量消耗。此外，我们表明BitNet遵循类似于全精度Transformers的缩放规律，表明它可以有效扩展到更大的语言模型，并在性能和效率方面具有潜在优势。

2 BitNet

如图2所示，BitNet使用与Transformers相同的布局，堆叠自注意力机制和前馈网络的块。与原始Transformer相比，BitNet使用BitLinear（公式11）代替传统的矩阵乘法，后者采用二值化（即1-bit）的模型权重。我们将其他组件保持高精度，例如在实验中使用8-bit。我们总结了以下原因。首先，残差连接和层归一化对大型语言模型的计算成本贡献可以忽略不计。其次，随着模型的增大，QKV转换的计算成本远小于参数投影的计算成本。第三，我们保留输入/输出嵌入的精度，因为语言模型必须使用高精度的概率来进行采样。

2.1

2.1 BitLinear

我们首先使用符号函数将权重二值化为+1或-1。根据[LOP+22]的方法，在二值化之前，我们将权重中心化为零均值，以在有限的数值范围内增加容量。在二值化之后使用一个缩放因子β，以减少实值权重和二值化权重之间的l2误差。一个权重矩阵W ∈ Rn×m的二值化可以表示为：

[ W = \text{Sign}(W - α), ]

[ \text{Sign}(W_{ij}) = \begin{cases}
+1, & \text{if } W_{ij} > 0 \
-1, & \text{if } W_{ij} \leq 0
\end{cases} ]

[ α = \frac{1}{nm} \sum_{ij} W_{ij} ]

我们进一步将激活值量化为b-bit精度。根据[DLBZ22]的方法，我们使用absmax量化，它通过将激活值乘以Qb并除以输入矩阵的绝对最大值，将激活值缩放到范围[-Qb, Qb] (Qb = 2^b−1)：

[ x = \text{Quant}(x) = \text{Clip}\left(\frac{x \times Qb}{γ}, -Qb + ϵ, Qb - ϵ \right), ]

[ \text{Clip}(x, a, b) = \max(a, \min(b, x)), ]

[ γ = ||x||_∞, ]

其中ϵ是防止溢出的小浮点数。

对于非线性函数（例如ReLU）之前的激活值，我们通过减去输入值的最小值将它们缩放到范围[0, Qb]，以确保所有值为非负：

[ x = \text{Quant}(x) = \text{Clip}\left(\frac{(x - η) \times Qb}{γ}, ϵ, Qb - ϵ \right), ]

[ η = \min_{ij} x_{ij}. ]

在本工作中，我们将激活值量化为8-bit，并在未来工作中探索更低的精度。此外，训练期间按张量进行量化，而推理期间按token进行量化，以确保稳定性和效率。

根据上述量化方程，矩阵乘法可以表示为：

[ y = Wx ]

我们假设W和x中的元素相互独立并共享相同的分布，并且W和x彼此独立。那么输出y的方差估计为：

[ \text{Var}(y) = n \text{Var}(wx) ]

[ = n E[w^2] E[x^2] ]

[ = n β^2 E[x^2] ≈ E[x^2] ]

对于全精度计算，输出的方差Var(y)在使用标准初始化方法（例如Kaiming初始化或Xavier初始化）时通常为1，这对训练稳定性有很大好处。为了在量化后保持方差，我们在激活量化之前引入了LayerNorm [BKH16]函数。通过这种方式，输出y的方差估计为Var(y) ≈ E[LN(x)^2] = 1，这与全精度的Var(y)具有相同的量级。在Transformers的上下文中，它与SubLN [WMH+22]具有相同的实现。使用SubLN和上述量化方法，我们得到BitLinear，其公式为：

[ y = Wx = W \text{Quant}(\text{LN}(x)) \times \frac{βγ}{Qb} ]

[ \text{LN}(x) = \frac{x - E(x)}{\sqrt{\text{Var}(x)} + ϵ}, β = \frac{1}{nm ||W||_1} ]

图2展示了BitLinear的计算流程。在SubLN操作之后，激活值使用absmax函数进行量化。然后在1-bit权重和量化激活值之间进行矩阵乘法。输出激活值通过{β,γ}重新缩放，以将它们反量化到原始精度。

模型并行化与组量化和归一化

扩展大型语言模型的一个基本技术是模型并行化[SPP+19]，它将矩阵乘法分割到多个设备上。现有模型并行化方法的前提是张量在分区维度上是独立的。然而，所有参数α, β, γ和η都是从整个张量计算出来的，破坏了独立性前提。一个解决方案是为每个参数引入一个all-reduce操作。然而，即使每个参数的通信量很小，随着模型变得更深，同步的数量也在增加，这显著减慢了前向传递。这个问题也存在于SubLN中，其中均值和方差需要在分区维度上估计。

为此，我们提出了一种使模型并行化更高效的简单方法。我们将权重和激活值划分为多个组，然后独立估计每个组的参数。这样，参数可以在本地计算，而无需额外的通信。这种方法称为组量化，其公式如下：

对于一个权重矩阵W ∈ Rn×m，我们沿着分区维度将其分为G组，每组大小为(\frac{n}{G} × m)。然后我们独立估计每组的参数：

[ α_g = \frac{G}{nm} \sum_{ij} W(g)_{ij}, β_g = \frac{G}{nm ||W(g)||_1} ]

其中W(g)表示权重矩阵的第g组。同样，对于激活值，我们可以将输入矩阵x∈Rn×m分为G组，并计算每组的参数：

[ γ_g = ||x(g)||∞, η_g = \min{ij} x(g)_{ij} ]

对于LN，我们可以应用组归一化技术[WH20]，独立计算每组的均值和方差：

[ \text{LN}(x(g)) = \frac{x(g) - E(x(g))}{\sqrt{\text{Var}(x(g))} + ϵ} ]

通过这种方式，我们可以高效地实现具有组量化和归一化的模型并行化，无需额外的通信，并且可以扩展到大型语言模型。

2.2 模型训练

直接估计器

为了训练我们的1-bit模型，我们采用了直接估计器（STE）[BLC13]在反向传播过程中近似梯度。这种方法在反向传递过程中绕过了不可微函数，如符号函数（公式2）和剪辑函数（公式5）。STE允许梯度在网络中流动而不受这些不可微函数的影响，从而使得训练我们的量化模型成为可能。

混合精度训练

虽然权重和激活值被量化到低精度，但梯度和优化器状态被存储为高精度，以确保训练的稳定性和准确性。根据之前的工作[LSL+21]，我们维护了一个高精度格式的潜在权重，以便累积可学习参数的更新。潜在权重会在前向传递过程中动态二值化，并且永远不会用于推理过程。

大学习率

优化的一大挑战在于，潜在权重上的小更新常常对1-bit权重没有影响。这会导致基于1-bit权重估计的有偏梯度和更新。在训练开始阶段，这个问题尤为严重，因为模型需要尽快收敛。为了解决这个问题，我们探索了各种方法，最后得出结论增加学习率是加速优化的最简单和最有效的方法。我们的实验表明，BitNet在收敛速度方面受益于大学习率，而FP16 Transformer在训练开始时使用相同的学习率会发散。更多细节可以在第3节中找到。

2.3 计算效率

我们从算术操作能量和内存占用的角度估计了BitNet的计算效率。我们主要关注矩阵乘法的计算，因为它对大型语言模型的成本贡献最大。

算术操作能量

根据[Hor14, ZZL22]中的能量模型，不同算术操作的能量消耗可以估算如下：

位数	加法能量估算 ( \hat{E}_{add} ) (pJ)	乘法能量估算 ( \hat{E}_{mul} ) (pJ)
	45nm	7nm
FP32	0.9	0.4
FP16	0.03	0.16
INT8	0.007	0.38

表2：不同位数表示在45nm和7nm工艺节点上的加法和乘法能量消耗[Hor14, ZZL22]。

在普通Transformers中，对于维度为 (m \times n) 和 (n \times p) 的矩阵乘法，能量消耗可以计算如下：

[ E_{add} = m \times (n - 1) \times p \times \hat{E}_{add} ]

[ E_{mul} = m \times n \times p \times \hat{E}_{mul} ]

对于BitNet，由于权重是1-bit，矩阵乘法的能量消耗主要由加法操作主导。乘法操作只用于用标量 (\beta) 和 (\frac{\gamma}{Q_b}) 缩放输出，所以乘法的能量消耗可以计算如下：

[ E_{mul} = (m \times p + m \times n) \times \hat{E}_{mul} ]

这比Transformers中的能量消耗要小得多。表1展示了W1A8 BitNet与全精度（32-32）和半精度（16-16）Transformer相比的能量节省情况。可以看到，BitNet提供了显著的能量节省，特别是在乘法操作方面，这是矩阵乘法能量消耗的主要组成部分。

3 与 FP16 Transformers 的比较

3.1 设置

我们训练了一系列不同规模的自回归语言模型，BitNet的参数量从125M到30B不等。模型在一个英语语料库上进行训练，该语料库包括Pile数据集、CommonCrawl快照、RealNews和CC-Stories数据集。我们使用SentencePiece分词器对数据进行预处理，词汇表大小为16K。除了BitNet，我们还使用相同的数据集和设置训练了Transformer基线模型，以进行公平的比较。更多细节可以在附录中找到。

3.2 推理优化的缩放法则

神经语言模型已经证明可以按照可预测的方式扩展[KMH+20]，特别是在普通Transformer架构下。损失与用于训练的计算量成幂律关系。这使我们能够确定计算预算的最佳分配，并从较小的模型中预测大型语言模型的性能。

为了研究二值化Transformer的缩放法则，我们首先绘制了BitNet和FP16 Transformer基线模型的缩放曲线，比较它们的参数量。我们固定训练token的数量并改变模型的大小。图3显示，BitNet的损失缩放与FP16 Transformer相似，都遵循幂律。我们随后使用一个不可约损失项来拟合缩放法则：

[ L(N) = aN^b + c ]

为了评估缩放法则是否能够准确预测损失，我们选择了从125M到6.7B的模型来拟合幂律中的参数，并使用该法则来预测13B和30B模型的损失。结果表明，拟合的缩放法则能够高度准确地预测BitNet的损失。此外，随着模型规模的增长，BitNet和FP16 Transformer之间的差距变得更小。

虽然上述幂律度量了BitNet的缩放趋势，但它并不能准确地建模损失和实际计算量之间的关系。之前的工作 [KMH+20, HKK+20, HBM+22] 通过计算FLOPs来估算计算量。然而，这不适用于主要由整数计算主导的1-bit模型。此外，它主要衡量训练计算，而不是推理计算。为了更好地理解神经语言模型的缩放效率，我们引入了推理优化的缩放法则。它根据能量消耗预测损失。我们关注推理能量成本，因为它随着模型的使用而扩展，而训练成本只有一次。我们按照第2.3节中的方法估算能量消耗。图3展示了在7nm工艺节点上推理能量成本的缩放曲线。结果证明，BitNet具有更高的缩放效率。在固定的计算预算下，BitNet实现了显著更好的损失表现。同时，为了获得与FP16模型相同的性能，BitNet的推理成本更低。

3.3 下游任务的结果

除了损失，我们还关注BitNet在扩展时的能力。相比于损失，由于神经语言模型的涌现性质，能力更难预测。为了用可解释的指标评估能力，我们测试了四个下游任务上的0-shot和4-shot结果，包括Hellaswag [ZHB+19]、Winogrande [SBBC20]、Winograd [LDM12]和Storycloze [MCH+16]。图4报告了不同规模的BitNet和FP16 Transformer的平均结果。类似于损失缩放曲线，下游任务的表现可以随着计算预算的增长而扩展。此外，无论是零样本还是少样本性能，BitNet的能力扩展效率都远高于FP16 Transformer基线。

3.4 稳定性测试

训练低比特Transformer的主要挑战是优化的稳定性。因此，我们通过训练一系列不同峰值学习率的模型，进行了BitNet和FP16基线的稳定性测试。图5a展示了稳定性测试的结果。结果显示，BitNet可以在大学习率下收敛，而FP16 Transformer则不能，展示了BitNet更好的训练稳定性。这种优化优势使得可以使用更大的学习率进行训练。图5b显示，BitNet可以从学习率的增加中受益，在PPL（困惑度）方面实现更好的收敛。

4 与训练后量化的比较

4.1 设置

我们使用与第3.1节中描述的相同设置训练 BitNet。我们将 BitNet 与最新的量化方法进行比较，包括 Absmax [DLBZ22]、SmoothQuant [XLS+23]、GPTQ [FAHA23] 和 QuIP [CCKS23]。这些方法是在 FP16 Transformer 模型上进行训练后量化，遵循与 BitNet 相同的训练设置和数据。其中，Absmax 和 SmoothQuant 量化了权重和激活值，而 GPTQ 和 QuIP 只降低了权重的精度。我们将这些方法应用于各种量化级别。对于仅量化权重的方法（即 GPTQ 和 QuIP），我们尝试了 W4A16 和 W2A16。对于量化权重和激活值的方法（即 Absmax 和 SmoothQuant），我们使用它们将 FP16 Transformers 量化为 W8A8、W4A4 和 W1A8。我们实现的 BitNet 是二进制权重 8 位激活（W1A8），其位数低于或等于基线方法。

4.2 结果

表3详细比较了我们提出的方法 BitNet 与各种基线方法在四个基准数据集（即 Winogrande、Winograd、Storycloze 和 Hellaswag）上的零样本性能。所有模型的大小均为 6.7B，以确保公平比较。这些方法在多个权重位级别上进行了评估，范围从 16 位到 1 位。除了下游任务的零样本准确性外，评估指标还包括验证集上的语言模型困惑度（perplexity），这为每种方法的性能提供了综合理解。

结果表明，BitNet 在实现与基线方法相当的性能水平方面表现出色，特别是在较低位级别时。BitNet 的零样本得分与8位模型相当，而推理成本则低得多。对于4位模型，仅量化权重的方法优于量化权重和激活值的方法，主要原因是激活值更难量化。作为1位模型，BitNet 显著优于量化权重和激活值的方法以及仅量化权重的方法。对于更低位的模型，BitNet 在各个基线方法中表现出一致的优越性。这证明了量化感知训练方法优于训练后量化方法。图6总结了我们的方法与基线方法从1.3B到6.7B模型规模的零样本准确性和少样本准确性。这证明了这种优势在不同规模上是一致的。

5 消融研究

在表4中，我们展示了与几种替代方法的消融研究。我们消融了不同激活量化方法以及模型训练稳定化技术的影响。BitNet 实现了 absmax 来量化激活值，并使用 SubLN 来稳定训练。一种量化替代方法是弹性函数 [LOP+22]，它通过可学习参数动态调整比例。在我们的实验中，我们发现 absmax 比弹性函数性能更好。此外，absmax 函数导致更稳定的训练，从而使得 BitNet 可以使用更大的学习率。我们还将 SubLN 与 Pre-LN 和 BMT 架构 [ZGC+23] 进行了比较。Pre-LN 是 GPT 预训练的默认架构，而 BMT 被证明可以提高二值化模型的稳定性。我们的实验表明，SubLN 优于 Pre-LN 和 BMT。因此，我们选择 absmax 和 SubLN 作为 BitNet 的实现。

6 结论与未来工作

我们提出了 BitNet，一种新颖的 1-bit Transformer 架构，用于大型语言模型。我们的方法旨在具有可扩展性和稳定性，能够高效处理大型语言模型。实验结果表明，BitNet 在困惑度和下游任务性能方面实现了具有竞争力的表现，同时显著减少了与基线方法相比的内存占用和能耗。此外，BitNet 遵循了类似于全精度 Transformers 的扩展法则，这表明它可以有效地扩展到更大的语言模型，并在性能和效率方面带来潜在的好处。未来，我们希望在模型规模和训练步骤方面扩大 BitNet 的应用。我们还对将 BitNet 应用于其它架构（例如 RetNet [SDH+23]）以训练大型语言模型感兴趣。