Oh! We Freeze: Improving Quantized Knowledge Distillation via Signal Propagation Analysis for Large Language Models

📄 arXiv: 2403.18159v2 📥 PDF

作者: Kartikeya Bhardwaj, Nilesh Prasad Pandey, Sweta Priyadarshi, Kyunggeun Lee, Jun Ma, Harris Teague

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-03-26 (更新: 2024-03-28)

备注: Accepted at Practical ML for Low Resource Settings Workshop at ICLR 2024

💡 一句话要点

提出KD-QAT方法以解决低比特量化模型的稳定性问题

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

关键词: 量化感知微调 知识蒸馏 低比特量化 梯度传播分析 自然语言处理 边缘计算 模型稳定性

📋 核心要点

  1. 现有的低比特量化模型在推理时存在稳定性差和易受量化误差影响的问题，限制了其实际应用。
  2. 论文提出了一种基于知识蒸馏的量化感知微调技术（KD-QAT），并通过分析梯度传播来提高其稳定性。
  3. 实验结果显示，使用ov-freeze技术的7B LLaMAv2-Chat模型在4位量化下的性能接近浮点精度，准确率损失小于0.7%。

📝 摘要（中文）

大型生成模型如大型语言模型（LLMs）和扩散模型在自然语言处理和计算机视觉领域取得了革命性进展。然而，它们的推理速度慢、计算和内存需求高，使得在边缘设备上部署变得困难。本研究提出了一种轻量级的量化感知微调技术（KD-QAT），旨在通过知识蒸馏提高4位权重量化LLMs的性能，尤其是在设备聊天应用场景中。我们通过实证研究训练过程中的梯度传播，深入分析KD-QAT的稳定性，并提出了ov-freeze技术以稳定该过程。实验表明，在4位量化水平下，使用ov-freeze的7B LLaMAv2-Chat模型在常识推理基准测试中仅损失0.7%的准确率，接近浮点精度表现。

🔬 方法详解

问题定义：本研究旨在解决低比特量化模型在训练和推理过程中稳定性差的问题，现有方法在面对量化误差时表现脆弱，影响了模型的实际应用效果。

核心思路：提出了一种量化感知微调技术（KD-QAT），通过对训练过程中的梯度传播进行深入分析，识别并缓解低比特量化带来的不稳定性，从而提升模型的整体性能。

技术框架：整体框架包括数据准备、KD-QAT训练过程和ov-freeze稳定化步骤。首先使用常见数据集进行微调，然后通过分析梯度传播来优化训练过程，最后应用ov-freeze技术以增强稳定性。

关键创新：最重要的创新在于通过梯度传播分析识别KD-QAT的脆弱性，并提出ov-freeze技术来有效稳定训练过程，与传统方法相比，显著提高了低比特量化模型的性能。

关键设计：在模型训练中，采用特定的损失函数来平衡知识蒸馏和量化误差的影响，同时在网络结构上进行调整，以适应4位量化的需求，确保模型在推理时的高效性和准确性。

📊 实验亮点

实验结果表明，使用ov-freeze技术的7B LLaMAv2-Chat模型在4位量化下的准确率损失小于0.7%，接近浮点精度，显著优于传统量化方法，展示了KD-QAT在实际应用中的有效性。

🎯 应用场景

该研究的潜在应用领域包括边缘设备上的自然语言处理任务，如智能助手和聊天机器人。通过提高低比特量化模型的性能，能够在资源受限的环境中实现高效的语言理解和生成，具有重要的实际价值和广泛的市场前景。

📄 摘要（原文）

Large generative models such as large language models (LLMs) and diffusion models have revolutionized the fields of NLP and computer vision respectively. However, their slow inference, high computation and memory requirement makes it challenging to deploy them on edge devices. In this study, we propose a light-weight quantization aware fine tuning technique using knowledge distillation (KD-QAT) to improve the performance of 4-bit weight quantized LLMs using commonly available datasets to realize a popular language use case, on device chat applications. To improve this paradigm of finetuning, as main contributions, we provide insights into stability of KD-QAT by empirically studying the gradient propagation during training to better understand the vulnerabilities of KD-QAT based approaches to low-bit quantization errors. Based on our insights, we propose ov-freeze, a simple technique to stabilize the KD-QAT process. Finally, we experiment with the popular 7B LLaMAv2-Chat model at 4-bit quantization level and demonstrate that ov-freeze results in near floating point precision performance, i.e., less than 0.7% loss of accuracy on Commonsense Reasoning benchmarks.