探针空间预条件处理：快速稳定的零阶训练
Probe-Space Preconditioning for Fast and Stable Zero-Order Training
查看机构详情
- Stanford University（斯坦福大学）
机构由 AI 辅助整理，请以论文原文为准。
浏览论文内容
中文总结 AI 辅助
本文提出1.5-SPSA零阶优化方法，通过探针空间对角预条件器，在更少步骤下超越MeZO和BP，实现快速稳定的模型训练。
中文摘要 AI 辅助
反向传播（BP）主导深度学习，但会带来巨大的内存开销。例如，使用Adam训练OPT-30B需要约600GB的GPU内存（假设批大小为8，序列长度为2048）。另一种选择是零阶优化（ZOO），它在推理模式下训练（同一模型仅需约60GB）：无需存储激活值、无需梯度、也无需优化器状态。然而，ZOO的收敛速度一直落后于BP。在这项工作中，我们评估了两种方法来缩小这一差距。首先，我们表明将训练计算预算从多步重新分配到具有大量扰动（或探针）的大有效批大小但更少的步骤，使得1SPSA（Spall, 1992）能够以更少的训练计算量超越MeZO（Malladi等，2023）等零阶方法。接下来，我们引入1.5-SPSA，在1SPSA的每一步中添加一次“干净”的前向传播，以计算探针空间中的廉价对角预条件器，通过降低高曲率方向的权重来改善收敛速度和收敛性。在Qwen3和OPT模型家族的6个训练后数据集上进行基准测试，我们表明1.5-SPSA以更少的优化步骤取得了优于先前ZOO求解器的最先进结果。例如，我们训练OPT-13B（用于与MeZO直接比较），发现1.5-SPSA仅用70步就在SST-2上比MeZO和BP均提高了+3.1%的准确率，而MeZO需要100,000步。最后，我们结合8位打包随机生成器、Triton融合解包/应用内核以及分布式并行，实现了在商用GPU（如A100）上就地快速稳定地训练高达OPT-30B规模的模型。
英文摘要
Backpropagation (BP) dominates deep learning but imposes a massive memory tax. For example, training OPT-30B with Adam requires $\approx$ 600GB of GPU memory (assuming batch size 8 and sequence length 2048). Alternatively, zero-order optimization (ZOO) trains in inference-mode (requiring only $\approx$ 60GB for the same model): no stored activations, no gradients, and no optimizer states. However, ZOO convergence has lagged behind BP. In this work, we evaluate two methods to close this gap. First, we show that reallocating training compute budget from many steps to large effective batch sizes with many perturbations (or probes) but fewer steps, allows 1SPSA (Spall, 1992) to outperform zero order methods like MeZO (Malladi et al., 2023) with less training compute. Next, we introduce 1.5-SPSA, adding a single "clean" forward-pass per step to 1SPSA to calculate a cheap diagonal preconditioner in probe-space, which improves convergence rate and convergence by down-weighting high curvature directions. Benchmarking on 6 post-training datasets on both Qwen3 and OPT model families, we show that 1.5-SPSA achieves State-of-the-Art results over previous ZOO solvers with much less optimization steps. For example, we train OPT-13B (for direct comparison to MeZO) and find 1.5-SPSA achieves +3.1% accuracy on SST-2 over both MeZO and BP in only 70 steps vs. MeZO's 100,000 steps. Finally, we combine an 8-bit-packing random generator, triton fused unpack/apply kernels, and distributed parallelism to achieve fast and stable training of models as large as OPT-30B in-place on commodity GPUs (e.g. A100).
