Extreme Compression of Large Language Models via Additive Quantization阅读

Abstract

随着高精度的大型语言模型(LLMs)的出现,促使了对高性能量化技术的竞赛,以便在终端用户设备上执行这些模型。本文重新审视了“极端”LLM压缩的问题——定义为目标极低比特数,如每个参数2到3比特——从多码本量化(MCQ)经典方法的角度出发。我们的算法AQLM,通过两项创新,推广了信息检索中的经典加法量化(AQ)方法,以提升LLM压缩的最新技术水平:1)学习输入自适应的权重矩阵加法量化,2)跨每个Transformer块的码本参数联合优化。总体而言,AQLM是首个在压缩至每个参数少于3比特时,在准确性与模型大小方面帕累托最优的方案,并在极端压缩(2比特)情况下显著超越了所有已知方案。此外,AQLM很实用:我们为AQLM的令牌生成提供了快速的GPU和CPU实现,使我们能够在速度上匹配或超越优化的FP16实现,同时在更小的内存占用下运行。

什么是经典加法量化?
经典加法量化通过多个码本向量的和来表示数据点,虽然可能需要存储多个码本,但这些码本通常是共享的,并且每个数据点仅需存储引用码本的索引。因此,整体上可以显著减少内存使用。通过优化码本的设计和索引方式,经典加法量化在压缩率和精度之间能取得良好平衡。


什么是帕累托最优?
帕累托最优本身没有具体的计算公式,因为它是一个概念状态。然而,在多目标优化中,通常使用以下方法来寻找帕累托最优解:
1. **目标函数集合**:定义多个需要优化的目标函数。
2. **约束条件**:确定优化问题的限制条件。
3. **支配关系**:对于两个解 \( A \)\( B \),如果 \( A \) 至少在一个目标上优于 \( B \),且在其他目标上不逊于 \( B \),则称 \( A \) 支配 \( B \)4. **帕累托前沿**:通过比较支配关系,识别出无法被其他解支配的一组解。

1. Introduction

生成型大型语言模型(LLMs)的快速发展,引发了广泛的工业和大众兴趣,部分原因是高精度开放LLMs的可用性,如LLAMA 1和2、Falcon、BLOOM、OPT和NeoX/Pythia等。开放模型的一个关键优势是,在假设计算和内存成本可以降低到在普通硬件上可管理的情况下,终端用户可以本地进行推理或微调。这导致了几种用于压缩LLMs的推理和微调的方法。目前,LLMs准确后训练压缩的主要方法是量化,它通过减少存储模型权重(和可能的激活)的比特宽度来改进模型的占用空间和内存传输。大体上,LLM权重通过“直接”量化进行压缩,也就是为每个矩阵子组件选择合适的量化网格和归一化,然后通过直接舍入或更复杂的分配将权重映射到网格上。量化引发了自然的压缩与准确性折衷,通常通过模型大小与模型困惑度(PPL)来衡量。 现有方法可以在3-4比特每元素时达到较低的准确性损失,甚至可以稳定地将超大型模型压缩至2比特或更少。然而,在大多数情况下,低比特数会导致显著的准确性下降、更高的实现复杂性和运行时开销。从实际角度来看,使用当前技术进行2比特范围的“极端”量化不如简单地使用更小的基础模型并将其量化到更高的比特宽度,如每个参数3-4比特,因为后者在给定相同字节的模型大小时能提供更高的准确性。
贡献:在这项工作中,我们通过首次展示多码本量化(MCQ)技术可以扩展到LLM权重压缩,提升了LLM压缩的最新技术水平。总体而言,MCQ是一组信息检索方法,包括专门的量化算法来压缩向量数据库,以实现高效搜索。与直接量化不同,MCQ通过利用量化值的互信息来共同压缩多个值。

具体来说,我们将加法量化(AQ)扩展到压缩LLM权重的任务,使每层和Transformer块的输出大致保留。我们的扩展重新制定了经典的AQ优化问题,以减少LLM层输出在输入标记分布下的误差,并联合优化层块上的码,而不仅仅是保留标准AQ中的权重。我们将这种结果过程称为语言模型的加法量化(AQLM)。与一些需要混合稀疏-量化格式的极端LLM量化方法不同,AQLM以简单的同质格式量化模型,易于在实践中支持。我们的主要贡献如下:

