迈向极低位大语言模型——1bit量化

摘要

当量化技术中的位宽（用于表示每种权重的位数）急剧减少时，通常会导致模型性能的大幅下降。

为了缓解这种问题，作者提出了一种开创性的方法—— OneBit，该方法将 LLM 权重矩阵量化为 1 位值。这代表了一种激进的量化形式，旨在尽可能降低位宽。

OneBit 包含：

• 量化感知训练 (QAT) 框架，该框架包括一种以 1 位形式表示参数的独特方法，可增强 LLM 的量化能力。
• 采用基于矩阵分解的参数初始化策略，加速了 QAT 过程的收敛。

For this target, we introduce a 1-bit quantization-aware training (QAT) framework named OneBit, including a novel 1-bit parameter representation method to better quantize LLMs as well as an effective parameter initialization method based on matrix decomposition to improve the convergence speed of the QAT framework.

问题和挑战

1. 但是，基于 Transformer 的 LLM 的部署充满了困难，这主要是由于它们的高计算需求和大量的内存占用量。随着 LLM 参数数量的增加，操作它们所需的资源也会相应增加。

   Despite their popularity, deploying transformer-based LLMs presents significant challenges due to their computational intensity and considerable memory requirements as the parameters of LLMs become more and more.

2. 许多人（Dettmers 等人，2022 年；Frantar 等人，2022 年；Frantar 和 Alistarh，2023 年）致力于减少 LLM 的计算和内存开销，同时仍保留其大部分原始模型功能。在这些努力中，量化技术得到了广泛关注，尤其是训练后量化（Post-Training PTQ），因为它的迁移(transfering)成本更低。GPTQ（Frantar 等人，2022 年）、SpQR（Dettmers 等人，2023 年b）和 AWQ（Lin 等人，2023 年）等开创性研究成功地将 LLM 的权重矩阵压缩为 4 bit值，同时保持了 LLM 的主要能力。高效量化实现了时间和空间效率以及模型性能之间的平衡，是 LLM 优化的重大进展。

   Among these efforts, quantization has gained widespread attention, particularly Post-Training Quantization (PTQ), benefitted from its lower transferring costs. Seminal studies such as GPTQ (Frantar et al., 2022), SpQR (Dettmers et al., 2023b), and AWQ (Lin et al., 2023) successfully compress the weight matrices of LLMs to 4-bit values while maintaining the main abilities of LLMs. Efficient quantization represents significant advances in LLM optimization, by achieving a balance between time and space efficiency as well as model performance.

3. 尽管 PTQ 可以有效地将权重矩阵压缩为至少 3 位表示形式，但随着位宽的进一步降低，其有效性会大大降低。如下图所示：

   Unfortunately, the efficacy of PTQ rapidly diminishes when the quantization bit-width is extremely low, as shown in Figure 1.

   

4. 最近的研究希望利用量化感知训练（QuantizationAware Training，QAT）来克服 PTQ 面临的瓶颈。LLM-QAT （Liu 等人，2023 年）在量化过程中引入了一些可学习的参数，取得了显著的效果。OmniQuant （Shao 等人，2023 年）整合了可学习的等价变换，在 2 bit量化方面取得了可喜的成果。然而，现有的方法在将模型权重压缩到 1 位时效果下降，难以保持有效性。

   主要问题是在如此低的位宽下会出现严重的精度损失。这种精度的损失极大地影响了线性投影运算（表示为 WX，其中 W 是权重矩阵，X 是输入），这是 LLM 的基本组成部分。精度的急剧降低导致在此操作期间丢失的信息增加，从而影响模型的整体性能。

   Recent researches hope to leverage Quantization Aware Training (QAT) to overcome the bottlenecks faced by PTQ. LLM-QAT (Liu et al., 2023) introduces a few learnable parameters into the quantization process, achieving notable results. OmniQuant (Shao et al., 2023), integrating learnable equivalent transformation, presents promising results in 2-bit quantization. However, existing methods decline when compressing model weights to 1 bit, struggling to maintain effectiveness.

方法介绍

