PaperlayerView exact revision on arXiv

Source-linked AI summary

Feature Learning in Infinite-Width Neural Networks

Greg Yang, Edward J. Hu

arXiv:2011.14522v3cs.LGcond-mat.dis-nncs.NE

TL;DR

The paper asks whether infinite-width neural networks can retain feature learning, which standard and NTK parametrizations do not. It classifies parametrizations and uses Tensor Programs to derive feature-learning limits, finding that µP performs well on Word2Vec and few-shot Omniglot tasks. These results are bounded by a scope focused on training times independent of width and by specific initialization considerations.

  • Problem

    Standard and NTK parametrizations lack feature learning in their infinite-width limits, despite feature learning being central to pretraining and transfer learning.

  • Method

    The paper classifies abc-parametrizations and uses Tensor Programs to derive explicit infinite-width training limits, including the Maximal Update Parametrization.

  • Results

    µP admits feature learning and performs well on Word2Vec and few-shot learning through MAML on Omniglot, while the parametrization classification separates feature-learning and kernel regimes.

  • Takeaways & Limitations

    Tensor Programs provide a systematic way to derive infinite-width limits for computations expressible as Tensor Programs, including feature-learning neural-network regimes.

  • Takeaways & Limitations

    The results concern training times independent of width, and the transfer-learning comment does not directly apply to the few-shot setting because readout weights carry over from pretraining.

Abstract

from arXiv · show

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.

1 Introduction

The paper argues that standard and NTK parametrizations cannot learn features at infinite width, then proposes a parametrization and Tensor Programs analysis that recover feature learning limits. It classifies stable, nontrivial parametrizations into either feature-learning or kernel regimes and evaluates the resulting limits on feature-dependent tasks.

  • NTK and standard parametrizations lack feature learning in their infinite-width limits, despite feature learning’s importance for pretraining and transfer learning.
  • Tensor Programs derive infinite-width training dynamics for µP and, more generally, any abc-parametrization.
  • The µP limit outperforms NNGP/NTK baselines and finite networks on Word2Vec and Omniglot few-shot learning via first-order MAML.
  • Any nontrivial stable abc-parametrization yields either a feature-learning limit or kernel gradient descent, but not both.
  • The Maximal Update Parametrization (µP) modifies standard scaling to permit maximal feature learning without blowup.
  • The paper defines abc-parametrizations as a natural space encompassing standard, NTK, and Mean Field parametrizations.

2 Related Works

Prior work developed NTK and mean-field limits, but this paper distinguishes its discrete-time large-width analysis and addresses broader methodological and scope concerns.

  • Mean-field limits for 1-hidden-layer MLPs are equivalent to the µP limit up to the symmetry of Eq. (5).
  • Unlike most prior analyses, this paper takes the large-width limit while keeping gradient descent discrete-time.Continuous-time results can be recovered by taking a subsequent continuous-time limit.
  • The paper emphasizes that Gaussian N(0, 1/n) initialization is empirically crucial for performance relative to mean-field-style initialization.
  • The analysis concerns training times independent of width, rather than explicit convergence-time dependence that allows wider networks to remain near the limit longer.
  • Earlier work characterized NTK training as lazy because weights move little, contrasting it with substantially moving weights in the mean-field regime.
  • The paper notes that prior approaches loosely mimicked feature learning with deep kernel processes instead of taking exact infinite-width SGD limits.

3 Feature Learning vs Kernel Behavior

The paper distinguishes infinite-width feature learning from kernel behavior and characterizes which abc-parametrizations produce each regime. It shows that standard and NTK parametrizations yield kernel limits, whereas feature-learning limits evolve representations and support transfer-learning benefits.

  • Key Observations: In NTK, embedding changes generically have coordinate size Θ(1/√n), while Mean Field changes have size Θ(1) for fixed t ≥ 1.Both parametrizations start with embedding coordinates of Θ(1) size.
  • Key Observations: NTK feature kernels do not change in the infinite-width limit, whereas Mean Field feature kernels do.Feature-kernel evolution is linked to feature learning and transfer learning.
  • Key Observations: NTK training follows kernel gradient descent in function space, while Mean Field updates can have higher-order dependence on the learning rate.For identity activation, the cited example indicates a quadratic rather than linear dependence on η.
  • abc-Parametrizations: The stability conditions for abc-parametrizations are a1 + b1 = 0, al + bl = 1/2 for l ∈ [2,L], and aL+1 + bL+1 ≥ 1/2.These conditions characterize when the parametrization remains stable as width grows.
  • Dynamical Dichotomy: A nontrivial stable abc-parametrization admits feature learning iff it evolves the feature kernel iff r = 0; it is in kernel regime iff r > 0.Thus, every nontrivial stable parametrization falls into one regime rather than both.
  • Consequences: Any nontrivial stable feature-learning parametrization has f0(ξ) → 0 almost surely for every input, and feature learning in one layer propagates through later layers.The general classification also says the infinite-width function becomes deterministic in feature-learning parametrizations.
  • Standard and Maximal Update Parametrizations: Standard parametrization can allow only an O(1/n) learning rate while keeping finite outputs, and then remains in the kernel regime without feature learning.The paper proposes µP as a modification that admits maximal feature learning while remaining stable.

