[论文解读] On the Generalization of Stochastic Gradient Descent with Momentum

Ali Ramezani-Kebrya, Antonakopoulos, Kimon+3 more|arXiv (Cornell University)|Sep 12, 2018

Stochastic Gradient Optimization Techniques参考文献 37被引用 11

一句话总结

本文提出SGDEM（早期动量随机梯度下降），一种改进的随机梯度下降变体，仅在训练初期的若干轮次中应用动量，以提升泛化性能。理论证明标准SGDM在凸损失函数下可能存在无界的稳定性差距，而SGDEM可确保对光滑Lipschitz函数和强凸函数的泛化性，并推导出预期真实风险的理论边界。实验验证在ResNet-18和逻辑回归模型上表现良好。

ABSTRACT

While momentum-based accelerated variants of stochastic gradient descent (SGD) are widely used when training machine learning models, there is little theoretical understanding on the generalization error of such methods. In this work, we first show that there exists a convex loss function for which the stability gap for multiple epochs of SGD with standard heavy-ball momentum (SGDM) becomes unbounded. Then, for smooth Lipschitz loss functions, we analyze a modified momentum-based update rule, i.e., SGD with early momentum (SGDEM) under a broad range of step-sizes, and show that it can train machine learning models for multiple epochs with a guarantee for generalization. Finally, for the special case of strongly convex loss functions, we find a range of momentum such that multiple epochs of standard SGDM, as a special form of SGDEM, also generalizes. Extending our results on generalization, we also develop an upper bound on the expected true risk, in terms of the number of training steps, sample size, and momentum. Our experimental evaluations verify the consistency between the numerical results and our theoretical bounds. SGDEM improves the generalization error of SGDM when training ResNet-18 on ImageNet in practical distributed settings.

研究动机与目标

• 为解决基于动量的SGD变体（如SGDM）在泛化性方面缺乏理论理解的问题。
• 表明即使在凸损失函数下，标准SGDM在多轮训练中也可能出现无界的稳定性差距。
• 提出SGDEM——一种修改后的动量更新规则，仅在训练初期阶段应用动量，以确保泛化性保证。
• 推导出以训练步数、样本量和动量参数表示的预期真实风险的上界。
• 在ResNet-18（ImageNet）和逻辑回归（notMNIST/MNIST）上通过实验验证理论结果，表明在调整动量持续时间后泛化性能得到提升。

提出的方法

• 提出SGDEM，即SGDM的一种变体，其中动量仅在前$ t_d $个周期内应用，之后禁用，以控制泛化行为。
• 采用统一稳定性框架，分析SGDEM在光滑Lipschitz函数和强凸损失函数下的泛化误差。
• 通过稳定性分析推导出预期真实风险的上界，纳入学习率、动量、训练步数和样本量等参数。
• 证明对于光滑Lipschitz函数，SGDEM在广泛的学习率条件下无论动量值如何均能实现泛化。
• 建立对于强凸函数的结论：当动量处于特定范围内时，标准SGDM（作为SGDEM的特例）可实现泛化。
• 通过在ResNet-18（ImageNet）和逻辑回归（notMNIST/MNIST）上的实验验证理论边界，比较不同动量调度下的泛化误差与测试准确率。

实验结果

研究问题

• RQ1 标准SGDM在使用Heavy-ball动量时，是否在多轮训练中对凸损失函数表现出无界的稳定性差距？
• RQ2 SGDEM——一种修改后的动量更新规则——是否能在广泛学习率条件下，对光滑Lipschitz损失函数实现泛化保证？
• RQ3 在强凸损失函数下，标准SGDM的动量参数应处于何种范围才能实现泛化？
• RQ4 动量应用时长（$ t_d $）在实际训练场景中如何影响泛化误差？
• RQ5 能否为SGDEM推导出预期真实风险的理论边界，并通过实证方法验证？

主要发现

• 对于特定凸损失函数，标准SGDM在多轮训练中稳定性差距趋于无界，表明其可能损害泛化性能。
• SGDEM在广泛学习率条件下，对光滑Lipschitz损失函数可实现泛化，且不依赖于动量值大小。
• 对于强凸函数，当动量参数位于特定范围内时，标准SGDM可实现泛化，这是新的理论贡献。
• 推导出预期真实风险的理论边界，并与数值结果保持一致，最优$ \tilde{t}^* $通过求解非解析方程获得。
• 在分布式设置下于ImageNet上训练ResNet-18时，当动量仅在前50轮（共90轮）内应用时，SGDEM的泛化误差低于SGDM。
• 在notMNIST和MNIST上的逻辑回归实验中，测试准确率在动量仅应用于有限轮次时达到峰值，证实过长的动量应用可能降低泛化性能。