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

TL;DR

This paper introduces a prior-informed zeroth-order optimization method with adaptive direction alignment that significantly improves memory efficiency and convergence speed in fine-tuning large language models, outperforming traditional approaches.

Contribution

The paper presents a novel prior-informed perturbation technique for zeroth-order optimization that enhances gradient estimation accuracy and convergence in large language model fine-tuning.

Findings

Outperforms traditional ZO methods on OPT-13B across 11 benchmarks.

Achieves better results than gradient-based methods on 9 out of 11 tasks.

Demonstrates faster convergence and improved efficiency in large-scale LLM fine-tuning.

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.