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

Infinite-width NTK limits simplify training dynamics, but they leave learned features essentially like random initialization despite feature learning's importance to deep learning. This paper classifies parametrizations and derives their limits with Tensor Programs, finding that µP models feature-learning behavior in experiments.

  • Problem

    In the NTK limit, learned features remain essentially like random initialization, despite feature learning's importance to deep learning.

    Source: §1 · p. 1
  • Method

    The authors classify neural-network parametrizations and use Tensor Programs to derive infinite-width limits, proposing µP as a parametrization that admits feature learning.

    Source: §7 · p. 25, §10 · p. 30
  • Results

    The classification separates parametrizations into feature-learning and kernel regimes, and experiments show µP models feature-learning behavior.

    Source: §1 · p. 2, §10 · p. 30
  • Takeaways & Limitations

    Tensor Programs provide a systematic way to derive infinite-width limits for computations expressible in the framework.

    Source: §7 · p. 22, §10 · p. 30

    For general nonlinearities, the derived expectations lack closed-form formulas, and Monte Carlo evaluation errors can compound over training time.

    Source: §8 · p. 26

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 develops infinite-width parametrizations that can learn features, addressing a limitation of standard and NTK parametrizations. It proves a dichotomy between feature-learning and kernel limits and derives these limits using Tensor Programs.

1 Introduction

  • In the NTK limit, features learned during pretraining are essentially unchanged from random initialization, a limitation for applications relying on feature learning.The paper verifies this behavior empirically on Word2Vec.Source: §1 · p. 1
  • The proposed modification to standard parametrization enables maximal feature learning, and the paper computes its infinite-width limit rigorously using Tensor Programs.The limits are explicitly calculated for Word2Vec and Omniglot few-shot learning via MAML.Source: §1 · pp. 1–2
  • Any nontrivial stable abc-parametrization yields either evolving embeddings in a feature-learning limit or kernel gradient descent in a kernel limit, but not both.In any feature-learning limit, the network function is identically zero at initialization.Source: §1 · p. 2, Our Contributions · p. 4
  • The paper’s Glorot-initialized TensorFlow default makes first-layer preactivations converge to zero as width grows, yielding pathological infinite-width behavior.The caveat concerns initialization scaling with width.Source: §1 · p. 2

Our Contributions

  • µP admits maximal feature learning while allowing a Θ(1) maximum learning rate; standard parametrization requires O(1/width), and NTK parametrization O(1), to avoid blowup.In a one-hidden-layer MLP, µP specializes to mean-field parametrization up to symmetry.Source: §1 · p. 3
  • Tensor Programs derive infinite-width training limits by expressing the computation as a program and recursively calculating coordinate distributions with the Master Theorem.The method applies to any abc-parametrization and can extend to architectures such as ResNets and Transformers, though the paper focuses on MLPs.Source: §1 · p. 3, Our Contributions · p. 4
  • The paper formulates abc-parametrizations as a natural space that includes standard, NTK, and mean-field parametrizations.It shows standard and NTK parametrizations yield kernel limits rather than feature-learning limits.Source: §1 · p. 2, Our Contributions · p. 4
  • On Word2Vec and Omniglot few-shot learning, µP limits outperform NTK baselines and finite networks, with finite networks approaching infinite-width performance as width increases.The Omniglot models use first-order MAML.Source: §1 · p. 2, Our Contributions · p. 4

2 Related Works

The related works distinguish this paper’s discrete-time, feature-learning limits from prior kernel and mean-field frameworks. They also position its parametrization classification, assumptions, and training-time scope against earlier studies.

Comparison with Mean Field Limits

  • For one-hidden-layer MLPs, the mean-field limit is equivalent to the µP limit up to Eq. (5)’s symmetry, while this work uses typical Gaussian N(0, 1/n) initialization [11] [43] [45].The paper contrasts this central-limit-effect initialization with prior deeper mean-field frameworks, which often use different initializations and integral-kernel effects; the framework of [15] tracks coordinate distributions recursively but assumes atypical ℓ2 regression initialization.Source: §2 · p. 4

Discrete- vs Continuous-Time Gradient Descent

  • Unlike prior infinite-width studies that jointly take continuous time, this work takes only the large-width limit, keeping gradient descent discrete-time [11] [26] [43] [45].The authors note that continuous-time limits can obscure practical step size, task, and batch-size effects, and introduce ODE/PDE well-posedness questions absent from discrete-time evolution.Source: §2 · p. 4

Technical Assumptions

  • The rigorous limit requires only a polynomially bounded weak second derivative for φ and a continuous loss derivative with respect to predictions, unlike many earlier regularity assumptions [11, 15, 26, 30, 35, 43, 45].In the one-hidden-layer discrete-time mean-field case, the result covers Gaussian initialization and heavy-tailed distributions expressible as Gaussian images under pseudo-Lipschitz functions.Source: §2 · p. 5

