通过部分正交化加速零阶优化,提升大模型微调速度
Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration
- 用幂迭代替代牛顿-舒尔兹,仅放大主要谱方向
- 在OPT-13B上比ZO-Muon快1.5至4倍,收敛更快
- 适合边缘设备上大模型微调,尤其对低方差梯度有效
零阶(ZO)优化在无需内存密集型反向传播的情况下,成为边缘设备上微调大语言模型的重要方法。尽管近期工作通过低维子空间搜索降低ZO方差,但子空间限制仍未能充分挖掘优化几何结构,导致加速不足。本文聚焦隐藏层训练问题,发现谱优化器如Muon因正交化能力优于AdamW,但全正交化在零阶设置中表现不佳,因梯度估计噪声大且不可靠。为此,我们提出部分谱正交化策略:将Muon中的牛顿-舒尔兹过程替换为更快速、更集中的幂迭代方法,仅放大主导谱方向;同时采用流式幂迭代,通过动量投影的子空间约束降低梯度方差,提升效率与泛化性。在OPT-13B模型的SuperGlue数据集上,本方法相比当前最优的ZO-Muon实现1.5至4倍的收敛加速;在不同模型上,也以更少时间达到与MeZO、LOZO和ZO-Muon相当甚至更优的最终精度。代码已开源。
原文摘要 · Abstract (English)
Zeroth-order (ZO) optimization has become increasingly popular and important in fine-tuning large language models (LLMs), especially on edge devices due to its ability to adjust the model to local data without the need for memory-intensive back-propagation. Recent works try to reduce ZO variance through low-dimensional subspace search, but subspace restriction alone leaves key optimization geometry under-exploited, motivating additional acceleration. In this work, we focus on the hidden layer training problem in which spectral optimizers like Muon outperform AdamW due to its ability to exploit weak spectral directions by orthogonalization. However, we have discovered that unlike in the first-order setting, full orthogonalization works poorly in the ZO setting since the gradient estimates are highly noisy and unreliable. To address this issue, we propose applying partial spectral orthogonalization to accelerate ZO optimization. To do so, we replace the iconic Newton-Schulz procedure in Muon with the faster, more concentrated power-iteration method so that it only amplifies dominant spectral directions. Furthermore, to improve the efficiency and generalization of the algorithm, we adopted a streaming variant of power-iteration that requires low variance in gradients, which was achieved through constraining our search inside a subspace obtained through the projection of momentum, echoing recent advances. Experiments on LLM fine-tuning show that our method can achieve from 1.5x to 4x the convergence speed of ZO-Muon, the current SOTA algorithm, across SuperGlue datasets in the OPT-13B model. Across different models, we also reach competitive final accuracies with less time in most cases compared with strong ZO baselines such as MeZO, LOZO and ZO-Muon. Code is available at https://github.com/MOFA-LAB/ZO-MOPI.git.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。