可扩展的 Levenberg-Marquardt 深度学习训练框架

arXiv:2610.22925v1 公告类型：新论文 摘要：Levenberg-Marquardt（LMA）算法是一种广泛使用的非线性最小二乘求解器，以快速收敛而闻名。对于小规模神经网络，与 Adam 等一阶优化器相比，LMA 能够提供更好的收敛性和架构鲁棒性。然而，LMA 会带来高昂且难以承受的内存和计算开销。这些开销主要来自雅可比矩阵，其维度同时随样本数量和参数数量增长。这一瓶颈阻碍了 LMA 在大规模神经网络中的应用。近期研究尝试使用预条件共轭梯度（PCG）求解器扩展 LMA，以避免显式构造雅可比矩阵。然而，PCG 只能近似求解，因此需要进行多次网络传播，并且需要针对具体任务设计预条件器。其他研究则通过优化 LMA 来降低开销，但将其方法限制在特定任务或网络架构上。因此，LMA 对于训练通用的大规模神经网络仍不实用。在本文中，我们提出了 LeMA，这是一种面向多 GPU 平台大规模神经网络训练的分布式 LMA 算法。我们的贡献有两点：第一，LMA 的精确对偶形式消除了对参数维度的内存依赖。第二，雅可比分片策略能够动态并行计算雅可比分片，从而避免显式构造完整的雅可比矩阵。对于大规模神经网络，LeMA 将 LMA 的内存占用从关于参数数量的二次复杂度降低为关于批次大小的二次复杂度。这使得无需使用近似的 PCG 求解，就能够对十亿参数规模的模型执行精确的 LMA 步骤。此外，在一台配备 8 个 GPU 的节点上，对于一个包含 1,000,022,632 个参数的卷积神经网络，LeMA 相较于单 GPU 实现了 7.97 倍加速，对应 99.7% 的强扩展效率。