4 Standard Parametrization

The standard parametrization forces a trade-off between finite logits and feature learning: learning rates large enough to learn features cause blowup, while stable rates yield kernel dynamics.

  • Learning rates above O(1/n) make standard-parametrized logits blow up after one SGD step.
  • The resulting infinite-width dynamics do not admit feature learning because parameter-induced activation changes vanish with width.
  • In a two-hidden-layer linear MLP, the first-step logit change converges to −χK(ξ, ξ)=−χ∥ξ∥2, the linear-kernel gradient-descent update.
  • With learning rate O(1/n), standard parametrization remains stable and nontrivial but enters the kernel regime.

5 Maximal Update Parametrization

The Maximal Update Parametrization modifies standard scaling by shrinking logits and rescaling the first layer, enabling stable, nonvanishing feature learning throughout the network.

  • 5 Maximal Update Parametrization: Dividing logits by √n and using Θ(1) learning rate enables feature learning, while rescaling the first layer lets every layer learn features.
  • 5.1 Dividing Logits by √n: Under the initial modification, logits converge to 0 because their variance is Θ(1/n), yet the loss derivative remains Θ(1) when y≠0.
  • 5.1 Dividing Logits by √n: After one SGD step, hidden activations and logits both change by Θ(1) coordinatewise when c=0, directly demonstrating feature learning.
  • 5.2 First Layer Parametrization: Setting a1=−1/2 and b1=1/2 scales the first-layer gradient by n, making its preactivation change Θ(1) coordinatewise.
  • 5.3 What is µP Maximal In?: µP uniquely makes every weight matrix update maximally while also initializing the final matrix maximally within stable abc-parametrizations.

6 Deriving Feature Learning Infinite-Width Limit: Intuition and Examples

The paper develops intuition and exact formulas for µP infinite-width training by tracking coordinate distributions, correlations, and scalar limits across SGD steps.

  • Tensor Programs automate rigorous infinite-width limit derivations by recursively calculating coordinate distributions and the resulting network evolution.
  • 6.1 1-Hidden-Layer MLP: At initialization, a one-hidden-layer µP network has deterministic zero logits, while its loss derivative remains deterministic and generally nonzero.
  • 6.1 1-Hidden-Layer MLP: Theorem 6.1 gives almost-sure convergence of the finite network output to an expectation over recursively updated Gaussian variables.
  • 6.1 1-Hidden-Layer MLP: For a one-hidden-layer linear MLP, the limit reduces to four scalar coefficients whose updates determine the infinite-width output.
  • 6.2 2-Hidden-Layer MLP: In deeper networks, Gaussian matrix multiplication creates both central-limit behavior and correlations between a matrix and its transpose.
  • 6.2 2-Hidden-Layer MLP: Using a decoupled backpropagation matrix removes the copied contribution that otherwise prevents Gaussian coordinate behavior.

7 Tensor Programs Framework

The Tensor Programs framework represents training as a computation graph and recursively converts its vectors and scalars into limiting random variables and deterministic quantities.

  • Tensor Programs express SGD through matrix multiplication, coordinatewise nonlinearities, and coordinatewise averages.
  • The framework represents vectors by coordinate-distribution random variables and scalars by deterministic limits computed recursively from program operations.
  • The framework handles transpose-induced correlations through the ZDot rule, whereas unused transposes leave matrix products Gaussian under the ZHat rule.
  • The Tensor Program workflow is to encode the computation graph, calculate each vector and scalar limit, then read the limiting logits at inference time.
  • The Master Theorem guarantees almost-sure convergence of scalar program quantities and pseudo-Lipschitz statistics under the stated assumptions.
  • Tensor Programs extend across architectures and training procedures, covering both feature-learning and kernel limits.