Training Time

  • The results cover training times independent of width; a quantitative Master Theorem could extend the analysis to training time increasing with width.This O(1)-time scope does not contradict [27], which studies leaving the kernel regime in Ω(log(width)) time; the paper argues its limit mode is practically relevant under fixed computational budgets [29].Source: §2 · p. 5

Classification of Parametrizations

  • Earlier work described lazy/active or kernel/rich regimes, but did not formally classify natural parametrizations; this paper studies arbitrary-depth MLPs under usual SGD [19] [31] [48].Prior work also found that mean-field limits, unlike NTK limits, can learn low-dimensional linear structure [12].Source: §2 · p. 5

Other Related Works

  • Related work studied unstable standard parametrization [47] and deep-kernel processes that mimic feature learning [2] [3], while this paper takes the exact infinite-width limit of SGD training.The abc framework also captures subtle parametrization differences [4] [14] [57], and prior work observed wider networks can improve downstream linear-transfer performance [17].Source: §2 · pp. 5–6

3 Feature Learning vs Kernel Behavior

The section contrasts feature learning with kernel behavior and classifies stable abc-parametrizations by their infinite-width training dynamics. It shows that feature learning and kernel gradient descent are mutually exclusive in this setting.

3.1 Motivating Examples: Neural Tangent Kernel and Mean Field Limits

  • The shallow-network examples use NTP and an equivalent MFP formulation; the GP limit is a special case of NTK when only the final layer is trained.The section introduces these parametrizations to motivate the broader classification.Source: §3.1 · p. 6

3.1.1 Key Observations

  • For fixed t ≥1, embedding changes are Θ(1/√n) in NTP versus Θ(1) in MFP, so the feature kernel stays fixed in the NTK limit but evolves in the MF limit.In this shallow-network comparison, NTP exemplifies kernel behavior and MFP exemplifies feature learning.Source: §3.1.1 · p. 7
  • NTK pretraining gives no linear-transfer benefit and is no better than random initialization even with whole-network finetuning; adaptation is not automatically trivial, but NTK underperforms feature learning.The section identifies Word2Vec as a linear transfer task and few-shot learning as an adaptation setting.Source: §3.1.1 · p. 8

3.2 abc-Parametrizations and Dynamical Dichotomy

  • For nontrivial stable abc-parametrizations under tanh or smooth σ-gelu, feature learning and feature-kernel evolution occur iff r = 0, while kernel regime occurs iff r > 0.Thus these parametrizations admit feature learning or kernel gradient descent, but not both.Source: §3.2 · pp. 8, 10
  • Any nontrivial stable feature-learning parametrization has f0(ξ) → 0; feature learning in a layer implies it in all later layers, and its limiting function is deterministic.The deterministic limit is incompatible with a Bayesian perspective.Source: §3.2 · p. 10
  • The standard parametrization cannot admit stable feature learning, whereas the proposed µP modification does and achieves it maximally in a suitable sense.The section motivates the µP proposal by analyzing standard-parameterization dynamics in the next section.Source: §3.2 · p. 11

4 Standard Parametrization

In standard parametrization, keeping logits bounded requires an SGD learning rate of order O(1/n), and at that rate the infinite-width network is in the kernel regime and does not learn features.

After First SGD Step

  • In the two-hidden-layer linear example, a learning rate n^-c gives f1 = f − Θ(n^(1−c)), so c ≥ 1 is required to keep the logits O(1).The initialization has roughly iid Gaussian coordinates of variance Θ(1), and the argument considers a single SGD step.Source: §4 · pp. 11–12

Kernel Regime and Lack of Feature Learning

  • At c = 1, hidden-representation changes have coordinates of size O(n^-1/2), so there is no feature learning, while f1 − f → −χ∥ξ∥², kernel gradient descent with the linear kernel.The same kernel-gradient-descent conclusion holds across all SGD steps.Source: §4 · p. 12

Theorem 4.1

  • For an L-hidden-layer MLP, bounded logits for all training routines, times, and inputs require an SGD learning rate O(1/n); at this rate, the network is in the kernel regime and does not admit feature learning.This is the general result for standard parametrization.Source: §4 · p. 12

5 Maximal Update Parametrization

This section introduces the Maximal Update Parametrization (µP), a modification of the standard parametrization that enables feature learning while maintaining stability. It scales updates so every layer can contribute maximally without logits or activations blowing up.

5 Maximal Update Parametrization

  • µP divides logits by √n (aL+1=1/2, c=0) and additionally sets a1=−1/2, b1=1/2 to enable feature learning with maximal stable updates.The authors call this parametrization MUP or µP.Source: §5 · p. 12, §5.3 · p. 15

