ConMeZO: Adaptive Descent-Direction Sampling for Gradient-Free Finetuning of Large Language Models

TL;DR

ConMeZO tackles the slow convergence of gradient-free finetuning for large language models by introducing cone-based sampling around a momentum direction, which reduces variance in zeroth-order gradient estimates while preserving memory advantages. The method merges momentum with a geometrically constrained perturbation distribution, yielding a Descent Lemma that can provide faster per-iteration progress when momentum aligns with the true gradient, while maintaining the same worst-case rate as MeZO. Theoretical guarantees show a convergence bound of $\frac{1}{T}\sum_{t=0}^{T-1} \mathbb{E}[\|\nabla f(x_t)\|^2] \le \frac{2 \ell d (f(x_0) - f^*)}{T}$, and empirical results demonstrate up to $2\times$ speedups on RoBERTa-large and OPT-1.3B/13B with modest memory overhead. Overall, ConMeZO enables faster, memory-efficient zeroth-order finetuning of billion-parameter LLMs with practical improvements in early convergence and final performance.

Abstract

Zeroth-order or derivative-free optimization (MeZO) is an attractive strategy for finetuning large language models (LLMs) because it eliminates the memory overhead of backpropagation. However, it converges slowly due to the inherent curse of dimensionality when searching for descent directions in the high-dimensional parameter space of billion-scale LLMs. We propose ConMeZO, a novel zeroth-order optimizer that accelerates convergence by adaptive directional sampling. Instead of drawing the direction uniformly at random, ConMeZO restricts the sampling to a cone centered around a momentum estimate. This concentrates the search in directions where the true gradient is more likely to lie and thus reduces the effect of high dimensions. We prove that ConMeZO achieves the same worst-case convergence rate as MeZO. Empirically, when finetuning LLMs on natural language tasks, ConMeZO is up to 2X faster than MeZO while retaining the low-memory footprint of zeroth-order methods.

Figures (7)

• Figure 1: ConMeZO achieves 2$\times$ speedup over MeZO when finetuning OPT-1.3B on SQuAD dataset.
• Figure 2: 2D- and 3D-representation of the cone-sampling approach. (a) Sphere with radius $\sqrt{d}$ and (gray) search space cone of half-angle $\theta$ around promising search direction $\hat{m}_t$. We can set random direction $z_t=z_t^\parallel + z_t^\perp$ with angle $\gamma$ to $\hat{m}_t$. (b) 3D representation of cone sampling in red area.
• Figure 3: Synthetic Optimization Problem: ConMeZO achieves 2.45$\times$ speedup over MeZO on the synthetic quadratic problem.
• Figure 4: Peak GPU memory usage (MiB) increase of ConMeZO over MeZO is negligible when compared to the memory usage of first-order methods like AdamW: Top: RoBERTa-Large on SST2 (batch size 64). Bottom: OPT-1.3B on BoolQ (batch size 16).
• Figure 5: Heatmaps of Test Accuracy of ConMeZO on TREC dataset for different $\theta$ and $\beta$ values and fixed learning rate $\eta = 10^{-6}$.

Theorems & Definitions (11)

• Definition 1
• Lemma 1
• Proposition 1
• Proposition 2
• Lemma 2
• Theorem 1: Descent Lemma
• Corollary 1
• Proof
• Proof
• Proof