SKIM: Any-bit Quantization Pushing The Limits of Post-Training Quantization

📄 arXiv: 2412.04180v2 📥 PDF

作者: Runsheng Bai, Bo Liu, Qiang Liu

分类: cs.LG

发布日期: 2024-12-05 (更新: 2024-12-07)

💡 一句话要点

SKIM：提出任意比特量化方法，突破后训练量化的性能极限

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 后训练量化 低比特量化 K-means聚类 贪婪算法 比特分配 大型语言模型 模型压缩

📋 核心要点

  1. 现有量化方法在低比特位下性能显著下降，且仅提供有限的、手动调整的特定比特位解决方案。
  2. SKIM方法通过贪婪算法进行比特分配，并引入可训练缩放向量优化K-means聚类，提升量化性能。
  3. 实验表明，SKIM方法显著缩小了3比特量化LLaMA模型与全精度模型之间的性能差距，平均降低了16.3%的困惑度。

📝 摘要（中文）

大型语言模型（LLMs）在各种任务中表现出令人印象深刻的性能，但部署它们进行推理带来了挑战。它们对资源的高需求通常需要复杂、昂贵的多GPU流水线，或者使用更小、能力较弱的模型。量化提供了一种有希望的解决方案，它利用较低的精度进行模型存储，但现有方法通常在较低的精度水平上经历显著的性能下降。此外，它们通常只提供一组有限的、特定比特级别的解决方案，其中许多都经过了广泛的手动调整。为了应对这些挑战，我们提出了一种名为SKIM的新方法：具有混合精度的缩放K-means聚类。我们的方法引入了两项新技术：1. 一种贪婪算法，用于解决权重通道上近似最优的比特分配；2. 一个用于不可微K-means聚类的可训练缩放向量。这些技术显著提高了性能，并且可以适应任何给定的比特。值得注意的是，在模型困惑度方面，我们的方法将3比特量化的LLaMA模型与其全精度模型之间的差距平均缩小了16.3%。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）在低比特量化时性能显著下降的问题，以及现有量化方法依赖手动调整、仅支持特定比特位的问题。现有方法的痛点在于无法在保证性能的同时，实现任意比特的灵活量化。

核心思路：论文的核心思路是通过优化比特分配和改进K-means聚类算法，实现在任意比特位下都能保持较高性能的量化。通过贪婪算法进行比特分配，使得不同权重通道可以根据其重要性分配不同的比特数。引入可训练的缩放向量，优化K-means聚类的效果，从而提升量化模型的精度。

技术框架：SKIM方法主要包含两个阶段：比特分配和量化。首先，使用贪婪算法对权重通道进行比特分配，确定每个通道的量化比特数。然后，使用带有可训练缩放向量的K-means聚类算法对权重进行量化。可训练缩放向量用于调整聚类中心，从而优化量化效果。整个过程无需手动调整，可以自适应地选择最佳的量化方案。

关键创新：SKIM方法的关键创新在于：1. 提出了一种贪婪算法，用于解决权重通道上近似最优的比特分配问题，实现了自适应的比特分配。2. 引入了一个可训练的缩放向量，用于优化不可微的K-means聚类算法，提升了量化精度。这两个创新点使得SKIM方法能够在任意比特位下实现高性能的量化。

关键设计：贪婪比特分配算法通过迭代的方式，每次选择降低困惑度最大的通道进行比特分配。可训练缩放向量通过梯度下降进行优化，损失函数为量化后的模型与原始模型输出之间的差异。K-means聚类的初始化采用K-means++算法，以提高聚类的稳定性和准确性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SKIM方法在LLaMA模型上取得了显著的性能提升。实验结果表明，SKIM方法将3比特量化的LLaMA模型与其全精度模型之间的困惑度差距平均缩小了16.3%。这一结果表明，SKIM方法能够有效降低低比特量化带来的性能损失，使得量化模型在保持较高性能的同时，显著降低了资源消耗。

🎯 应用场景

SKIM方法可应用于各种需要低精度推理的大型语言模型部署场景，例如移动设备、边缘计算设备等资源受限的环境。通过降低模型存储空间和计算复杂度，SKIM方法可以有效降低部署成本，提高推理速度，并为更广泛的用户提供高质量的AI服务。该方法还有助于推动LLM在资源受限场景下的应用，例如智能家居、自动驾驶等。

📄 摘要（原文）

Large Language Models (LLMs) exhibit impressive performance across various tasks, but deploying them for inference poses challenges. Their high resource demands often necessitate complex, costly multi-GPU pipelines, or the use of smaller, less capable models. While quantization offers a promising solution utilizing lower precision for model storage, existing methods frequently experience significant performance drops at lower precision levels. Additionally, they typically provide only a limited set of solutions at specific bit levels, many of which are extensively manually tuned. To address these challenges, we propose a new method called SKIM: Scaled K-means clustering wIth Mixed precision. Our approach introduces two novel techniques: 1. A greedy algorithm to solve approximately optimal bit allocation across weight channels, and 2. A trainable scaling vector for non-differentiable K-means clustering. These techniques substantially improve performance and can be adapted to any given bit. Notably, in terms of model perplexity, our method narrows the gap between 3-bit quantized LLaMA models and their full precision counterparts by 16.3% on average.