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

📄 arXiv: 2604.18264v1 📥 PDF

作者: Fei Wang, Li Shen, Liang Ding, Chao Xue, Ye Liu, Changxing Ding

分类: cs.LG

发布日期: 2026-04-20

💡 一句话要点

提出AdaLeZO，通过自适应层采样优化零阶优化，加速大语言模型微调。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 零阶优化 大语言模型 自适应采样 多臂老虎机 逆概率加权

📋 核心要点

  1. 现有零阶优化方法在微调大模型时存在收敛速度慢、方差大的问题，扰动生成和参数更新成为性能瓶颈。
  2. AdaLeZO将层选择建模为多臂老虎机问题，动态分配扰动预算给最敏感的层，提高采样效率。
  3. AdaLeZO采用逆概率加权机制，保证梯度估计的无偏性，并起到时间去噪的作用，加速收敛。

📝 摘要（中文）

零阶优化（ZO）为大语言模型微调提供了一种内存高效的范式，仅依赖前向传播。然而，其缓慢的收敛速度和高估计方差严重限制了实际应用。本文分析了ZO算法的运行时特征，发现扰动生成和参数更新占据了超过40%的训练延迟，是关键的系统瓶颈。标准的一致探索策略存在根本缺陷，因为它未能考虑深度网络中各层异构的敏感性，导致计算资源的浪费。为此，我们提出了自适应层ZO优化框架AdaLeZO。通过将层选择过程建模为非平稳多臂老虎机问题，AdaLeZO动态地将有限的扰动预算分配给最敏感的参数。我们进一步引入了基于有放回采样的逆概率加权机制，保证无偏梯度估计，同时有效地充当时间去噪器以减少方差。在参数规模从67亿到300亿的LLaMA和OPT模型上的大量实验表明，与最先进的方法相比，AdaLeZO实现了1.7倍到3.0倍的加速。重要的是，AdaLeZO作为一个通用的即插即用模块，可以无缝地提高现有ZO优化器的效率，而不会产生额外的内存开销。

🔬 方法详解

问题定义：论文旨在解决零阶优化（ZO）在大语言模型微调中效率低下的问题。现有的ZO方法，如基于均匀采样的算法，在深度网络中对所有层进行同等程度的扰动，忽略了不同层对模型性能的异构影响。这种盲目的搜索方式导致大量的计算资源浪费在不敏感的层上，从而降低了训练效率。此外，ZO方法固有的梯度估计方差也进一步阻碍了收敛速度。

核心思路：AdaLeZO的核心思路是自适应地调整每一层的采样概率，将更多的计算资源分配给对模型性能影响更大的层。通过将层选择过程建模为一个非平稳的多臂老虎机（MAB）问题，AdaLeZO能够动态地学习每一层的敏感度，并根据其敏感度分配扰动预算。这种自适应的采样策略能够更有效地探索参数空间，从而提高训练效率。

技术框架：AdaLeZO的整体框架包含以下几个主要模块：1) 敏感度估计模块：该模块负责估计每一层参数的敏感度。论文将层选择建模为MAB问题，并使用某种策略（例如，UCB或Thompson Sampling）来选择要扰动的层。2) 扰动生成模块：该模块根据选定的层生成扰动。3) 梯度估计模块：该模块使用ZO方法估计梯度。4) 参数更新模块：该模块使用估计的梯度更新模型参数。5) 逆概率加权模块：该模块使用逆概率加权来保证梯度估计的无偏性，并降低方差。

关键创新：AdaLeZO的关键创新在于其自适应的层采样策略和逆概率加权机制。与传统的均匀采样方法相比，AdaLeZO能够更有效地利用计算资源，从而提高训练效率。逆概率加权机制则保证了梯度估计的无偏性，并降低了方差，从而加速了收敛。此外，AdaLeZO作为一个通用的即插即用模块，可以方便地集成到现有的ZO优化器中。

关键设计：AdaLeZO的关键设计包括：1) 多臂老虎机策略：论文需要选择一种合适的MAB策略来平衡探索和利用。2) 奖励函数：论文需要设计一个合适的奖励函数来衡量每一层的敏感度。3) 逆概率加权：论文使用有放回采样，并根据采样概率计算逆概率权重，以保证梯度估计的无偏性。4) ZO梯度估计器：论文可以使用不同的ZO梯度估计器，例如，两点估计器或随机梯度估计器。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AdaLeZO在LLaMA和OPT模型（6.7B-30B参数）上实现了显著的加速。与最先进的ZO优化器相比，AdaLeZO实现了1.7倍到3.0倍的wall-clock加速。这些结果表明，AdaLeZO能够有效地提高ZO优化的效率，并降低大语言模型微调的计算成本。

🎯 应用场景

AdaLeZO可广泛应用于大语言模型的微调，尤其是在计算资源受限的场景下。该方法能够显著提高微调效率，降低训练成本，加速模型的迭代和部署。此外，AdaLeZO的自适应采样思想也可以推广到其他优化问题中，例如，神经网络架构搜索和超参数优化。

📄 摘要（原文）

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.