8 Computational Considerations

The Tensor Programs feature-learning limits are analytically tractable in special cases but become increasingly expensive with nonlinearities, training time, and depth.

  • General nonlinearities lack closed-form expectation formulas, so Monte Carlo errors compound quickly with training time.Polynomial activations are an exception because Gaussian polynomial moments can be evaluated analytically.
  • With polynomial nonlinearities, expectation integrands can reach degree Ω(2t), causing super-exponential evaluation costs as training time t increases.This bottleneck arises from nested activations and derivatives over correlated Gaussian variables.
  • Deeper networks add covariance-storage and backward-pass costs, with t SGD steps requiring Ω(t^2) space and time before further expectation costs.Intermediate n × n Gaussian matrices create additional dependencies across earlier training steps.
  • µP-limit computation ranges from easiest for one-hidden-layer linear networks to hardest for nonlinear MLPs with nonpolynomial activations.One-hidden-layer linear networks suffice for the paper’s Word2Vec and MAML demonstrations.

9 Experiments

The experiments validate the theoretical µP limits and evaluate feature learning on Omniglot few-shot adaptation and Word2Vec transfer, where µP outperforms kernel baselines.

  • Experiments: The experiments divide into theory verification and realistic-dataset studies, while computational limits restrict infinite-width training duration.The verification experiments stress-test theoretical predictions across multiple scenarios.
  • 9.1 Verifying the Theory: Finite-width simulations verify the analytically computed µP limits for quadratic and linear activations.Figure 4 compares widths 1024 and 4096 across 100 random seeds; quadratic networks converge more slowly with width.
  • 9.2 Few-Shot Learning on Omniglot via First Order MAML: First Order MAML adapts parameters on a task’s training set, then updates the initialization using loss on a separate test set.The procedure omits the Hessian-based second-order update used by Second Order MAML, which is slower but generally slightly better.
  • 9.2 Few-Shot Learning on Omniglot via First Order MAML: The µP feature-learning limit significantly outperforms kernel limits, while finite-width µP networks approach it from below as width increases.This pattern holds in the Omniglot MAML findings.
  • 9.3 Word2Vec: Word2Vec learns word feature vectors from context using CBOW, negative sampling, and sigmoid loss before evaluation on word analogy.The experiments compare µP against finite-width µP, NTK, and GP limits.
  • 9.3 Word2Vec: In Word2Vec, µP and finite-width µP networks significantly outperform the NTK limit, with finite-width performance converging from below as width increases.The NTK and GP features remain fixed at initialization, making the NTK analogy accuracy practically 0 for the reported vocabularies.

10 Conclusion

The paper unifies major infinite-width parametrizations, classifies their dynamics, proposes µP for feature learning, and demonstrates its behavior experimentally using Tensor Programs.

  • The abc-parametrization and Tensor Programs framework unify the NTK and Mean Field limits of large-width neural networks.
  • The Dynamical Dichotomy classifies abc-parametrizations into feature-learning and kernel regimes, identifying absent feature learning as a weakness of NTK and standard parametrizations.
  • The proposed Maximal Update Parametrization, µP, admits feature learning in the infinite-width limit.Experiments on Word2Vec and few-shot learning support µP as a model of feature-learning behavior.
  • Tensor Programs provide a systematic route to deriving infinite-width limits for computations expressible in the framework, including entire training procedures.The framework was reorganized into a user-friendly Tensor Programs paper series before this feature-learning work.

B Further Discussions on the Shallow NTK and MF Examples

In the NTK limit, function evolution comes from separate parameter-feature interactions, whereas the Mean Field limit also retains their joint interaction.

  • In the NTK limit, ∆f_t(ξ) is predominantly driven by V_0∆x_t(ξ) and ∆V_tx_0(ξ), while ∆V_t∆x_t(ξ) vanishes.
  • The Mean Field limit differs by retaining a nontrivial joint interaction between output-weight changes and feature changes.
  • For the NTK parameterization, V_0∆x_1(ξ) remains Θ(1) despite individual feature changes being O(1/√n).The nonvanishing sum results from interactions between initialized output weights and feature changes.
  • Tensor Programs extend this shallow-network derivation to general training time t.

C abc-Parametrization for General Neural Architectures

