[论文解读] Gram-Gauss-Newton Method: Learning Overparameterized Neural Networks for Regression Problems

Tianle Cai, Ruiqi Gao+6 more|arXiv (Cornell University)|May 28, 2019

Neural Networks and Applications参考文献 53被引用 31

一句话总结

本文提出Gram-Gauss-Newton（GGN）方法，一种用于在平方损失回归任务中训练过参数化神经网络的新型二阶优化算法。通过利用神经正切核（NTK）并使用雅可比矩阵的格拉姆矩阵（J J⊤）求解核回归问题，GGN在仅增加少量计算开销的情况下实现了二次收敛，使其在大规模模型中既具备理论高效性又具备实际可行性。

ABSTRACT

First-order methods such as stochastic gradient descent (SGD) are currently the standard algorithm for training deep neural networks. Second-order methods, despite their better convergence rate, are rarely used in practice due to the prohibitive computational cost in calculating the second-order information. In this paper, we propose a novel Gram-Gauss-Newton (GGN) algorithm to train deep neural networks for regression problems with square loss. Our method draws inspiration from the connection between neural network optimization and kernel regression of neural tangent kernel (NTK). Different from typical second-order methods that have heavy computational cost in each iteration, GGN only has minor overhead compared to first-order methods such as SGD. We also give theoretical results to show that for sufficiently wide neural networks, the convergence rate of GGN is \emph{quadratic}. Furthermore, we provide convergence guarantee for mini-batch GGN algorithm, which is, to our knowledge, the first convergence result for the mini-batch version of a second-order method on overparameterized neural networks. Preliminary experiments on regression tasks demonstrate that for training standard networks, our GGN algorithm converges much faster and achieves better performance than SGD.

研究动机与目标

• 为解决由于计算成本过高而导致二阶优化方法在深度学习中实际应用受限的问题。
• 开发一种计算高效的二阶方法，同时保留基于海森矩阵方法的快速收敛特性。
• 为过参数化神经网络上的二阶方法的随机小批量变体建立理论收敛保证。
• 利用神经网络优化与核回归之间的联系，通过神经正切核（NTK）实现。
• 通过在每一步直接求解基于NTK的回归问题，实现更快的训练速度和更优的性能表现，以提升回归任务的性能。

提出的方法

• GGN方法通过使用格拉姆矩阵 J J⊤ 代替海森矩阵近似 J⊤J，重新表述了高斯-牛顿更新，将计算成本从 O(m²) 降低至 O(n²)，其中 n 为训练样本数量。
• 在每次迭代中，利用NTK求解一个核回归问题，其中核函数由网络输出相对于其参数的雅可比矩阵导出。
• 该方法采用小批量训练以引入随机性，从而提升泛化能力并降低每轮迭代的计算成本，格拉姆矩阵的规模随小批量大小 b 而缩放为 b×b。
• 该算法通过使用格拉姆矩阵的逆矩阵隐式执行牛顿型更新，即使在海森矩阵不可逆时，该格拉姆矩阵仍具有良好的条件性和可逆性。
• 理论分析表明，GGN在足够宽的网络中可实现二次收敛，优于SGD的线性收敛速度。
• 通过使用谱范数和特征值界，该方法设计为在紧凑区域 B(R) 内限制参数更新，以保持稳定性和收敛性。

实验结果

研究问题

• RQ1 能否通过降低计算开销，使二阶优化方法在大规模过参数化神经网络中具备实际可行性？
• RQ2 在每一步求解基于NTK的核回归是否能带来比标准一阶方法（如SGD）更快的收敛速度？
• RQ3 在过参数化设置下，二阶方法的小批量版本能否实现理论收敛保证？
• RQ4 格拉姆矩阵 J J⊤ 是否可作为深度学习中二阶优化的海森矩阵近似的一种可行且稳定的替代方案？
• RQ5 在NTK框架下，将二阶方法应用于过参数化网络时，其收敛速率如何？

主要发现

• GGN在足够宽的过参数化神经网络中实现了二次收敛，显著优于SGD的线性收敛速率。
• 与SGD相比，GGN每轮迭代的计算开销极低，仅需计算大小为 n×n 或 b×b 的 J J⊤ 及其逆矩阵。
• 首次为过参数化神经网络上的小批量二阶方法建立了理论收敛保证，收敛速率为 O((1−Ω(λ₀²/n²))ᵗ)。
• 实验结果表明，GGN在标准回归任务中收敛速度远快于SGD，并取得了更优的性能表现。
• 该方法在整个训练过程中保持稳定，因为参数更新被限制在紧凑区域 B(R) 内，其中 R=Θ(n⁵/λ₀⁴)。
• 理论分析证实，格拉姆矩阵 Gₜᵢ,ᵢ 的逆矩阵具有良好的条件性和有界性，从而确保了数值稳定性。