🤖 machine learning

Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration

This paper proposes ZO-MOPI, a zeroth-order optimization method for fine-tuning large language models that accelerates convergence by replacing full orthogonalization with a partial orthogonalization strategy using power iteration and momentum-based subspace projection, achieving 1.5x to 4x faster convergence than state-of-the-art ZO-Muon while maintaining competitive accuracy.

Original authors: Jiahe Chen, Ziye Ma

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiahe Chen, Ziye Ma

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Tuning a Giant Radio Without a Manual

Imagine you have a massive, complex radio (a Large Language Model or LLM) with billions of knobs. You want to tune it to play a specific song perfectly (fine-tuning it for a new task).

Usually, engineers use a "First-Order" method: they look at the knobs, feel the vibration, and know exactly which way to turn them to get a better sound. This is fast, but it requires a huge, expensive control room (lots of computer memory) to hold all the instructions. If you are trying to tune this radio on a small device like a smartphone or a car computer, you don't have that much memory.

Zeroth-Order (ZO) Optimization is the "blind" method. Instead of feeling the vibration, you just guess a direction, turn the knobs, and listen to see if the music got better or worse. You don't need the big control room, so it fits on small devices. However, because you are guessing, it's incredibly slow and noisy. It's like trying to find the perfect station by randomly spinning the dial and hoping you land on the right frequency.

The Problem: The "Noisy" Signal

The researchers noticed that while this "blind" method saves memory, the information it gathers is very messy.

The Real Signal: In a perfect world, the "best" directions to turn the knobs are clear and strong (like a loud, clear voice).
The ZO Reality: Because the method relies on random guessing, the signal is filled with static and noise. It's like trying to hear a whisper in a hurricane.

There was a new tool called Muon that tried to help. Muon is like a smart assistant that looks at all the knobs and says, "Okay, let's organize these directions so we don't waste energy on the weak ones." It tries to make all the directions equally strong.

The Catch: Muon was designed for the "perfect world" (First-Order). When the researchers tried to use Muon on the "noisy" Zeroth-Order data, it made things worse. Because the data was so full of static, Muon tried to boost everything, including the noise. It was like turning up the volume on a radio that's just playing static; the noise became louder than the music.

The Solution: "Partial Orthogonalization" (The Selective Filter)

The authors, Jiahe Chen and Ziye Ma, came up with a new approach called ZO-MOPI. Their secret sauce is Partial Orthogonalization.

Instead of trying to organize every direction (which includes the noisy, useless ones), they decided to only focus on the top few directions that are actually strong and clear. They ignore the rest of the noise.

Think of it like this:

Old Way (Muon): You have a bucket of water mixed with sand, dirt, and gold dust. You try to separate everything perfectly. You end up wasting a lot of time sorting through the dirt.
New Way (ZO-MOPI): You realize the gold dust is heavy and sinks to the bottom. You just scoop up the bottom layer (the strong signals) and ignore the floating dirt (the noise). You get the gold much faster.

How They Did It: The "Streaming Power" Trick

To make this selective filtering work, they used a technique called Streaming Power Iteration (SPI).

Imagine you are trying to find the strongest wind in a stormy field.

The Old Method: You stand still and measure every single gust of wind, then try to calculate the average direction. It takes forever and the wind is too chaotic.
The New Method (SPI): You hold a flag. You only pay attention to the gusts that push the flag the hardest. You ignore the tiny breezes that barely move it. By focusing only on the strong pushes, you can quickly figure out which way the wind is really blowing.

They also added a "Momentum" feature. This is like remembering where the wind was blowing a few seconds ago. Even if one gust is a fluke (noise), your memory of the previous gusts helps you stay steady and not get confused by the chaos.

The Results: Faster Tuning, Same Memory

The researchers tested this on giant AI models (like OPT-13B and LLaMA3-8B) using standard language tests (SuperGLUE).

Speed: Their new method was 1.5 to 4 times faster than the current best "blind" tuning methods. It reached the same level of accuracy in much less time.
Accuracy: It ended up with the same (or slightly better) final quality as the other methods.
Memory: It kept the same low memory usage, meaning it can still run on small devices.

Summary

The paper introduces a smarter way to tune AI models on devices with limited memory. Instead of trying to fix every single noisy guess, the new method (ZO-MOPI) acts like a filter, focusing only on the strongest, most reliable signals and ignoring the static. This allows the AI to learn much faster without needing a supercomputer.

Technical Summary: Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration

Problem Statement
Zeroth-order (ZO) optimization has emerged as a critical technique for fine-tuning Large Language Models (LLMs) on memory-constrained edge devices, as it avoids the memory-intensive backpropagation required by first-order (FO) methods. However, ZO methods suffer from high variance in gradient estimates due to Monte Carlo perturbations, leading to significantly slower convergence compared to FO methods.