什么是经典的AQ量化?
经典的AQ优化(AQ Optimization)通常指的是一种归纳学习方法,主要用于生成规则和决策树。这种方法起源于机器学习领域,尤其是规则学习和数据挖掘。AQ算法通过不断调整规则来优化分类精度。以下是一些关键特性:
1. **规则生成**:AQ算法从训练数据中生成一组规则,用于分类任务。每个规则通常包括一个条件(前提)和一个结论(分类结果)。
2. **覆盖性**:算法尝试生成覆盖尽可能多正例(正向样本)且不包含负例(负向样本)的规则。
3. **迭代优化**:通过迭代过程,逐步改进规则的准确性和覆盖范围,以提高整体分类性能。
4. **应用领域**:广泛应用于分类问题,如医学诊断、市场营销和风险评估等领域。
  1. 我们提出了AQLM算法,将AQ扩展到LLM权重的后训练压缩,通过两项创新:(1)将AQ背后的MAP-MRF优化问题调整为实例感知,考虑层校准输入和输出激活;(2)结合高效的层内调优技术,利用校准数据联合优化多个层的量化参数。
MAP-MRF优化问题:

MAP(最大后验概率):这是一种统计方法,用于在给定数据下找到最可能的参数设置。
MRF(马尔可夫随机场):这是一个用于建模上下文依赖关系的数学模型。
在这里,MAP-MRF被用于调整量化过程,使其更加“实例感知”,也就是说,更好地适应每一层的输入和输出特性。
  1. 我们评估了该算法在压缩LLAMA 2系列的高精度开放LLM时的有效性,压缩率为每个参数2-4比特。我们发现AQLM在标准2-4比特压缩范围内优于之前的最新技术,特别是在极端2比特量化中有显著提升。我们提供了详尽的各种算法参数影响的消融实验,如码宽和码本数量,并将我们的分析扩展到最近的Mixtral模型。我们还与后续工作的改进微调算法一起评估了AQLM,进一步提高了2和3比特模型的准确性。

  2. 我们展示了AQLM的实用性,通过为特定编码提供高效的GPU和CPU内核实现,以及端到端生成。结果表明,我们的方法在减少内存占用最多8倍的同时,可以在速度上匹配或甚至超越浮点基线。具体而言,AQLM可以在GPU上实现约30%的分层加速,在CPU推理中实现高达4倍的加速。

2. Background & Related Work

2.1. LLM量化

早期的后训练量化(PTQ)方法(如ZeroQuant、LLM.int8()和nuQmm)通过直接的四舍五入(RTN)投影,并调整量化粒度来平衡内存效率和准确性。GPTQ提出了一种更精确的数据感知方法,通过大型求解器来最小化逐层的ℓ2误差。Dettmers和Zettlemoyer研究了这些早期方法的精度压缩权衡,建议4比特量化可能是RTN量化的最佳选择,并观察到像GPTQ这样的数据感知方法允许更高的压缩,即严格低于每权重4比特,同时保持帕累托最优性。我们的工作首次将这一帕累托前沿推进到每权重低于3比特。

并行工作中,Dettmers等人研究了同时量化权重和激活到8比特,指出大型LLM中的“异常值特征”会导致大量错误,从而引发了各种缓解策略。最近,一些改进技术专注于量化高影响输出误差的权重异常值的困难。SpQR通过将异常值保存为高度稀疏的高精度矩阵来解决这个问题。AWQ通过每通道缩放来减少重要权重的误差。SqueezeLLM利用对角费舍尔作为Hessian的代理,并通过K-means聚类实现非均匀量化。发表的最新方法是QuIP,同时我们的工作中推出了改进版本QuIP。

大致而言,它们首先通过旋转矩阵“平滑”权重,然后将其映射到晶格上。QuIP和QuIP#旨在最小化给定初始权重和校准数据情况下每层的“最坏情况”误差。例如,在QuIP#中,旋转后的权重分布近似为高斯分布,而编码晶格(E8P)被选择以最小化“舍入”误差。相比之下,我们的方法使用不同的权重编码(码本是加法的),并学习码本而不是固定码本。我们的见解是,通过直接优化校准集上的码本,可以获得更高的准确性,而不需要旋转。此外,我们展示了不同层的码本可以通过联合微调校准数据进行共同训练。

码本是怎么学习的?

2.2. 最近邻搜索的量化

