Leveraging Coordinate Momentum in SignSGD and Muon: Memory-Optimized Zero-Order

TL;DR

The paper targets memory-efficient fine-tuning of large language models by replacing costly first-order updates with zero-order optimization. It introduces JAGUAR SignSGD and JAGUAR Muon, both leveraging a momentum-based zero-order gradient estimator that uses only $2d+1$ parameters and $\mathcal{O}(1)$ function evaluations per iteration, and it extends these ideas to matrix-parameter optimization. The authors provide rigorous convergence guarantees in the stochastic non-convex setting and validate the methods empirically on challenging LLM fine-tuning tasks, demonstrating competitive convergence and substantial memory savings compared to traditional FO methods. The results indicate that zero-order optimization with memory- and computation-efficient momentum can enable effective resource-constrained LLM adaptation, with practical implications for edge devices and distributed settings.

Abstract

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

Figures (1)

• Figure 1: Test accuracy of JAGUAR SignSGD on SST-2 for RoBERTa-large with LoRA for different values of $\beta$.

Theorems & Definitions (13)

• lemma 1
• theorem 1
• corollary 1
• theorem 2
• corollary 2
• remark 1
• proof
• proof
• lemma 2
• proof