MeZO-A$^{3}$dam: Memory-efficient Zeroth-order Adam with Adaptivity Adjustments for Fine-tuning LLMs
Keywords: Optimization, Zeroth-Order Optimization, Large Language Models, Fine-tuning
Abstract: Recently, fine-tuning of language models (LMs) via zeroth-order (ZO) optimization have gained significant traction due to their ability of memory-efficient deployment, significantly reducing memory cost over first-order methods. However, the existing studies on ZO optimization for LM fine-tuning often exhibit slow convergence and the reliance on the hand-crafted prompts. Towards mitigating these limitations, in this paper, we first investigate on the importance of adaptive gradient based ZO optimization method. Toward this, we revisit memory-efficient zeroth-order Adam (MeZO-Adam) and make important findings that merely considering adaptivity can enable faster convergence while improving the generalization ability compared to previous studies. Interestingly, we further observe that decreasing the level of adaptivity might be recommended in ZO optimization potentially due to the high variance of ZO gradient estimate, hypothesized as \emph{weak adaptivity hypothesis}. Based upon our hypothesis, we propose MeZO-A$^3$dam, MeZO-Adam with Adaptivity Adjustments according to the parameter dimension. We provide the dimension-free theoretical guarantee on both the convergence and the generalization of MeZO-A$^3$dam, providing strong evidence for our hypothesis. Extensive experiments show that MeZO-A$^3$dam can achieve faster convergence and better generalization over several baselines across LMs of various sizes on diverse datasets. By adaptivity adjustments, MeZO-A$^3$dam outperforms MeZO, MeZO-SVRG, and MeZO-Adam, with up to an average of $36.6\\%$, $16.9\\%$, $6.8\\%$ improvements in performance and up to an average of $\times 12.6$ and $\times1.8$ faster convergence, respectively. Furthermore, by leveraging an off-the-shelf low-bit optimizer, MeZO-A$^3$dam achieves an average of $40.3\\%$ and $43.6\\%$ memory reduction from MeZO-SVRG and MeZO-Adam.
Supplementary Material: zip
Primary Area: optimization
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 9024
Loading