我们的工作建立在近似最近邻搜索(ANN)算法的基础上。与PTQ不同,ANN量化旨在压缩向量数据库,使用户能够高效计算相似性并找到与查询点集的最近邻。为了实现高压缩,现代ANN搜索算法采用矢量量化(VQ)——这是一种联合量化多个向量维度的方法。其通过学习“码本”,即一组可学习的候选向量来编码数据。为了对给定的数据库向量进行编码,VQ将其分割为子组,然后通过从学习到的码本中选择向量来编码每个组。该算法通过利用点积的线性性来高效计算相似性搜索的距离或点积。用于ANN搜索的量化方法推广了矢量量化,被称为多码本量化(MCQ)。MCQ方法通常在查询端不涉及信息损失,因此是内存高效ANN的领先方法。我们简要回顾MCQ。

近似最近邻搜索(ANN)和K-means是两种不同的算法,用于解决不同的问题。

### 近似最近邻搜索(ANN)
- **用途**:用于快速查找高维空间中与查询点最近的点。
- **特点**:重点在于效率和速度,而非精确性,适用于大规模数据集。
- **应用**:推荐系统、图像检索、自然语言处理等。
#### 举例说明
假设有一个包含数百万个高维向量的数据集,ANN可以快速找到与给定查询向量最接近的几个向量,而不需要遍历所有数据,节省计算时间。

### K-means
- **用途**:用于将数据集划分为 \(k\) 个簇,每个簇由一个质心表示,目的是最小化簇内的方差。
- **特点**:是典型的聚类算法,主要用于数据分组。
- **应用**:市场细分、图像压缩、模式识别等。
#### 举例说明
假设我们有客户的购买数据,K-means可以将客户分成几个组(簇),每个组代表一种消费行为模式。

### 区别总结
- **目的不同**:ANN侧重于查找最近邻,K-means侧重于数据聚类。
- **算法性质**:ANN是搜索算法,K-means是聚类算法。
- **结果输出**:ANN输出与查询点相近的点,K-means输出数据分组和簇中心。

产品量化(PQ)是MCQ的早期版本,将每个向量编码为来自M个码本的M个码字的串联。PQ将向量分解为M个独立的子向量,并对每个子向量应用矢量量化,同时使用单独的码本。因此,每个向量由码字索引的元组编码,并被近似为x ≈ [c1i1,…,cMiM]。快速欧氏距离计算可以通过查找表实现:

[ M ∥q − x∥² ≈ ∥q − [c1i1,...,cMiM]∥² = \sum_{m=1}^M ∥qm − cmim∥², ]

其中qm是查询q的第m个子向量。如果查询子向量到码字的距离预先计算好,则可以通过M次加法和查找计算该和。由于基于产品的近似在D/M维组件独立分布时效果更好,后续工作研究了寻找更好变换的方法。

对于其他相似性函数,(Guo等)提出了一种最大内积搜索(MIPS)的量化程序。该方法通过解决约束优化问题来最小化数据库和查询向量之间内积的量化误差。与上述公式类似,该过程通过预先计算查询q与所有码本中的编码之间的点积,然后合并这些部分点积来恢复完整的相似性分数,以实现高效的内积搜索。

非正交量化的后续工作推广了产品量化的思想,通过将每个向量近似为M个码字的和来代替串联。这种方法仍然高效,同时提高了近似精度。
为此,残差矢量量化首先对原始向量进行量化,然后迭代量化上一次迭代的近似残差。加法量化(AQ)更为通用,因为它对不同码本的码字没有施加约束。通常情况下,AQ提供了最小的压缩误差,但在处理大M时训练更加复杂。

最后,一些最近的工作详细阐述了加法量化的思想,提出了更有效的码本学习程序。复合量化(CQ)学习具有固定内积值的不同码本之间的码本。目前,LSQ方法在压缩精度上达到了最先进的水平。

矢量量化在模型压缩中的应用也有显著进展。例如,Zhou等人利用多码本量化压缩深度学习模型中的词嵌入。另一类工作则探索了矢量量化在线性模型或深度模型中的线性层中的应用。与上面的PQ类似,这些技术预先计算输入与所有码的内积,然后通过查找计算线性层,从而加速推理。然而,这些算法引入了显著的预测误差,不允许它们压缩深度模型。因此,我们相信我们是首次成功地将MCQ适应并扩展到LLM的研究者。

3.AQLM:Additive Quantization for LLMs

3.1. 概述

我们从以下观察开始:加法量化(AQ)解决的问题与后训练量化(PTQ)相关:两种设置假设存在一组“输入”向量,即AQ的输入数据和PTQ的权重矩阵行。目标是在保留点积相似性的同时压缩这些输入,针对查询向量(对于AQ)和层输入嵌入(对于PTQ)。两者的区别在于,AQ假设查询分布未知,而PTQ方法表明优化一组校准数据中的样本输入嵌入就足够。

