理解随机伪梯度上的外层 Nesterov 动量稳定性

arXiv:2610.22623v1 公告类型：新 摘要：大型神经网络通常由多个工作节点协同训练，节点之间的通信是一项主要成本。因此，工作节点会先执行若干次本地优化步骤，再将其参数的总位移（即伪梯度）发送给外层优化器，而外层优化器通常会应用 Nesterov 动量。我们研究外层动量、本地步骤数以及工作节点间的平均操作如何共同影响训练动态，尤其是其稳定性。对于随机二次函数，我们在任意维度下推导出了精确的均方稳定性判据。该判据揭示，内层批大小决定了增加本地步骤数最终会使外层迭代趋于稳定还是失稳。此外，尽管对独立工作节点的伪梯度进行平均可以允许采用更大的外层动量值，但实现稳定性所需的工作节点数量可能会随本地步骤数呈指数增长。最后，我们表明，可以使用共享小批量扰动，在神经网络训练过程中的中间检查点直接探测决定稳定性的均方动态。我们在不同网络上检验了理论发现，并测量了它们对训练损失的影响。