Search papers, labs, and topics across Lattice.
This paper introduces AdaLeZO, an Adaptive Layer-wise Zeroth-Order optimization framework designed to improve the efficiency of fine-tuning Large Language Models by addressing the significant runtime bottleneck caused by uniform exploration strategies. By treating the layer selection process as a non-stationary Multi-Armed Bandit problem, AdaLeZO dynamically allocates perturbation budgets to the most sensitive parameters, leading to substantial reductions in training latency. Experimental results show that AdaLeZO accelerates wall-clock time by 1.7x to 3.0x compared to existing state-of-the-art methods, while maintaining unbiased gradient estimation and requiring no additional memory overhead.
A single adaptive framework can boost the efficiency of zeroth-order optimization by up to 3x without increasing memory usage.
Zeroth-Order optimization presents a promising memory-efficient paradigm for fine-tuning Large Language Models by relying solely on forward passes. However, its practical adoption is severely constrained by slow wall-clock convergence and high estimation variance. In this work, we dissect the runtime characteristics of ZO algorithms and identify a critical system bottleneck where the generation of perturbations and parameter updates accounts for over 40% of the training latency. We argue that the standard uniform exploration strategy is fundamentally flawed as it fails to account for the heterogeneous sensitivity of layers in deep networks, resulting in computationally wasteful blind searches. To address this structural mismatch, we propose AdaLeZO, an Adaptive Layer-wise ZO optimization framework. By formulating the layer selection process as a non-stationary Multi-Armed Bandit problem, AdaLeZO dynamically allocates the limited perturbation budget to the most sensitive parameters. We further introduce an Inverse Probability Weighting mechanism based on sampling with replacement, which guarantees unbiased gradient estimation while effectively acting as a temporal denoiser to reduce variance. Extensive experiments on LLaMA and OPT models ranging from 6.7B to 30B parameters demonstrate that AdaLeZO achieves 1.7x to 3.0x wall-clock acceleration compared to state-of-the-art methods. Crucially, AdaLeZO functions as a universal plug-and-play module that seamlessly enhances the efficiency of existing ZO optimizers without incurring additional memory overhead.
