通过随机子空间降维,实现大模型零阶微调的高效低内存训练。
Zeroth-Order Fine-Tuning of LLMs in Random Subspaces
- 在低秩子空间中进行零阶优化,降低梯度估计方差。
- 相比传统方法,收敛更快且性能提升显著,实验覆盖多种语言任务。
- 适合资源受限场景下的大模型微调,尤其关注内存效率的实践者。
大语言模型(LLMs)的微调在下游任务中表现优异,但随着模型规模增大,反向传播带来的内存开销日益严峻。零阶(ZO)优化通过前向传播估计梯度,提供了内存友好的替代方案,但其梯度估计方差通常随参数维度线性增长,对大模型构成挑战。本文提出随机子空间零阶(SubZero)优化方法,设计适用于大模型的低秩扰动机制,在显著降低内存消耗的同时提升训练性能。理论证明,该方法的梯度估计接近反向传播梯度,方差低于传统零阶方法,并与SGD结合可保证收敛。实验表明,SubZero在多种语言建模任务中优于标准零阶方法(如MeZO),实现更优的微调性能和更快的收敛速度。代码已开源:https://github.com/zimingyy/SubZero。
原文摘要 · Abstract (English)
Fine-tuning Large Language Models (LLMs) has proven effective for a variety of downstream tasks. However, as LLMs grow in size, the memory demands for backpropagation become increasingly prohibitive. Zeroth-order (ZO) optimization methods offer a memory-efficient alternative by using forward passes to estimate gradients, but the variance of gradient estimates typically scales linearly with the model's parameter dimension$\unicode{x2013}$a significant issue for LLMs. In this paper, we propose the random Subspace Zeroth-order (SubZero) optimization to address the challenges posed by LLMs' high dimensionality. We introduce a low-rank perturbation tailored for LLMs that significantly reduces memory consumption while improving training performance. Additionally, we prove that our gradient estimation closely approximates the backpropagation gradient, exhibits lower variance than traditional ZO methods, and ensures convergence when combined with SGD. Experimental results show that SubZero enhances fine-tuning performance and achieves faster convergence compared to standard ZO approaches like MeZO across various language modeling tasks. Code is available at https://github.com/zimingyy/SubZero.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。