🤖 AI Summary
This work addresses the high memory cost of backpropagation in large language model fine-tuning and the slow convergence of existing zeroth-order optimizers like MeZO, which struggle to capture loss landscape structure. The authors propose AdaMeZO, the first zeroth-order optimizer incorporating an Adam-like adaptive mechanism that estimates first- and second-order moments using only forward passes, eliminating the need to store momentum or incur additional memory overhead. Theoretical analysis establishes its convergence guarantees, while experiments demonstrate that AdaMeZO reduces forward computation by up to 70% compared to MeZO while achieving superior performance. Trajectory visualizations further confirm its ability to adapt effectively to diverse loss landscapes.
📝 Abstract
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 AdaMeZO's performance, showing that AdaMeZO can outperform MeZO while requiring up to $70\%$ fewer forward passes. Trajectory visualizations affirm AdaMeZO's ability to adapt to diverse loss landscapes.