5.1 Dividing Logits by √n

  • Although initialized logits converge to 0, with Θ(1) learning rate (c=0), hidden activations and logits change Θ(1) coordinatewise after one SGD step, yielding feature learning.For general MLPs, the parametrization admits feature learning by Theorem 3.6.Source: §5.1 · p. 13
  • The µP dynamics need not have a kernel limit: a tanh network or a linear MLP training its middle and last layers can exhibit non-kernel dynamics.Training only the middle layer in the linear MLP example is still in the kernel regime; for tanh, bounded activations conflict with Θ(1) hidden-state changes, while training middle and last layers introduces quadratic weight interactions.Source: §5.1 · p. 13
  • After one step, the modified parametrization keeps logits Θ(1), whereas standard parametrization makes them Θ(√n), despite similar weight updates.The difference arises because updates create correlations between the final-layer weights and activations.Source: §5.1 · p. 13

5.2 First Layer Parametrization

  • To make first-layer preactivations learn, µP sets a1=−1/2, b1=1/2, leaving the forward pass unchanged but scaling U's gradient by n so h changes Θ(1) coordinatewise.Without this change, h changes by at most O(1/n) coordinatewise during training.Source: §5.2 · pp. 13–14

5.3 What is µP Maximal In?

  • In µP, every W^l, including W^(L+1), is updated maximally, the output layer is initialized maximally, and µP is the unique stable abc-parametrization with these properties.For stable abc-parametrizations, the output layer is updated maximally iff 2aL+1+c=1 and initialized maximally iff aL+1+bL+1+r=1; µP has r_l=0 for every hidden layer.Source: §5.3 · pp. 14–15

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

This section builds intuition for Tensor Programs and derives infinite-width training dynamics for abc- and µP-parametrized networks, including normal SGD and arbitrary-depth MLPs. It also explains which weight matrices learn in feature-learning parametrizations.

6.1 1-Hidden-Layer MLP

  • For 1-hidden-layer µP, logits converge almost surely to an explicit scalar-variable recursion, the discrete-time mean-field limit [11] [43] [45].For linear networks, the limit can be reduced to a width-(d + do) linear network with diagonal initialization.Source: §6.1 · p. 16

6.2 2-Hidden-Layer MLP: SGD with Partially Decoupled Backpropagation

  • With partially decoupled backpropagation, an independent backward matrix isolates the Gaussian central-limit effect and yields Gaussian coordinate processes in the infinite-width dynamics.The Gaussian heuristic for Wx1 depends on using the independent backward matrix rather than W⊤; normal backpropagation introduces additional correlations.Source: §6.2 · pp. 17–19

6.3 2-Hidden-Layer MLP: Normal SGD

  • In normal SGD, correlations between the updated activations and W add a copied-gradient term to Wx1, which is generally non-Gaussian for nonlinear φ.For φ = id, Wx1 remains Gaussian, but with a variance different from the partially decoupled calculation.Source: §6.3 · pp. 19–20
  • For 2-hidden-layer MLPs in µP, any learning-rate-1 training routine with pseudo-Lipschitz φ′ has an almost-sure infinite-width limit whose Gaussian terms have covariance determined by the corresponding limiting variables.The normal-SGD dynamics modify the Gaussian terms and their covariances relative to the earlier equations; the Gaussian sets are mutually independent alongside the initialized variables.Source: §6.3 · p. 20

6.4 MLP of Arbitrary Depth

  • In feature-learning parametrizations, any weight matrix that is not maximally updated remains effectively fixed at initialization in the infinite-width limit.The µP limit for deeper MLPs can be derived by similar reasoning, with a rigorous Tensor Programs treatment covering stable abc-parametrizations.Source: §6.4 · p. 21

6.5 Summary of Main Intuitions for Deriving the µP Limit

  • The law of large numbers makes coordinates of a vector approximately iid with their pairwise products converging to expectations, supporting derivations for one-hidden-layer dynamics and weight updates.The passage applies this intuition to the limits of ΔWtxt in Sections 6.2 and 6.3.Source: §6.5 · p. 21
  • When Wᵀ is involved, Wz includes a Gaussian component plus a linear combination of previously defined variables from W–Wᵀ correlation, needed for normal-SGD dynamics.This correlation calculation is intricate, motivating the earlier treatment of decoupled SGD, which does not require it.Source: §6.5 · p. 21

7 Tensor Programs Framework

Tensor Programs automate and rigorously formalize infinite-width limit derivations by representing training and inference computations as programs. The framework supplies convergence guarantees and applies to feature-learning and kernel limits across architectures and tasks.

  • Tensor Programs generate computations from matrix multiplication, coordinatewise nonlinearities, and coordinatewise averages, representing the core operations underlying SGD.The formal program definition includes Gaussian matrix-vector products, nonlinear vector operations, and scalar moments.Source: §7 · pp. 22–23
  • The ZMatMul rule combines a Gaussian component with a derivative-based correction, which is generally nonzero when both a matrix and its transpose are used.This correction generalizes the calculation in Theorem 6.4; it vanishes when the transpose is unused.Source: §7 · p. 24
  • Algorithm 1 computes an infinite-width limit by encoding training and inference as a Tensor Program, deriving symbolic vector and scalar limits, and reading inference logits from those limits.The procedure generalizes the manual derivations and is used for the paper’s MLP experiments.Source: §7 · p. 25
  • Theorem 7.4 shows that vector-coordinate distributions converge to their symbolic limits and every program scalar converges almost surely under the stated setup and regularity conditions.For fixed k and pseudo-Lipschitz ψ, coordinate slices of k vectors behave asymptotically as iid copies of their symbolic variables.Source: §7 · p. 25
  • Tensor Programs can express practically any training and inference procedure involving common neural-network architectures, including both feature-learning and kernel limits.The framework is used to derive and compute µP and kernel limits for metalearning and Word2Vec.Source: §7 · p. 25

