Sample adaptive zeroth order method speeds up language model tuning
Sample-Adaptive Zeroth-Order Post-Training
Machine Learning
Summary
Fine-tuning large language models usually requires lots of memory and time because it uses backpropagation. The authors point out that current alternatives that don't use backpropagation treat all training examples the same, even though some are more important during training. They propose SAZO, which learns which examples are more important and focuses more effort on those, speeding up training. This method keeps track of example importance efficiently without extra model checks and proves to converge faster than previous approaches. Their experiments show that SAZO can reach good results much faster on different models.
What this means in practice
- •For machine learning engineers: Speed up fine-tuning of large language models by adaptively focusing on more important training samples.
- •For automated machine learning teams: Reduce memory usage and accelerate model updates in black-box optimization settings with limited gradient access.
Authors
Haiyuan Wan, Yanyi Li, Weicheng Lin, Cong Fang
Abstract
Zeroth-order (ZO) optimization offers a memory-efficient alternative to backpropagation for fine-tuning large language models. However, existing ZO methods typically sample training examples uniformly, overlooking their evolving importance and leaving avoidable estimation variance. To address this limitation, we propose Sample-Adaptive Zeroth-Order Optimization (SAZO), which learns sample importance online and adaptively allocates queries to accelerate convergence. SAZO extracts gradient-norm information from existing directional differences to estimate sample importance without access to true gradients. Furthermore, to track evolving importance from sparse minibatch feedback, SAZO maintains per-example exponential moving averages that aggregate observations across updates. The method requires no additional model evaluations and only $O(N)$ additional scalar state compared with MeZO. We prove that SAZO asymptotically matches the leading convergence term of the minimum-variance sampling oracle, and achieves a strictly faster convergence rate than standard methods. Extensive experiments across four model families demonstrate improved average downstream performance and up to $6\times$ faster convergence to the same training loss in terms of query count.