从高层次上看,我们从解决以下问题开始:对于一个具有(d_{in})输入特征和(d_{out})输出特征的线性层,给定其权重(W \in \mathbb{R}^{d_{out} \times d_{in}})和一组校准输入(X \in \mathbb{R}^{d_{in} \times n}),寻找一种量化权重的配置(\hat{W}),使原始层和压缩层的输出之间的平方误差最小化:

[ \text{argmin}_{W} |WX - \hat{W}X|^2_2. ]
在这里插入图片描述

在下面的内容中,我们假设W使用AQ进行量化,并采用标准符号。AQ将权重行分为(g)个连续元素的组,并将每组权重表示为从多个学习到的码本(C_1, …, C_M)中选择的(M)个向量之和,每个码本包含(2^B)个向量(对于B位码)。一个权重通过从每个码本中选择一个码并将其相加来编码。我们将此选择表示为一个one-hot向量(b_m),从而得到一组的表示:
在这里插入图片描述

在这里插入图片描述

这类似于PTQ算法,只是每组使用更复杂的编码。为了表示完整的权重,我们简单地连接:
[ \sum_{m=1}^{M}C_m b_{ijm}. ]

其中(\oplus)表示连接,(b_{ijm} \in \mathbb{R}{2B})是一个用于第i个输出单元、第j个输入维度组和第m个码本的one-hot码。

我们的算法将学习码本(C_m \in \mathbb{R}^{g \times 2^B})以及由one-hot表示的离散码(b \in \mathbb{R}^{d_{out} \times d_{in}/g \times M \times 2^B})。结果方案使用(M \cdot B)位编码每组(g)权重,并进一步需要(g \cdot 2^B \cdot 16)位用于FP16码本。误差变为:

在这里插入图片描述

为了学习这种权重表示,我们通过运行残差K-means来初始化码本C和码b。具体地,初始化算法首先对权重组进行K-means聚类并保存结果的聚类索引。接下来,通过从每个权重中减去最近的聚类来计算量化误差。最后,算法再次运行K-means聚类,但这次是在量化误差而不是权重上进行。因此,每个后续码本被初始化以补偿来自之前码本的量化误差。初始化后,我们在更新码(b_{i,j,m})和码本(C_m)之间交替,直到损失函数(3)停止在指定的容差内改善。由于码是离散的而码本是连续的,并且我们正在优化多个交互层,我们的方法分为三个阶段,在算法1中描述并在下面详细说明。

3.1.0 补充

这段文字描述了一种用于权重量化的算法,涉及学习码本和离散编码。目的是通过最小化误差来表示权重。以下是对其关键步骤的解释和举例:

步骤说明
  1. 学习码本和编码

    • 码本 (C_m) 大小为 (g \times 2^B)。
    • 编码 (b) 为一位热向量,维度为 (d_{out} \times \frac{d_{in}}{g} \times M \times 2^B)。
  2. 压缩表示

    • 每组 (g) 个权重使用 (M \cdot B) 位表示。
    • 码本存储使用 (g \cdot 2^B \cdot 16) 位(FP16格式)。
  3. 误差最小化

    • 目标是选择码本和编码,使得原始输出与量化输出之间的误差最小化。
  4. K-means初始化

    • 通过残差K-means初始化码本和编码。
    • 初始步骤包括对权重组进行K-means聚类,确定初始码本。
    • 计算量化误差,并对误差进行再次聚类,更新码本。
  5. 优化循环

    • 交替更新编码和码本,直到误差不再显著减少。
举例说明

假设我们需要对一个简单的权重行进行量化:

权重行及初始设置

  • 权重行:[0.6, -0.4, 0.8, 0.2]
  • 每组包含 2 个元素((g=2)),使用 2 个码本((M=2)),每个码本有 2 个向量((B=1))。

初始化步骤

  1. K-means聚类

    • 初始聚类权重组 ([0.6, -0.4]) 和 ([0.8, 0.2]),得到两个初始码本,比如:
      • 码本1:([0.5, -0.3])
      • 码本2:([0.7, 0.1])
  2. 误差计算和修正

    • 计算误差:([0.1, -0.1]) 和 ([0.1, 0.1])
    • 对误差进行K-means聚类,得到更新的码本用于修正。
  3. 迭代优化

    • 更新编码选择和码本向量。
    • 反复进行,直到误差收敛。

通过这个过程,可以有效地用较少的存储表示权重,同时保持模型性能。

3.2. 阶段1:代码的波束搜索

