[论文解读] Squat: Quant Small Language Models on the Edge

Xuan Shen, Peiyan Dong+10 more|arXiv (Cornell University)|Feb 16, 2024

Advanced Neural Network Applications被引用 4

一句话总结

EdgeQAT 提出了一种新颖的量化感知训练框架，通过熵和分布引导联合优化轻量级大语言模型的权重与激活，以最小化自注意力机制中的信息失真。该方法在树莓派5等边缘设备上实现了高达2.37倍的推理加速，使用4-bit权重和激活，优于FP16基线模型，同时保持了模型精度。

ABSTRACT

A growing trend has emerged in designing high-quality Small Language Models (SLMs) with a few million parameters. This trend is driven by the increasing concerns over cloud costs, privacy, and latency. Considering that full parameter training is feasible for SLMs on mobile devices, Quantization-Aware Training (QAT) is employed to improve efficiency by reducing computational overhead and memory footprint. However, previous QAT works adopt fine-grained quantization methods to compress models with billions of parameters on GPUs, incompatible with current commodity hardware, such as mobile and edge devices, which relies on Single Instruction Multiple Data (SIMD) instructions. Thus, the generalization of these methods to SLMs on mobile devices is limited. In this paper, we propose Squat method, an effective QAT framework with deployable quantization for SLMs on mobile devices. Specifically, we propose entropy-guided and distribution-aligned distillation to mitigate the distortion of attention information from quantization. Besides, we employ sub-8-bit token adaptive quantization, assigning varying bit widths to different tokens based on their importance. Furthermore, we develop a SIMD-based Multi-Kernel Mixed-Precision (MKMP) multiplier to support sub-8-bit mixed-precision MAC on mobile devices. Our extensive experiments verify the substantial improvements of our method compared to other QAT methods across various datasets. Furthermore, we achieve an on-device speedup of up to 2.37x compared with its FP16 counterparts, signaling a great advancement. Code: https://github.com/shawnricecake/squant

研究动机与目标

• 解决在边缘设备上对轻量级大语言模型进行低比特权重与激活量化时导致的性能下降问题。
• 识别出量化注意力图中的信息失真——特别是查询和键的分布失真——是精度下降的主要原因。
• 开发一种联合优化权重与激活量化的QAT框架，以实现在边缘硬件上充分利用INT8运算。
• 引入基于标记重要性的自适应量化方法，进一步降低比特宽度而不牺牲精度。
• 实现在智能手机、树莓派和FPGA等边缘平台上的高速、低资源推理。

提出的方法

• 在查询和键张量上提出熵最大化，以最小化量化误差并恢复信息内容。
• 引入一种分布引导损失，最小化FP16与量化后注意力图之间的余弦相似度差异。
• 设计一种基于标记重要性的机制，根据激活模式动态分配4-bit或8-bit精度。
• 对权重和激活均应用QAT，使边缘设备能够通过纯整数运算实现完整加速。
• 采用混合量化策略，高重要性标记使用8-bit精度，低重要性标记使用4-bit，以优化效率。
• 使用统一的训练流程，将熵、分布和自适应比特宽度优化整合到一个端到端的QAT框架中。

实验结果

研究问题

• RQ1 为何在8-bit以下联合量化权重、激活和KV缓存时，后训练量化（PTQ）会失效？
• RQ2 大语言模型在低比特量化中，尤其是注意力机制中，导致主要性能下降的原因是什么？
• RQ3 熵最大化与分布对齐能否共同减少注意力图中的量化误差？
• RQ4 基于重要性的自适应激活量化在降低平均比特宽度的同时，保持精度的有效性如何？
• RQ5 在树莓派和FPGA等真实边缘设备上，对权重和激活进行完整QAT能否实现显著加速？

主要发现

• 在树莓派5上，EdgeQAT 使用W4A4量化实现了高达2.37倍的推理加速，显著优于FP16推理。
• 在OnePlus 11智能手机上实现了1.91倍的本地推理加速，在FPGA平台上实现了高达1.9倍的加速，均采用W4A4设置。
• 在SuperGLUE基准上，EdgeQAT 使用4-bit权重和激活实现了最先进精度，优于PACT和LSQ等QAT基线模型。
• 采用25%的标记使用8-bit精度的自适应量化方法，其精度高于等效的非自适应5-bit量化。
• 分布引导损失有效减少了FP16输出与量化输出之间注意力图的差异，直接缓解了精度损失。
• 基于熵的查询和键分布优化能有效降低量化误差，尤其在低比特环境下表现更优。