本文提出Quantized Zeroth-order Optimization (QZO)，通过扰动量化尺度参数并结合方向导数裁剪，在量化神经网络上实现零阶优化微调，将内存使用减少18倍以上，并在LLMs和Stable Diffusion上展示出显著的内存效率和一定的性能提升。
Large Language Model, Generative AI, Efficiency, Pre-training, Fine-tuning
Sifeng Shang, Jiayi Zhou, Chenyu Lin, Minxian Li, Kaiyang Zhou
Hong Kong Baptist University, Nanjing University of Science and Technology
Generated by grok-3
Background Problem
随着大型语言模型（LLMs）规模的指数级增长，GPU内存成为适配这些模型到下游任务的瓶颈。传统微调方法需要大量内存来存储模型权重、梯度和优化器状态，例如微调一个Llama-7B模型需要56GB GPU内存。本文旨在通过最小化模型权重、梯度和优化器状态的内存使用，探索内存高效训练的极限，解决资源受限环境下LLMs微调的难题。
Method
本文提出了一种新颖的方法——Quantized Zeroth-order Optimization (QZO)，用于在量化神经网络上应用零阶优化（ZO），以实现内存消耗的最大化减少。
- 核心思想：通过ZO消除梯度和优化器状态的内存需求，同时通过模型量化（如从bfloat16转为int4）减少权重内存占用。QZO通过扰动量化的连续尺度参数（而非离散权重）来估计梯度，避免了量化和去量化的复杂操作。
- 具体实现：
- Quantized Simultaneous Perturbation Stochastic Approximation (Q-SPSA)：扩展传统SPSA方法，仅对量化尺度参数进行扰动，保持离散权重不变，通过两次前向传递估计梯度方向（公式如 ）。
- Directional Derivative Clipping (DDC)：通过裁剪方向导数（公式如 ）减少梯度估计的方差，稳定训练过程，并提供理论证明其有效性。
- 关键问题与批评：虽然QZO在概念上创新，但其依赖于量化方法的质量，若量化误差较大，前向传递中的噪声可能显著影响梯度估计精度。此外，论文未充分探讨如何在不同量化方法间优化Q-SPSA的扰动策略，可能导致性能不稳定。
Experiment
实验在多个LLMs（如OPT-6.7B、Llama-2-7B、Llama-3.1-8B、Llama-2-13B）和Stable Diffusion 3.5 Large上进行，覆盖NLP分类和生成任务以及文本到图像生成任务。
- 数据集与设置：NLP任务使用SST-2、SuperGLUE子集和SQuAD数据集，Stable Diffusion使用Styled Image Dataset。量化方法包括GPTQ（4-bit）和AQLM（2-bit），实验在单张Nvidia RTX 4090 GPU（24GB）上完成，内存分析基于SST-2数据集。
- 结果分析：
- 在4-bit LLMs上，QZO显著优于Zero-Shot-Q（量化零样本模型），与MeZO（未量化ZO方法）性能相当，但内存使用减少约3倍（如Llama-2-7B上QZO仅用4.99GB vs MeZO的14.79GB）。
- 在2-bit极端量化下，QZO仍显著优于Zero-Shot-Q，显示出在边缘设备上的潜力。
- 在Stable Diffusion上，QZO将内存需求从86.43GB降至12.4GB，但生成图像质量与LLM任务相比提升有限，论文归因于噪声调度冲突。
- 评价与批评：实验设置较为全面，覆盖多种模型和量化方法，但结果显示QZO与全精度微调（如LoRA）仍有较大性能差距，特别是在扩散模型上，改进效果不明显。实验未充分探讨超参数（如扰动尺度、裁剪阈值）对结果的影响，缺乏调优细节。此外，Stable Diffusion实验仅提供定性结果，未给出量化指标，难以评估实际效果。总体而言，内存效率提升显著，但性能提升的适用性有限。
Further Thoughts
QZO的内存效率优势令人印象深刻，尤其是在资源受限环境下的潜力。然而，其性能与全精度微调的差距提示我们需要在梯度估计精度上进一步突破。未来的研究可以探索结合更先进的量化方法（如自适应量化）来减少初始量化误差对ZO的影响。此外，针对扩散模型的噪声调度冲突问题，可以尝试设计与扩散过程兼容的扰动策略，例如调整ZO引入的噪声分布以匹配扩散模型的预定噪声调度。另一个有趣的方向是将QZO与其他内存高效方法（如LoRA）结合，探索是否能在保持低内存占用的同时进一步缩小性能差距。最后，QZO在边缘设备上的应用潜力值得深入研究，尤其是在联邦学习或隐私保护场景中，结合隐私保护机器学习技术可能带来新的突破。