Universally Empowering Zeroth-Order Optimization via Adaptive Layer-wise Sampling

2026-04-20

AI Summary

This work addresses the inefficiency of zeroth-order optimization in large language model fine-tuning, where uniform exploration strategies ignore layer-wise sensitivity differences, leading to slow convergence and high gradient variance. To overcome this, the authors propose AdaLeZO, a novel framework that formulates layer selection as a non-stationary multi-armed bandit problem, dynamically allocating perturbation budgets to the most sensitive layers. By integrating sampling with replacement and inverse probability weighting, AdaLeZO yields unbiased, low-variance gradient estimates. The method features a plug-and-play adaptive inter-layer sampling module that incurs no additional memory overhead. Experiments on LLaMA and OPT models ranging from 6.7B to 30B parameters demonstrate 1.7–3.0× faster training compared to state-of-the-art baselines, substantially enhancing zeroth-order optimization efficiency.

Technology Category

• Machine Learning: Online Learning & Bandits
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Zeroth-Order Optimization
• Large Language Models
• Convergence Speed
• Estimation Variance
• Layer-wise Sensitivity

Innovation

• Zeroth-Order Optimization
• Adaptive Layer-wise Sampling
• Multi-Armed Bandit
• Inverse Probability Weighting
• Memory-Efficient Fine-Tuning