Low-rank evolution strategies improve large language model fine-tuning accuracy
EGGROLL, Unrolled: Understanding and Improving Low-Rank Evolution Strategies at Scale
Machine LearningArtificial Intelligence
Summary
Fine-tuning large language models is hard because changing many parameters at once requires a lot of computation. The paper studies EGGROLL, a technique that changes model weights in simpler, low-rank ways that are easier to handle. The authors analyze the math behind these changes and identify potential problems and benefits. They also create LOO-ROLL, an improved estimator that reduces errors and speeds up tuning, showing big accuracy gains on language tasks.
What this means in practice
- •For machine learning engineers: Speed up fine-tuning of large language models by using low-rank perturbations that keep quality high while requiring fewer computations and less variance in gradient estimates.
- •For ai infrastructure teams: Improve tuning workflows for transformer-based models by integrating LOO-ROLL to halve estimator error at similar evaluation cost, leading to better model performance in less time.
Authors
Ege C. Kaya, Abolfazl Hashemi
Abstract
EGGROLL makes evolution strategies (ES) practical for LLMs by replacing dense Gaussian weight perturbations with low-rank Gaussian products, often of rank one. This choice is computationally attractive but geometrically severe: each rank-one perturbation lies in a zero-volume subset of the ambient matrix space, despite having identity covariance. We characterize the mean EGGROLL update field at finite rank and nonzero perturbation radii, then analyze the error of its finite-population estimator. The population field is obtained by applying an explicit resolvent to the gradient of the objective smoothed by the perturbations. We show that the resolvent can introduce a nonconservative component and can reverse the local stability of an optimum. EGGROLL is nevertheless exact on every quadratic objective at every rank and radius. For smooth objectives, its first local finite-rank correction is $O(σ^2/r)$, and nonasymptotic bounds control the resulting field error under smoothness assumptions. Under a local affine model, rank-one perturbations increase the variance of the gradient estimator by only $\frac{2(m+n+1)}{mn+1}$ relative to dense Gaussian ES, or $0.098\%$ for a $4096\times4096$ matrix. We then introduce LOO-ROLL, a leave-one-out estimator that preserves the finite-rank population field while replacing EGGROLL's two antithetic evaluations per direction by one. At equal evaluation cost, LOO-ROLL halves estimator MSE in transformer blocks. At matched wall time across ten post-training settings and models up to 8B parameters, LOO-ROLL improves seven outcomes in individual paired tests, with no significant loss. On the GSM8K test set, accuracy increases from $38.1\%$ to $63.0\%$ at 0.6B and from $65.9\%$ to $80.0\%$ at 8B. Transformer measurements recover the predicted finite-rank variance, while the rank comparisons show no reproducible reward-based advantage for rank eight.