8 Computational Considerations

Analytic evaluation of feature-learning limits is computationally intensive beyond special cases: general nonlinear expectations lack closed forms, while polynomial and deeper-network cases still face rapidly growing costs.

Computational limitations

  • The convergence-in-distribution result [52] is strictly weaker than the formulation in Theorem 7.4.The passage describes the latter formulation as generally more useful.Source: §8 · p. 25
  • General nonlinearities lack closed-form expectation formulas, and Monte Carlo evaluation can accumulate error quickly over training.The example expectation nests φ′ inside φ, preventing a general closed-form solution.Source: §8 · p. 26
  • Polynomial activations permit analytical Gaussian-moment evaluation, but the required computation grows super-exponentially with training time.The integrand’s degree and monomial count grow exponentially in t, making moment evaluation increasingly costly.Source: §8 · p. 26
  • In deeper networks, Gaussian-matrix dependencies make t SGD steps cost Ω(t^2) in space and time, with backward-pass terms adding costly expectations.The backward-pass ˙Z is generally a linear combination of Ω(t) terms whose coefficients require expectations that encounter the exponential bottleneck.Source: §8 · p. 26
  • Computational cost rises from one-hidden-layer linear networks to deeper linear, polynomial-activation, and nonpolynomial-activation MLPs; one-hidden-layer linear networks suffice for the Word2Vec and MAML demonstrations.The stated order is one-hidden-layer linear, L-hidden-layer linear, nonlinear polynomial, then nonlinear nonpolynomial MLPs.Source: §8 · pp. 25–26

9 Experiments

The experiments test whether predicted µP limits match finite-width networks, then evaluate feature learning on Omniglot and Word2Vec tasks where the paper says GP and NTK limits are inadequate.

9 Experiments

  • The experiments split between stress-testing the theory and scaling to feature-learning tasks, with simpler architectures used for the latter because of computational constraints.The µP networks could not be trained for very long in the theory-verification experiments, but were trained long enough to verify the theory.Source: §9 · p. 26

9.1 Verifying the Theory

  • Figure 4 compares analytically computed µP limits for linear and quadratic activations with finite-width networks and reports empirical agreement.The finite widths are 1024 and 4096; quadratic networks converge more slowly to the limit, and training is limited to a few steps by computational difficulties.Source: §9.1 · p. 26, Figure 4 · p. 27

9.2 Few-Shot Learning on Omniglot via First Order MAML

  • On Omniglot few-shot learning, the feature-learning µP limit significantly outperforms kernel limits, feature learning benefits more than nonlinearities, and finite µP networks approach the limit from below as width increases.The evaluation uses 1-shot 5-way classification with First Order MAML.Source: §9.2 · pp. 27–28, Table 2 · p. 28

9.3 Word2Vec

  • On Word2Vec analogy evaluation, infinite- and finite-width µP networks significantly outperform the NTK limit, and finite-width performance approaches the µP limit from below as width increases.The comparisons use CBOW Word2Vec; the GP and NTK limits have fixed features and select answers uniformly from the vocabulary.Source: §9.3 · p. 29, Table 3 · p. 30

10 Conclusion · A A Short Origin Story of the Tensor Programs Paper Series

The paper unifies infinite-width limits and distinguishes feature-learning from kernel regimes, proposing µP as a feature-learning parametrization. The Tensor Programs series first made its mathematical foundation more accessible before this paper presented its payoffs.

10 Conclusion

  • The framework unifies NTK and Mean Field limits through abc-parametrizations and Tensor Programs, while the Dynamical Dichotomy classifies parametrizations into feature-learning and kernel regimes.Source: §10 · p. 30
  • The NTK and standard parametrizations both lack feature learning, which the paper identifies as a fatal weakness of NTK as a model for real neural networks.Source: §10 · p. 30
  • µP is proposed with a derived infinite-width limit that admits feature learning, and Word2Vec and few-shot learning experiments show it models feature-learning behavior.Source: §10 · p. 30
  • Any computation expressible as a Tensor Program has a derivable infinite-width limit; Tensor Programs' universality for deep-learning computation [49] [51] makes this a systematic approach to the limit problem.Source: §10 · p. 30
  • Tensor Programs' flexibility makes feature-learning analyses of overparametrized neural networks in reinforcement learning, self-supervised learning, and deep generative models ripe for further work.Source: §10 · p. 30

