Prior-Informed Zeroth-Order Optimization with Adaptive Direction Alignment for Memory-Efficient LLM Fine-Tuning

2026-01-08

arXiv.org

AI Summary

This work addresses the high memory overhead of backpropagation in large language model (LLM) fine-tuning and the inefficiency of existing zeroth-order optimization methods, which suffer from high-variance gradient estimates due to random perturbations, leading to slow convergence. To overcome these limitations, the authors propose a plug-and-play zeroth-order optimization method that adaptively constructs perturbation directions using prior information to improve alignment between estimated and true gradients. A greedy perturbation strategy is further introduced to enhance optimization efficiency. The approach significantly accelerates convergence and improves performance across LLMs of varying scales and architectures. On the OPT-13B model, it outperforms conventional zeroth-order methods on all 11 evaluated tasks and surpasses gradient-based baselines on 9 of them.

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
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Fine-tuning large language models (LLMs) has achieved remarkable success across various NLP tasks, but the substantial memory overhead during backpropagation remains a critical bottleneck, especially as model scales grow. Zeroth-order (ZO) optimization alleviates this issue by estimating gradients through forward passes and Gaussian sampling, avoiding the need for backpropagation. However, conventional ZO methods suffer from high variance in gradient estimation due to their reliance on random perturbations, leading to slow convergence and suboptimal performance. We propose a simple plug-and-play method that incorporates prior-informed perturbations to refine gradient estimation. Our method dynamically computes a guiding vector from Gaussian samples, which directs perturbations toward more informative directions, significantly accelerating convergence compared to standard ZO approaches. We further investigate a greedy perturbation strategy to explore the impact of prior knowledge on gradient estimation. Theoretically, we prove that our gradient estimator achieves stronger alignment with the true gradient direction, enhancing optimization efficiency. Extensive experiments across LLMs of varying scales and architectures demonstrate that our proposed method could seamlessly integrate into existing optimization methods, delivering faster convergence and superior performance. Notably, on the OPT-13B model, our method outperforms traditional ZO optimization across all 11 benchmark tasks and surpasses gradient-based baselines on 9 out of 11 tasks, establishing a robust balance between efficiency and accuracy.

Problem

• Large Language Models
• Zeroth-Order Optimization
• Memory Efficiency
• Gradient Estimation
• Fine-Tuning

Innovation

• Zeroth-Order Optimization
• Prior-Informed Perturbation
• Memory-Efficient Fine-Tuning
• Gradient Estimation
• Large Language Models