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

2025-06-04

AI Summary

To address the excessive memory and computational overhead in parameter-efficient fine-tuning (PEFT) of large language models (LLMs), this paper proposes two zeroth-order optimization algorithms: JAGUAR SignSGD and JAGUAR Muon. First, it establishes the first rigorous convergence guarantee for stochastic zeroth-order SignSGD. Second, it introduces the first Muon variant incorporating structural priors of parameter matrices—specifically, low-rank structure—and provides a universal convergence rate analysis under arbitrary noise. Both methods integrate sign-based gradient updates, coordinate-wise momentum, and LoRA adaptation. Empirically, they match or surpass the convergence quality of Adam/SGD on LLM fine-tuning tasks while reducing memory complexity to O(1) and requiring only a single function evaluation per iteration—significantly lowering resource consumption.

Technology Category

• Machine Learning: Optimization
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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 $ exttt{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 $ exttt{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

Problem

• Memory-efficient fine-tuning for large language models
• Reducing computational costs in zero-order optimization
• Convergence guarantees for SignSGD in stochastic cases

Innovation

• ZO momentum-based algorithm for SignSGD
• Matrix-structured ZO extension of Muon
• Memory-efficient zero-order LLM fine-tuning