The paper extends abc-parametrizations to arbitrary neural architectures by assigning scaling exponents to each parameter tensor. In this framework, µP uses width-independent learning rates and is uniquely maximally stable and updated.

  • Generalization: Each parameter tensor W receives exponents aW and bW, with W = n^-aW w, Gaussian initialization variance n^-2bW, and learning rate ηn^-c.This generalizes the parametrization beyond multilayer perceptrons.
  • µP rules: For general architectures expressible as Tensor Programs, µP sets bW = 1/2 for every parameter and keeps c = 0.The prescription includes architectures such as ResNets and Transformers.
  • µP rules: Non-output parameter exponents aW depend on tensor shape, while output matrices use aW = 1 and output biases use aW = 0.The shape cases distinguish vectors, one-fixed-dimension matrices, and matrices whose two sides scale with width.
  • Optimality: µP is the unique stable abc-parametrization in which all parameter tensors update maximally and all readout weights initialize maximally.The optimality statement applies in the generalized setting.

D Experimental Details

The experiments compare finite- and infinite-width linear MLPs, kernel models, and MAML procedures using explicit parametrizations and coefficient-based infinite-width tracking. Word2Vec analysis shows that NTK predictions are random at infinite width because initialization determines a fixed argmax.

  • Models: The basic experimental model is a 1-hidden-layer linear MLP with input dimension d, output dimension do, and width n.More advanced variants are also considered.
  • Infinite-width limit: In µP, the infinite-width network output converges almost surely to a deterministic trajectory defined recursively for every input.The recursion updates coefficient objects using the loss derivative and learning rate η.
  • Implementation: The implementation tracks deterministic coefficients because trained coordinates remain linear combinations of initialized Gaussian vectors and converge deterministically as width grows.This reparametrization avoids tracking all width-n coordinates directly.
  • Training procedures: Finite- and infinite-width training algorithms are essentially identical apart from initialization and hidden size, with forward, backward, clipping, and update stages shared.The infinite-width hidden size is d + do rather than n.
  • MAML: MAML modifies these procedures by disabling weight decay and clipping gradients against each task’s validation loss during meta-training.During adaptation, gradients are not clipped.
  • Word2Vec: NTK Word2Vec outputs are random because the argmax varies across initializations but remains fixed throughout training at large width.The associated inner products are random under the stated central-limit behavior.

F Nuances of the Master Theorem

The Master Theorem tracks limiting tensor-program variables through recursively defined Gaussian quantities and partial-derivative expectations. Its assumptions require continuity for parameter-only functions and pseudo-Lipschitz behavior otherwise.

  • Limiting variables: Each program variable’s limiting quantity is represented as a deterministic function of Gaussian base variables and limit scalars.This representation defines the derivatives used in the theorem.
  • Partial derivatives: Partial derivatives are taken with respect to the Gaussian components appearing in the program’s recursive representation, with other components assigned zero derivative.The definition depends on the precise written program.
  • Derivative expectations: The derivative expectation can be defined through a covariance matrix and pseudoinverse, extending the construction beyond directly differentiable nonlinearities.Stein’s lemma recovers the usual expectation when that expectation is well defined.
  • Regularity: A pseudo-Lipschitz function has polynomially controlled growth, and compositions remain pseudo-Lipschitz with degree equal to the sum of component degrees.Such functions are also Lipschitz on compact sets.
  • Assumptions: The theorem assumes continuity for functions with only parameter arguments and pseudo-Lipschitz dependence for other nonlinear and moment functions.Continuous loss derivatives are therefore permitted in the stated setting.
  • Scope: The paper uses pseudo-Lipschitz conditions instead of weaker polynomial-boundedness assumptions because the alternative rank-stability condition is subtle to verify rigorously.Other Master Theorem versions trade smoothness restrictions for rank stability.

G A Rough Sketch of the Geometry of abc-Parametrizations

Stable abc-parametrizations form an unbounded polyhedron whose nontrivial regions lie on two facets. The µP-to-NTK family forms an edge, with µP and NTK at its endpoints under the paper’s uniform-update characterization.

  • Polyhedral geometry: Stable abc-parametrizations form an unbounded polyhedron because increasing a layer’s a_l while decreasing b_l reduces that layer’s learning rate without losing stability.As the shift grows, that layer eventually is not trained.
  • Polyhedral geometry: Nontrivial stable parametrizations lie on two unbounded facets of the stable polyhedron.The same learning-rate rescaling accounts for the facets’ unboundedness.
  • Endpoints and edge: NTP lies at the intersection of the two facets, while NTP and µP are connected by an edge.NTP uniquely combines uniform body updates with maximal initialization and updating of the output matrix.
  • Uniform parametrizations: Uniform body updates mean every layer changes at the same Θ(n^-r) scale during training.The exponent r is shared across all body layers.
  • Uniform parametrizations: For every r in [0, 1/2], there is a unique stable uniformly updated parametrization with that r and maximal output-weight initialization.These parametrizations are denoted UP_r.
  • Endpoints and edge: UP_0 is µP and UP_1/2 is NTP, while r > 1/2 is unstable because the required initial weights would make the initial Gaussian process diverge.Thus the UP_r family exactly traces the stable edge between the two endpoints.