A A Short Origin Story of the Tensor Programs Paper Series

B Further Discussions on the Shallow NTK and MF Examples

The NTK and mean-field limits differ in how parameter and feature changes interact: NTK evolution comes from cross-interactions, while mean-field evolution also includes a nontrivial interaction between both changes.

C abc-Parametrization for General Neural Architectures · C.1 Maximal Update Parametrization

The appendix extends abc-parametrizations to arbitrary neural architectures and specifies µP rules for biases and parameter tensors. It also characterizes µP as uniquely stable under maximal-update and maximal-readout-initialization conditions.

C abc-Parametrization for General Neural Architectures

  • Each parameter tensor W in an arbitrary architecture receives its own exponents aW and bW, with W = n^-aW w, w initialized with variance n^-2bW, and learning rate ηn^-c.The trainable parameter is w, initialized as wαβ ∼ N(0, n^-2bW).Source: C abc-Parametrization for General… · p. 35

C.1 Maximal Update Parametrization

  • For MLP biases in µP, hidden-layer biases use abl = -1/2 and bbl = 1/2, while the output bias uses abL+1 = bbL+1 = 0.The output bias is added to the logits.Source: §C.1 · p. 35
  • µP extends to architectures expressible as Tensor Programs, including ResNets and Transformers, and uses a learning rate independent of width (c = 0).For any parameter tensor W, bW = 1/2.Source: §C.1 · p. 35
  • For non-output weights, µP sets aW = -1 + 1/(2pW), where pW is the width-growth exponent of the number of entries in W.The passage assigns pW = 0 when both sides stay fixed, 1 for a vector or one fixed-dimensional side, and 2 when both matrix sides scale with width.Source: §C.1 · pp. 35–36
  • For output weight matrices with fixed output dimension, µP sets aW = 1/2; for output biases, it sets aW = 0.These output-parameter rules are exceptions to the non-output-weight rule.Source: §C.1 · p. 36
  • µP is the unique stable abc-parametrization in this general context whose parameter tensors are all updated maximally and whose readout weights are all initialized maximally.The result formalizes stability, maximal parameter updates, and maximal readout initialization for general architectures.Source: §C.1 · p. 36

D Experimental Details … E More Detailed Comparison with Deep Mean Field Limits

The appendices specify experimental procedures and coefficient-space implementations for the infinite-width µP models, then contrast their Word2Vec behavior with the NTK limit and deep mean-field scaling. The supplied passages give Omniglot evaluation procedures and a qualitative NTK Word2Vec result, but no final benchmark comparison or accuracy value.

D Experimental Details

  • The appendix gives an almost-sure infinite-width output limit for a one-hidden-layer linear µP network, expressed through a recurrence for its state matrices.The recurrence starts from identity matrices for A and D and zero matrices for B and C.Source: D Experimental Details · p. 36

D.1.1 Linear 1-Hidden-Layer µP Network

  • The linear µP implementation tracks evolving parameters through coefficients in initialized Gaussian vectors, making its infinite-width algorithm essentially match finite-width SGD apart from initialization and hidden size.For MAML, adaptation omits gradient clipping, validation-loss gradients are clipped, and weight decay is disabled.Source: §D.1.1 · pp. 36–37

D.1 Few-Shot Learning on Omniglot via MAML

  • Omniglot evaluation uses 1-shot, 5-way tasks, 100 training epochs over three seeds, and a 15-seed test rerun for the selected hyperparameters.Validation accuracy is averaged over the last 10 epochs; NTK and GP additionally tune bias initialization, while finite and µP models set it to zero.Source: §D.1.1 · p. 38

D.1.2 NNGP and NTK for Relu Networks

  • For ReLU NNGP and NTK MAML, the trained linear layer can be represented by a dynamic stack of input–coefficient pairs, avoiding explicit kernel-feature embeddings.The method uses only the kernel, and the passages describe five-epoch hyperparameter sweeps followed by a 15-seed test rerun.Source: §D.1.2 · pp. 38–39

D.2 Word2Vec Experimental Details

  • The Word2Vec experimental setup initializes µP input and output weights with variances scaled by 1/n and uses the same selected hyperparameters across widths.The reported choices are σu = σv = 1, η = 0.05, and γ = 0.001.Source: §D.2.1 · pp. 39–40

D.2.1 µP Limit

  • µP Word2Vec uses σu = σv = 1, η = 0.05, and γ = 0.001, and its word vectors are read from the trained embedding rows.The limit is trained in coefficient space with hidden size 2|V|, one-hot embedding initialization, and no biases.Source: §D.2.1 · p. 40

D.2.2 NTK Limit

E More Detailed Comparison with Deep Mean Field Limits

  • With iid zero-mean 1/n weight scaling, deep mean-field preactivations vanish and gradients vanish beyond two hidden layers, unlike the 1/√n scaling considered here.The 1/√n scaling also requires handling central-limit effects, whereas the cited mean-field calculation proceeds through the law of large numbers.Source: E More Detailed Comparison with Deep… · p. 40

