Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent
Jaehoon Lee , Lechao Xiao , Samuel S. Schoenholz, Yasaman Bahri Roman Novak, Jascha Sohl-Dickstein, Jeffrey Pennington Google Brain {jaehlee, xlc, schsam, yasamanb, romann, jaschasd, Both authors contributed equally to this work. Work done as a member of the Google AI Residency program (https://g.co/airesidency).
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 have made a theory of learning dynamics elusive. In this work, we show that for wide neural networks 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 neural networks and Gaussian processes, gradient-based training of wide neural networks with a squared loss produces test set predictions drawn from a Gaussian process 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.
中文速览
宽神经网络（wide neural network）在训练过程中究竟发生了什么，一直缺乏清晰的理论解释。本文发现，当网络宽度趋于无穷时，梯度下降的训练动态可以被大幅简化：网络等价于其参数在初始值处的一阶泰勒展开所得到的线性模型，对应的训练过程存在解析闭式解。进一步地，在平方损失下，这一无限宽网络的预测分布始终是一个由神经切核（Neural Tangent Kernel, NTK）决定的高斯过程，且作者推导出了该高斯过程随训练时间演化的显式表达式。实验表明，即使对有限宽度的实际网络（包括宽残差网络等常见架构），线性化近似与原始网络的预测也高度吻合，且这一结论对不同优化方法和损失函数均保持稳健。这项工作为理解深度神经网络的训练机理提供了严格的理论基础，揭示了宽网络"懒惰训练"的本质，并将梯度下降与贝叶斯后验采样之间的联系与差异明确量化。
原文 arXiv:1902.06720;中英对照 + 大白话阅读 https://aha.fim.ai/paper/1902.06720v4