Low-Rank Curvature for Zeroth-Order Optimization in LLM Fine-Tuning

AI summary

Overview

Research area: Machine learning — specifically zeroth-order (gradient-free) optimization for fine-tuning large language models.

Technical level: Advanced. The paper assumes familiarity with gradient-based optimization, Hessian/curvature approximations, Kronecker factorization, and evolution strategies.

Scope: The paper introduces LOREN, a curvature-aware, variance-reduced zeroth-order optimizer for LLM fine-tuning, and evaluates it against leading ZO baselines on GLUE and SuperGLUE tasks.

What This Paper Is About

Zeroth-order (ZO) optimizers like MeZO can fine-tune LLMs with far less memory than standard first-order methods such as SGD or AdamW, because they estimate gradients using only forward passes rather than backpropagation. However, existing ZO methods suffer from two problems: their finite-difference gradient estimates are very noisy (high variance), and they are blind to the differing curvature of the loss landscape across weights and layers. This paper's goal is to fix both problems at once — producing better-conditioned and lower-variance search directions — while keeping the memory advantage that makes ZO methods attractive in the first place.

Key Contributions

1. LOREN, described by the authors as the first ZO optimizer to simultaneously adapt to curvature information and apply variance reduction, enabling stable updates even in high-dimensional and ill-conditioned settings.
2. A reformulation of gradient preconditioning in ZO optimization as the problem of adaptively estimating an anisotropic perturbation distribution, connected to evolution strategies so that preconditioned ZO gradients are estimated directly via a score-function estimator without any additional forward passes.
3. A damped low-rank (rank-1) Kronecker-factored covariance structure that preserves memory efficiency — the authors state this is the first use of a block-diagonal Hessian approximation in ZO optimization, capturing richer curvature information than a pure diagonal approximation.
4. Extensive experiments on standard LLM benchmarks (GLUE and SuperGLUE) comparing LOREN with MeZO, MeZO-Adam, MeZO-SVRG, LOZO, and HiZOO, with LOREN reported as delivering higher test accuracy at lower memory footprint than other state-of-the-art preconditioned or adaptive ZO methods.

Main Findings

