🤖 machine learning

Learning a Zeroth-Order Optimizer for Fine-Tuning LLMs

This paper introduces ZO-Finetuner, a learning-based zeroth-order optimizer that automatically learns efficient perturbation strategies through a memory-efficient design, enabling one-time per-model training and superior performance across diverse downstream tasks compared to existing static zeroth-order methods.

Original authors: Kairun Zhang, Haoyu Li, Yanjun Zhao, Yifan Sun, Huan Zhang

Published 2026-06-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kairun Zhang, Haoyu Li, Yanjun Zhao, Yifan Sun, Huan Zhang

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

Imagine you have a massive, incredibly complex library (a Large Language Model, or LLM) that knows how to write, reason, and chat. You want to teach this library a new, specific skill, like writing legal contracts or solving math problems. This process is called "fine-tuning."

Usually, teaching this library involves a very expensive method called "backpropagation." Think of this like a teacher who, after every sentence the student writes, has to walk through the entire library, check every single book, and calculate exactly how to adjust the student's brain. It's accurate, but it requires a huge amount of memory and energy, often making it impossible to do on standard computers.

To fix this, researchers developed a "Zeroth-Order" method (like MeZO). Instead of walking through the whole library, this method is like a blindfolded explorer. The explorer makes a tiny guess (a perturbation), sees if the result got better or worse, and then makes another tiny guess in the opposite direction. By comparing the two results, they can figure out which way to go without ever seeing the whole picture. This saves a massive amount of memory.

The Problem: The "One-Size-Fits-All" Map
The existing blindfolded explorers use a fixed, random strategy. They guess directions based on a standard, boring rule (like rolling a fair die). The paper argues this is inefficient. Sometimes, the library has specific "rooms" (blocks of parameters) where you need to be very precise, and other rooms where you can be more wild. A fixed strategy doesn't know this; it treats every part of the library the same.

The Solution: ZO Fine-tuner (The Smart Explorer)
The authors created ZO Fine-tuner, a "learning-to-learn" system. Instead of using a fixed rule, they trained a tiny, super-efficient "coach" (a small neural network) to teach the blindfolded explorer how to guess better.

Here is how it works, using a few analogies:

The Coach (ZO Fine-tuner): Imagine a coach who watches the explorer. The coach doesn't need to see the whole library. Instead, the coach looks at a few simple stats: "How messy is this room right now?" and "Did the last guess help or hurt?" Based on this, the coach tells the explorer: "In this specific room, take a big, bold step. In that other room, take a tiny, careful step."
The Block Strategy: The library is huge (billions of parameters). If the coach tried to give instructions for every single book individually, it would be too slow and memory-heavy. However, the researchers noticed that the library is organized into "blocks" (like sections of books). The coach learns to give one instruction per block. This is like a coach telling a whole team of players how to move together, rather than shouting at every single player individually. This keeps the memory usage tiny.
The "Train Once, Reuse Widely" Magic: Usually, you have to train a new coach for every new task. But the authors found something amazing: The "shape" of the library doesn't change much, whether you are teaching it math or writing stories. So, they trained the coach once on a single dataset (COPA). Then, they took that same coach and sent it to teach the library on completely different tasks (like sentiment analysis or reading comprehension) and even on different versions of the library.
- The Result: The coach trained on one task worked surprisingly well on all the others. It's like training a driver on a specific track, and then having them drive perfectly on a totally different highway without any new practice.

The Results
The paper tested this on four different large language models and seven different tasks.

Performance: In about 82% of the cases, this new "Smart Explorer" reached a better result faster than the old "Fixed Rule" explorers.
Efficiency: It didn't require much extra memory. The "coach" itself is so small (less than 2MB for a 30-billion parameter model) that it's like adding a single page of notes to a 60GB encyclopedia.
Stability: The new method was also less sensitive to the "learning rate" (how big the steps are). Even if you told it to take bigger steps, it didn't crash as easily as the old methods.

In Summary
The paper introduces a way to teach giant AI models new skills without needing a supercomputer. They replaced a rigid, random guessing strategy with a tiny, smart coach that learns how to guess the right directions for different parts of the model. Once this coach is trained on one task, it can be reused to teach the model many other tasks efficiently, saving both time and computer memory.

Technical Summary: Learning a Zeroth-Order Optimizer for Fine-Tuning LLMs

Problem Statement

Fine-tuning large language models (LLMs) using standard first-order optimizers (e.g., Adam) incurs prohibitive memory overhead, often requiring up to 12 times the memory of inference due to the storage of optimizer states and backpropagation activations. While zeroth-order (ZO) optimizers like MeZO have emerged as a memory-efficient alternative by relying on forward passes only, existing methods typically employ hand-crafted, static sampling strategies (e.g., standard normal distributions). These static strategies are not adaptable to the specific structural properties of different models or tasks, often requiring extensive hyperparameter tuning and failing to fully exploit the local geometry of the loss landscape.

