BitNet：面向大型语言模型的 1 位 Transformer 扩展架构
微软研究院
中国科学院大学
清华大学
摘要
大型语言模型规模的不断扩大给部署带来了挑战，同时高能耗也引发了对环境影响的担忧。在本研究中，我们提出了 BitNet------ 一种面向大型语言模型的可扩展且稳定的 1 位 Transformer 架构。具体而言，我们引入 BitLinear 作为 nn.Linear 层的即插即用替代方案，以便从零开始训练 1 位权重。语言建模任务的实验结果表明，与最先进的 8 位量化方法和 FP16 Transformer 基线模型相比，BitNet 在实现具有竞争力性能的同时，大幅降低了内存占用和能耗。此外，BitNet 展现出与全精度 Transformer 相似的缩放定律，这表明其在保持效率和性能优势的前提下，有望有效扩展到更大规模的语言模型。
图 1：BitNet 从零开始训练 1 位 Transformer，以高能效方式取得了具有竞争力的结果。BitNet 显著优于最先进的量化方法。随着模型规模的扩大，成本节省效果更加显著，同时与采用 FP16 训练的模型相比，性能依然具有竞争力。
（注：作者贡献说明：同等贡献。⋄通讯作者。arXiv:2310.11453v1 [cs.CL] 2023 年 10 月 17 日）
1 引言
大型语言模型的快速发展 [BMR⁺20, Ope23, CND+22, ADF⁺23, TLI⁺23, TMS⁺23] 已在各类任务中取得了显著进展。然而，由于推理成本高、能耗大，大型语言模型的部署和运行成本高昂。随着模型规模的增长，访问和处理模型参数所需的内存带宽成为主要瓶颈，限制了整体推理性能。此外，当在分布式系统或多设备平台上部署这些模型时，设备间的通信开销会显著影响推理延迟和能耗。模型量化 [FAHA23, CCKS23, XLS+23] 已成为一种极具潜力的解决方案，它能在保持模型竞争力的同时，大幅降低大规模模型的内存占用和计算成本。
目前大多数针对大型语言模型的量化方法都是训练后量化。这类方法简单易用，无需修改训练流程或重新训练模型。然而，尤其是在精度较低的情况下，训练后量化会导致更显著的精度损失，因为模型在训练过程中并未针对量化表示进行优化。
深度神经网络量化的另一类方法是量化感知训练。与训练后量化相比，量化感知训练通常能获得更高的精度，因为模型从训练之初就考虑了精度降低的影响。此外，它支持模型的持续训练或微调，这对大型语言模型至关重要。量化感知训练的核心挑战在于优化过程 ------ 随着精度降低，模型的收敛难度会增大。此外，量化感知训练是否遵循神经语言模型的缩放定律尚不明确。
本研究聚焦于二值化（即 1 位量化）这一量化的极端情况，并将其应用于大型语言模型。以往关于二值化神经网络的研究 [RORF16, BT19] 大多围绕卷积神经网络展开。近年来，已有部分研究关注二值化 Transformer，但这些研究主要集中在机器翻译或 BERT 预训练任务上，与大型语言模型存在显著差异。例如，机器翻译采用编码器 - 解码器架构，BERT 预训练使用双向编码器，而大型语言模型则采用单向解码器。此外，大型语言模型通常会扩展到更大的规模，而 BERT 和机器翻译模型并未进行如此大规模的扩展。
据我们所知，本研究是首个探索 1 位大型语言模型量化感知训练的工作。我们提出的 BitNet 是一种面向大型语言模型的 1 位 Transformer 架构，旨在实现内存和计算效率的高效扩展。BitNet 采用低精度二进制权重和量化激活，同时在训练过程中为优化器状态和梯度保留高精度。我们的方法设计兼具可扩展性和稳定性，能够高效处理大型语言模型。BitNet 架构的实现非常简单，只需替换 Transformer 中的线性投影层（即 PyTorch 中的 nn.Linear）即可。此外，它还能与其他大型语言模型加速方法互补，例如 PagedAttention [KLZ⁺23]、FlashAttention [DFE⁺22, Dao23] 和推测解码 [LKM23]。
我们在一系列语言建模基准测试中对 BitNet 进行了评估，并与最先进的量化方法和 FP16 Transformer 进行了对比。实验结果表明，BitNet 在困惑度和下游任务精度方面均取得了具有竞争力的性能。更重要的是，与基线模型相比，BitNet 显著降低了内存占用和能耗。此外，我们还发现 BitNet 遵循与全精度 Transformer 相似的缩放定律，这表明它可以有效扩展到更大规模的语言模型，并在性能和效率方面带来潜在优势。
2 BitNet
如图 2 所示，BitNet 采用与 Transformer 相同的架构布局，由自注意力机制和前馈网络块堆叠而成。与标准 Transformer 不同的是，BitNet 使用 BitLinear（公式 11）替代传统的矩阵乘法，采用二值化（即 1 位）模型权重。我们将其他组件保留为高精度（例如，实验中采用 8 位），原因如下：首先，残差连接和层归一化对大型语言模型的计算成本贡献微乎其微；其次，随着模型规模扩大，QKV 转换的计算成本远小于参数投影；第三，输入 / 输出嵌入需要保留高精度，因为语言模型需要使用高精度概率进行采样。
图 2：(a) BitLinear 的计算流程。(b) BitNet 的架构，由注意力机制和前馈网络（FFN）堆叠而成，其中矩阵乘法通过 BitLinear 实现。
2.1 BitLinear
我们首先使用符号函数将权重二值化为 + 1 或 - 1。借鉴 [LOP⁺22] 的方法，在二值化之前对权重进行零均值中心化处理，以在有限的数值范围内提升模型容量。二值化后引入缩放因子 β，以减小实值权重与二值化权重之间的 L2 误差。权重W∈Rn×m的二值化可表示为：W~=Sign(W−α),(1)Sign(Wij)={+1,−1,ifWij>0,ifWij≤0,α=nm1∑ijWij(3)
我们进一步将激活量化到 b 位精度。借鉴 [DLBZ22] 的方法，采用绝对最大值（absmax）量化，通过乘以Qb并除以输入矩阵的绝对最大值，将激活缩放到[−Qb,Qb]范围（其中Qb=2b−1）：x~=Quant(x)=Clip(x×γQb,−Qb+ϵ,Qb−ϵ),其中，ε 是一个小浮点数，用于避免裁剪时发生溢出。
对于非线性函数（如 ReLU）之前的激活，我们通过减去输入的最小值将其缩放到[0,Qb]范围，确保所有值非负：x~=Quant(x)=Clip((x−η)×γQb,ϵ,Qb−ϵ),η=minijxij.(6)
本研究中，我们将激活量化为 8 位，更低精度的量化将留待未来研究。此外，为兼顾稳定性和效率，训练过程中按张量进行量化，推理过程中按令牌进行量化。
基于上述量化公式，矩阵乘法可表示为：y=W~x~(7)
假设 W 和 x 中的元素相互独立且服从相同分布，且 W 和 x 彼此独立，则输出 y 的方差可估计为：
在全精度计算中，使用标准初始化方法（如 Kaiming 初始化或 Xavier 初始化）时，输出的方差Var(y)量级为 1，这对训练稳定性非常有利。为了在量化后保留这一方差特性，我们在激活量化之前引入层归一化（LayerNorm）[BKH16] 函数。这样，输出 y 的方差可估计为Var(y)≈E[LN(x~)2]=1，与全精度计算的方差量级一致。在 Transformer 中，其具体实现与 SubLN [WMH⁺22] 完全相同。结合 SubLN 和上述量化方法，我们得到 BitLinear，其表达式为：y=W~x~=W~Quant(LN(x))×Qbβγ(11)LN(x)=Var(x)+ϵx−E(x),β=nm1∥W∥1(12)
图 2 展示了 BitLinear 的计算流程：经过 SubLN 操作后，激活通过绝对最大值函数进行量化；1 位权重与量化后的激活进行矩阵乘法；输出激活通过 β 和 γ 进行缩放，反量化恢复到原始精度。
基于分组量化与归一化的模型并行
模型并行 [SPP⁺19] 是大型语言模型扩展的关键技术，它将矩阵乘法分配到多个设备上执行。现有模型并行方法的前提是张量在划分维度上相互独立。然而，参数 α、β、γ 和 η 均基于整个张量计算，这破坏了独立性前提。一种解决方案是为每个参数引入全归约（all-reduce）操作，但尽管每个参数的通信量较小，随着模型深度增加，同步次数会显著增多，从而大幅减慢前向传播速度。SubLN 中也存在类似问题，其均值和方差需要跨划分维度估计。
为此，我们提出了一种简化方法来提升模型并行效率：将权重和激活分组，然后独立估计每个组的参数。这样，参数可以在本地计算，无需额外通信。这种方法称为分组量化，其表达式如下：
对于权重矩阵W∈Rn×m，沿划分维度将其分为 G 组，每组大小为Gn×m，然后独立估计每组的参数：αg=nmG∑ijWij(g),βg=nmGW(g)1,其中W(g)表示权重矩阵的第 g 组。类似地，对于激活，可将输入矩阵x∈Rn×m分为 G 组，并独立计算每组的参数：γg=x(g)∞,ηg=minijxij(g)(14)
对于层归一化（LN），可采用分组归一化技术 [WH20]，独立计算每组的均值和方差：LN(x(g))=Var(x(g))+ϵx(g)−E(x(g))(15)
通过这种方式，我们可以高效实现模型并行，无需额外通信，且能够扩展到大型语言模型。
表 1：不同模型规模下 BitNet 和 Transformer 的能耗对比。输入长度设为 512 时的结果报告。
2.2 模型训练
直通估计器
为了训练 1 位模型，我们采用直通估计器（STE）[BLC13] 在反向传播过程中近似梯度。该方法在反向传播时绕过符号函数（公式 2）和裁剪函数（公式 5）等非可微函数，使梯度能够不受这些函数影响地在网络中传播，从而实现量化模型的训练。
混合精度训练
虽然权重和激活被量化为低精度，但梯度和优化器状态仍以高精度存储，以确保训练稳定性和精度。借鉴先前的研究 [LSL⁺21]，我们为可学习参数维护一个高精度的潜在权重，用于累积参数更新。潜在权重在正向传播时动态二值化，且不用于推理过程。
大学习率
优化过程中的一个挑战是：潜在权重的微小更新往往不会改变 1 位权重，导致基于 1 位权重估计的梯度和更新存在偏差。这种问题在训练初期尤为严重，因为此时模型需要快速收敛。我们尝试了多种方法来解决这一问题，最终发现增大学习率是最简单有效的优化加速方式。实验表明，BitNet 能够在大学习率下收敛，而 FP16 Transformer 在相同学习率下会在训练初期发散（详见第 3 节）。
2.3 计算效率
我们从算术运算能耗和内存占用两方面评估 BitNet 的计算效率。重点关注矩阵乘法的计算成本，因为它是大型语言模型的主要计算开销来源。
算术运算能耗
根据 [Hor14, ZZL22] 中的能耗模型，不同算术运算的能耗估计如下表所示：
表 2：45 纳米和 7 纳米工艺节点下不同比特表示的加法（ADD）和乘法（MUL）运算能耗 [Hor14, ZZL22]。
在标准 Transformer 中，对于维度为m×n和n×p的矩阵乘法，能耗计算如下：Eadd=m×(n−1)×p×E^add(16)Emul=m×n×p×E^mul(17)
对于 BitNet，由于权重为 1 位，矩阵乘法的能耗主要由加法运算主导。乘法运算仅用于通过标量 β 和Qbγ对输出进行缩放，因此乘法能耗可计算为：Emul=(m×p+m×n)×E^mul这一能耗显著低于 Transformer 中的乘法能耗。表 1 展示了 1 位权重 8 位激活（W1A8）的 BitNet 与全精度（32-32）和半精度（16-16）Transformer 的能耗对比。结果表明，BitNet 实现了显著的能耗节省，尤其是在矩阵乘法能耗的主要组成部分 ------ 乘法运算上。
3 与 FP16 Transformer 的对比
3.1 实验设置
我们训练了一系列不同规模（125M 至 30B 参数）的 BitNet 自回归语言模型。模型在英语语料库上训练，该语料库包括 Pile 数据集、Common Crawl 快照、RealNews 和 CC-Stories 数据集。我们使用 SentencePiece 分词器对数据进行预处理，词汇表大小为 16K。为保证对比公平性，我们还在相同数据集和设置下训练了 Transformer 基线模型。更多细节详见附录。
3.2 推理最优缩放定律
研究表明，采用标准 Transformer 架构的神经语言模型具有可预测的缩放特性 [KMH⁺20]，其损失随训练计算量呈幂律变化。这使得我们能够确定计算预算的最优分配，并从小型模型预测大型语言模型的性能。
为了研究二值化 Transformer 的缩放定律，我们首先绘制了 BitNet 和 FP16 Transformer 基线模型的损失随参数数量变化的缩放曲线。我们固定训练令牌数量，改变模型规模。图 3 显示，BitNet 的损失缩放与 FP16 Transformer 相似，均遵循幂律。随后，我们用一个不可约损失项拟合缩放定律：L(N)=aNb+c(19)
为了评估该缩放定律对损失的预测准确性，我们选择 125M 至 6.7B 参数的模型拟合幂律参数，并用于预测 13B 和 30B 参数模型的损失。结果表明，拟合后的缩放定律能够高精度预测 BitNet 的损失。此外，随着模型规模增大，BitNet 与 FP16 Transformer 的性能差距逐渐缩小。
上述幂律仅衡量了 BitNet 的缩放趋势，但未能很好地建模损失与实际计算量之间的关系。先前的研究 [KMH⁺20, HKK⁺20, HBM⁺22] 通过计算浮点运算次数（FLOPs）来估计计算量，但这不适用于以整数计算为主的 1 位模型，且主要衡量训练计算量而非推理计算量。为了更深入理解神经语言模型的缩放效率，我们提出了 "推理最优缩放定律"，该定律预测损失与能耗之间的关系。我们重点关注推理能耗，因为它随模型使用次数增长，而训练成本是一次性的。能耗估计方法详见第 2.3 节。图 3 展示了 7 纳米工艺节点下损失随推理能耗变化的缩放曲线，结果表明 BitNet 具有更高的缩放效率：在固定计算预算下，BitNet 的损失显著更低；同时，要达到与 FP16 模型相同的性能，BitNet 所需的推理成本远小于 FP16 模型。
图 3：BitNet 和 FP16 Transformer 的缩放曲线。
3.3 下游任务结果
除损失外，我们还关注 BitNet 随规模扩展的能力表现。与损失相比，由于神经语言模型的涌现特性，能力更难预测。为了通过可解释指标评估模型能力，我们在四个下游任务上测试了零样本和 4 样本性能，包括 Hellaswag [ZHB⁺19]、Winogrande [SBBC20]、Winograd [LDM12] 和 Storycloze [MCH⁺16]。图 4 报告了不同规模 BitNet 和 FP16 Transformer 的平均结果。与损失缩放曲线类似，下游任务性能随计算预算增加而提升。此外，无论是零样本还是少样本性能，BitNet 的能力缩放效率都远高于 FP16 Transformer 基线模型。
图 4：BitNet 和 FP16 Transformer 的零样本（左）和少样本（右）性能与推理成本的关系。
3.4 稳定性测试
低比特 Transformer 训练的主要挑战是优化稳定性。因此，我们通过在不同峰值学习率下训练一系列模型，对 BitNet 和 FP16 基线模型进行了稳定性测试。图 5 (a) 展示了稳定性测试结果：BitNet 能够在大学习率下收敛，而 FP16 Transformer 则不能，这表明 BitNet 具有更好的训练稳定性。这种优化优势使 BitNet 能够采用更大的学习率，图 5 (b) 显示，更大的学习率使 BitNet 实现了更好的收敛（以困惑度 PPL 衡量）。
图 5：相同学习率下，BitNet 比 FP16 Transformer 更稳定（左）。训练稳定性使 BitNet 能够采用更大的学习率，从而实现更好的收敛（右）。
4 与训练后量化的对比
4.1 实验设置
我们采用第 3.1 节所述的相同设置训练 BitNet，并将其与最先进的量化方法进行对比，包括 Absmax [DLBZ22]、SmoothQuant [XLS⁺23]、GPTQ [FAHA23] 和 QuIP [CCKS23]。这些方法均基于与 BitNet 相同训练设置和数据的 FP16 Transformer 模型进行训练后量化。其中，Absmax 和 SmoothQuant 同时量化权重和激活，而 GPTQ 和 QuIP 仅降低权重精度。我们在不同量化级别上应用这些方法：对于仅权重量化（GPTQ 和 QuIP），实验采用 W4A16（4 位权重 16 位激活）和 W2A16（2 位权重 16 位激活）；对于权重 - 激活量化（Absmax 和 SmoothQuant），将 FP16 Transformer 量化为 W8A8（8 位权重 8 位激活）、W4A4（4 位权重 4 位激活）和 W1A8（1 位权重 8 位激活）。我们的 BitNet 实现为 1 位权重 8 位激活（W1A8），其比特数低于或等于所有基线模型。
4.2 结果
表 3 详细对比了我们提出的 BitNet 与各类基线方法在四个基准数据集（Winogrande、Winograd、Storycloze 和 Hellaswag）上的零样本性能。所有模型的规模均为 6.7B 参数，以保证对比公平性。评估涵盖了从 16 位到 1 位的多个权重比特级别，除了下游任务的零样本精度外，还包括验证集上的语言模型困惑度，以全面了解每种方法的性能。
结果表明，BitNet 在低比特级别上表现出显著的有效性，与基线方法相比具有竞争力。BitNet 的零样本得分与 8 位模型相当，而推理成本远低于 8 位模型。对于 4 位模型，仅权重量化方法优于权重 - 激活量化方法，主要原因是激活更难量化。作为 1 位模型，BitNet 的结果显著优于所有权重 - 激活量化方法和仅权重量化方法。在低比特模型中，BitNet 在所有基线上均保持了优越的得分，这证明了量化感知训练方法相比训练后量化方法的优势。图 6 总结了模型规模从 1.3B 扩展到 6.7B 时，我们的方法与基线模型的零样本和少样本精度，结果表明 BitNet 的优势在不同规模下均保持一致。
图 6：BitNet 与训练后量化基线模型在下游任务上的零样本（左）和少样本（右）结果。
表 3：BitNet 与基线模型的零样本结果（PTQ：训练后量化，WGe：Winogrande 数据集，WG：Winograd 数据集，SC：Storycloze 数据集，HS：Hellaswag 数据集）。
5 消融实验
表 4 展示了 BitNet 与多种替代方法的消融实验结果。我们分析了激活量化方法选择和模型训练稳定性技术对性能的影响。BitNet 采用绝对最大值（absmax）进行激活量化，并使用 SubLN 保证训练稳定性。另一种量化替代方案是弹性函数 [LOP⁺22]，其通过可学习参数动态调整缩放比例。实验发现，绝对最大值量化的性能优于弹性函数，且能带来更稳定的训练，使 BitNet 能够采用更大的学习率。我们进一步将 SubLN 与 Pre-LN 和 BMT 架构 [ZGC⁺23] 进行对比：Pre-LN 是 GPT 预训练的默认架构，而 BMT 已被证明能提升二值化模型的稳定性。实验结果表明，SubLN 的性能优于 Pre-LN 和 BMT。因此，我们在 BitNet 中采用绝对最大值量化和 SubLN 的组合方案。
表 4：BitNet 的消融实验结果（WGe：Winogrande 数据集，WG：Winograd 数据集，SC：Storycloze 数据集，HS：Hellaswag 数据集）。Elastic 是来自 [LOP⁺22] 的一种激活量化方法，而 BMT 是来自 [ZGC⁺23] 的一种架构，用于稳定低比特模型的训练。
6 结论与未来工作
我们提出了 BitNet，一种面向大型语言模型的新型 1 位 Transformer 架构。该方法设计兼具可扩展性和稳定性，能够高效处理大型语言模型。实验结果表明，与基线模型相比，BitNet 在困惑度和下游任务性能方面均取得了具有竞争力的表现，同时大幅降低了内存占用和能耗。此外，BitNet 遵循与全精度 Transformer 相似的缩放定律，这表明它可以有效扩展到更大规模的语言模型，并在性能和效率方面带来潜在优势。未来，我们计划扩大 BitNet 的模型规模和训练步数，并探索将 BitNet 应用于其他架构（如 RetNet [SDH⁺23]）以训练大型语言模型。
附录 A 超参数
表 5：缩放实验中 BitNet 的模型配置。
表 6：缩放实验中 BitNet 和 FP16 Transformer 的超参数。对于 13B 和 30B 模型，为保证训练稳定性，权重衰减设为 0.05。
表 7：BitNet 和 FP16 Transformer 稳定性测试的超参数。
表 8：BitNet 消融实验的超参数。