LOZO+: Provably Efficient Zeroth-Order Fine-Tuning via Greedy Low-Rank Subspace Selection
Abstract
Zeroth-order (ZO) optimization offers a more memory-efficient alternative to first-order methods for fine-tuning large language models (LLMs). Recent ZO methods, exemplified by LOZO, estimate gradients within low-rank subspaces to align with the low-rank structure of LLM gradients. However, these methods rely on randomly generated subspaces of a fixed rank, which provides no guarantee of alignment with the actual dominant subspaces of the gradients; essentially, they remain ZO gradient descent with stochastic subspace sampling. To more effectively exploit the low-rank nature of LLM gradients, we propose \textbf{LOZO+}, an efficient \textbf{ZO} fine-tuning algorithm for LLMs that incorporates greedy \textbf{Lo}w-Rank subspace selection. Specifically, LOZO+ leverages loss-based feedback to assess alignment between candidate directions and the dominant low-rank gradient subspaces, and employs an adaptive thresholding criterion to retain only directions yielding substantial gradient descent, thereby steering ZO optimization toward more effective convergence. Importantly, we establish a theoretical framework that characterizes the convergence behavior of LOZO+, formally prove its superiority over existing methods. Extensive experiments demonstrate that LOZO+ consistently outperforms existing ZO methods and achieves performance competitive with FO algorithm, while retaining the memory efficiency inherent to ZO optimization.