Variance-reduced Zeroth-Order Methods for Fine-Tuning Language Models
Tanmay Gautam, Youngsuk Park, Hao Zhou, Parameswaran Raman, Wooseok Ha
TL;DR
This work tackles the memory bottlenecks of fine-tuning large language models by introducing MeZO-SVRG, a variance-reduced zeroth-order optimization method that preserves memory efficiency while improving convergence without task prompts. By integrating full-batch and minibatch SPSA gradient estimators in a memory-efficient SVRG framework, it achieves up to 20% gains in test accuracy over the prior MeZO method and reduces GPU-hours by about a factor of 2, all while halving the memory footprint compared to first-order SGD on autoregressive models. Theoretical guarantees show a convergence bound of the form $\mathbb{E}\left[\|\nabla f(\boldsymbol{\theta}^{(T)})\|_2^2\right] = O\left(\frac{d}{T} + \frac{\mathbf{1}(b<n)}{b}\right)$ under suitable choices, supporting the practical results. Overall, MeZO-SVRG enables scalable, prompt-free fine-tuning of large LMs with reduced memory and computation, broadening access to high-capacity models in resource-constrained settings.
Abstract
Fine-tuning language models (LMs) has demonstrated success in a wide array of downstream tasks. However, as LMs are scaled up, the memory requirements for backpropagation become prohibitively high. Zeroth-order (ZO) optimization methods can leverage memory-efficient forward passes to estimate gradients. More recently, MeZO, an adaptation of ZO-SGD, has been shown to consistently outperform zero-shot and in-context learning when combined with suitable task prompts. In this work, we couple ZO methods with variance reduction techniques to enhance stability and convergence for inference-based LM fine-tuning. We introduce Memory-Efficient Zeroth-Order Stochastic Variance-Reduced Gradient (MeZO-SVRG) and demonstrate its efficacy across multiple LM fine-tuning tasks, eliminating the reliance on task-specific prompts. Evaluated across a range of both masked and autoregressive LMs on benchmark GLUE tasks, MeZO-SVRG outperforms MeZO with up to 20% increase in test accuracies in both full- and partial-parameter fine-tuning settings. MeZO-SVRG benefits from reduced computation time as it often surpasses MeZO's peak test accuracy with a $2\times$ reduction in GPU-hours. MeZO-SVRG significantly reduces the required memory footprint compared to first-order SGD, i.e. by $2\times$ for autoregressive models. Our experiments highlight that MeZO-SVRG's memory savings progressively improve compared to SGD with larger batch sizes.