1. 作者提出了一种新颖的线性层与符号值无关分解法（Sign-Value-Independent Decomoisition, SVID），以使用约 1 bit值来表示 LLM。
   • 每个原始高位(high-bit)权重矩阵被分解为一个符号矩阵（±1）和两个值向量(value vectors)。
     • 值向量以较低的成本提供了线性投影所需的浮点精度，并有助于轻松训练模型。
     • 符号矩阵以较小的空间成本保持了原始权重矩阵的高秩，从而保留了较高的信息容量。
   • SVID 为 1 bit模型提供了更好的参数初始化，作者采用量化感知知识蒸馏法(quantization-aware knowledge distillation)，将原始模型的能力迁移到所提出的 1 bit对应模型中。
2. 作者提出用 SVID 将高位矩阵分解为低位矩阵，这对作者的 1 位架构的初始化至关重要。实验证明，基于 SVID 的初始化可以提高模型性能和收敛速度。
3. 广泛的实验证明，作者的方法在 OPT、LLaMA 和 LLaMA2 中的模型规模从 1.3B 到 13B 不等时都能很好地发挥作用，这也展示了它的通用性。

方法实现

背景

模型量化的主要思路是将 FP32 或 FP16 格式模型中的每个权重矩阵 \boldsymbol{W} 压缩成低比特对应格式。具体来说，作者通常将transformer中线性层的权重矩阵量化为 8 位、4 位甚至 2 位。

大多数量化研究主要采用四舍五入法（Round-To-Nearest, RTN），即将权重 w 四舍五入到量化网格中最接近的值。它可以表述为：
$$
\hat{w} = \text{Clip}(\lfloor \frac{w}{s} \rceil + z,\ 0,\ 2^N -1) \tag{1}
$$
其中：

• s 表示量化缩放参数，
• z 表示零点参数，
• N 表示量化位宽。
• $\text{Clip}(\cdot)$将结果截断在 0 到 2^N - 1 的范围内。

随着位宽越来越小，量化网格也越来越稀疏。当作者将 LLM 量化为 1bit 值时，量化模型中只有 2 个可用的数字可供选择。

现有研究（Dettmers 和 Zettlemoyer，2023 年）指出，基于 RTN 方法的量化可以在 4 位水平上获得最佳性能。如之前的图示所示，进一步压缩（如将模型量化为 2 位值）会导致性能大幅下降（Shao 等人，2023 年）。

此外，当 N 等于 1 时，基于 RTN 方法的量化本质上等同于设置一个阈值，阈值两侧的权重 w 被转换为相应的整数值 $\hat{w}$ 。在这种情况下，公式 (1) 中的参数 s 和 z 实际上就失去了实际意义。因此，当将权重量化为 1 bit时，元素向 RTN 运算会大大降低权重矩阵 W 的精度，从而导致量化模型的性能不佳。

当 N 等于 1 时，基于舍入到最近 (RTN) 方法的量化过程简化为每个权重的二进制决策，将其分配给 0 或 1，在考虑符号后相当于 -1 或 +1。这是因为 1 位量化值只有两种可能的状态。在这种情况下，方程 (1) 中的比例参数和零点参数 (z) 变得多余，因为量化网格是最稀疏的，不会在两种状态之间提供渐变。参数 s 和 z 旨在调整量化范围和中心，但由于只有两种可能的结果，它们在微调量化过程中的作用被否定了，因此它们实际上毫无意义！

1 bit 线性架构

由于 1 bit权重量化的精度损失严重，将线性层中的权重矩阵直接从 FP32/16 转换为基于 RTN 的 1 bit格式具有挑战性。Wang 等人（2023 年）通过研究纯 1 bit权重矩阵的能力，从头开始训练 1 bit模型，探索了这种可能性。在 W1A16 设置中，他们的线性层设计为：
$$
\begin{align} W_{\pm} = \text{Sign}[\boldsymbol{W} - \text{Mean}(\boldsymbol{W})]\newline
\ \eta = \text{Mean}[\text{Abs}(\boldsymbol{W} - \text{Mean}(\boldsymbol{W}))]\newline
\ \boldsymbol{Y} = \eta \cdot \text{LayerNorm}(\boldsymbol{X}) \boldsymbol{W}_{\pm}^T \end{align} \tag{2}
$$
其中：

• $\boldsymbol{W}$ 表示 m×n 形的量化权重矩阵
• $W_{\pm 1}$ 表示 1 bit量化矩阵
• $\boldsymbol{X}$ 是线性层的输入
• $\boldsymbol{Y}$ 是输出
• $\text{Sign}(\cdot) $、 $\text{Mean}(\cdot)$ 和 $\text{Abs}(\cdot)$ 函数返回符号矩阵、平均值和绝对值矩阵。

