Wide neural networks of any depth evolve as linear models under gradient descent*This article is an updated version of a paper presented at 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada. (21st December 2020)

• Record Type: Journal Article
• Authors: Lee, Jaehoon; Xiao, Lechao; Schoenholz, Samuel S; Bahri, Yasaman; Novak, Roman; Sohl-Dickstein, Jascha; Pennington, Jeffrey
• Is Part Of: Journal of statistical mechanics. (2020:Dec.)
• Publication Date: 2020-12-21
• Subjects: machine learning; Statistical mechanics -- Periodicals; Mechanics -- Statistical methods -- Periodicals; 530.1305
• DOI: 10.1088/1742-5468/abc62b
• Languages: English
• ISSNs: 1742-5468

Abstract

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.