Improving large language model tuning with multiscale space optimization
ENGINE: Endogenous Variational MultiScale Optimization for Zeroth-Order LLM Fine-Tuning
Machine Learning
Summary
Fine-tuning large language models can be tricky because using all model parameters at once is noisy and slow, but focusing on just a few parameters loses important details. The authors solve this by switching between a low-detail and a high-detail view during tuning, gaining the best of both worlds. This approach, called ENGINE, makes tuning faster and more accurate without using more memory. They tested this method on common benchmarks and found it speeds up the process significantly.
What this means in practice
- •For machine learning engineers: Speed up fine-tuning of large language models with lower memory use and better accuracy using multiscale optimization techniques.
- •For ai infrastructure teams: Enable more efficient resource use in deploying fine-tuning pipelines for large models by reducing variance and computational overhead.
Authors
Zhuoli Ouyang, Changxi Chi, Siyuan Li, Tailin Wu
Abstract
Zeroth-order (ZO) optimization enables memory-efficient large language model (LLM) fine-tuning, and recent methods incorporate curvature-aware strategies. However, in the high-dimensional parameter space of LLMs, these methods inherently suffer from exploding gradient variance and prohibitive costs. Conversely, confining updates to a low-dimensional subspace reduces variance but discards critical high-dimensional information, degrading performance. This creates a fundamental dilemma: full-space optimization is informative but has high variance, while subspace optimization has low variance but loses information. To break this trade-off, we propose ENGINE (Endogenous Variational MultiScale). Inspired by computational mechanics, we adapt static Variational MultiScale (VMS) theory to dynamic optimization, treating low- and high-dimensional spaces as resolved coarse and unresolved fine scales. Through a multigrid-style V-cycle, ENGINE alternates between spaces, dynamically probing fine-grained information in the high-dimensional space and injecting it to correct the efficient low-dimensional sampling process. Under appropriate assumptions, statistical analysis characterizes variance reduction and correction of curvature-induced bias, and convergence analysis establishes a stationarity bound for the corrected coarse updates. Evaluations on standard LLM benchmarks show that ENGINE achieves higher accuracy and faster convergence than existing ZO methods, achieving up to a $2.61\times$ speed-up in forward passes with zero additional peak memory overhead. Code is available at https://github.com/AI4Science-WestlakeU/ENGINE.