敏感度感知混合精度量化 V1

问题： 统一的超低精度量化会导致严重的精度损失，因此需要逐层方法。
解决方案： 智能混合精度量化逐层调整位宽——保护敏感部分，同时积极压缩其余部分。

工作原理

• 🔍 敏感度评分 – 衡量每层从量化到较低位数的退化程度
• ⚖️ 精度分配 – 自动为脆弱层分配更高的位数，为鲁棒层分配更低的位数
• 🔄 优化 – 迭代调整以达到精度目标

重要性： 在不显著降低性能的情况下，提供更精简、更快的模型。
（深入了解方法论的细节！🧠）

高效部署大型神经网络（如基于Transformer的模型或扩散模型）仍然是一个主要挑战，因为它们需要高内存和计算资源。量化通过降低精度帮助解决内存瓶颈，但将相同的低精度统一应用于所有层通常会降低模型性能。

在这项工作中，我们探索了一种敏感度感知混合精度量化（MPQ）方法，该方法根据每个块对量化的敏感程度分配精度级别。更敏感的块保持更高的精度，而更鲁棒的块则进行更积极的量化。这种方法适用于不同的架构——从LLM到扩散模型——旨在与统一量化相比，在效率和精度之间取得更好的权衡。

这种敏感度感知的自动量化框架提出了一系列优化问题，其仔细的解决对于实现最高水平的量化并最大限度地减少模型性能的降级至关重要。核心思想是根据每个层对精度损失的敏感度来调整量化策略。

我们首先遍历模型以识别所有 `nn.Linear` 层（除了 `lm_head`）。在计算每层（或块）的敏感度分数后，我们分配适当的量化策略。可用选项包括

• 未更改（`nn.Linear`）：保留原始的FP32权重。
• W16A16LinearLayer：权重量化为BF16（16位），激活仍为FP16。
• W8A16LinearLayer：权重量化为INT8，激活仍为FP16。
• W4A16LinearLayer：权重量化为INT4（4位），激活仍为FP16。

每个选定的量化方案都会将标准 `nn.Linear` 层替换为相应的自定义层类，以支持所选的精度级别。这些自定义类处理量化和反量化过程，包括权重打包（针对W4A16）、尺度和零点计算以及缓冲区管理。

原始模型

MPQ模型版本

图1：EleutherAI/gpt-neo-125M模型在MPQ前后的概述

如背景部分所述，有多种方法可以估计层或块对量化的敏感度。然而，其中大多数方法要么计算成本高昂，要么特定于模型。在这项工作中，我们的目标是开发一种既有效又计算高效的方法来估计这些敏感度。一个简单而有前景的想法是基于散度的敏感度估计。核心思想很简单：要估计给定层或块的敏感度，我们计算原始模型输出与仅对该特定层或块进行量化的模型输出之间的Jensen-Shannon散度（JSD）。高散度表示高度敏感的层，对应于高敏感度分数，该层应在混合精度量化阶段分配更多位。我们还根据激活幅度对敏感度分数进行小幅提升，以进一步完善估计。

基于散度的敏感度分数是使用Jensen-Shannon散度（JSD）计算的，用于比较基线和量化输出概率分布。我们选择JSD进行敏感度估计，因为它具有对称性和有界性，非常适合比较概率分布，而不会偏向原始或量化输出。与KL散度不同，JSD能够很好地处理零概率，并提供一种平滑的分布偏移度量，这与感知模型退化良好相关。然而，JSD也有限制：它假定输出分布是任务性能的可靠代理，这可能不适用于所有架构（例如，回归任务）。此外，它需要足够的校准数据来稳定估计，并且其计算成本随词汇量增长。未来的工作将探索互补度量（例如，基于Hessian的分数）来解决这些限制。在所有批次$BB$上平均并由$\log 2\backslash log\; 2$归一化的最终敏感度分数为

最终的敏感度分数使用归一化激活幅度进行调整

变量 $\lambda \backslash lambda$ 是一个可以在验证数据集上调整的超参数。

用于计算这些分数的样本数据集是 `wikitext` 数据集的子集。所使用的样本数量也需要研究，以确定提供代表性子集并产生可靠敏感度估计的最小尺寸。

以下是三个不同模型的层敏感度分析

图2：EleutherAI/gpt-neo-125M的层敏感度

图3：facebook/opt-125m的层敏感度

图4：TinyLlama/TinyLlama-1.1B-Chat-v1.0的层敏感度

