论文

通过功率迭代的部分正交化加速零阶谱优化

Accelerating Zeroth-Order Spectral Optimization with Partial Orthogonalization from Power Iteration

模型训练参数高效训练

摘要

零阶 (ZO) 优化在 微调 大语言模型 (LLM) 中变得越来越流行和重要,特别是在边缘设备上,因为它能够将模型调整为本地数据,而不需要内存密集型反向传播。最近的工作尝试通过低维子空间搜索来减少 ZO 方差,但仅子空间限制就导致关键优化几何体未得到充分利用,从而激发了额外的加速。在这项工作中,我们重点关注隐藏层训练问题,其中 Muon 等光谱优化器由于能够通过正交化利用弱光谱方向而优于 AdamW。然而,我们发现,与一阶设置不同,完全正交化在 ZO 设置中效果不佳,因为梯度估计噪声很大且不可靠。为了解决这个问题,我们建议应用部分光谱正交化来加速 ZO 优化。为此,我们用更快、更集中的功率迭代方法取代了 Muon 中标志性的牛顿-舒尔茨过程,这样它只放大主要的光谱方向。此外,为了提高算法的效率和泛化性,我们采用了功率迭代的流式变体,它需要梯度的低方差,这是通过将我们的搜索限制在通过动量投影获得的子空间内来实现的,这与最近的进展相呼应。 LLM 微调 上的实验表明,我们的方法可以在 OPT-13B 模型中的 SuperGlue 数据集上实现当前 SOTA 算法 ZO-Muon 1.5 倍到 4 倍的收敛速度。在不同的模型中,与 MeZO、LOZO 和 ZO-Muon 等强大的 ZO 基线相比,我们在大多数情况下都用更少的时间达到了有竞争力的最终精度。代码可在 https://github.com/MOFA-LAB/ZO-MOPI.git 获取。