遗憾的是，这种方法降低了计算需求，但也导致性能明显下降（Wang 等人，2023 年）。

受 Wang 等人（2023）的启发，作者还使用函数 $\text{Sign}(\cdot)$ 对权重矩阵进行量化，量化后的矩阵元素被设置为 +1 或 -1。

此外，作者还注意到，虽然 $\boldsymbol{W}_{\pm 1}$ 保持了 $\boldsymbol{W}$ 的高秩，但遗漏的浮点精度仍然会破坏模型性能。

因此，与之前的研究不同，作者引入了两个 FP16 格式的值向量，以减少量化过程中的精度损失。作者提出的线性层设计如下：
$$
\begin{align} W_{\pm 1} = \text{Sign}(\boldsymbol{W})\newline
\ \boldsymbol{Y} = [(\boldsymbol{X}\odot \boldsymbol{g}) \boldsymbol{W}_{\pm 1}^T]\odot \boldsymbol{h}\newline
\ \boldsymbol{Z} = \text{LayerNorm}(\boldsymbol{Y}) \end{align} \tag{3}
$$
其中：

• g 和 h 是两个 FP16 值向量。

注意，作者在公式 (3) 中使用括号指定了计算顺序，以尽量减少时间和空间成本。

Wang 等人（2023 年）与 OneBit 的主要区别在于额外的参数 g 和 h 。例如，当作者量化一个形状为 4096×4096 的权重矩阵时，量化结果的平均位宽为 1.0073。

在作者提出的 1 bit架构中，权重矩阵 $\boldsymbol{W}$ 在数学上分为两个部分：

• 一个是 INT1 格式的符号矩阵 $W_{\pm 1}$ ，
• 一个是 FP16 格式的值向量 $\boldsymbol{g/h}$ 。

为了借助完全训练过的权重初始化 1 位模型，作者引入了权重矩阵 W 的符号-值无关分解（SVID），可表述为：
$$
\boldsymbol{W}=\boldsymbol{W}_{sign} \odot \boldsymbol{W}_{value}
$$
其中：

• $\boldsymbol{W}_{value} = |\boldsymbol{W}|$
• $\boldsymbol{W}_{sign} = \text{Sign}(\boldsymbol{W})$

对于 $\boldsymbol{W}_{value}$ ，作者进一步将其近似分解为两个向量 a 和 b 的外积，这也被称为 rank-1 近似。

因此，作者提出的矩阵分解方法可以表示为：
$$
\boldsymbol{W} \approx \boldsymbol{W}_{sign} \cdot (\boldsymbol{ab}^T) \tag{4}
$$

作者可以采用一些广泛使用的矩阵分解方法来进行 rank-1 近似，如 SVD（Beltrami，1990 年）和 NMF（Paatero 和 Tapper，1994 年）。

命题 1

给定权重矩阵 $\boldsymbol{W}$ 和输入 $\boldsymbol{X}$ 后，线性层可根据 SVID 重构如下：
$$
\boldsymbol{XW}^T \approx [(\boldsymbol{X} \odot \boldsymbol{b}^T)\boldsymbol{W}_{sign}^{T}]\odot \boldsymbol{a}^T \tag{5}
$$
这弥补了量化模型的结构与原始权重之间的差距。

它表明，如果作者将$\boldsymbol{W}_{sign}$赋值为$\boldsymbol{W}_{\pm 1}$，将$\boldsymbol{a}^T$赋值为 h ，将$\boldsymbol{b}^T$赋值为 g ，那么量化模型就是原始模型的近似初始化。

此外，与先还原原始矩阵$\boldsymbol{W}$相比（如公式 (4)），公式 (5) 中的计算顺序由于无需还原 FP16W，因此可以在内存中节省大约一个 FP16 格式的矩阵$\boldsymbol{W}$。

SVID 的主要目的是让符号矩阵参与矩阵的近似，而不是仅仅依赖 FP16 格式的值向量。为了证明符号矩阵在矩阵逼近中的作用，作者提出了以下命题。

命题2

给定矩阵$\boldsymbol{W}$和$|\boldsymbol{W}|$， $\boldsymbol{W}=\boldsymbol{W}_{sign}\cdot |\boldsymbol{W}|$。

