[论文解读] On the infinite width limit of neural networks with a standard parameterization

Jascha Sohl‐Dickstein, Roman Novak+2 more|arXiv (Cornell University)|Jan 21, 2020

Stochastic Gradient Optimization Techniques参考文献 26被引用 14

一句话总结

本文提出了一种改进的标准参数化方法，用于神经网络，该方法保持了学习率的稳定缩放，保留了各层宽度的相对影响，并在宽度趋于无穷时确保梯度流动的稳定性。与NTK或朴素的标准参数化不同，该方法可得到一个明确定义的神经正切核，同时更好地匹配有限宽度下的训练动态，并在实验中实现了相当或更优的性能。

ABSTRACT

There are currently two parameterizations used to derive fixed kernels corresponding to infinite width neural networks, the NTK (Neural Tangent Kernel) parameterization and the naive standard parameterization. However, the extrapolation of both of these parameterizations to infinite width is problematic. The standard parameterization leads to a divergent neural tangent kernel while the NTK parameterization fails to capture crucial aspects of finite width networks such as: the dependence of training dynamics on relative layer widths, the relative training dynamics of weights and biases, and overall learning rate scale. Here we propose an improved extrapolation of the standard parameterization that preserves all of these properties as width is taken to infinity and yields a well-defined neural tangent kernel. We show experimentally that the resulting kernels typically achieve similar accuracy to those resulting from an NTK parameterization, but with better correspondence to the parameterization of typical finite width networks. Additionally, with careful tuning of width parameters, the improved standard parameterization kernels can outperform those stemming from an NTK parameterization. We release code implementing this improved standard parameterization as part of the Neural Tangents library at https://github.com/google/neural-tangents.

研究动机与目标

• 解决现有无限宽度神经网络参数化中因无法保持关键有限宽度训练动态而产生的不一致问题。
• 解决朴素标准参数化在宽度增加时出现的神经正切核发散问题。
• 在无限宽度极限下保持稳定且有限的学习率缩放，与实际训练实践相匹配。
• 确保权重和偏置的相对训练动态得以保留，这与NTK参数化不同。
• 开发一种参数化方法，使其产生明确定义的神经正切核，同时紧密匹配典型有限宽度网络的行为。

提出的方法

• 引入一个宽度缩放因子 $ s $，其中每层的宽度为 $ n^l = sN^l $，并在取 $ s \to \infty $ 时得到无限宽度极限。
• 将仿射层重写为 $ z^{l+1} = \frac{1}{\sqrt{s}} W^l y^l + b^l $，其中 $ W^l \sim \mathcal{N}(0, \sigma_w^2 / N^l) $，以保持方差缩放。
• 使用 $ b^l \sim \mathcal{N}(0, \sigma_b^2) $ 以在无限宽度极限下保持偏置贡献的恒定。
• 从该参数化推导出无限宽度下的NNGP和NTK核，证明其收敛到明确定义的形式。
• 通过蒙特卡洛模拟验证核的行为，并将其集成到Neural Tangents库中。
• 调整层宽 $ N^l $ 以优化泛化性能，证明在特定设置下优于NTK核。

实验结果

研究问题

• RQ1 能否在不导致神经正切核发散且不丢失有限宽度训练动态的前提下，将标准参数化扩展至无限宽度？
• RQ2 改进的标准参数化是否在无限宽度极限下保留了层宽对训练动态的相对影响？
• RQ3 当宽度参数被仔细调优时，改进的标准参数化是否能在泛化性能上优于NTK参数化？
• RQ4 在无限宽度下，改进的标准参数化与NTK参数化相比，权重和偏置的学习动态有何差异？
• RQ5 改进的标准参数化是否产生一个更能反映有限宽度网络实际训练行为的神经正切核？

主要发现

• 改进的标准参数化产生了一个不会发散的明确定义的神经正切核，与朴素标准参数化不同。
• 所得核在全连接ReLU网络中的测试准确率与NTK参数化相当。
• 通过仔细调优层宽 $ N^l $，改进的标准参数化在全连接网络上的泛化性能可优于NTK参数化。
• 改进的参数化保留了层宽对训练动态的相对影响，而NTK参数化在宽度增加时会失去这一特性。
• 权重和偏置的相对训练动态得以保留，而NTK参数化中偏置被过度强调。
• 随着宽度增加，偏置对神经正切核的贡献保持恒定，而权重贡献随 $ N^l $ 增大，表明在宽网络中偏置的影响变小。