首先,AQLM更新代码(b_{i,j,m})以最小化MSE目标(3)。类似于Babenko和Lempitsky(2014);Martinez等人(2016;2018),我们将目标重新表述为完全连接的离散马尔可夫随机场(MRF),以利用MRF求解器。

什么是离散马尔科夫随机场?

完全连接的离散马尔可夫随机场(MRF)是一种概率模型,用于描述一组随机变量之间的依赖关系。其特点如下:
### **特点**
1. **节点和边**:
   - 每个节点代表一个离散随机变量。
   - 在完全连接的MRF中,每对节点之间都有一条边,表示变量之间可能的相互依赖。
2. **马尔可夫性**:
   - 每个节点的概率只依赖于与其直接相连的节点(邻域)。
   - 即一个节点的状态与其他非邻域节点的状态条件独立。
3. **势函数**:
   - 边和节点关联的函数,用于定义状态组合的优先级。
   - 势函数用于计算联合概率分布。
4. **联合概率分布**:
   - 由势函数决定,通常表示为规范化的乘积形式。

### **应用**
- 图像处理:如去噪、分割。
- 计算机视觉:如目标识别。
- 自然语言处理:如词性标注。

### **示例**
假设有三个变量 \(X, Y, Z\),每个取值为0或1。在完全连接的MRF中:
- 存在边 \(X-Y\), \(Y-Z\), \(X-Z\)。
- 每个节点的状态依赖于另两个节点。
通过势函数定义每种状态组合的权重,计算出联合概率分布。常用的方法包括最大似然估计和吉布斯抽样来推断和学习参数。
这种模型强大在于能够处理复杂的依赖关系,但计算代价较高,尤其是在变量数量很大时。

为了简化推导,我们首先考虑一个特殊情况,即单个输出单元((d_{out}=1))和单个量化组(即(g=d_{in})),以去除连接运算符:

[ ||WX - \sum_{m=1}{M}C_mb_mX||2_2 ]

我们通过展开平方差来重写这个目标:
[ ||WX - \sum_{m=1}{M}C_mb_mX||2_2 = ||WX||^2_2 - 2 \langle WX, \sum_{m=1}^{M}C_mb_mX \rangle_F + ||\sum_{m=1}{M}C_mb_mX||2_2 ] (4)

上式中,(\langle ·,· \rangle_F)表示两个矩阵的Frobenius内积。接下来,让我们分别考虑方程(4)的三个部分。首先注意到,(||WX||^2_2)在b中是常数,可以忽略。第三部分可以进一步展开为成对的点积:
在这里插入图片描述


请注意,第二部分和第三部分都依赖于类似(C_mb_mX)的矩阵的Frobenius积。这些矩阵在实践中可能不方便:因为(X \in \mathbb{R}^{d_{in} \times n})
,每个矩阵的大小将随着校准数据集n的大小而变化。为了解决这个问题,我们将积重写为:

[ \langle C_ib_iX, C_jb_jX \rangle_F = \langle C_ib_iXX^T, C_jb_j \rangle_F. ] (6)

在这里插入图片描述

因此,可以预先计算(XX^T \in \mathbb{R}^{d_{in} \times d_{in}})
。在以后的推导中,我们将这种类型的积表示为(\langle A,B \rangle_{XX^T} \equiv \langle AXX^T, B \rangle_F)
。然后,方程(4)变为:

[ ||WX - \sum_{m=1}{M}C_mb_mX||2_2 = ||WX||^2_2 - 2 \sum_{m=1}^{M} \langle W, C_mb_m \rangle_{XX^T} + \sum_{i=1}^{M} \sum_{j=1}^{M} \langle C_ib_i, C_jb_j \rangle_{XX^T}. ] (7)

最后,我们将该方程推广到多个输出单元((d_{out} > 1))和量化组((g=d_{in}))。对于(d_{out} > 1),请注意,原始目标(3)对输出单元是可加的:因此,我们可以将(7)独立应用到每个输出维度并汇总结果。为了支持多个输入组((g=d_{in})),我们可以将每个组视为一个单独的码本,其中只有活动组的代码为非零。因此,我们需要重复每个码本(d_{in}/g)次,并根据活动组用零填充。
现在很明显,最小化(4)等同于在马尔可夫随机场中进行MAP推断,在这里插入图片描述
为一元势,(\langle C_ib_i, C_jb_j \rangle_{XX^T})
为成对势。虽然找到精确的最优解是不可行的,但先前的工作表明,这种类型的MRF可以通过波束搜索或ICM(Besag, 1986)来近似求解。