F Nuances of the Master Theorem

The appendix clarifies how Master Theorem partial derivatives are defined from a program’s representation and broadens derivative expectations using Stein’s lemma. It also specifies regularity assumptions on functions used in the theorem.

Partial derivative

  • The partial derivative treats each program variable’s limit as a deterministic function of earlier inputs’ limits, differentiating that function with respect to its inputs.Derivatives with respect to variables outside that input set are defined as zero.Source: F Nuances of the Master Theorem · p. 41
  • The partial derivative depends on how the program is written, even when the underlying mathematics is equivalent; the resulting expression for ZDot can nevertheless agree across forms.For the two forms shown, ZDot_W⊤x = (Zy + Zz) E φ′(Ẑ_W(y+z)).Source: F Nuances of the Master Theorem · p. 41

Partial derivative expectation

  • The partial derivative expectation is extended to cases where the nonlinearity is not differentiable by using Stein’s lemma, and Theorem 7.4 holds with this broader definition.The passage gives the Heaviside step function as an example and describes a covariance-pseudoinverse construction.Source: F Nuances of the Master Theorem · p. 41

Pseudo-Lipschitz functions

  • Pseudo-Lipschitz functions are polynomially bounded, remain Lipschitz on compact sets, and compose with degrees that add.Their definition bounds function differences by input distance times a polynomial-growth factor.Source: F Nuances of the Master Theorem · p. 41

Master Theorem assumptions

  • Functions with inputs used in Nonlin or Moment must be pseudo-Lipschitz in inputs and parameters; parameter-only functions in Moment need only be continuous.The continuity condition permits any loss function with a continuous derivative.Source: F Nuances of the Master Theorem · p. 41

Alternative theorem conditions

  • Other Master Theorem versions [52] relax smoothness requirements in exchange for rank stability, which is generically expected but subtle to verify rigorously.This paper instead adopts the pseudo-Lipschitz condition.Source: F Nuances of the Master Theorem · pp. 41–42

G A Rough Sketch of the Geometry of abc-Parametrizations · H Proofs of Main Results

The stable abc-parametrizations form an unbounded polyhedron whose nontrivial parametrizations lie on two facets. Uniformly updated parametrizations form an edge connecting µP and NTP, with a unique parametrization for each r in [0, 1/2].

G A Rough Sketch of the Geometry of abc-Parametrizations

  • The stable polyhedron is unbounded: shifting a_l up and b_l down by θ ≥ 0 preserves stability while decreasing layer l’s learning rate until W^l is not trained.This shift generates another stable parametrization for every θ ≥ 0.Source: G A Rough Sketch of the Geometry of… · p. 42
  • Nontrivial parametrizations lie on two unbounded facets, whose intersection contains NTP and µP as vertices connected by an edge.Source: G A Rough Sketch of the Geometry of… · p. 42
  • A stable parametrization updates the MLP body uniformly exactly when r_l = r for every layer l.Uniform updating means that for some training routine, time t ≥ 1, and input ξ, every layer has ΔW^l_t(ξ) = Θ(n^-r).Source: G A Rough Sketch of the Geometry of… · p. 42
  • NTP uniquely combines uniform body updates with maximally initialized and updated W^L+1; its initial output converges in distribution to a Gaussian Process with nonzero kernel.Uniqueness is modulo Eq. (5).Source: G A Rough Sketch of the Geometry of… · p. 42
  • For each r in [0, 1/2], a unique stable Uniform Parametrization UPr has a_l = 2I(l = 1) + r, a_L+1 = 1/2, b_l = 1/2 − r, and c = 0.The parametrization is unique modulo Eq. (5).Source: G A Rough Sketch of the Geometry of… · p. 42
  • UP0 is µP and UP1/2 is NTP; no such uniform parametrization is stable for r > 1/2 because W0 would scale as Θ(n^r−1), making the initial GP blow up.Thus the UPr family for r in [0, 1/2] forms an edge of the stable polyhedron.Source: G A Rough Sketch of the Geometry of… · p. 42

H.1 Rigorous Statements of Main Results … H.5 The Infinite-Width Limit

This section rigorously characterizes stability and feature-learning behavior in infinite-width training, including the conditions under which stable dynamics learn features or follow kernel gradient descent. It also constructs a Tensor Program for MLP training and establishes convergence properties of the resulting infinite-width limit.

