Powering Up Zeroth-Order Training via Subspace Gradient Orthogonalization

TL;DR

This paper introduces ZO-Muon, a novel zeroth-order optimization method that leverages subspace projection and spectral gradient orthogonalization to improve efficiency and accuracy in large-scale model fine-tuning.

Contribution

The paper proposes a unified framework of subspace gradient orthogonalization and introduces ZO-Muon, a new method that enhances zeroth-order optimization for large models.

Findings

ZO-Muon accelerates convergence on LLMs and ViTs.

ZO-Muon reduces query complexity by over 75% compared to MeZO.

ZO-Muon improves accuracy significantly in fine-tuning tasks.

Abstract

Zeroth-order (ZO) optimization provides a gradient-free alternative to first-order (FO) methods by estimating gradients via finite differences of function evaluations, and has recently emerged as a memory-efficient paradigm for fine-tuning large-scale models by avoiding backpropagation. However, ZO optimization has a fundamental tension between accuracy and query efficiency. In this work, we show that ZO optimization can be substantially improved by unifying two complementary principles: (i) a projection-based subspace view that reduces gradient estimation variance by exploiting the intrinsic low-rank structure of model updates, and (ii) Muon-style spectral optimization that applies gradient orthogonalization to extract informative spectral structure from noisy ZO gradients. These findings form a unified framework of subspace gradient orthogonalization, which we instantiate in a new method, ZO-Muon, admitting a natural interpretation as a low-rank Muon optimizer in the ZO setting. Extensive experiments on large language models (LLMs) and vision transformers (ViTs) demonstrate that ZO-Muon significantly accelerates convergence and achieves a win-win improvement in accuracy and query/runtime efficiency. Notably, compared to the popular MeZO baseline, ZO-Muon requires only 24.7% of the queries to reach the same SST-2 performance for LLM fine-tuning, and improves accuracy by 25.1% on ViT-B fine-tuning on CIFAR-100.