为了解决这个问题,我们选择改编Babenko和Lempitsky(2014)的波束搜索算法。该算法保持k(波束大小)个最佳代码配置,从前一个解开始。在每一步中,算法尝试通过尝试所有(2^Bk) 个替代方案并根据MSE(7)选择最佳k个来替换一个代码。由于损失函数是可加的,改变一个代码只会影响一小部分损失组件。因此,我们可以通过从先前的损失函数(在代码替换之前)开始,然后添加和减去在此迭代中更改的组件来有效地计算损失函数。这些少量的损失组件可以通过在波束搜索前与(XX^T)相乘来有效地计算。波束搜索在所有(d_{out})输出单元上并行运行。这是可能的,因为编码一个输出单元不会影响其他单元的目标(7)。请注意,波束搜索不一定是解决此问题的最佳方案。用于检索的AQ变体(Martinez等人,2016;2018)使用随机ICM更快地找到解决方案。在本研究中,我们选择波束搜索是因为它更容易在PyTorch/JAX等ML框架中实现。

3.3. 阶段2:码本更新

在第二阶段,我们寻找优化的码本向量 (C_1,…,C_M),以最小化与波束搜索相同的平方误差。如果将代码 (b) 视为常量,最小化(3)就成为一个关于 (C_m) 的最小二乘问题。原始AQ算法通过闭式形式解决该问题,依赖于每个向量维度可以独立优化这一事实。由于存在 (XX^T),问题变得复杂:一个码本坐标的最优值依赖于所有其他坐标的值。原则上,我们可以以闭式形式优化 (C_m),但这需要求解一个大矩阵,或使用专门的迭代最小二乘求解器(如共轭梯度法)。

为了简化,我们当前的实现默认使用Adam(Kingma & Ba, 2015)来近似解决这个最小化问题。实际上,这个码本调优阶段只占总计算时间的一小部分。我们的目标计算如下:

[ ||WX - \hat{W}X||^2_2 = ||(W - \hat{W})X||^2_2 = \langle (W - \hat{W})XX^T, (W - \hat{W}) \rangle_F, ]

其中 (\hat{W}) 是来自步骤2的量化权重矩阵,(XX^T) 矩阵是预先计算的。我们通过迭代(非随机的)全批量梯度下降来优化这一目标。在每个更新阶段,我们的实现运行100次Adam的迭代,学习率为 (10^{-4})。然而,我们发现最终结果对这些参数不敏感:使用较少的步骤或较小的学习率进行训练可以实现相同的损失,但需要更长时间才能收敛。在未来的工作中,可以通过使用专用的码本最小二乘求解器来消除这些超参数。与其他算法类似,我们还学习每个单元的比例 (s \in \mathbb{R}^{d_{out}}),初始化为 (s_i := ||W_i||_2) 并通过相同的优化器与码本一起更新(算法1中的第10行)。

3.4. 阶段3:层内一致性微调

到目前为止,我们的算法是独立于模型的其他部分压缩每个权重矩阵。然而,在实际中,量化误差在不同矩阵间有不同的交互。这一问题在极端(2位)压缩的情况下尤其明显,因为量化误差较大。以往的研究通过量化感知训练(QAT)解决这一问题,例如(Gholami等,2021)。他们不是一次性压缩整个模型,而是逐步量化模型参数,并训练其余参数以补偿量化误差。不幸的是,在我们的情况下运行QAT是不切实际的,因为现代的大型语言模型(LLM)训练或微调成本极高。因此,大多数用于LLM的后训练量化(PTQ)算法仅在同一线性层内调整模型参数(Frantar等,2022a;Lin等,2023;Dettmers等,2023b)。

在此,我们选择在单个transformer块的层级进行优化,即由4-8个线性层组成的组,这些层形成一个多头自注意力机制,随后是一个MLP层。在对单个transformer块中的所有线性层进行量化后,我们通过反向传播调整其余参数,以更好地逼近该transformer块的原始输出。具体来说,我们使用PyTorch的自动微分引擎来微分 (||block(X_{block}) - Y_{block}||^2)
,其中 (X_{block}) 是该transformer块的输入激活,(Y_{block}) 是量化前记录的block(Xblock)的输出激活。我们训练码本 (C_m)、比例向量 (s) 和所有非量化参数(RMSNorm的比例和偏置),同时保持代码 (b_{i,j,m}) 不变。与第3.3节类似,我们使用Adam训练这些参数,以最小化与原始块输出(量化前)的均方误差。这一阶段使用与单层量化相同的校准数据。完整过程在算法1中总结。

