Feature Learning in Infinite-Width Neural Networks
arXiv:2011.14522
Abstract
As its width tends to infinity, a deep neural network's behavior under gradient descent can become simplified and predictable (e.g. given by the Neural Tangent Kernel (NTK)), if it is parametrized appropriately (e.g. the NTK parametrization). However, we show that the standard and NTK parametrizations of a neural network do not admit infinite-width limits that can learn features, which is crucial for pretraining and transfer learning such as with BERT. We propose simple modifications to the standard parametrization to allow for feature learning in the limit. Using the *Tensor Programs* technique, we derive explicit formulas for such limits. On Word2Vec and few-shot learning on Omniglot via MAML, two canonical tasks that rely crucially on feature learning, we compute these limits exactly. We find that they outperform both NTK baselines and finite-width networks, with the latter approaching the infinite-width feature learning performance as width increases. More generally, we classify a natural space of neural network parametrizations that generalizes standard, NTK, and Mean Field parametrizations. We show 1) any parametrization in this space either admits feature learning or has an infinite-width training dynamics given by kernel gradient descent, but not both; 2) any such infinite-width limit can be computed using the Tensor Programs technique. Code for our experiments can be found at github.com/edwardjhu/TP4.
arXiv comment4th paper in the Tensor Programs series. Appearing in ICML 2021
References in corpus (23)
- Efficient Estimation of Word Representations in Vector Space
- Distributed Representations of Words and Phrases and their Compositionality
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks
- Gaussian Error Linear Units (GELUs)
- Language Models are Few-Shot Learners
- Neural Tangent Kernel: Convergence and Generalization in Neural Networks
- A Convergence Theory for Deep Learning via Over-Parameterization
- A Mean Field View of the Landscape of Two-Layers Neural Networks
- Exponential expressivity in deep neural networks through transient chaos
- Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation
- Dynamical Isometry and a Mean Field Theory of RNNs: Gating Enables Signal Propagation in Recurrent Neural Networks
- The large learning rate phase of deep learning: the catapult mechanism
- Kernel and Rich Regimes in Overparametrized Models
- A Mean Field Theory of Batch Normalization
- Tensor Programs II: Neural Tangent Kernel for Any Architecture
- On the Selection of Initialization and Activation Function for Deep Neural Networks
- Dynamics of Deep Neural Networks and Neural Tangent Hierarchy
- Dynamical Isometry and a Mean Field Theory of LSTMs and GRUs
- On the infinite width limit of neural networks with a standard parameterization
- Learning Over-Parametrized Two-Layer ReLU Neural Networks beyond NTK
- Wider Networks Learn Better Features
- Dynamically Stable Infinite-Width Limits of Neural Classifiers