观察三种不同语言模型（EleutherAI GPT-Neo-125M、Facebook OPT-125M和TinyLlama 1.1B）的敏感度分析结果，出现了几个一致的模式，为量化行为提供了有价值的见解。最引人注目的观察是，所有模型的最终投影层（c_proj）都表现出持续高敏感度，在GPT-Neo和OPT模型中，敏感度分数范围为1.0e-04到1.4e-04，在TinyLlama的down_proj层中达到3.3e-05。这种模式反映了这些层在将高维表示投影回词汇空间中的关键作用，其中即使是很小的量化误差也可能显著影响最终输出概率。此外，值投影层（v_proj）表现出显著敏感度，尤其是在GPT-Neo（1.0e-04）和TinyLlama（1.1e-05）中，这表明注意力值计算特别容易受到精度损失的影响，因为它们在注意力机制中决定信息流。

分析揭示了模型之间敏感度模式的显著架构差异。Facebook OPT表现出独特的模式，其全连接层（fc1和fc2）显示出相对较高的敏感度分数（3.87e-04和5.76e-04），这与其他模型显著不同，可能反映了前馈网络实现中的架构差异。相反，键和查询投影层（k_proj和q_proj）在所有模型中始终表现出较低的敏感度，在GPT-Neo中得分约为2.7e-05，在OPT中甚至更低，分别为5.1e-06和9.5e-06。这种对量化的鲁棒性表明，注意力机制可以容忍键-查询相似度计算中的一些噪声，而不会导致显著的性能下降。

层级敏感度分布揭示了模型深度上的非均匀模式，清晰的峰值出现在特定层，而非渐进过渡。GPT-Neo和OPT模型在早期到中期层和模型末尾显示出最高的敏感度，而TinyLlama则表现出更分散的模式，反映了其不同的架构和训练方法。这些发现为混合精度量化策略提供了明确的指导：投影层，特别是c_proj/down_proj组件，应分配更高的比特以保持性能，而k_proj和q_proj层可以容忍更激进的量化。模型特定模式强调了逐架构敏感度分析的重要性，因为即使模型家族具有相似的参数数量和总体结构，最佳量化策略也可能因模型家族而异。

接下来，我们将实现更稳健的方法来估计层或块的敏感度。特别是，我们将探索基于Hessian的方法，如HAWQ论文中所介绍的。这些方法利用二阶信息，通过检查损失函数相对于权重的曲率，更准确地评估层对量化的敏感度。

虽然计算量比基于散度的方法更大，但基于Hessian的技术在估计敏感度方面提供了更高的保真度，并且可以更好地指导混合精度量化中的精度分配。我们计划研究Hessian的近似（例如，块对角或对角）以减少计算开销，并使这些方法更具可扩展性。

此外，我们将比较基于散度和基于Hessian的敏感度估算器在不同模型架构和任务中的有效性和效率，以评估权衡并指导实际中的方法选择。

一旦计算出逐层敏感度分数，接下来的挑战是根据这些分数确定适当的量化配置。这提出了一个非平凡的优化问题，取决于各种因素，包括模型架构、精度约束以及内存或延迟预算。

在没有敏感度估计的情况下，需要对所有可能的位宽分配进行穷举搜索。然而，对于深度神经网络来说，这是计算上不可行的，因为混合精度量化的搜索空间随层数呈指数级增长。为了解决这个问题，我们提出了一种启发式初始化策略，该策略利用敏感度分数来指导初始位分配，然后进行迭代优化以满足特定的性能约束。

我们考虑以下三种初始化策略

• adaptive_threshold：使用预定义阈值从敏感度统计数据中导出位分配。
• int8_only：统一将INT8精度分配给所有层。
• int4_only：将INT4精度分配给所有层（除非后续进行优化，否则通常不是最佳选择）。

在`adaptive_threshold`方法中，我们定义了三个阈值来将敏感度分布分割成四个量化区间：FP32、BF16、INT8和INT4。这些阈值计算如下：

$${\alpha }_1=\mu +0.5\sigma ,{\alpha }_2=\mu ,{\alpha }_3=\mu -0.5\sigma \backslash alpha\_1\; =\; \backslash mu\; +\; 0.5\backslash sigma,\; \backslash quad\; \backslash alpha\_2\; =\; \backslash mu,\; \backslash quad\; \backslash alpha\_3\; =\; \backslash mu\; -\; 0.5\backslash sigma$$

其中 $\mu \backslash mu$ 和 $\sigma \backslash sigma$ 表示所有层敏感度分数的均值和标准差。

这些策略充当了寻找最佳混合精度配置的起点——一种在最大化模型压缩的同时最小化精度降低的配置，通常以困惑度（越小越好）来衡量。与暴力搜索相比，这种方法更具可扩展性和实用性。

