[论文解读] FlexRound: Learnable Rounding based on Element-wise Division for Post-Training Quantization

Jung Hyun Lee, Jeonghoon Kim+2 more|arXiv (Cornell University)|Jun 1, 2023

Advanced Neural Network Applications被引用 5

一句话总结

FlexRound 提出了一种用于训练后量化（PTQ）的可学习量化机制，通过逐元素除法联合优化统一的量化网格大小和每权重缩放因子，实现灵活且基于权重幅度的量化。该方法在图像分类、自然语言理解与生成任务中均达到最先进性能，包括对大语言模型实现高效 4 位量化，且精度损失极小。

ABSTRACT

Post-training quantization (PTQ) has been gaining popularity for the deployment of deep neural networks on resource-limited devices since unlike quantization-aware training, neither a full training dataset nor end-to-end training is required at all. As PTQ schemes based on reconstructing each layer or block output turn out to be effective to enhance quantized model performance, recent works have developed algorithms to devise and learn a new weight-rounding scheme so as to better reconstruct each layer or block output. In this work, we propose a simple yet effective new weight-rounding mechanism for PTQ, coined \emph{FlexRound}, based on element-wise division instead of typical element-wise addition such that FlexRound enables jointly learning a common quantization grid size as well as a different scale for each pre-trained weight. Thanks to the reciprocal rule of derivatives induced by element-wise division, FlexRound is inherently able to exploit pre-trained weights when updating their corresponding scales, and thus, flexibly quantize pre-trained weights depending on their magnitudes. We empirically validate the efficacy of FlexRound on a wide range of models and tasks. To the best of our knowledge, our work is the first to carry out comprehensive experiments on not only image classification and natural language understanding but also natural language generation. Moreover, we demonstrate, for the first time, that large language models can be efficiently quantized, with only a negligible impact on performance compared to half-precision baselines, achieved by reconstructing the output in a block-by-block manner. Our code is available at \url{https://github.com/onliwad101/FlexRound_LRQ}.

研究动机与目标

• 为解决传统训练后量化（PTQ）中因僵化、统一的量化方案导致的性能下降问题。
• 通过基于权重幅度的自适应缩放，提升预训练模型的量化效率。
• 开发一种可学习的量化机制，联合优化量化网格大小与每权重缩放因子，无需完整微调。
• 在多样化任务中（包括自然语言生成与大语言模型）验证该方法的有效性。
• 证明通过分块输出重建，可实现大语言模型的 4 位量化，且性能接近半精度基准。

提出的方法

• FlexRound 将传统权重量化中的逐元素加法替换为逐元素除法，以实现尊重预训练权重幅度的梯度传播。
• 通过可微量化机制，联合学习共享的量化网格大小与每个预训练权重的独立缩放因子。
• 利用除法的导数链式法则，使预训练权重引导缩放因子的更新，实现基于幅度感知的量化。
• 通过在层/块输出上使用重建损失进行梯度下降优化缩放因子，最小化量化引起的性能退化。
• 该方法支持通道级与张量级量化，并采用分块方式重建输出，以保持模型性能。
• 与现有 PTQ 框架兼容，可仅通过最小修改集成至现有量化流程。

实验结果

研究问题

• RQ1 基于逐元素除法的可学习量化机制是否在准确率与灵活性方面优于现有 PTQ 方法？
• RQ2 通过除法实现的幅度感知缩放是否能提升量化性能，尤其在大权重情况下？
• RQ3 FlexRound 是否能有效实现大语言模型的低比特（如 4 位）量化，且性能退化极小？
• RQ4 在图像分类、自然语言理解与生成等多样化任务中，FlexRound 与 AdaRound 及其他最先进 PTQ 方法相比表现如何？
• RQ5 在低精度量化下，分块输出重建是否能有效维持 LLaMA-33B 等高参数量模型的性能？

主要发现

• FlexRound 在图像分类与自然语言理解基准上达到最先进性能，优于 AdaRound 与其他 PTQ 方法。
• 在 WikiText2 上，FlexRound 使用 8 位权重时，LLaMA-7B 的困惑度为 5.73，优于 AdaRound 的 6.19，相对提升 7.3%。
• 在 4 位量化下，FlexRound 将 LLaMA-7B 在 WikiText2 上的困惑度降低至 5.63，优于 AdaRound 的 5.72，且与半精度基线相比退化极小。
• 在 PTB 数据集上，FlexRound 使用 8 位权重时，LLaMA-7B 的困惑度为 9.28，优于 AdaRound 的 9.85，相对提升 6.1%。
• 使用 3 位权重时，FlexRound 在 WikiText2 上保持 LLaMA-7B 的困惑度为 5.88，显著优于 AdaRound 的 6.41，表明其具备强大的低比特量化能力。
• 对于 LLaMA-33B，FlexRound 使用 4 位权重时在 WikiText2 上的困惑度为 4.32，仅比半精度基线（4.06）高出 0.26，证明其接近半精度性能。