On the infinite width limit of neural networks with a standard parameterization
This work addresses a foundational issue in neural network theory for researchers, offering a more accurate infinite width limit that captures finite network dynamics, though it is incremental in refining existing parameterizations.
The authors tackled the problem of problematic infinite width extrapolations in neural networks by proposing an improved standard parameterization that preserves key finite width properties and yields a well-defined neural tangent kernel. They showed experimentally that this approach achieves similar accuracy to NTK parameterization kernels, with better correspondence to finite networks and potential outperformance through tuning.
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.