RoSTE: An Efficient Quantization-Aware Supervised Fine-Tuning Approach for Large Language Models

📄 arXiv: 2502.09003v3 📥 PDF

作者: Quan Wei, Chung-Yiu Yau, Hoi-To Wai, Yang Katie Zhao, Dongyeop Kang, Youngsuk Park, Mingyi Hong

分类: cs.LG, cs.AI

发布日期: 2025-02-13 (更新: 2025-06-06)

备注: accepted by ICML 2025

🔗 代码/项目: GITHUB

💡 一句话要点

提出RoSTE算法，实现高效的大语言模型量化感知监督微调

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

关键词: 大语言模型 量化 监督微调 量化感知训练 激活值旋转 低比特量化 模型压缩 RoSTE算法

📋 核心要点

  1. 现有方法先微调再量化，忽略了二者协同作用，导致量化后的大语言模型性能次优。
  2. RoSTE算法结合量化感知监督微调和自适应旋转策略，减少激活异常值，提升量化效果。
  3. 实验表明，RoSTE在多种模型和任务上优于现有后量化方法，证明了其有效性。

📝 摘要（中文）

监督微调是将预训练大语言模型(LLMs)适配到下游任务的标准方法。量化最近被研究作为一种高效的LLM部署后训练技术。为了获得量化的微调LLM，传统的流程通常首先微调预训练模型，然后进行后训练量化。由于未能利用微调和量化之间的协同作用，这通常会产生次优的性能。为了有效地实现LLM中权重、激活和KV缓存的低比特量化，我们提出了一种名为Rotated Straight-Through-Estimator (RoSTE)的算法，该算法结合了量化感知监督微调(QA-SFT)和自适应旋转策略，该策略识别有效的旋转配置以减少激活异常值。我们通过分析RoSTE应用于过度参数化的最小二乘量化训练问题时的预测误差，提供了关于RoSTE的理论见解。我们的发现表明，预测误差与收敛权重的量化误差成正比，这可以通过优化的旋转配置有效地管理。在不同规模的Pythia、Qwen和Llama模型上的实验证明了RoSTE的有效性。与现有的后SFT量化基线相比，我们的方法在各种任务和不同的LLM架构上始终如一地实现了卓越的性能。我们的代码可在https://github.com/OptimAI-Lab/RoSTE获得。

🔬 方法详解

问题定义：论文旨在解决大语言模型（LLM）在量化后性能下降的问题。传统的微调后量化方法未能充分利用微调和量化之间的协同效应，导致量化后的模型性能不佳。激活值中的异常值是导致量化误差增大的一个重要因素。

核心思路：论文的核心思路是提出一种量化感知的监督微调（QA-SFT）方法，并在微调过程中自适应地旋转激活空间，以减少激活值中的异常值。通过在微调阶段就考虑到量化的影响，并优化激活分布，从而提高量化后模型的性能。

技术框架：RoSTE算法主要包含两个核心部分：量化感知监督微调（QA-SFT）和自适应旋转策略。QA-SFT在微调过程中模拟量化操作，使模型适应量化的影响。自适应旋转策略则通过搜索最佳旋转配置，减少激活值中的异常值。整个流程是在标准的监督微调框架下进行的，但加入了量化感知和旋转优化的模块。

关键创新：RoSTE的关键创新在于将量化感知和自适应旋转策略结合起来，在微调阶段就优化模型的量化性能。与传统的微调后量化方法相比，RoSTE能够更好地利用微调和量化之间的协同效应，从而获得更高的量化模型性能。此外，论文还提供了RoSTE的理论分析，证明了其有效性。

关键设计：RoSTE的关键设计包括：1) 使用Straight-Through Estimator (STE)来近似量化操作的梯度，从而实现量化感知的微调；2) 设计了一种自适应旋转策略，通过搜索最佳旋转角度来最小化激活值的量化误差；3) 理论分析表明，预测误差与量化权重的量化误差成正比，通过优化旋转配置可以有效控制量化误差。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RoSTE在Pythia、Qwen和Llama等不同规模的模型上均取得了显著的性能提升。与现有的后量化方法相比，RoSTE在各种任务上始终如一地表现出更优的性能。例如，在某些任务上，RoSTE能够将量化模型的性能提升到接近甚至超过全精度模型的水平，证明了其在低比特量化方面的有效性。

🎯 应用场景

RoSTE算法可应用于各种需要高效部署的大语言模型场景，例如移动设备、边缘计算和资源受限的服务器。通过降低模型大小和计算复杂度，RoSTE能够使LLM在这些平台上更高效地运行，从而促进LLM在更广泛的应用领域中的部署和应用。该技术对于开发低功耗、高性能的AI应用具有重要意义。

📄 摘要（原文）

Supervised fine-tuning is a standard method for adapting pre-trained large language models (LLMs) to downstream tasks. Quantization has been recently studied as a post-training technique for efficient LLM deployment. To obtain quantized fine-tuned LLMs, conventional pipelines would first fine-tune the pre-trained models, followed by post-training quantization. This often yields suboptimal performance as it fails to leverage the synergy between fine-tuning and quantization. To effectively realize low-bit quantization of weights, activations and KV caches in LLMs, we propose an algorithm named Rotated Straight-Through-Estimator (RoSTE), which combines quantization-aware supervised fine-tuning (QA-SFT) with an adaptive rotation strategy that identifies an effective rotation configuration to reduce activation outliers. We provide theoretical insights on RoSTE by analyzing its prediction error when applied to an overparameterized least square quantized training problem. Our findings reveal that the prediction error is directly proportional to the quantization error of the converged weights, which can be effectively managed through an optimized rotation configuration. Experiments on Pythia, Qwen and Llama models of different sizes demonstrate the effectiveness of RoSTE. Compared to existing post-SFT quantization baselines, our method consistently achieves superior performances across various tasks and different LLM architectures. Our code is available at https://github.com/OptimAI-Lab/RoSTE.