虽然微调块比单个线性层更昂贵,但仍可以在合理的时间内在单个GPU上对数十亿参数的模型进行量化。此外,由于算法只修改少量可训练参数,优化器状态占用的显存较少。这种微调在几次迭代后收敛,因为它从一个良好的初始猜测开始。实际上,微调transformer层只占总校准时间的一小部分(10-30%或更少)。

4. Experiments

我们在现代大型语言模型(LLM)的后训练量化典型场景下评估AQLM算法。我们的评估主要集中在LLAMA 2模型系列,因为它是微调模型或一般LLM应用的流行基础,例如(Dettmers等,2023a),我们还展示了Mistral系列模型的结果(Jiang等,2024)。在4.1节中,我们对各种LLAMA 2模型和量化位宽进行了完整的AQ流程评估;4.3节展示了对单个AQ组件的消融分析和实现细节。

4.1. 现代LLM的压缩质量

我们在WikiText-2(Merity et al., 2016)和C4(Raffel et al., 2020)验证集上报告困惑度(perplexity)。同时,我们通过LMEval Harness(Gao et al., 2021)测量WinoGrande(Sakaguchi et al., 2021)、PiQA(Tata & Patel, 2003)、HellaSwag(Zellers et al., 2019)、ARC-easy和ARC-challenge(Clark et al., 2018)的零样本准确率。我们遵循GPTQ(Frantar et al., 2022a)的评估设置,并在附录C中提供AQLM和基准的配置。我们考虑了压缩范围的三个主要目标:每个模型参数2-2.8位、3-3.1位和4-4.1位。在下面的结果中,参数的平均位数仅考虑量化的权重,不包括以浮点精度保存的参数,类似于相关工作。关于模型大小估算的详细信息见附录H。我们将AQ与3和4位的GPTQ(Frantar等,2022a)、3和4位的SpQR(Dettmers等,2023b)、2、3和4位的QuIP(Chee等,2023)以及2和4位的QuIP#(Tseng等,2024)进行比较。虽然GPTQ和SpQR技术上支持2位量化,但在2-3位范围内表现不佳。对于QuIP,我们调整的实现显示LLAMA2 13B和70B的性能尚可,但对7B模型表现不佳。我们使用RedPajama数据集的子集(Computer, 2023)进行每个算法的校准,序列长度为4096。

每种方法的具体位宽由如码本数量和码宽等参数决定。我们在表1和2中分别报告2-2.8和3-3.1位宽范围的结果。更多关于4-4.1位的结果在附录F.2中。

结果显示,AQLM在所有设置中均优于先前的最佳PTQ算法,通常具有较大的优势,特别是在高压缩时。这在标准验证集(Wiki-Text2和C4)的困惑度和零样本任务的准确率方面都成立。具体来说,在"极端"的每个参数2-2.1位范围内,我们观察到最高的准确率提升,此时所有方法偏离未压缩模型都很大。

Mixtral量化。表3展示了Mixtral MoE的结果,与2位的QuIP#进行比较。(完整结果见附录F.1)在这种情况下,AQLM也优于QuIP#。虽然与LLAMA2模型相比,优势较小,但对于“更难”的任务,如Arc Challenge,仍然显著(+3分)。