作者按以下方式分解这些矩阵：
$$
\boldsymbol{W}= \boldsymbol{ab}^T+\boldsymbol{E}_1 和 |\boldsymbol{W}| = \boldsymbol{\widetilde{a}\widetilde{b}}^T+\boldsymbol{E}_2
$$
其中：

• $\boldsymbol{E}_i$表示误差矩阵。

从 Frobenius-norm 的角度来看，SVID 更接近原始矩阵 $\boldsymbol{W}$ ：
$$
\lVert \boldsymbol{W} - \boldsymbol{W}_{sign} \odot \boldsymbol{\widetilde{a}\widetilde{b}}^T \rVert_{F}^2 \le \lVert \boldsymbol{W} - \boldsymbol{ab}^T \rVert_F^2 \tag{6}
$$
它清楚地表明了符号矩阵$\boldsymbol{W}_{sign}$在矩阵逼近中的实际作用。

需要注意的是，由于大多数参数的精度都很低，因此要精确地近似权重矩阵$\boldsymbol{W}$是相当具有挑战性的。SVID 的目的不是精确复制原始模型的参数，而是利用原始模型的大量训练，为进一步训练提供一个有效的起点。

知识蒸馏

作者采用量化感知知识蒸馏(quantization-aware knowledge distillation)技术，将知识从原始模型（即教师模型）迁移到量化模型（即学生模型）。

在学生模型中，矩阵 $\boldsymbol{W}$ 中的元素和公式 (3) 中的向量$\boldsymbol{g/h}$将被训练。

作者使用基于交叉熵的对数和基于隐藏状态的均方误差全精度教师模型来指导量化的学生模型（Sun et al.） 不使用语言建模损失。交叉熵定义为：
$$
\mathcal{L}_{CE}=-\frac{1}{n_s}\sum_{i=1}^{n_s}{\sum_c{P_{c}^{\mathcal{T}}\left( \boldsymbol{o}_i \right) \log P_{c}^{\mathcal{S}}\left( \boldsymbol{o}_i \right)}} \tag{7}
$$
其中：

• c 表示类别数
• n_s 表示当前批次的训练样本数
• $\mathcal{T}$和$\mathcal{S}$分别是教师模型和学生模型

隐藏状态的误差定义为：
$$
\mathcal{L}_{MSE}=\sum_{i=1}^{n_s}{\sum_{j=1}^{n_l}{\lVert \frac{\boldsymbol{q}_{i,j}^{\mathcal{T}}}{\lVert \boldsymbol{q}_{i,j}^{\mathcal{T}} \rVert _2}-\frac{\boldsymbol{q}_{i,j}^{\mathcal{S}}}{\lVert \boldsymbol{q}_{i,j}^{\mathcal{S}} \rVert} \rVert}_{2}^{2}} \tag{8}
$$
其中：

• n_l 表示层数
• $\boldsymbol{q}$表示隐藏状态。

因此最终的目标函数可以表示为：
$$
\mathcal{L}_{KD} = \mathcal{L}_{CE} + \alpha \mathcal{L}_{MSE} \tag{9}
$$
其中：

• $ \alpha $是平衡交叉熵损失和中间层特征重要性的超参数。

主要结果

• 在各种模型大小下，作者的 1 bit权重量化方法在 W2A16 设置下明显优于其他方法。

  此外，随着模型规模的增大，基于 QAT 的方法的有效性持续提高，而 PTQ 方法、 GPTQ 的结果可能会随着模型规模的增大而下降（例如，在 LLaMA 上从 7B 到 13B）。这表明，基于 QAT 的方法可以在极低位量化时获得稳定的结果。

• 具体来说，随着模型大小的增加，作者的方法更接近 FP16 的性能。

  例如，当从 LLaMA-7B 扩展到 LLaMA-13B 时，FP16 模型的困惑度仅降低了 0.59，而作者的方法则降低了 1.20。

• 在困惑度方面，只有作者的方法取得了与最强 FP16 基准线相当的结果。

  例如，作者的方法在 LLaMA-13B 模型的 Wiki2 数据集上取得了 9.18 的成绩，而 FP16 基线为 5.09。其他方法的性能损失很大，尽管它们使用的是 2 bit量化，比作者的 1 bit量化要多。对于 GPTQ 和 LLM-QAT，量化后的性能下降相当严重。至于 OmniQuant，尽管在 W2A16 设置下它是最强的基线，但与作者的 W1A16 设置相比，它仍然遭受了更大的性能损失。