H.1 Rigorous Statements of Main Results

  • Stability holds iff a1 + b1 = 0, al + bl = 1/2 for l = 2,…,L, and aL+1 + bL+1 ≥ 1/2, alongside the stated initialization and training boundedness conditions.The theorem assumes tanh or sufficiently small-σ gelu; stability requires order-one initial activations and logits and bounded coordinates and logits during training.Source: §H.1 · p. 44
  • A stable parametrization is nontrivial iff aL+1 + bL+1 + r = 1 or 2aL+1 + c = 1.Nontriviality means that the function changes in the infinite-width limit for some training routine.Source: §H.1 · p. 44
  • The feature-learning definition only requires a feature embedding to change for some training scenario; it does not itself measure feature quality.The authors note that their proofs show the existence of such a scenario implies more general scenarios, while experiments address downstream task performance.Source: §H.1 · p. 45
  • For nontrivial stable parametrizations, r = 0 corresponds to feature learning and r > 0 to kernel regime with fixed features; these regimes are mutually exclusive.Feature learning and feature-kernel evolution are equivalent in this setting, and feature or prefeature learning in one layer entails both in that layer and all subsequent layers.Source: §H.1 · p. 46
  • In a stable kernel-regime limit, transfer learning is trivialized whether the new classifier alone or the whole network is trained.This includes linear or nonlinear transfer learning after replacing the classifier layer.Source: §H.1 · p. 47

H.2 Stability at Initialization

  • For any parametrization with a nonzero-almost-everywhere activation, the initialization bounds on preactivations, activations, and expected squared logits hold iff the stability exponent conditions hold.The bounds require order-one initial preactivations and activations in every layer and expected squared logits of order Oξ(1).Source: §H.2 · p. 47
  • For Θξ(1) hidden activations, stability holds iff a1 + b1 = 0 and al + bl = 1/2 for 2 ≤ l ≤ L; additionally, aL+1 + bL+1 ≥ 1/2.Source: §H.2 · p. 48

H.3 Program Setup

  • The program assumes initialization stability, d = 1, pseudo-Lipschitz φ′, and a non-lag output-scale condition; it represents activation and weight increments in Θ(1)-scaled coordinates.The assumed output-scale condition is aL+1 + bL+1 ≤ 2aL+1 + c; the reverse case entails a one-step lag and can be handled similarly.Source: §H.3 · pp. 48–49

H.4 Program Construction

  • The Tensor Program encodes SGD for an L-hidden-layer MLP through separate first and subsequent forward and backward passes.The setup specifies the program’s initial matrices, vectors, and scalars.Source: §H.3 · p. 48, §H.4 · p. 50

H.4.1 First Forward Pass

  • When aL+1 + bL+1 = 1/2, the initial output converges to a nontrivial Gaussian by the CLT, and the program conditions on its values over the input set.The first forward pass constructs the initial activations and output.Source: §H.4.1 · p. 50

H.4.2 First Backward Pass

  • The first backward pass defines the initial loss derivative and represents the normalized change in the output-layer weights after the first gradient step.The loss derivative convergence follows from continuity of the loss derivative.Source: §H.4.2 · p. 51

H.4.3 tth Forward Pass, t ≥1

  • Each subsequent forward pass constructs preactivations and activations from earlier states and represents output changes with increments δft(ξ) = ft(ξ) − ft−1(ξ).Activation increments are expressed through Ψ applied to the previous activation and its change.Source: §H.4.3 · pp. 51–52

H.4.4 tth Backward Pass, t ≥1

  • Each subsequent backward pass computes hidden-layer backpropagation quantities, the loss derivative for the current example, and the normalized output-layer weight change after the SGD update.Source: §H.4.4 · p. 52

H.5 The Infinite-Width Limit

  • Under pseudostability, the program’s θ limits are 0 or 1, and the Master Theorem gives roughly iid limiting coordinates represented by Z variables and program-specific recurrences.The activation-increment recurrence uses Ψ; the weight-change construction uses a ZHat-ZDot decomposition.Source: §H.5 · pp. 52–53
  • The limiting output ˚f_t(ξ) is the sum of its output increments, δ˚f_1(ξ) through δ˚f_t(ξ).Source: §H.5 · p. 53
  • The limiting loss derivative is defined as ˚χ_t = L′(˚f_t, y_0).Source: §H.5 · p. 53
  • A stated identity gives 0(ξ) = 0 and equates Zdxl 0(ξ) with ˆZdxl 0(ξ) for any ξ ∈ X.Source: §H.5 · p. 53
  • For pseudostable parametrizations, all θ• converge to 0 or 1, satisfying Setup 7.2; the Master Theorem then gives a.s. convergence to ˚χ_t for every ξ ∈ X and t ≥ 1.The theorem applies to collections of vectors for which Zvi is defined.Source: §H.5 · p. 54

H.5.1 Some Immediate Results … H.7.3 Applications to Tanh

This appendix section develops results on pseudostable parametrizations, distinguishing kernel dynamics from feature learning and deriving recurrence-based results for the r = 0 case. It also gives stability consequences and activation-specific sign and limit properties for σ-Gelu and tanh, alongside theorem settings for pseudostable parametrizations.

H.5.1 Some Immediate Results

  • Pseudostable parametrizations are stable; when r > 0 they fix all (pre)features and (pre)feature kernels, while an initialization-stable parametrization with r < 0 is not stable.For r > 0, a pseudostable parametrization is nontrivial iff aL+1 + bL+1 + r = 1 or 2aL+1 + c = 1.Source: §H.5.1 · p. 54, §H.6 · p. 57

