Zero-Order Optimization for LLM Fine-Tuning via Learnable Direction Sampling

TL;DR

This paper introduces a learnable sampling policy for zero-order optimization in large language model fine-tuning, significantly reducing variance and enabling scalable, memory-efficient training.

Contribution

It proposes a novel policy-driven zero-order framework with theoretical analysis and practical algorithms that improve gradient estimates for large-scale NLP models.

Findings

Enhanced fine-tuning performance on LLM benchmarks

Reduced variance in gradient estimation

Relaxed dependence on parameter dimensionality

Abstract

Fine-tuning large pretrained language models (LLMs) is a cornerstone of modern NLP, yet its growing memory demands (driven by backpropagation and large optimizer States) limit deployment in resource-constrained settings. Zero-order (ZO) methods bypass backpropagation by estimating directional derivatives from forward evaluations, offering substantial memory savings. However, classical ZO estimators suffer from high variance and an adverse dependence on the parameter dimensionality $d$, which has constrained their use to low-dimensional problems. In this work, we propose a policy-driven ZO framework that treats the sampling distribution over perturbation directions as a learnable policy and updates it to reduce the variance of directional estimates. We develop a practical algorithm implementing this idea and provide a theoretical analysis, showing that learned sampling distributions improve the quality of gradient information and relax the explicit dependence on $d$ in convergence bounds. Empirically, we validate the approach on challenging LLM fine-tuning benchmarks, demonstrating substantially improved performance compared to standard ZO baselines. Our results suggest that adaptive direction sampling is a promising route to make ZO fine-tuning viable at scale. The source code is available at https://github.com/brain-lab-research/zo_ldsd