Agile-Quant：面向大语言模型边缘端更快推理的激活引导量化框架

Shen X, Dong P, Lu L, et al. Agile-quant: Activation-guided quantization for faster inference of LLMs on the edge[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2024, 38(17): 18944-18951.

1. 引言与研究动机

大语言模型（LLMs）如LLaMA、OPT和BLOOM等基于Transformer架构，在自然语言处理任务中展现出了前所未有的性能。然而，这些模型的部署面临着严峻的计算和存储挑战。以GPT3-175B为例，即使采用紧凑的float16格式存储，其参数也需要326GB（以1024为倍数）的内存空间。这一需求不仅超出了当前最强大的单个GPU的容量限制，更使得在硬件资源受限的边缘设备上部署这些模型变得极其困难。

量化技术作为一种模型压缩方法，通过降低数值精度来减少模型的存储需求和计算复杂度。特别是在边缘设备上，量化可以充分利用CPU和树莓派等平台中普遍存在的SIMD（单指令多数据）单元所支持的高效8位定点（INT8）运算，从而大幅提升推理吞吐量和能效。

1.1 现有方法的局限性

当前的量化方法存在两个关键问题，严重制约了LLMs在边缘设备上的实际部署效果。

问题一：激活量化的困难性。现有工作如GPTQ、AWQ和SpQR等主要专注于权重的4位量化，而将激活保留在浮点（FP16）域中。这种做法的根本原因在于激活量化会严重影响模型性能，特别是当模型规模增大时，激活中会出现大量的异常值（outliers）。Dettmers等人的实验表明，直接将这些异常值设置为零会导致任务性能下降高达45%。此外，由于LLMs的模型规模巨大，有限的学术计算资源难以承担相关的训练成本，因此后训练量化（PTQ）成为主流方法，但PTQ在处理异常值引起的量化误差方面能力有限。

问题二：硬件支持的不匹配。主流边缘处理器如CPU和树莓派通过SIMD单元高效地并行执行多个操作。SIMD指令擅长利用字节级数据（8位整数）的并行性，并在常见的指令集架构（ISA）和深度神经网络处理框架中得到良好支持，例如TensorFlow Lite中的GEMMLOWP和PyTorch中的QNNPACK。然而，这些低精度内核仅通过零扩展将亚字节操作数对齐到字节边界，将其作为8位或16位操作数处理，无法真正实现亚字节精度的高效计算。

2. 激活特性的深入分析

2.1 LLMs中的Token重要性机制

在自然语言处理中，句子中往往存在大量非关键词汇，它们对整体理解贡献较小。这意味着我们可以使用更少的资源高效处理这些词汇，甚至可以将其排除以降低复杂度。由于词汇在语言模型中被嵌入为token，研究者通过分析注意力机制来评估每个token的重要性。

然而，LLMs中使用的因果注意力掩码确保在自注意力机制中，每个token只能与之前的token交互，而不能与后续的token交互。这种因果机制使得累积概率不适合评估token重要性，因为它对前面的token的累积概率不公平。

图1的详细分析：

• 图1(a)展示了LLaMA-7B模型第一层的注意力图，包含16个token，激活未被量化。可以观察到明显的三角形模式，表明token主要关注相邻位置，特别是前一个位置。
• 图1(b)展示了最后一层的注意力图，几乎所有token都呈现垂直条纹模式，表明所有token都与第一个起始token强相关。
• 图1©展示了激活量化后的注意力图，原本的垂直条纹模式退化为三角形模式，量化后的token倾向于关注局部相邻位置而非全局的起始token。
• 图1(d)展示了包含2048个token的一个激活中异常值的分布情况。深色的直线表明异常值倾向于集中在相同或相邻的通道中，通道索引保持不变。

2.2 激活量化对注意力机制的影响

激活量化对模型性能的影响主要体现在两个方面。首先，异常值的分布特性使得量化变得困难。如图1(d)所示，异常值分布在相邻甚至相同的通道中，这种聚集性使得传统的均匀量化方法难以有效处理。其次，注意力图由查询矩阵$Q$和键矩阵$K$生成，是输入依赖的，因此会受到激活量化的直接影响。

量化后的注意力图从全局模式退化为局部模式，这种变化意味着：

• 代表性特征的全局性严重恶化
• 模型捕获长距离依赖关系的能力下降
• 量化误差导致的原始注意力图信息损失会直接影响最终的任务性能

2.3 硬件延迟的详细分析

图2的硬件分析结果显示了LLaMA块在Snapdragon 870 CPU上的延迟分布：

• 在FP16精度下：FC1和FC2占总运行时延迟的54%，线性变换和投影占约30%，非线性操作（LayerNorm/Softmax/SwiGLU）仅占不到8%
• 在INT8精度下：FC1和FC2占49.5%，其他线性操作占约38%，非线性操作占不到12%
• 在INT4精度下：主要计算负载进一步向线性操作倾斜，非线性操作占比增加到约16%

基于这一分析，论文确定了量化优先级顺序：FC1和FC2（最高优先级）→ 线性变换 → 线性投影 → AttnV → QK（最低优先级）。

3. Agile-Quant方法论

3.1 基础量化器的数学定义

3.1.1 均匀量化器

均匀量化是最常用的量化方法，已被大多数硬件设备原生支持。对于输入张量$X \in \mathbb{R}^{m \times n}$和位宽$b$，均匀量化器定义为：

$$Q(X|b) = \text{CLIP}\left(\left\lfloor\frac{X}{s}\right\rceil + zp, 0, 2^b - 1\right)$$

其中缩放因子$s$和零点$zp$的计算方式为：

$$s = \frac{\max(X) - \min(X)}{2^b - 1}$$

$$zp = \text{CLIP}\left(\left\lfloor-\frac{\min(X)}{s}\right\rceil, 0, 2^b - 1\right)$$

3.1.2 Log2量化器

Log2量化器通过引入指数运算来处理具有长尾分布的激活值：

$$Q(X|b) = \text{Sign}(X) \cdot \text{CLIP}\left(\left\lfloor-\log_2\frac{|X|}{\max(|X|)}\right\rceil, 0, 2^{b-1} - 1\right)$$

这种量化器特别适合处理Softmax输出，因为Softmax的输出值通常呈现指数分布特性。

3.2 激活量化流水线

图3的详细描述展示了完整的激活量化流水线：

1. 嵌入的token首先经过LayerNorm标准化
2. 自注意力模块包含：
   • $Q \times K^T$计算（8x8精度）
   • Softmax激活（使用Log2量化）
   • Attention × V计算（4x8精度）
   • 线性投影（8x4或4x4精度）
3. MLP模块包含：
   • 全连接层FC1（8x4或4x4精度）
   • 激活函数
   • 全连接层FC2（8x4或4x4精度）
4. 残差连接贯穿整个流程

3.3 激活引导的Token剪枝优化

3.3.1 TRIP（Two-Refine Improved by Pruning）方法

图4的实现细节展示了激活量化与token剪枝的结合过程。该图说明了如何通过减少token数量来降低量化误差并改善注意力机制。

对于transformer模型中的激活矩阵$X \in \mathbb{R}^{m \times d}$，其中$m$是token数量，$d$是特征维度，token剪枝函数$F_P(\cdot)$根据token重要性级联地剪枝token：

$$X^P = F_P(X) \in \mathbb{R}^{n \times d}, \quad n < m$$

然后应用TRIP因子$\alpha = \{\alpha_1, \alpha_2, ..., \alpha_d\}$到不同通道：

$$X^P_Q = Q(X^P|b, \alpha) = \text{CLIP}\left(\left\lfloor\frac{X^P}{2^{\alpha}s}\right\rceil + zp, 0, 2^b - 1\right)$$

修正后的缩放因子和零点为：

$$s = \frac{\max(X^P) - \min(X^P)}{(2^b - 1) \cdot 2^{\alpha}}$$

$$zp = \text{CLIP}\left(\left\lfloor-\frac{\min(X^P)}{2^{\alpha} \cdot s}\right\rceil, 0, 2^b - 1\right)$$

对于每个通道$c \in [1, d]$，最优TRIP因子通过最小化量化误差获得：

$$\alpha_c = \arg\min_{\alpha_c \in \{0,1,...,K\}} \left\|X^P_c - \left\lfloor\frac{X^P_c}{2^{\alpha_c} \cdot s}\right\rceil \cdot 2^{\alpha_c} \cdot s\right\|_2^2$$

其中$K$是最大细化级别，依赖于包含相对较多异常值的通道。

3.3.2 渐进式Token剪枝策略

假设模型有$n$层$L = (l_1, l_2, ..., l_n)$，剪枝操作应用于层子集$L_p = (l_{p1}, l_{p2}, ..., l_{pm})$。设置最后一层$l_n$的剪枝比率为$\beta$，则对于$l_i \in L_p$的渐进比率$\gamma$计算为：

$$\gamma = 1 - (1 - \beta)^{\frac{1}{m}}$$

这种指数衰减策略确保剪枝强度随深度增加而逐渐增强。累积的token稀疏度$s$定义为：

$$s = 1 - \sum_{i=1}^{n} r_i$$

其中$r_i$表示层$l_i$中剩余token的比例。

3.4 边缘导向的硬件优化

3.4.1 SIMD基础的4位乘法器设计

图5的硬件实现细节展示了INT4乘法器的具体设计：

1. 点乘法阶段：将相邻的两个权重值$W_{i,j}$和$W_{i+1,j}$合并成一个INT16值，然后与共享的激活值相乘。结果的前8位分配给$W_{i,j}$的乘积，后8位分配给$W_{i+1,j}$的乘积。

2. 加法阶段：使用位移操作将16位输出扩展到32位。具体布局为：[0, $O_{i,j}$, 0, $O_{i+1,j}$]，每部分占8位。

3. 量化回INT4：最终将输出分割成两个INT16值，并量化回INT4精度。

这种设计减少了50%的乘法和加法操作，显著提高了计算效率。

3.4.2 高效TRIP矩阵乘法实现

图6的实现策略说明了如何在硬件上高效实现TRIP矩阵乘法：

• TRIP因子预先应用于权重矩阵的相应通道
• 修改后的权重在设备上预加载
• 不同通道使用不同的缩放因子（×2, ×1, ×8等）
• 所有计算在层级共享相同的量化参数

4. 实验结果与分析

4.1 量化性能评估

表1-3展示了在Wikitext-2和C4数据集上的详细实验结果。Agile-Quant在LLaMA、OPT和BLOOM模型上都取得了优异的性能：

• 对于LLaMA-7B：使用Agile-Quant-8（带token剪枝）的困惑度为6.09，仅比FP16的5.68略高，但远优于其他激活量化方法如MoFQ8的6.49
• 对于更大规模的模型，性能差距进一步缩小：LLaMA-65B的困惑度从FP16的3.53增加到3.92

4.2 Token稀疏度分析

图7的关键发现：

• 所有三种规模的LLaMA模型（7B、13B、30B）都在适当的token稀疏度下（红星标记处）实现了比密集模型更好的性能
• 最优稀疏度随模型规模增加而降低：7B模型约0.3，13B模型约0.25，30B模型约0.2
• 过高的稀疏度（>0.6）会导致性能急剧下降

图8的序列长度影响：

• 序列长度对token剪枝效果有决定性影响
• 当序列长度为256时，即使少量剪枝也会导致性能严重退化
• 只有当序列长度达到2048时，token剪枝才能有效发挥作用
• 这一发现解释了为什么token剪枝在长上下文场景中特别有效

4.3 端到端硬件加速

表4展示了在实际硬件设备上的推理性能：

• 在Android手机（Snapdragon 870）上：

  • OPT-125M：从1.03秒加速到0.44秒（2.3倍加速）
  • LLaMA-7B：从10.6秒加速到4.44秒（2.4倍加速）

• 在树莓派4B上：

  • OPT-1.3B：从754.25秒加速到296.50秒（2.5倍加速）
  • OPT-2.7B：从1150.9秒加速到455.64秒（2.5倍加速）

5. 消融研究

5.1 输入序列长度的影响

实验表明，输入序列长度是影响token剪枝效果的关键因素。当序列长度从2048减少到256时，即使是轻微的token剪枝（稀疏度0.1）也会导致困惑度从约6增加到超过15。这是因为短序列中每个token都携带重要信息，剪枝会导致关键信息丢失。

5.2 不同量化策略的比较

通过对比不同组件使用不同量化位宽的效果，研究发现：

• FC1和FC2使用4位量化可以获得最佳的速度-精度权衡
• 注意力机制中的QK计算对量化更敏感，保持8位精度更合适
• Softmax使用Log2量化比均匀量化效果更好

6. 结论与展望

Agile-Quant框架成功地解决了LLMs在边缘设备上部署的关键挑战。通过结合激活引导的量化策略、智能的token剪枝优化以及专门设计的硬件实现，该方法在保持模型性能的同时实现了显著的推理加速。主要贡献包括：

1. 提出了平衡延迟降低和任务性能的激活引导边缘导向量化策略
2. 设计了激活感知的token剪枝方法，有效减少异常值的负面影响
3. 实现了SIMD基础的4位乘法器和高效的TRIP矩阵乘法
4. 在多个边缘设备上实现了高达2.55倍的实际加速

附录：数学推导

A. TRIP因子优化

给定量化误差函数：

$$E(\alpha) = \|X - Q(X|\alpha)\|_2^2$$

其中$Q(X|\alpha)$是带有TRIP因子$\alpha$的量化函数。展开误差函数：

$$E(\alpha) = \sum_{i=1}^{n}\sum_{j=1}^{d} \left(x_{ij} - q_{ij}(\alpha_j)\right)^2$$

对于每个通道$j$，独立优化其TRIP因子：

$$\frac{\partial E}{\partial \alpha_j} = -2\sum_{i=1}^{n} \left(x_{ij} - q_{ij}(\alpha_j)\right) \frac{\partial q_{ij}}{\partial \alpha_j}$$

由于量化函数的阶梯性质，使用网格搜索找到最优$\alpha_j$：

$$\alpha_j^* = \arg\min_{\alpha_j \in \{0,1,...,K\}} \sum_{i=1}^{n} \left(x_{ij} - \left\lfloor\frac{x_{ij}}{2^{\alpha_j} \cdot s_j}\right\rceil \cdot 2^{\alpha_j} \cdot s_j\right)^2$$

B. Token重要性评分的理论基础

定义token $i$对起始token的注意力分数为：

$$a_i = \frac{1}{L}\sum_{l=1}^{L} A^{(l)}_{i,0}$$

其中$A^{(l)}$是第$l$层的注意力矩阵。考虑因果掩码的影响，实际的重要性分数需要归一化：

$$\text{Importance}_i = \frac{a_i}{\sum_{j=1}^{i} a_j} \cdot \log(i+1)$$

$\log(i+1)$项补偿了位置偏差，确保后面的token不会因为累积效应而被低估。

C. 渐进式剪枝比率的最优性证明

设总信息损失为$L(\beta)$，计算复杂度降低为$C(\beta)$。优化目标是：

$$\min_{\beta} L(\beta) \quad \text{s.t.} \quad C(\beta) \geq C_{\text{target}}$$

使用拉格朗日乘数法：

$$\mathcal{L}(\beta, \lambda) = L(\beta) + \lambda(C_{\text{target}} - C(\beta))$$

求解一阶条件：

$$\frac{\partial \mathcal{L}}{\partial \beta} = \frac{\partial L}{\partial \beta} - \lambda\frac{\partial C}{\partial \beta} = 0$$

假设信息损失随剪枝比率指数增长：$L(\beta) = L_0 e^{k\beta}$，
计算复杂度线性降低：$C(\beta) = C_0(1-\beta)$。

代入并求解得到最优剪枝比率：

$$\beta^* = \frac{1}{k}\log\left(\frac{\lambda C_0}{kL_0}\right)$$

这解释了为什么渐进式策略$\gamma = 1 - (1-\beta)^{1/m}$是有效的：它在各层之间分配剪枝负担，使得总体信息损失最小化。

D. 4位乘法器的计算复杂度分析

传统的$n \times n$矩阵乘法需要$O(n^3)$次乘法操作。使用4位打包技术后：

1. 权重打包：将两个4位权重合并为8位，复杂度$O(n^2)$
2. SIMD乘法：每次操作处理2个乘法，复杂度降为$O(n^3/2)$
3. 结果解包：分离高低位，复杂度$O(n^2)$

总体复杂度：

$$T_{\text{4-bit}} = O(n^2) + O(n^3/2) + O(n^2) = O(n^3/2)$$

相比标准8位计算，理论加速比为2倍，考虑到内存访问和缓存效率的提升，实际加速比可达2.3-2.5倍。