H.6 r > 0 Implies Kernel Regime

  • For pseudostable parametrizations, r > 0 yields kernel gradient descent, while r = 0 admits feature learning and is not in the kernel regime.For r > 0, δf_t(ξ) = −ηχ_{t−1}Σ(ξ_{t−1}, ξ).Source: §H.6 · p. 56, §H.7 · p. 57
  • If aL+1 + bL+1 + r > 1 and 2aL+1 + c = 1, the r > 0 limit is the NNGP kernel, even when the network body is trained slightly.The stated condition allows body changes that do not interact with W L+1 enough to change f_t.Source: §H.6 · p. 57

H.7 r = 0 Implies Feature Learning

  • When r = 0, pseudostable parametrizations admit (pre)feature learning and (pre)feature kernel evolution, but are not in the kernel regime.The analysis shows feature learning by examining η derivatives; analogous results for prefeatures and other layers can be derived similarly.Source: §H.7 · pp. 57–58

H.7.1 Deriving Recurrence Relations on ∂ηλl, ∂ηγl, ∂²ηλl

  • The recurrence analysis finds that the first η derivative does not establish feature evolution, so the authors use second derivatives and derive recurrences for the relevant quantities.The recurrences concern ∂ηλl, ∂ηγl, and ∂²ηλl.Source: §H.7.1 · pp. 58–60

H.7.2 Applications to σ-Gelu

  • For σ-gelu, the stated constants κl1, κl2, κl3, and γl02(0) are positive and converge to 0 with small σ, whereas κ̇l3 is negative and diverges to −∞.The σ-gelu results include theorems for pseudostable parametrizations with r = 0, including under the stated boundary conditions aL+1 + bL+1 + r = 1 or 2aL+1 + c = 1.Source: §H.7.2 · pp. 60–61
  • As σ → 0, dotκ_l3 and γ_l13 tend to −∞, γ_l22 tends to 0, and C tends to a positive constant; γ_l13 < 0 and γ_l22 > 0.The passage also states that ∂2 ηγL−1(0) and ∂2 ηλl(0) tend to 0 by Theorem H.45.Source: §H.7.2 · p. 61

H.7.3 Applications to Tanh

  • For tanh, γ_l02(0), dotκ_l3, and γ_l13 are negative; Theorem H.48 considers pseudostable parametrizations with r = 0 for all l ≥ ℓ.Source: §H.7.3 · p. 62
  • Theorem H.49 considers tanh with a pseudostable parametrization at r = 0, assuming aL+1 + bL+1 + r = 1 or 2aL+1 + c = 1.Source: §H.7.3 · p. 62

H.7.4 Main Results · H.8 Putting Everything Together

For tanh or sufficiently small-σ σ-gelu, stability is equivalent to pseudostability. The results classify stable parametrizations by whether they learn features or instead lie in a kernel regime with fixed features.

H.7.4 Main Results

  • For r = 0, a pseudostable parametrization is nontrivial iff aL+1 + bL+1 = 1 or 2aL+1 + c = 1.The proposition applies to tanh or σ-gelu with sufficiently small σ.Source: §H.7.4 · p. 62
  • For nontrivial pseudostable parametrizations with r = 0, the dynamics are not in the kernel regime and include feature learning in the Lth layer.These results assume tanh or σ-gelu with sufficiently small σ.Source: §H.7.4 · p. 62
  • Feature-kernel evolution and prefeature learning or kernel evolution in a layer imply feature learning and all four kinds of evolution or learning in that layer and every higher layer through L.The listed kinds are feature learning, feature-kernel evolution, prefeature learning, and prefeature-kernel evolution.Source: §H.7.4 · p. 62
  • An initialization-stable parametrization with r = 0 is not stable if aL+1 + bL+1 < 1 or 2aL+1 + c < 1.The proof establishes instability through divergence of f1(ξ0) in the stated cases.Source: §H.7.4 · p. 63

H.8 Putting Everything Together

  • Stable parametrizations are exactly the pseudostable parametrizations.This equivalence is stated for tanh or σ-gelu with sufficiently small σ.Source: §H.8 · p. 63
  • For nontrivial stable parametrizations, r = 0 entails feature learning, while r > 0 is equivalent to the kernel regime and fixed features and feature kernels.The r > 0 equivalences also include fixed prefeatures and prefeature kernels.Source: §H.8 · pp. 63–64
  • If r = 0, then for all ξ ∈ X, f0(ξ) converges almost surely to 0 and ft(ξ) converges almost surely to a deterministic limit; however, the converse is not true.Source: §H.8 · p. 64
  • The stated result is that, when r > 0, aL+1 + bL+1 + r > 1, and 2aL+1 + c = 1, the limit is a Neural Network-Gaussian Process limit.Source: §H.8 · p. 64
Loading 2011.14522v3…