On the infinite width limit of neural networks with a standard parameterization

By Jascha Narain Sohl-Dickstein, Roman Novak, Sam Schoenholz, Jaehoon Lee. 2020

56 citations recorded by the source corpus.

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 this https URL.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Statistical and Nonlinear Physics, Scaling & Learning Theory, Statistical Learning Theory

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. J. Edward Hu
2. Greg Yang
3. Timothy Nguyen
4. Jaehoon Lee
5. Zhourung Chen
6. Sam Schoenholz
7. Lechao Xiao
8. Jeffrey Pennington

In Computer Vision and Pattern Recognition

1. J. Edward Hu
2. Timothy Nguyen
3. Jaehoon Lee
4. Zhourung Chen
5. Greg Yang
6. Hakan Bilen
7. Bo Zhao
8. Sam Schoenholz

In Statistical and Nonlinear Physics

1. Greg Yang
2. Timothy Nguyen
3. Zhourung Chen
4. J. Edward Hu
5. Jaehoon Lee
6. Sam Schoenholz

In Scaling & Learning Theory

1. J. Edward Hu
2. Greg Yang
3. Jaehoon Lee
4. Timothy Nguyen
5. Zhourung Chen
6. Sam Schoenholz
7. Lechao Xiao
8. Jeffrey Pennington

In Statistical Learning Theory

1. J. Edward Hu
2. Jaehoon Lee
3. Zhourung Chen
4. Timothy Nguyen
5. Greg Yang
6. Etai Littwin

Citations in the corpus

• Feature Learning in Infinite-Width Neural Networks (2020, background): “[41] also noted that standard parametrization leads to unstable training dynamics.”
• : "Statistical Modeling: The Two Cultures" (Statistical Science, 2001, 16(3), 199-231) (unknown): “In support of this hypothesis are perhaps the many papers illustrating close connections between neural networks with an infinite number of parameters and specialized, complex kernel ridge regression methods (Lee et al., 2019; Sohl-Dickstein et al., 2020).”
• Dataset Meta-Learning from Kernel Ridge-Regression (2020, methodology): “7While with NTK parameterization Jacot et al. (2018), the kernel does not have any width dependence, with the improved standard parameterization of Sohl-Dickstein et al. (2020), one could encode width dependence into NTK. Table A1: Accuracy on random subsets of MNIST.”