• Variance reduction is demonstrably effective: On three 1,000-dimensional test functions (Sphere, Rastrigin, Rosenbrock), the authors generated 5,000 gradient estimates at a fixed point and computed mean squared error relative to the true gradient. ZO-SGD with RLOO consistently achieved lower MSE than ZO-SGD without it. Both used four perturbations per gradient estimate for fairness.
• Better saddle-point escape: On the monkey saddle function, with all optimizers initialized at (2.9, -0.01), ZO-SGD and ZO-Adam struggled due to noisy gradient estimates, HiZOO showed moderate improvement via its ZO-Hessian estimate, and LOREN followed the most efficient path, escaping the saddle region by combining low-rank curvature with low-variance gradient estimates.
• Memory reduction: LOREN cuts peak memory usage by up to 27.3% compared with MeZO-Adam.
• Memory complexity advantage: The additional memory requirement relative to MeZO is O(mn) for MeZO-Adam, O(mn) for MeZO-SVRG, O(nr) for LOZO (r is the low rank), O(mn) for HiZOO, and O(n) for LOREN. This lower cost enables the use of heavyball momentum, from which LOREN benefits through acceleration.
• Convergence guarantee: Under L-smoothness of the objective and bounded variance of gradient estimates, LOREN converges to a stationary point at a rate of O(1/√T), where T is the number of iterations.
• DistilBERT (66M), FP32 — average accuracy over MNLI, QNLI, SST-2, CoLA: LOREN 65.4, MeZO-Adam 64.6, HiZOO 62.6, MeZO-SVRG 61.6, LOZO 55.3, MeZO 54.4. Per-task, LOREN reached 73.0 on QNLI and 81.7 on SST-2 — the highest of all listed methods on both.
• RoBERTa-large (355M), FP32: LOREN reached the highest SST-2 score (86.1) and CoLA score (73.8) among the listed methods, but its average across the four tasks was 70.1, below MeZO-Adam's 73.6 according to the reported table.
• Training dynamics: Accuracy curves for GPT-2-XL fine-tuned on QNLI and OPT-13B fine-tuned on CB (with early stopping) show LOREN achieving the highest mean accuracy and markedly more stable convergence.
• Evaluation budget: LOREN uses six forward evaluations per iteration with RLOO variance reduction; the same budget was applied to all baselines for fair comparison. Other ZO optimizers typically use only two or three forward passes per step, and the authors note this default setting leads to degraded performance compared to using six passes (results under default settings are given in the paper's Appendix E).

Methodology in Plain English

The team starts from a simple intuition. A zeroth-order optimizer guesses the gradient direction by nudging the model's parameters in random directions and seeing how the loss changes. If those random nudges are drawn from a plain spherical distribution, the optimizer treats every direction in parameter space as equally important — which is wrong, because some directions are steep and others nearly flat.

The first move is a mathematical re-framing: instead of thinking of "preconditioning" (rescaling) the gradient separately, treat it as learning the shape of the random perturbation distribution itself. If the perturbation distribution is an anisotropic Gaussian whose covariance equals the inverse curvature, then drawing perturbations from it is mathematically equivalent to preconditioning the gradient. This means curvature-aware optimization becomes a problem of learning a distribution, which is exactly what evolution strategies (ES) do.

Because storing and updating a full covariance matrix would be prohibitive for LLMs, the authors approximate curvature per layer with a damped rank-1 block-diagonal form — a Kronecker product of an identity and a low-rank-plus-diagonal term. This small structure has closed-form expressions for both the inverse and the inverse square root, so the optimizer can cheaply scale perturbations. The only extra learnable object is a single vector per layer, which is why the extra memory cost drops to O(n).

Finally, instead of the usual two-point finite-difference estimator, LOREN evaluates the loss at multiple perturbed points and combines them with the REINFORCE leave-one-out (RLOO) estimator: each sample's function value is compared against the average of the other samples in the batch. This baseline means each evaluation is used as a variance-reduction reference for the others, squeezing more signal out of the same number of forward passes. The whole procedure is presented as Algorithm 1, which alternately updates the parameters (the mean of the search distribution) and the small curvature parameter (its shape).

Why This Matters

Impact on research. The paper challenges the assumption that ZO methods must choose between memory efficiency and optimization quality. By showing that curvature adaptation and variance reduction can be combined cheaply via a Kronecker-factored rank-1 structure, it opens a middle path between pure first-order methods (accurate but memory-hungry) and simple ZO methods (cheap but noisy). It also provides a convergence guarantee at O(1/√T), lending theoretical grounding to a family of methods that has largely been empirical.

Real-world applications:

• Fine-tuning large language models on consumer or single-GPU hardware where backpropagation memory costs are prohibitive.
• On-device or edge personalization of LLMs, where memory budgets are tight and only forward passes are practical.
• Domain adaptation of foundation models in regulated industries where models must be tuned locally without large-scale infrastructure.
• Parameter-efficient and privacy-sensitive fine-tuning pipelines where intermediate activations cannot be stored or shared.

Industry relevance. Memory is often the binding constraint on LLM fine-tuning in production. A method that reduces peak memory by up to 27.3% versus MeZO-Adam while improving accuracy on several benchmark tasks has direct implications for cost, hardware requirements, and the scale of models that can be adapted by smaller teams. The released code at https://github.com/hseung88/loren lowers the barrier to reproduction and adoption.

Future Directions

• Scale beyond the tested models. The reported benchmark models are DistilBERT (66M) and RoBERTa-large (355M) in Table 2, with GPT-2-XL and OPT-13B shown only in accuracy curves. Whether the O(n) memory advantage and accuracy gains hold for much larger LLMs is not fully established in the provided content.
• Reconciling the RoBERTa-large results. LOREN leads on SST-2 and CoLA for RoBERTa-large but trails MeZO-Adam on MNLI and on the reported four-task average (70.1 vs 73.6). Understanding when curvature preconditioning helps versus when a plain adaptive ZO method wins is an open question.
• Hyperparameter sensitivity. LOREN relies on a damping factor ρ, a smoothing parameter ε, a rank-1 curvature structure per layer, and a forward-pass count K (set to six here). How sensitive performance is to these choices, and whether adaptive schemes could set them automatically, is not addressed in the provided content.
• Combining with parameter-efficient methods. The experiments use full-parameter fine-tuning without prompts; whether LOREN composes with LoRA-style adapters or prompt-based tuning is not explored.

Target Audience

This paper is aimed at machine learning researchers and engineers working on LLM fine-tuning efficiency, zeroth-order and derivative-free optimization, and memory-constrained training. It is most useful to readers who already understand gradient estimators, Hessian approximations, and Kronecker factorization, and who want to know how curvature information can be injected into ZO methods without sacrificing their memory advantages. Practitioners seeking a drop-in memory-efficient fine-tuning method will also benefit, though the dense mathematical derivations in Sections 3 and 4 assume a strong optimization background.

Authors’ abstract

We introduce LOREN, a curvature-aware zeroth-order (ZO) optimization method for fine-tuning large language models (LLMs). Existing ZO methods, which estimate gradients via finite differences using random perturbations, often suffer from high variance and suboptimal search directions. Our approach addresses these challenges by: (i) reformulating the problem of gradient preconditioning as that of adaptively estimating an anisotropic perturbation distribution for gradient estimation, (ii) capturing curvature through a low-rank block diagonal preconditioner using the framework of natural evolution strategies, and (iii) applying a REINFORCE leave-one-out (RLOO) gradient estimator to reduce variance. Experiments on standard LLM benchmarks show that our method outperforms state-of-the-art ZO methods by achieving higher accuracy and faster convergence, while cutting peak memory usage by up to 27.3% compared with MeZO-Adam.