[论文解读] Extreme Compression of Large Language Models via Additive Quantization

Vage Egiazarian, Andrei Panferov+4 more|arXiv (Cornell University)|Jan 11, 2024

Topic Modeling被引用 4

一句话总结

该论文提出AQLM，一种新颖的加法量化方法，可将大语言模型（LLMs）极端压缩至每参数2–3比特。通过采用实例感知、输入自适应的码书优化及跨层块联合调优，AQLM在2比特设置下实现了最先进（SOTA）的准确率，同时支持高达8倍的内存减少和GPU上3.05倍、CPU上4.07倍的加速，实现快速GPU/CPU推理。

ABSTRACT

The emergence of accurate open large language models (LLMs) has led to a race towards performant quantization techniques which can enable their execution on end-user devices. In this paper, we revisit the problem of "extreme" LLM compression-defined as targeting extremely low bit counts, such as 2 to 3 bits per parameter-from the point of view of classic methods in Multi-Codebook Quantization (MCQ). Our algorithm, called AQLM, generalizes the classic Additive Quantization (AQ) approach for information retrieval to advance the state-of-the-art in LLM compression, via two innovations: 1) learned additive quantization of weight matrices in input-adaptive fashion, and 2) joint optimization of codebook parameters across each transformer blocks. Broadly, AQLM is the first scheme that is Pareto optimal in terms of accuracy-vs-model-size when compressing to less than 3 bits per parameter, and significantly improves upon all known schemes in the extreme compression (2bit) regime. In addition, AQLM is practical: we provide fast GPU and CPU implementations of AQLM for token generation, which enable us to match or outperform optimized FP16 implementations for speed, while executing in a much smaller memory footprint.

研究动机与目标

• 为解决在每参数2–3比特的极端大语言模型压缩中，现有方法因准确率显著下降而面临的挑战。
• 将多码书量化（MCQ）技术，特别是加法量化（AQ），扩展至大语言模型权重的后训练量化。
• 通过引入输入感知校准与联合块级优化，提升低比特设置下的准确率-大小帕累托最优性。
• 通过优化的GPU和CPU推理内核，实现压缩大语言模型在实际部署中的可行性，且延迟开销极低。

提出的方法

• 提出AQLM，一种学习型加法量化框架，重新表述经典AQ优化问题，以在输入标记分布下最小化层输出误差。
• 通过联合优化整个层块的码书参数，引入实例感知、输入自适应的码书学习，利用校准数据。
• 采用残差K均值初始化，加速码书训练过程中的收敛。
• 采用两阶段训练流程：先进行码书学习，再通过校准进行码书参数微调。
• 实现高效的GPU和CPU内核用于AQLM推理，利用更小的8比特码书提升缓存利用率并减少内存占用。
• 支持同质量化格式，无需混合稀疏-量化组件，简化部署。

实验结果

研究问题

• RQ1 多码书量化（MCQ）技术能否有效适配至大语言模型权重的后训练量化，以实现极端压缩？
• RQ2 与标准直接量化相比，实例感知、输入自适应的码书优化是否能显著提升2比特大语言模型量化中的准确率？
• RQ3 跨层块的联合优化是否能在保持低比特宽度压缩的同时提升模型性能？
• RQ4 考虑到计算复杂度，AQLM在通用硬件上是否具备实时推理的实际可行性？
• RQ5 AQLM的性能如何随校准集大小和码书配置（如码宽、码书数量）而变化？

主要发现

• AQLM在2比特设置下达到最先进性能，优于先前方法如QuIP#，在Llama 2 7B、13B和70B模型上表现更优。
• 在Llama 2 70B模型上，AQLM在Wiki2数据集上以每参数2.07比特实现困惑度7.98，显著提升了极端压缩下的准确率。
• 在GPU（Nvidia RTX 3090）上，AQLM相比FP16实现最高3.05倍加速；在CPU（Intel i9，8核）上实现最高4.07倍加速，适用于70B模型推理。
• 与FP16相比，该方法将内存占用减少高达8倍，同时保持或超过FP16的推理速度。
• 微调码书参数对准确率影响最大，当校准集大小达到4096个序列时，超过256个序列后收益递减。
• 使用多个较小码书（如2×8-bit）可提升GPU缓存利用率并实现更高加速，尽管存在轻微准确率损失。