提出零阶优化新方法,高效微调大模型且节省内存。
Leveraging Coordinate Momentum in SignSGD and Muon: Memory-Optimized Zero-Order
- 用坐标动量思想改进零阶优化,仅需常数次函数评估。
- 在多个大模型微调任务中达到或超过一阶方法的收敛精度。
- 适合资源受限场景下的大模型轻量化微调,尤其适配LoRA等技术。
微调大语言模型对下游任务适配至关重要,但传统一阶优化器如SGD和Adam存在随模型规模增长而急剧上升的内存与计算开销。本文研究零阶(ZO)优化方法作为更高效替代方案,尤其适用于参数高效微调技术如LoRA。提出$ exttt{JAGUAR SignSGD}$,一种基于零阶动量的算法,扩展了零阶SignSGD,参数量与标准零阶SGD相同,每轮迭代仅需$/mathcal{O}(1)$次函数评估。据我们所知,这是首个在随机零阶情形下为SignSGD建立严格收敛保证的研究。进一步提出$ exttt{JAGUAR Muon}$,一种新颖的零阶版Muon优化器,利用模型参数的矩阵结构,并在任意随机噪声下给出其收敛速率。在多个具有挑战性的大模型微调基准上进行大量实验,结果表明所提算法在收敛质量上达到或超过标准一阶方法,实现显著内存压缩。理论与实证结果确立了新型零阶优化方法在资源受限的大模型适应中的实用性与理论基础。代码开源于https://github.com/brain-mmo-lab/ZO_LLM。
原文摘要 · Abstract (English)
Fine-tuning Large Language Models (LLMs) is essential for adapting pre-trained models to downstream tasks. Yet traditional first-order optimizers such as Stochastic Gradient Descent (SGD) and Adam incur prohibitive memory and computational costs that scale poorly with model size. In this paper, we investigate zero-order (ZO) optimization methods as a memory- and compute-efficient alternative, particularly in the context of parameter-efficient fine-tuning techniques like LoRA. We propose $\texttt{JAGUAR SignSGD}$, a ZO momentum-based algorithm that extends ZO SignSGD, requiring the same number of parameters as the standard ZO SGD and only $\mathcal{O}(1)$ function evaluations per iteration. To the best of our knowledge, this is the first study to establish rigorous convergence guarantees for SignSGD in the stochastic ZO case. We further propose $\texttt{JAGUAR Muon}$, a novel ZO extension of the Muon optimizer that leverages the matrix structure of model parameters, and we provide its convergence rate under arbitrary stochastic noise. Through extensive experiments on challenging LLM fine-tuning benchmarks, we demonstrate that the proposed algorithms meet or exceed the convergence quality of standard first-order methods, achieving significant memory reduction. Our theoretical and empirical results establish new ZO optimization methods as a practical and theoretically grounded approach for resource-constrained LLM adaptation. Our code is available at https://github.com/brain-mmo-lab/ZO_LLM
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。