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

TL;DR

ConMeZO is a new zeroth-order optimizer that speeds up large language model finetuning by adaptively sampling descent directions, reducing convergence time while maintaining low memory usage.

Contribution

It introduces an adaptive directional sampling strategy for zeroth-order optimization, improving convergence speed in high-dimensional LLM finetuning.

Findings

ConMeZO is up to 2X faster than MeZO in finetuning LLMs.

ConMeZO maintains the low-memory footprint of zeroth-order methods.

Theoretical analysis shows ConMeZO has the same worst-case convergence rate as MeZO.

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.