[论文解读] Momentum-based variance-reduced proximal stochastic gradient method for composite nonconvex stochastic optimization

Yangyang Xu, Xu, Yibo|arXiv (Cornell University)|May 31, 2020

Stochastic Gradient Optimization Techniques参考文献 36被引用 6

一句话总结

该论文提出PStorm，一种用于复合非凸随机优化的基于动量的方差缩减近端随机梯度方法，仅需每次更新使用一个或$O(1)$个样本，即可实现$O(\varepsilon^{-3})$的最优复杂度。通过利用动量和基于Bregman散度的镜像-近端更新，PStorm实现了高效的小批量训练并提升了泛化性能，在稀疏神经网络上优于现有最优方法和原始SGD。

ABSTRACT

Stochastic gradient methods (SGMs) have been extensively used for solving stochastic problems or large-scale machine learning problems. Recent works employ various techniques to improve the convergence rate of SGMs for both convex and nonconvex cases. Most of them require a large number of samples in some or all iterations of the improved SGMs. In this paper, we propose a new SGM, named PStorm, for solving nonconvex nonsmooth stochastic problems. With a momentum-based variance reduction technique, PStorm can achieve the optimal complexity result $O(\varepsilon^{-3})$ to produce a stochastic $\varepsilon$-stationary solution, if a mean-squared smoothness condition holds. Different from existing optimal methods, PStorm can achieve the ${O}(\varepsilon^{-3})$ result by using only one or $O(1)$ samples in every update. With this property, PStorm can be applied to online learning problems that favor real-time decisions based on one or $O(1)$ new observations. In addition, for large-scale machine learning problems, PStorm can generalize better by small-batch training than other optimal methods that require large-batch training and the vanilla SGM, as we demonstrate on training a sparse fully-connected neural network and a sparse convolutional neural network.

研究动机与目标

• 设计一种随机梯度方法，实现在非凸非光滑问题上仅使用最少采样次数的最优收敛复杂度。
• 通过每次迭代仅使用一个或$O(1)$个样本，实现高效的小批量和在线学习，提升机器学习中的泛化性能。
• 将方差缩减技术扩展至具有非光滑正则项的复合非凸问题，同时保持最优收敛速率。
• 在稀疏神经网络训练中，相比原始SGD和其他最优方法，展示出更优的性能。

提出的方法

• PStorm采用镜像-近端框架，使用1-强凸势函数$w$来定义Bregman散度$V(\mathbf{x}, \mathbf{z})$。
• 该算法采用基于动量的随机梯度更新，利用序列$\{\beta_k\}$降低梯度方差并加速收敛。
• 在每次迭代中，仅使用一个或少数几个$\xi$样本，计算当前和前一迭代点的随机梯度。
• 该方法结合近端更新与基于动量的方差缩减，可在均方光滑性条件下实现$O(\varepsilon^{-3})$复杂度。
• 该算法是STORM方法的近端变体，命名为PStorm，专为在线学习和小批量学习设计。
• 其采用步长规则与自适应动量，以在非凸设置下平衡收敛性与稳定性。

实验结果

研究问题

• RQ1 能否设计一种随机梯度方法，在仅使用$O(1)$个样本每步更新的条件下，实现非凸非光滑问题的最优$O(\varepsilon^{-3})$复杂度？
• RQ2 基于动量的方差缩减如何提升稀疏神经网络小批量训练中的收敛性与泛化性能？
• RQ3 近端随机梯度方法是否能在不依赖大批次的情况下，于均方光滑性条件下保持最优复杂度？
• RQ4 在测试准确率、平稳性违反程度以及神经网络训练中的稀疏性方面，PStorm与原始SGD及其他最优方法相比表现如何？

主要发现

• 在均方光滑性条件下，PStorm仅使用每次更新一个样本，即可实现寻找$\varepsilon$-平稳解的最优$O(\varepsilon^{-3})$复杂度。
• 在稀疏全连接神经网络上，PStorm在$\lambda = 5 \times 10^{-4}$时达到88.17%的测试准确率，优于原始SGD和Hybrid-SGD。
• 在Cifar10数据集上使用AllCNN，PStorm在$\lambda = 5 \times 10^{-4}$时达到88.53%的测试准确率，与或超过Hybrid-SGD，显著优于Spiderboost。
• 在$\lambda = 0$时，PStorm的平稳性违反程度最低（0.19），在$\lambda = 2 \times 10^{-4}$时仍保持高稀疏性（72.78%非零参数）。
• 在所有测试超参数和数据集上，PStorm在测试准确率和平稳性违反程度方面始终表现最佳或第二佳。
• 由于具备小批量训练能力和方差缩减机制，该方法的泛化性能优于原始SGD和其他最优方法。