Wide neural networks of any depth evolve as linear models under gradient descent ^*

Abstract

A longstanding goal in deep learning research has been to precisely characterize training and generalization. However, the often complex loss landscapes of neural networks (NNs) have made a theory of learning dynamics elusive. In this work, we show that for wide NNs the learning dynamics simplify considerably and that, in the infinite width limit, they are governed by a linear model obtained from the first-order Taylor expansion of the network around its initial parameters. Furthermore, mirroring the correspondence between wide Bayesian NNs and Gaussian processes (GPs), gradient-based training of wide NNs with a squared loss produces test set predictions drawn from a GP with a particular compositional kernel. While these theoretical results are only exact in the infinite width limit, we nevertheless find excellent empirical agreement between the predictions of the original network and those of the linearized version even for finite practically-sized networks. This agreement is robust across different architectures, optimization methods, and loss functions.

Pathological Spectra of the Fisher Information Metric and Its Variants in Deep Neural Networks

“…Q depends on the softmax function g at each time step. One can numerically solve the training dynamics of g and obtain the theoretical value of the training loss (Lee et al, 2019). In Figure 5 (left), we confirmed that the theoretical line coincided well with the experimental results of gradient descent training.…”

“…We used is necessary for the steepest gradient method to converge (Karakida et al, 2019b). In fact, this 2/λ max acts as a boundary of neural tangent kernel regime (Lee et al, 2019;Lewkowycz, Bahri, Dyer, Sohl-Dickstein, & Gur-Ari, 2020). Because λ max increases depending on the width and depth, we need to carefully choose an appropriately scaled learning rate to train the DNNs.…”

“…See section 4 for more details. In the case of the cross-entropy loss, its functional gradient is given by Q (y − g) (Lee et al, 2019). is the NTK at random initialization described in equation 4.1.…”

“…Specifically, NTK's eigenvalues determine the speed of convergence of the training dynamics. Moreover, one can predict the network output on the test samples by using the gaussian process with the NTK (Jacot et al, 2018;Lee et al, 2019). The NTK and empirical FIM share essentially the same nonzero eigenvalues.…”

“…Such DNNs with random weights are substantially connected to gaussian process and kernel methods (Daniely, Frostig, & Singer, 2016;Lee et al, 2018;Matthews, Rowland, Hron, Turner, & Ghahramani, 2018;Jacot, Gabriel, & Hongler, 2018). Furthermore, the theory of the neural tangent kernel (NTK) explains that even trained parameters are close enough to the random initialization in sufficiently wide DNNs, and the performance of trained DNNs is determined by the NTK on the initialization (Jacot et al, 2018;Lee et al, 2019;Arora et al, 2019). Karakida, Akaho, and Amari (2019b) focused on the FIM corresponding to the mean square error (MSE) loss and proposed a framework to express certain eigenvalue statistics by using order parameters.…”