Methodology: ZO Fine-tuner

The authors propose ZO Fine-tuner, a learning-based zeroth-order optimizer designed to automatically learn efficient perturbation strategies through a "Learning to Learn" (L2L) framework. The core innovation lies in moving from static, uniform perturbation to adaptive, block-wise perturbation.

1. Block-Wise Perturbation Architecture

Motivated by recent analyses suggesting that Transformer Hessians exhibit an approximately block-diagonal structure, the method avoids coordinate-wise parameterization (which is memory-prohibitive at scale) in favor of block-wise adaptation.

PertNN: For each parameter block $i$ of the LLM, a lightweight auxiliary neural network (PertNN) predicts a shared perturbation variance $\sigma_t^{(i)}$ .
Inputs: PertNN takes as input compact summarizing statistics of the current parameters (Mean and Variance of the block), the previous perturbation variance, and the loss values from the previous step.
Output: The network outputs an updated variance $\sigma_t^{(i)}$ , which is used to construct a diagonal variance matrix $\Sigma_t = \text{diag}(\sigma_t^{(1)}I, \dots, \sigma_t^{(n)}I)$ .
Normalization: To decouple the learned variance from the effective learning rate, the method normalizes $\Sigma_t$ such that its Frobenius norm remains constant ( $\|\Sigma_t\|_F = \sqrt{d}$ ). This ensures the perturbation magnitude remains stable while allowing the optimizer to learn relative block-wise variances.

2. Training via Learning to Learn (L2L)

The optimizer is trained using a meta-objective where the fine-tuning trajectory of the LLM itself serves as supervision.

Meta-Objective: The goal is to minimize the post-update loss $L(\theta_1)$ after a single ZO step. The gradient of this loss with respect to the PertNN parameters $\omega$ is computed using the reparameterization trick (sampling $z \sim N(0, I)$ and setting $u = \Sigma z$ ), making the perturbation process differentiable.
Training Strategy: The PertNN is trained on a single dataset (e.g., COPA) using a first-order optimizer to generate model states. To prevent overfitting to low-loss regions, the authors introduce a periodic reset mechanism, re-initializing the LLM parameters to their pre-finetuning state after certain intervals.
Generalization: Once trained for a specific base model, the ZO Fine-tuner is reused across diverse downstream tasks and derivative checkpoints without further training.

Key Contributions

Extension of L2L to LLMs: The paper demonstrates that a single learned optimizer, trained on a base model, can generalize effectively across downstream tasks and derivative checkpoints, validating a "train once, reuse widely" workflow.
Block-Wise Variance Learning: The authors propose a memory-efficient block-wise parameterization with shared variance per block. This design is theoretically grounded in the block-diagonal structure of Transformer Hessians, allowing for L2L at the scale of foundation models without prohibitive memory costs.
Empirical Performance: Across four LLMs (LLaMA-3.2-1B, LLaMA-3.1-8B, Qwen2.5-14B, OPT-30B) and seven datasets, ZO Fine-tuner outperforms the strongest zeroth-order baselines (MeZO, MeZO-AdamU, HIZOO, LOZO) in 82.1% of task-model combinations regarding converged loss, with an average accuracy improvement of 2.5%.

Experimental Results

Convergence and Accuracy: ZO Fine-tuner consistently reaches lower loss faster than baselines. It demonstrates superior robustness to learning rate choices, often achieving comparable or better performance at lower learning rates (e.g., $10^{-8}$ ) where baselines fail to converge or become unstable.
Transferability: An optimizer trained on the COPA dataset successfully transfers to fine-tune models on SST-2, CB, SQuAD, WSC, BoolQ, and DROP, as well as to math-reasoning tasks (MetaMathQA).
Derivative Checkpoints: The method generalizes from a base model (LLaMA-3.1-8B) to its instruction-tuned variant (LLaMA-3.1-8B-Instruct), outperforming MeZO in both loss and accuracy.
Efficiency:
- Memory: The memory overhead is negligible compared to the model size (e.g., <2MB for auxiliary networks on OPT-30B). The total memory footprint remains comparable to MeZO and is significantly lower than first-order methods (e.g., 62GB vs. 312GB for OPT-30B).
- Time: The time overhead introduced by querying the PertNN is minimal (<3.4% for smaller models, decreasing for larger models), as the dominant cost remains the forward passes required for zeroth-order estimation.

Significance and Claims

The paper claims that ZO Fine-tuner addresses the critical bottleneck of memory efficiency in LLM fine-tuning while overcoming the rigidity of static zeroth-order strategies. By leveraging the "train once, reuse widely" paradigm, the authors argue that model creators could ship a pretrained learned finetuner alongside base models, enabling downstream users to perform efficient, high-performance fine-tuning with near-inference memory costs. The work establishes that learning adaptive perturbation distributions is feasible at the foundation model scale and offers a practical path toward scalable, memory-efficient LLM adaptation.