Recent attempts to accelerate ZO optimization have focused on low-dimensional subspace search (e.g., LOZO, SubZero) and spectral optimizers like Muon, which exploit matrix geometry through orthogonalization. While Muon outperforms AdamW in FO settings by orthogonalizing gradients to exploit weak spectral directions, the authors observe that applying Muon directly to ZO settings yields limited gains. The core issue identified is that ZO gradient estimates are inherently noisy and unreliable. Unlike FO gradients, which are concentrated in dominant singular directions, ZO estimates exhibit a "heavier and flatter spectral tail" (Figure 1). Consequently, Muon's standard Newton-Schulz (NS) orthogonalization, which equalizes energy across the entire active spectrum, inadvertently amplifies these noisy, inaccurate directions, turning estimation errors into harmful update components.

Methodology: ZO-MOPI
The authors propose ZO-MOPI (Zeroth-Order Momentum with Partial Orthogonalization), a method designed to accelerate ZO optimization by selectively extracting dominant spectral directions while suppressing noise. The approach integrates three key components:

Partial Orthogonalization via Streaming Power Iteration (SPI):
Instead of performing full-spectrum orthogonalization (like NS iterations), ZO-MOPI employs a "partial orthogonalization" strategy. It focuses exclusively on the top- $k$ dominant singular directions, discarding the noisy spectral tail.
- Mechanism: The method replaces the NS procedure with a faster, more concentrated Streaming Power Iteration (SPI). SPI tracks the leading $k$ -dimensional subspace over training steps.
- Efficiency: By caching the right singular subspace ( $V_{t-1}$ ) and using it as a warm start for the next iteration, SPI avoids full SVD or expensive matrix multiplications. The computational cost is reduced to $O(rnk + nk^2)$ , where $k \ll r \ll \min(m, n)$ , compared to the higher cost of full-spectrum methods.
- Stability: The method relies on the observation that dominant spectral directions evolve smoothly under momentum smoothing, making the previous subspace a reliable initialization for the current step.
Subspace ZO Estimation:
To further reduce variance, the gradient estimation is constrained to a low-rank subspace. The random perturbation matrix $Z$ is decomposed into $AB$, where $A$ is a fixed subspace matrix and $B$ is a random perturbation within that subspace. This reduces the effective search dimension from $mn$ to $rn$, significantly lowering the variance of the gradient estimator.
Projected Momentum:
Momentum is maintained within the reduced $r \times n$ coordinate system rather than the full parameter space to save memory. A critical component of the method is the projection of momentum when the subspace $A$ is periodically refreshed (every $\nu$ iterations). The momentum $M_{old}$ is projected onto the new subspace $A_{new}$ via $M_{new} \leftarrow A_{new}^\top A_{old} M_{old}$ .
- Necessity: The authors demonstrate that momentum is not merely for variance reduction but is essential for the stability of SPI. Without momentum, the dominant subspace fluctuates drastically due to noisy ZO estimates, causing the power iteration to become unstable.

Key Contributions

Identification of Spectral Mismatch: The paper highlights that full-spectrum orthogonalization (Newton-Schulz) is detrimental in ZO settings because it amplifies the noise inherent in ZO gradient estimates, whereas FO gradients are naturally concentrated.
Partial Orthogonalization: The introduction of a selective orthogonalization strategy that preserves only the most reliable spectral components (top- $k$ ) while ignoring the noisy tail.
Streaming Power Iteration (SPI): A lightweight, efficient algorithm for tracking dominant singular directions that leverages temporal continuity and warm-starting to avoid expensive computations.
Projected Momentum Integration: A mechanism to maintain momentum stability across subspace refreshes, which is shown to be critical for the efficacy of spectral tracking in noisy environments.

Experimental Results
The authors evaluate ZO-MOPI on LLM fine-tuning tasks using OPT-13B, LLaMA3-8B, and Gemma2-2B models across SuperGLUE benchmarks (SST-2, RTE, BoolQ, SQuAD).

Convergence Speed: ZO-MOPI achieves 1.5x to 4x faster convergence compared to the current state-of-the-art ZO-Muon. For instance, on the RTE and SQuAD tasks, ZO-Muon requires approximately 3x the wall-clock time to reach the same accuracy as ZO-MOPI.
Final Accuracy: The method achieves competitive or superior final accuracies compared to strong baselines like MeZO, LOZO, and ZO-Muon. In several cases (e.g., OPT-13B on SST-2 and RTE), ZO-MOPI sets new best ZO accuracy records.
Efficiency: The method maintains the memory efficiency of ZO fine-tuning, with only negligible additional memory overhead for storing the subspace and power iteration states.
Ablation Studies: Experiments confirm that removing momentum or using a subspace refresh interval that is too short significantly degrades performance, validating the necessity of the proposed momentum projection and lazy subspace sampling strategies.

Significance and Claims
The paper claims that ZO-MOPI successfully bridges the performance gap between ZO and FO optimization for LLM fine-tuning by addressing the specific spectral characteristics of noisy gradient estimates. The authors argue that the key to accelerating ZO optimization is not merely reducing variance through subspace projection, but also adopting a selective spectral alignment that avoids amplifying noise. By replacing aggressive full-spectrum orthogonalization with partial orthogonalization via power iteration, the method enables spectral optimizers to function effectively in the high-variance ZO regime. The work suggests that spectral structure can be better exploited even under noisy estimates, offering a path toward efficient, memory-constrained fine-tuning of large models without relying on backpropagation.