LOZO+: Provably Efficient Zeroth-Order Fine-Tuning via Greedy Low-Rank Subspace Selection

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.

Adapting LLMs to new tasks usually requires a lot of computer memory, which can make the process expensive or inaccessible. A promising alternative is to update the model by testing small changes and keeping the ones that improve performance, rather than storing all the detailed information used by standard training methods. However, existing versions of this approach often choose these changes at random, so they may miss the most useful directions for improving the model. We propose LOZO+, a new method that makes this process more selective. Instead of relying only on random choices, LOZO+ uses feedback from the model’s performance to identify which update directions are likely to be helpful and keeps only those that lead to meaningful improvement. This helps the model learn more efficiently while still using much less memory than standard methods. We provide both theoretical analysis and experiments showing that LOZO+ improves over existing memory-efficient methods and can achieve performance close to standard fine-tuning approaches.