初始化后，我们通过监控量化模型与原始模型之间的困惑度差距来迭代地优化位分配。如果差距超过预定义的容差（例如，`max_perplexity_increase`），我们会系统地逐级升级层：首先将所有INT4层升级到INT8，然后将所有INT8层升级到BF16，最后将所有BF16层升级到FP32。这种逐级方法是受量化中收益递减原则的启发——从INT4升级到INT8通常比从INT8升级到BF16或从BF16升级到FP32提供更大的精度改进，这使得在转向更高精度升级之前耗尽低精度改进更有效率。在每个位级别内，我们以批次（由`layers_per_iteration`定义）升级最敏感的层，直到该级别所有层都已升级、满足困惑度约束或达到最大迭代限制。这确保在转向更高精度级别之前充分探索每个量化级别。

接下来，我们计划探索超越启发式阈值化的替代位分配策略。特别是，我们将研究基于梯度的优化方法，通过直接最小化平衡模型精度和压缩的损失函数来动态调整位宽分配。这种方法可以利用敏感度分数作为初始化，但通过反向传播或强化学习技术来优化位分配。

使用100个样本进行敏感度估计和模型评估（困惑度），并使用50次迭代进行迭代优化，每次迭代更新3层，允许困惑度增加5%，我们获得了以下结果：

有无优化的策略比较

策略	无优化	有优化
`adaptive_threshold`
`int8_only`
`int4_only`

adaptive_threshold 策略展示了**保守的方法**，优先保留模型质量。这种方法在最小困惑度退化（在测试模型中为0.1-2.3%）的情况下保持模型性能，同时对于较小模型实现约32-37%的大小缩减，对于TinyLlama则达到62.7%。这种保守行为源于该策略依赖于从敏感度分布中导出的统计阈值，这些阈值为高于平均敏感度分数的层分配更高的精度。优化版本和未优化版本之间没有差异，这表明初始位分配已经满足了困惑度约束，证明了adaptive_threshold在质量保存方面的内置安全裕度。

混合精度方法揭示了关于逐层敏感度模式的重要见解。尽管统一的 INT8 量化在压缩率方面显示出具有竞争力的结果，但**混合精度框架为处理多样化的模型架构和不同的敏感度要求提供了关键的灵活性**。adaptive_threshold 策略能够根据经验敏感度测量自动分配精度，这代表了固定量化方案的重大进步，特别是对于层重要性差异很大的模型。这种方法可以实现有针对性的精度分配，并根据特定的性能要求和计算约束进行微调。

仅使用 INT4 策略提供了**迭代优化机制恢复能力**的有力证据。最初激进的 4 位精度量化导致性能下降，困惑度增加了 17.4% 到 416.8%（在测试模型中）。优化过程实现了显著恢复，将 EleutherAI GPT-Neo 模型从 416.8% 的困惑度增加恢复到 3.6%，同时保持了显著压缩（模型大小减少了 39.8%）。这种恢复证明了优化算法能够系统地识别和升级关键层，验证了敏感度估计方法的层重要性排名以及迭代升级策略在满足性能约束同时保留压缩优势的有效性。如果我们更关注性能，我们可以选择更严格的困惑度增加约束（例如，1%）。以下是获得的结果（`int4_only` 经过优化）：

图5：比较结果（策略：仅int4并优化，最大允许PPL增加：1%）

这在压缩和性能之间取得了良好的平衡，因为它实现了比 `int8_only` 模型更大的尺寸缩小，同时几乎没有性能下降。

在这项工作中，我们提出了一种敏感度感知的混合精度量化方法，该方法动态地为基于Transformer的语言模型的不同块分配精度级别。我们的方法通过利用内部敏感度信号，在不重新训练的情况下改善了压缩和性能之间的权衡，在严格的比特约束下，在精度方面取得了可喜的成果。

到目前为止，我们已经在仅解码器模型上评估了我们的量化策略，用于因果语言建模。作为未来的工作，我们计划将我们的方法扩展到其他架构家族（例如，扩散模型）和更多样化的任务，以评估其通用性。我们还旨在将我们的评估扩展到更大规模的网络，以在更具挑战性的设置下测试我们方法的极限。同时，我们打算通过将其与替代度量进行比较来验证我们的敏感度指标，并通过在MMLU和HellaSwag等标准下游数据集上对量化模型进行基准测试来S进行更全面的评估。这将有助于确认我们方法的稳健性及其在实际部署场景中的适用性。