AdaMeZO: Adam-style Zeroth-Order Optimizer for LLM Fine-tuning Without Maintaining the Moments

Fine-tuning LLMs is necessary for various dedicated downstream tasks, but classic backpropagation-based fine-tuning methods require substantial GPU memory. To this end, a recent work, MeZO, which relies solely on forward passes to fine-tune LLMs, significantly reduces GPU requirements at the cost of slower convergence due to its indifference to loss landscapes. Standard solutions, such as Adam, explore loss landscapes by estimating the first- and second-order moments and storing them in memory to guide the model's movement through dimensions with lower curvature and vice versa. However, directly applying Adam negates MeZO's advantage as it will triple the memory requirement. In light of this, we propose AdaMeZO, a zeroth-order optimizer that leverages Adam-style first- and second-moment estimates without maintaining them in memory. We present a theoretical analysis of AdaMeZO, corroborated by extensive experiments demonstrating its performance, showing that it can outperform MeZO while requiring up to 70% fewer forward passes. Trajectory visualizations affirm AdaMeZO's ability to adapt to diverse loss landscapes.

Large language models—the technology behind many chatbots and writing tools—often need to be fine-tuned before they work well on a specific task, but the usual fine-tuning process can require a lot of GPU memory. A cheaper alternative is to adjust the model by making small random changes and keeping those that improve performance, but this approach can learn slowly due to the model's size. We introduce AdaMeZO, a new fine-tuning method that gives this low-memory approach some of the guidance used by Adam, a popular optimizer, without storing Adam’s large extra memory buffers. AdaMeZO does this by briefly reconstructing useful recent information from random seeds and updating the model in small blocks, rather than keeping full histories in memory. In experiments on several language tasks and model families, AdaMeZO usually matched or improved the previous low-memory method while using far fewer model evaluations. This could make it easier for researchers and developers with limited hardware to adapt large language models for specialized applications, while keeping memory costs close to simply running the model.