[论文解读] Apollo: An Adaptive Parameter-wise Diagonal Quasi-Newton Method for Nonconvex Stochastic Optimization

Xuezhe Ma|arXiv (Cornell University)|Sep 28, 2020

Stochastic Gradient Optimization Techniques被引用 14

一句话总结

Apollo 是一种用于非凸随机优化的参数化对角拟牛顿方法，通过自适应地使用对角矩阵近似 Hessian 矩阵，并对其绝对值进行修正以确保正定性。该方法实现了线性时间与内存复杂度，在视觉和语言任务中均优于 SGD 和 Adam 变体，展现出更快的收敛速度与更优的泛化性能。

ABSTRACT

In this paper, we introduce Apollo, a quasi-Newton method for nonconvex stochastic optimization, which dynamically incorporates the curvature of the loss function by approximating the Hessian via a diagonal matrix. Importantly, the update and storage of the diagonal approximation of Hessian is as efficient as adaptive first-order optimization methods with linear complexity for both time and memory. To handle nonconvexity, we replace the Hessian with its rectified absolute value, which is guaranteed to be positive-definite. Experiments on three tasks of vision and language show that Apollo achieves significant improvements over other stochastic optimization methods, including SGD and variants of Adam, in term of both convergence speed and generalization performance. The implementation of the algorithm is available at https://github.com/XuezheMax/apollo.

研究动机与目标

• 解决如 Adam 等一阶自适应方法的局限性，这些方法可能收敛至较差的局部极小值且泛化性能不佳。
• 克服经典拟牛顿方法在高维、非凸、随机设置（如深度学习）中效率低下的问题。
• 开发一种可扩展的拟牛顿方法，高效地整合曲率信息，同时避免高昂的计算或内存开销。
• 通过将 Hessian 近似修正为正定矩阵（利用绝对值的修正操作），确保在非凸设置下的稳定性和收敛性。

提出的方法

• 基于梯度差值迭代更新一个对角矩阵来近似 Hessian 矩阵，满足参数化弱 secant 条件。
• 对对角 Hessian 近似应用修正的绝对值操作，以保证正定性并处理非凸性。
• 采用线性预热的学习率调度策略，以提升初始阶段的收敛稳定性。
• 通过仅存储对角元素并增量式更新，保持线性时间与内存复杂度。
• 通过将修正阈值与梯度差值的无穷范数耦合，引入一种尺度不变的变体，无需额外超参数。
• 支持参数化梯度裁剪作为自然扩展，提升训练的鲁棒性。

实验结果

研究问题

• RQ1 能否使一种对角拟牛顿方法在深度学习的非凸随机优化中兼具高效性与有效性？
• RQ2 通过使用对角 Hessian 近似引入曲率信息，相较于一阶方法，其对收敛速度与泛化性能的影响如何？
• RQ3 能否在保持线性复杂度的前提下，使 Hessian 近似对非凸性具有鲁棒性？
• RQ4 Hessian 近似的修正绝对值是否能在实际训练中提升稳定性和泛化性能？
• RQ5 能否在不引入新超参数的前提下使该方法具备尺度不变性？

主要发现

• 在 CIFAR-10 上，Apollo 的训练速度比 SGD 快 1.42 倍，在 ImageNet 上快 1.49 倍，且在 ResNeXt-50 上内存消耗降低 1.62 倍。
• 在 One Billion Words 语言建模任务中，Apollo 的测试损失为 31.75 ± 0.10，优于标准梯度裁剪方法（31.94 ± 0.09）。
• 在 WMT-14 神经机器翻译任务中，Apollo 的 BLEU 得分为 28.39 ± 0.11，略高于标准梯度裁剪方法（28.34 ± 0.10）。
• 尽管每次迭代快 10 倍以上且内存使用显著更低，Apollo 在小型 CNN 模型上的收敛速度仍快于 SdLBFGS。
• Apollo 的尺度不变变体消除了对额外超参数的需求，同时保持了性能。
• 与 AdaHessian（一种二阶方法）相比，Apollo 的计算成本显著更低——在 CIFAR-10 和 ImageNet 上分别快 5.76 倍和 11.78 倍，凸显了 Apollo 的高效性。