• 在零样本精度方面，虽然所有方法都不可避免地会出现一些性能下降，但在大多数模型中，作者的方法取得了最接近 FP16 基准线的性能。

  在 OPT-1.3B/2.7B 模型上，作者的方法在 PIQA 和 ARC-e 等大多数任务上的性能损失较小。此外，与排名第二的基线 OmniQuant 相比，其他任务的性能损失可以忽略不计。在 LLaMA-7B 模型上，除 ARC-e/ARCc 外，作者的方法在所有任务上的表现都明显优于 OmniQuant，平均整体提高约 4%。

在实际问题中的表现

重点专注：常识推理(commonsense reasoning)和世界知识(world knowledge)。

• 常识推理：虽然其他模型比作者的模型参数更多，训练更全面，但作者的模型在常识推理方面仍有优势。这反映了从更大的 7B 模型中继承的优势。
• 世界知识：在世界知识方面，尽管作者的模型在社会科学领域损失较大，但在其他领域却优于经过全面训练的 Pythia-1B。这些结果证明了 OneBit 的实用性。

分析和讨论

效率

显然，将模型权重量化为极低位可以大幅度减少其内存占用。实际上，如表3所示，随着模型规模的增大，压缩比也在不断提高。这一点对于大型模型来说尤为重要，因为它使得将这类模型部署到单个GPU上成为可能，从而极大地提高了模型的实际应用性和灵活性。

虽然该方法在实施过程中可能会导致一定的性能损失，但如图4所示，它在空间占用和模型性能之间达到了良好的平衡。

比如只需使用原始模型空间的0.2倍，即可实现与FP16相当的性能表现。

此外，通过将权重量化为±1，能够进一步加速CPU上的矩阵乘法操作。

这是因为在这种量化方式下，原本需要进行的浮点乘法可以转换为速度更快的位操作，显著提升计算效率。因此，该方法的内存占用显著减少，更符合个人电脑和智能手机等设备的部署需求。

鲁棒性

极低比特量化使训练过程对学习率高度敏感，当学习率过小或过大时，使模型难以收敛。这主要是由于权重元素在+1和-1之间波动时产生的梯度幅度较大，导致Linear层输出显著波动。

下图显示了另一种1位量化架构BitNet在使用不同学习率训练时的表现，可以看到非常不稳定：

而OneBit通过引入两个FP16格式的值向量来补充必要的浮点数值精度，同时限制了量化后矩阵乘法结果的波动范围，从而提高了模型的鲁棒性，展现出更稳定的训练过程，并且对学习率不敏感。

不同组件的影响

OneBit主要的可变组件包括Post-LayerNorm、参数初始化等。

• Post-LayerNorm：在量化感知训练期间，模型存在浮点数溢出的问题，尤其在模型深度增加时，激活值容易急剧增长。为应对此问题，本文选择了Post-LayerNorm策略，而非Pre-LayerNorm。因为在实际应用中，Pre-LayerNorm有时可能无法有效地控制激活值的增长。
• 参数初始化：在使用SVID做参数初始化时，非负矩阵分解和奇异值分解均可用于分解|W|，但推荐使用奇异值分解。它能加快训练收敛速度，如下图所示，NMF初始化有助于获得更好的性能。

总结

OneBit框架提出了一种大胆的1位量化方法，极大地减少了模型部署时的存储和计算开销。尽管存在一些性能损失，但其在模型压缩和部署效率方面的潜力是显而易见的，为研究极低位量化模型提供了一条新的探索途径。

限制

尽管作者提出的方法大大减少了 LLM 的内存占用，为高效部署 LLM 带来了希望，但仍然存在一些局限性。

1. 首先，与原始模型相比，作者的超低位量化不可避免地会造成性能损失。
2. 此外，作者还不了解 1 位量化模型最佳参数背后的数学原理，因此只能通过代价高昂的 KD 过程来实现能力转移。幸运的是，这种代价是一次性的。
3. 此外，由于 1 位量化的特殊性，作者的方法无法自然扩展到更高的位宽。
4. 最后，作者没有考虑激活量化的问题，将其作为未来的工作。