Hi-ZFO: Hierarchical Zeroth- and First-Order LLM Fine-Tuning via Importance-Guided Tensor Selection
Feihu Jin, Ying Tan
TL;DR
Hi-ZFO addresses the instability and inefficiency of pure zeroth-order fine-tuning for large language models in generative tasks by introducing a hierarchical hybrid optimization that assigns first-order updates to high-importance layers and zeroth-order updates to lower-importance layers. The method leverages cost-aware, DP-based tensor selection to partition parameters and uses a dual-stream forward pass with a coordinated hybrid objective, L_total = L_FO + alpha L_ZO, to achieve stable, exploratory optimization. Theoretical convergence is established under standard smoothness and variance assumptions, and extensive experiments across generative, reasoning, and coding benchmarks show improved generalization with reduced training time and memory footprint compared with full fine-tuning and existing baselines. Overall, Hi-ZFO demonstrates that structured, stochastic perturbations in non-critical layers can significantly enhance optimization dynamics for large models in resource-constrained fine-tuning scenarios.
Abstract
Fine-tuning large language models (LLMs) using standard first-order (FO) optimization often drives training toward sharp, poorly generalizing minima. Conversely, zeroth-order (ZO) methods offer stronger exploratory behavior without relying on explicit gradients, yet suffer from slow convergence. More critically, our analysis reveals that in generative tasks, the vast output and search space significantly amplify estimation variance, rendering ZO methods both noisy and inefficient. To address these challenges, we propose \textbf{Hi-ZFO} (\textbf{Hi}erarchical \textbf{Z}eroth- and \textbf{F}irst-\textbf{O}rder optimization), a hybrid framework designed to synergize the precision of FO gradients with the exploratory capability of ZO estimation. Hi-ZFO adaptively partitions the model through layer-wise importance profiling, applying precise FO updates to critical layers while leveraging ZO optimization for less sensitive ones. Notably, ZO in Hi-ZFO is not merely a memory-saving surrogate; it is intentionally introduced as a source of "beneficial stochasticity" to help the model escape the local minima where pure FO optimization tends to stagnate. Validated across diverse generative, mathematical, and code reasoning tasks, Hi-ZFO consistently achieves superior performance while significantly reducing the training time. These results demonstrate the effectiveness of hierarchical hybrid optimization for LLM fine-tuning.