AQLM的帕累托最优性。显著的误差改进引发了在特定内存预算内选择“最优”模型变体以最大化准确率的问题。为此,我们遵循Dettmers & Zettlemoyer(2022):若量化模型在相同或更小总大小(字节数)下最大化准确率,则称其为帕累托最优。尽管取得了快速进展,先前的方法在2位时并非帕累托最优:例如,之前最佳的2位LLAMA2 13B(QuIP#,表1)在Wiki2上达到困惑度6.06,但使用4位量化的7B模型可以获得更低的5.21困惑度,且更小(见附录表10)。AQLM在相同模型上严格压缩到2位也低于帕累托最优,因为LLAMA2 7B的4位AQLM压缩优于其5.21对5.59。为了找到帕累托最优的量化位宽,我们在每个参数2-3位之间进行实验,并在表1的水平线下报告。因此,AQLM的帕累托最优位宽似乎在每个参数2.5位左右(表1),此时我们与LLAMA2 7B的5位AQLM相当(附录表10)。反过来,13B的2.76位AQLM优于未压缩的7B模型。因此,AQLM是第一个在每个参数少于3位时实现帕累托最优的算法。

4.2. 端到端微调实验

在后续研究中,QuIP#(Tseng等,2024)通过微调整个模型以最小化KL散度,改进了我们的分块协议(第3.4节)。在此,我们分析这种端到端微调如何转化为AQLM。我们遵循QuIP#(Tseng等,2024)的设置,使用默认参数进行端到端微调(见附录A)。表4报告了使用AQLM和QuIP#在端到端微调下的2位量化结果。我们在补充材料的表6、13和15中报告了该设置下的其他结果。为了区分两个版本,我们用⋆标记经过端到端微调的量化模型。总体而言,端到端微调提升了QuIP#和AQLM的性能,使两种方法达到相当的准确率。此外,我们注意到端到端微调对2位量化模型的提升更显著,而对3位及以上的模型效果递减。最后,我们可以看到,经过端到端微调的13B的2.19位AQLM与未压缩的7B模型相当,在零样本任务上达到了帕累托最优。

4.3. 消融分析

在附录E中,**我们考察了关于初始化、交替优化、微调影响以及对超参数敏感性的重要设计选择。**简而言之,我们首先发现残差K-means初始化对于算法的快速收敛至关重要:与随机初始化相比,它需要显著更少的训练迭代。我们还比较了相同位宽下的不同超参数配置,改变码本数量和组大小。其次,为验证我们的校准微调过程,我们将其与以下情况进行比较:1)不进行微调,2)仅微调非线性层(如RMSNorm),不微调码本参数,以及3)仅微调码本(不微调其他层)。在附录E中完整展示的结果表明,微调码本参数对准确率的影响最大,而仅微调RMSNorm的影响较小。这验证了我们利用校准集学习码本的选择。

此外,我们观察到,将样本序列数量从128增加到4096会逐渐改善PPL,但效果递减。这对于初始AQLM校准和微调都是如此。在这方面,AQLM从更大的校准集(类似于QuIP#)中获益更多,而不像GPTQ等直接方法,在大约256个输入序列时精度就趋于饱和。最后,我们研究了在给定位预算下的各种投资选项,例如比较较长的编码(如1x15)与具有较短编码的多个码本(如2x8)。

4.4. 推理速度

虽然我们的主要目标是针对给定模型大小最大化准确性,但AQLM在推理延迟方面也可以很实用。为此,我们为几种硬件友好的AQLM配置实现了高效的GPU和CPU内核。结果见表5。对于GPU推理,我们针对16位码本量化的LLAMA2模型,这对应于LLAMA2 70B的2.07位,13B的2.19位,以及7B模型的2.29位(见表1和4),还有一个在Wiki2上困惑度为6.57的2x8位码本模型(见表12)。对于每个模型,我们在标准层上基准测试矩阵-向量乘法子程序的性能。结果显示,AQLM可以以与FP16相当或更好的速度执行。通过HuggingFace集成的端到端生成结果见附录I:例如,在这种设置下,我们可以在LLAMA2 70B上实现≈14个tokens/s。我们观察到,多个较小的码本可以有效利用GPU缓存,从而实现更大的加速,但代价是略微降低准确性。
接下来,我们探讨如何利用AQLM加速CPU推理。正如在第2.2节中讨论的,如果码本大小较小,加性量化可以高效地计算点积。对于AQLM,可以将每个16位码本替换为多个较小的8位码本。这样会导致更高的量化误差,但在准确性方面仍优于基线(见附录表9)。表5中的结果显示,这也使得在CPU上相对于FP32实现高达4倍的推理速度提升。

5.Conclusion and FutureWork

我们介绍了AQLM,这是一种针对大型语言模型(LLM)压缩的新型加性量化方法(AQ),在每权重2位和3位的量化条件下显著提升了最先进的LLM量化结果。在局限性方面,AQLM比直接后训练量化方法(如RTN或GPTQ)计算量更大,尤其是因为使用了更复杂的编码表示。然而,尽管编码和解码更复杂,我们已展示AQLM在CPU和GPU上的高效实现。总体而言,我们认为使用AQLM可以在低内存消耗下准确高效地执行大规模LLM是非常值得注意的。虽然AQLM在低位量化方面已经取得了显著改进,但我们没有在这项工作中探索的一些有前景的方向可以进一步提升。其中一个方向是更好的微调策略。在第4.2节中,我们发现更好的微调算法(Tseng et al., 2024; Malinovskii et al., 2024)可以显著提高量化模型的准确性。我们相信,通过系统地探索微调算法,AQLM可以在未来的工作中受益。另一个有前景的方向是将AQLM推广到其他量化场景。尽管我们的工作集中在LLM量化上,底层算法可能适用于其他问题,例如量化计算机视觉模型、压缩长序列的LLM注意力缓存等。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值