Feature Learning in Infinite-Width Neural Networks

By Greg Yang, J. Edward Hu. 2020

209 citations recorded by the source corpus.

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.

Subfields: 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 Statistical Learning Theory

Citations in the corpus

Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer (2022, methodology): “Recently, [45] showed that different neural network parametrizations induce different infinitewidth limits and proposed the Maximal Update Parametrization (abbreviated μP) (summarized in Table 3) that enables “maximal” feature learning in the limit.”
Adaptive Budget Allocation for Parameter-Efficient Fine-Tuning (2023, background): “Typically α is set as 16 or 32 and never tuned (Hu et al., 2022; Yang & Hu, 2020).”
Unit Scaling: Out-of-the-Box Low-Precision Training (2023, methodology): “The reparameterisation implied by unit scaling is also used by Jacot et al. (2018), later broadened by Yang and Hu (2020) and exploited by Yang et al. (2022) in their work analysing the training behaviour of deep networks.”
AdaLoRA: Adaptive Budget Allocation for Parameter-Efficient Fine-Tuning (2023, background): “Typically α is set as 16 or 32 and never tuned (Hu et al., 2022; Yang & Hu, 2020).”