BitDistiller: Unleashing the Potential of Sub-4-Bit LLMs via Self-Distillation

📄 arXiv: 2402.10631v1 📥 PDF

作者: Dayou Du, Yijia Zhang, Shijie Cao, Jiaqi Guo, Ting Cao, Xiaowen Chu, Ningyi Xu

分类: cs.CL

发布日期: 2024-02-16

🔗 代码/项目: GITHUB

💡 一句话要点

提出BitDistiller以提升超低精度LLM的性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 权重量化 知识蒸馏 量化感知训练 超低精度 模型优化 自然语言处理

📋 核心要点

  1. 现有大型语言模型在超低精度下的性能提升面临重大挑战，尤其是在权重量化过程中，模型的保真度和计算效率难以兼顾。
  2. BitDistiller框架通过结合量化感知训练和知识蒸馏，采用非对称量化和置信度感知Kullback-Leibler散度目标，提升了超低精度模型的性能。
  3. 实验结果显示，BitDistiller在3位和2位配置下的通用语言理解和复杂推理任务上，性能显著优于现有方法，且训练成本更低。

📝 摘要（中文）

大型语言模型（LLMs）的扩展在自然语言处理领域取得了显著进展，但也带来了部署挑战。权重量化成为降低内存和计算需求的广泛解决方案。本文提出了BitDistiller框架，结合了量化感知训练（QAT）和知识蒸馏（KD），以提升超低精度（低于4位）的LLM性能。具体而言，BitDistiller首先采用定制的非对称量化和裁剪技术，以最大程度地保留量化权重的保真度，然后提出了一种新颖的置信度感知Kullback-Leibler散度（CAKLD）目标，以自蒸馏的方式实现更快的收敛和更优的模型性能。实证评估表明，BitDistiller在3位和2位配置下显著超越现有方法，在通用语言理解和复杂推理基准上表现优异。值得注意的是，BitDistiller在数据和训练资源需求上更具成本效益。

🔬 方法详解

问题定义：本文旨在解决超低精度（低于4位）大型语言模型在权重量化过程中面临的性能下降问题。现有方法在保持模型精度与降低计算需求之间存在矛盾，导致模型在实际应用中的有效性受限。

核心思路：论文提出的BitDistiller框架通过结合量化感知训练（QAT）与知识蒸馏（KD），利用非对称量化和置信度感知Kullback-Leibler散度（CAKLD）目标，旨在最大程度地保留量化权重的保真度，同时加速模型收敛。

技术框架：BitDistiller的整体架构包括两个主要模块：首先是量化模块，采用非对称量化和裁剪技术；其次是自蒸馏模块，利用CAKLD目标进行模型训练。整个流程通过迭代优化实现模型性能的提升。

关键创新：BitDistiller的核心创新在于引入了置信度感知Kullback-Leibler散度目标，这一设计使得模型在自蒸馏过程中能够更有效地利用知识，从而在超低精度下实现更好的性能。与传统方法相比，BitDistiller在保真度和收敛速度上具有显著优势。

关键设计：在技术细节方面，BitDistiller采用了定制的非对称量化策略，以适应不同的权重分布。此外，CAKLD目标的设计考虑了模型输出的置信度，使得训练过程更加高效。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，BitDistiller在3位和2位配置下的性能显著提升，尤其在通用语言理解和复杂推理基准上，超越了现有的主流方法。具体而言，BitDistiller在相同数据和训练资源下，性能提升幅度达到XX%（具体数据未知），展现出更高的成本效益。

🎯 应用场景

BitDistiller的研究成果在多个领域具有潜在应用价值，尤其是在资源受限的环境中，如移动设备和边缘计算。通过提升超低精度模型的性能，该框架能够使得大型语言模型在实际应用中更加高效，降低计算和存储成本，推动自然语言处理技术的普及与应用。

📄 摘要（原文）

The upscaling of Large Language Models (LLMs) has yielded impressive advances in natural language processing, yet it also poses significant deployment challenges. Weight quantization has emerged as a widely embraced solution to reduce memory and computational demands. This paper introduces BitDistiller, a framework that synergizes Quantization-Aware Training (QAT) with Knowledge Distillation (KD) to boost the performance of LLMs at ultra-low precisions (sub-4-bit). Specifically, BitDistiller first incorporates a tailored asymmetric quantization and clipping technique to maximally preserve the fidelity of quantized weights, and then proposes a novel Confidence-Aware Kullback-Leibler Divergence (CAKLD) objective, which is employed in a self-distillation manner to enable faster convergence and superior model performance. Empirical evaluations demonstrate that BitDistiller significantly surpasses existing methods in both 3-bit and 2-bit configurations on general language understanding and complex reasoning benchmarks. Notably, BitDistiller is shown to be more cost-effective, demanding fewer data and training resources. The code is available at https://github.com/DD-DuDa/BitDistiller.