Momentum method fine-tunes large language models without slow trial and error

MpSub: A Momentum $p$-Dimensional Subspace Trust-Region Method for Derivative-Free Fine-Tuning of Large Language Models

Machine Learning

Summary

Fine-tuning large language models usually needs a lot of memory and careful tuning of learning rates, which can be tricky. The authors created MpSub, a method that tweaks these models by looking at a small set of directions for improvement, one of which uses past momentum and others that try new random directions. This approach estimates how to update the model without needing gradients or learning rate adjustments and adapts its step size based on how well predicted improvements match actual results. They tested MpSub on large language models and found that it performed just as well as other methods that require tuning but without needing that extra tuning effort.

Fine-tuningLarge language modelsDerivative-free optimizationTrust-region methodsMomentumZeroth-order optimizationCentral differencesGradient estimationLearning rateForward pass

Authors

Yuyang Wang, Haoyu Yao, Pengcheng Xie

Abstract

Full-parameter fine-tuning of large language models has substantial memory costs because backpropagation stores activations and gradients. Zeroth-order optimization avoids this by estimating update directions from loss evaluations, but existing methods require tuning a sensitive learning rate for each model and task. We propose the momentum $p$-dimensional subspace trust-region method (MpSub). At each iteration, MpSub searches within a $p$-dimensional subspace: one direction preserves historical momentum from the most recent accepted step, while the remaining directions explore via fresh random sampling. The subspace gradient is estimated by central differences, a trial step is computed from a linear trust-region model, and the trust-region radius adapts according to the agreement between predicted and observed loss reduction, eliminating the learning rate. For LLM fine-tuning, evaluations within an iteration share a minibatch, and directions are regenerated in place from seeds, using forward passes alone. For smooth deterministic objectives under unorthogonalized Gaussian directions, we bound the finite-difference error, quantify gradient energy captured by the subspace, and prove that $\lim_{k\to\infty} \|\nabla f(x_k)\|_2 = 0$ almost surely under a safeguarded radius update. Under a matched budget of 8,400 training-objective forward passes, we fine-tune OPT-125M and OPT-350M on CommitmentBank. With the same preset parameters at both model sizes, MpSub attains mean test accuracies of 0.673 and 0.690 over three seeds, matching tuned MeZO (0.685) without any learning-rate search.