H.1 Rigorous Statements of Main Results

For tanh or sufficiently small-σ σ-gelu, the paper characterizes stable, nontrivial abc-parametrizations and proves that infinite-width limits separate feature learning from kernel behavior. Feature learning occurs exactly when r = 0, whereas r > 0 corresponds to fixed features or kernels and can yield an NNGP limit.

  • Scope: The explicit classification and feature-learning construction rely on tanh or sufficiently small-σ σ-gelu; generic nonlinearities remain outside the paper’s proved scope.
  • Stability: Stability requires a1 + b1 = 0, al + bl = 1/2 for l ∈ [2, L], and aL+1 + bL+1 ≥ 1/2.These conditions keep hidden preactivations and activations at Θ(1) scale at initialization.
  • Nontriviality: A stable abc-parametrization is nontrivial iff aL+1 + bL+1 + r = 1 or 2aL+1 + c = 1.
  • Classification: For nontrivial stable parametrizations, feature learning, feature-kernel evolution, and r = 0 are equivalent.The characterization also links prefeature evolution to corresponding behavior through the network layers.
  • Classification: When r > 0, the kernel regime is equivalent to fixing all features and feature kernels, including their Lth-layer versions and prefeature counterparts.
  • Layerwise propagation: If feature or feature-kernel evolution occurs in layer l, then all four corresponding forms of evolution occur in layers l through L.
  • Infinite-width limits: If r = 0, f0(ξ) converges almost surely to 0 while ft(ξ) converges almost surely to a deterministic limit; if r > 0 and additional equalities hold, the limit is NNGP.Feature learning is therefore incompatible with the Bayesian, distributional NNGP perspective in this setting.
  • Transfer learning: In a stable kernel regime, both linear and nonlinear transfer learning are trivialized, even when the entire network is fine-tuned.The result applies whether only the classifier or the whole network is trained.

H.5.1 Some Immediate Results

The appendix separates stable parametrizations into feature-learning and kernel regimes, then derives their corresponding infinite-width behavior and stability conditions.

  • A pseudostable parametrization is stable, while initialization-stable parametrizations with r < 0 are not stable.
  • r > 0: For r > 0, the parametrization fixes all features and feature kernels, yielding kernel gradient descent dynamics.The kernel dynamics use δ˚f_t(ξ) = −η˚χ_t−1Σ(ξ_t−1, ξ).
  • r > 0: Nontriviality for r > 0 requires a_L+1 + b_L+1 + r = 1 or 2a_L+1 + c = 1.
  • r > 0: For r > 0, a_L+1 + b_L+1 + r > 1 and 2a_L+1 + c = 1 produce the NNGP limit.This limit can also arise when the network body is trained slightly, provided its changes do not affect the output dynamics.

H.7.4 Main Results

The main results classify stable abc-parametrizations: nontrivial cases either support feature learning when r = 0 or kernel behavior when r > 0.

  • For tanh or sufficiently small-σ σ-gelu, a pseudostable parametrization with r = 0 is nontrivial iff a_L+1 + b_L+1 = 1 or 2a_L+1 + c = 1.
  • Any nontrivial pseudostable parametrization with r = 0 is outside the kernel regime and supports feature learning and feature-kernel evolution.
  • Stable abc-parametrizations are nontrivial exactly when a_L+1 + b_L+1 + r = 1 or 2a_L+1 + c = 1.The classification assumes tanh or σ-gelu with sufficiently small σ.
  • For r > 0, the parametrization is in the kernel regime and fixes features, feature kernels, prefeatures, and prefeature kernels.
  • If feature learning or related kernel evolution occurs at layer l, it also occurs in layers l through L.
  • When r = 0, the infinite-width initialization output converges to zero, while later outputs converge almost surely to deterministic values.
Loading 2011.14522v3…