Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark

Enhancing Memory Efficiency in Fine-Tuning LLMs through Zeroth-Order Optimization

Overview

Fine-tuning pre-trained LLMs is a pervasive practice in natural language processing tasks. However, the substantial memory overhead associated with gradient computation through back-propagation remains a significant barrier, particularly for computational platforms with limited memory. This challenge has motivated a shift towards memory-efficient approaches, such as Zeroth-Order (ZO) optimization which eliminates the necessity for explicitly computing gradients. Building on the concept introduced by Malladi et al. (2023), this paper proposes a comprehensive analysis of Zeroth-Order Optimization for memory-efficient LLM fine-tuning, unveiling previously overlooked optimization principles and introducing novel enhancements.

Related Work and Theoretical Background

Previous efforts in Parameter-Efficient Fine-Tuning (PEFT) strategies and zeroth-order optimization have laid the groundwork for memory-efficient model training. Traditional approaches like Adapter-based methods, Low-Rank Adaptation (LoRA), and prompt tuning significantly reduce the number of parameters required for fine-tuning but still require considerable memory for gradient computation. In contrast, Zeroth-Order (ZO) optimization utilizes function value-based gradient estimation, thereby circumventing the need for back-propagation and subsequently reducing memory usage. Despite its promise, the exploration of ZO optimization techniques beyond basic ZO-Stochastic Gradient Descent (ZO-SGD) is scant, prompting this study.

Methodology and Key Contributions

1. Benchmark Creation: The study creates the first benchmark for ZO optimization in LLM fine-tuning, evaluating six BP-free ZO optimization methods across five LLM families, three task complexities, and five fine-tuning schemes.
2. Insights on Optimization Principles: The benchmark study reveals critical insights including the importance of task alignment, the utility of the forward gradient method as a baseline for ZO optimization, and the balance between algorithm complexity and fine-tuning performance.
3. Enhancements to ZO Optimization: Drawing from the underline insights, the study proposes techniques of block-wise descent, hybrid ZO and FO (First-Order) training, and gradient sparsity to improve ZO optimization-based LLM fine-tuning.

Theoretical and Practical Implications

From a theoretical standpoint, this work advances understanding of the optimization landscape for LLM fine-tuning, particularly under resource constraints. Practically, the introduced benchmark and ensuing insights offer a structured foundation for future research and development in memory-efficient fine-tuning methods. The proposed enhancements to ZO optimization—block-wise descent, hybrid training, and gradient sparsity—not only improve fine-tuning accuracy but also maintain memory efficiency. These advancements possess the potential to facilitate on-device training and deployment of sophisticated language models in memory-constrained environments.

Future Directions

Looking ahead, the exploration of further ZO optimization methods and their combinations with established PEFT strategies presents a promising avenue for research. Additionally, investigating the applicability of these memory-efficient fine-tuning techniques beyond language models to other domains of deep learning could broaden their utility.

Concluding Thoughts

This study's comprehensive benchmarking and innovative enhancements to ZO optimization mark significant steps towards overcoming the memory limitations in fine-tuning LLMs. By elucidating the trade-offs between algorithm complexity, accuracy, and memory efficiency, it lays the groundwork for more sustainable and accessible AI models, pushing the boundaries of what's possible within constrained computational environments.