PaperlayerView exact revision on arXiv

Source-linked AI summary

Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks

Blake Bordelon, Cengiz Pehlevan

arXiv:2205.09653v3stat.MLcond-mat.dis-nncs.LG

TL;DR

Standard infinite-width kernel limits keep data representations fixed and thus cannot explain feature learning. The paper derives self-consistent dynamical kernel equations from a field theory and solves them with a polynomial-time procedure for deep networks. At γ0 = 1, the saddle-point equations recover the Tensor Programs stochastic process, and numerical solutions predict training across feature-learning strengths, widths, and depths.

  • Problem

    Standard infinite-width kernel limits keep data representations fixed and thus cannot explain feature learning.

    Source: §1 · p. 1

  • Method

    The paper derives self-consistent dynamical kernel equations from a field theory and solves them with a polynomial-time procedure for deep networks.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    At γ0 = 1, the saddle-point equations recover the Tensor Programs stochastic process, and numerical solutions predict training across feature-learning strengths, widths, and depths.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The framework relates stochastic hidden-layer fields to evolving kernels and predictions, and its self-consistent solution remains accurate over a wider range of γ0 than the compared approximation schemes.

    Source: §4.2 · p. 7, §7 · p. 9, §5.2 · p. 9

    The gradient-independence ansatz begins to fail at larger γ0 or later times, while perturbation theory breaks down at large γ0 and depth.

    Source: §5.1 · p. 8, §5.2 · p. 8

Abstract

from arXiv · show

We analyze feature learning in infinite-width neural networks trained with gradient flow through a self-consistent dynamical field theory. We construct a collection of deterministic dynamical order parameters which are inner-product kernels for hidden unit activations and gradients in each layer at pairs of time points, providing a reduced description of network activity through training. These kernel order parameters collectively define the hidden layer activation distribution, the evolution of the neural tangent kernel, and consequently output predictions. We show that the field theory derivation recovers the recursive stochastic process of infinite-width feature learning networks obtained from Yang and Hu (2021) with Tensor Programs . For deep linear networks, these kernels satisfy a set of algebraic matrix equations. For nonlinear networks, we provide an alternating sampling procedure to self-consistently solve for the kernel order parameters. We provide comparisons of the self-consistent solution to various approximation schemes including the static NTK approximation, gradient independence assumption, and leading order perturbation theory, showing that each of these approximations can break down in regimes where general self-consistent solutions still provide an accurate description. Lastly, we provide experiments in more realistic settings which demonstrate that the loss and kernel dynamics of CNNs at fixed feature learning strength is preserved across different widths on a CIFAR classification task.

1 Introduction

The introduction motivates a tractable theory of feature learning in wide networks, beyond infinite-width kernel limits with fixed representations. It presents self-consistent DMFT as an alternative that models feature dynamics and connects to prior infinite-width analyses.

1 Introduction

  • Fixed representations in standard infinite-width kernel limits leave feature learning unexplained, despite its importance to deep learning practice [21].The introduction frames a useful theory as analytically tractable, efficiently computable, predictive of performance and learned features, and reducible to interpretable quantities.Source: §1 · pp. 1–2
  • The paper develops a DMFT for infinite-width feature learning, recovers the Tensor Programs stochastic process at γ0 = 1, and provides a polynomial-time solver.The parameter γ0 interpolates between lazy and rich behavior; the numerical method is reported to predict training across feature-learning strengths, widths, and depths.Source: §1 · p. 2

1.1 Related Works

  • Earlier finite-width correction methods are approximate and can give incorrect predictions in very rich feature-learning regimes [37] [38].The paper contrasts these methods with a parameterization that permits feature learning at infinite width.Source: §1.1 · p. 2
  • The work builds on Tensor Programs analyses of infinite-width gradient descent and develops self-consistent DMFT for deep feature learning, a setting not previously treated with that framework.It also proposes alternating sampling to avoid the exponential-time cost of computing the exact infinite-width gradient-descent limit.Source: §1.1 · p. 2

2 Problem Setup and Definitions

The setup considers infinite-width fully connected and convolutional networks and introduces dynamical quantities for describing training. Its parameterization controls the feature-learning regime, while the temporal NTK trajectory connects to network predictions.

2 Problem Setup and Definitions

  • The framework applies to infinite-width fully connected and convolutional networks, with the main text focusing on equal widths and NTK parameterization for simplicity.The paper states that the chosen width scaling is the only one allowing feature learning as N → ∞, and that γ0 = 0 recovers the static NTK limit [7].Source: §2 · p. 3
  • The parameter γ controls the richness or laziness of training dynamics, and the learning-rate choice γ^2 makes the initial loss derivative independent of γ.The setup introduces backpropagation variables and identifies the pre-gradient signal, though their full definitions are not present in these passages.Source: §2 · p. 3
  • Feature and gradient kernels are introduced for each hidden layer, and the temporal NTK trajectory on the t = s diagonal determines network prediction dynamics.The setup also places the analysis in the mean-field limit.Source: §2 · p. 3

3 Self-consistent DMFT

The self-consistent DMFT describes training through dynamical kernel order parameters and stochastic fields for hidden-layer preactivations and pre-gradients. A path-integral action yields saddle-point equations, while extensions cover several architectures and training settings.

3 Self-consistent DMFT

  • The dynamical order parameters concentrate across random initializations while adapting to task structure, and characterize the distribution of each layer’s preactivations and pre-gradients.The kernels are defined as self-consistent averages over the resulting neuron distribution.Source: §3 · p. 3

3.1 Path Integral Construction

  • The path-integral construction separates initial-weight contributions from gradient updates and represents the stochastic fields through a moment-generating functional.Integrating over initial Gaussian weights yields an action expressed using order parameters; its single-site moment-generating function defines the field distribution, with A and B encoding feedforward-feedback correlations.Source: §3.1 · p. 4

3.2 Deriving the DMFT Equations from the Path Integral Saddle Point

  • At infinite width, the moment-generating function is dominated by the action’s saddle point, whose stationarity conditions define the DMFT kernels.The stochastic fields are induced by the single-site process, with mean-zero Gaussian processes entering through their covariances.Source: §3.2 · pp. 4–5
  • The DMFT extends to varying widths, trainable biases, convolutional networks, weight decay, Langevin sampling, momentum, and discrete-time training.The derived stochastic process is stated to be equivalent to the μP scheme of Yang and Hu.Source: §3.2 · p. 5

4 Solving the Self-Consistent DMFT

The section solves the self-consistent DMFT kernels by alternating sampling for nonlinear networks and closed-form equations for deep linear networks. It shows kernel evolution across depths and a nontrivial regularized fixed point when feature learning is present.

4 Solving the Self-Consistent DMFT

  • For nonlinear networks, alternating Monte Carlo sampling updates kernels and activation-gradient distributions until they are self-consistent.The resulting dynamical kernels can be used to compute network observables.Source: §4 · p. 5

4.1 Deep Linear Networks: Closed Form Self-Consistent Equations

  • Deep linear networks reduce the saddle-point equations to algebraic formulas in the kernels, avoiding the sampling procedure required for nonlinear networks.The equations use causal integral operators determined by kernels in adjacent layers.Source: §4.1 · p. 5
  • The self-consistent equations describe deep linear kernel evolution across depths, with larger feature-kernel changes at greater depth.The DMFT kernels also approach finite-width network estimates as width increases.Source: §4.1 · pp. 5–6, Figure 2 · p. 7
  • The deep linear DMFT applies to arbitrary input kernel Kx and targets y, without the input-structure restrictions of many prior results [69] [70].For whitened inputs, the dynamics interpolate from NTK-like linear convergence at γ0 → 0 toward logistic dynamics as γ0 increases.Source: §4.1 · pp. 5–6

4.2 Feature Learning with L2 Regularization

  • With weight decay, homogeneous networks can retain a nonzero infinite-width fixed point for the predictor and final NTK when feature learning strength γ0 is positive.The fixed point depends on λ, κ, depth L, and γ0.Source: §4.2 · pp. 6–7, Figure 3 · p. 7

5 Approximation Schemes

The section compares the self-consistent DMFT with gradient-independence and perturbative approximations, identifying where each approximation works and where it breaks down. DMFT remains accurate across a wider range of feature-learning strengths.

5.1 Gradient Independence Ansatz

  • The gradient-independence ansatz treats forward- and backward-pass initial weight matrices as independent Gaussian matrices, working near γ0 ≈ 0 or t ≈ 0 but failing at larger γ0 or t.It entails conditional Gaussianity and statistical independence of χℓ and ξℓ given {Φℓ, Gℓ}.Source: §5.1 · p. 8

5.2 Perturbation theory in γ0 at infinite-width

  • Kernel corrections at O(γ0) and O(γ0^3) vanish, leaving leading corrections at O(γ0^2); the leading-order relative NTK change scales linearly with depth.The leading-order correction also has temporal dependencies controlled by P(P + 1) functions.Source: §5.2 · p. 8
  • The perturbative theory begins to break down at large γ0 and large depth, despite predicting more significant kernel evolution for deeper networks.The predicted depth trend agrees with other perturbative studies and the paper’s non-perturbative results.Source: §5.2 · p. 8
  • Only non-perturbative DMFT remains accurate over a wide range of γ0, while the compared approximations coincide at γ0 = 0 and deviate in the feature-learning regime.Figure 4 compares loss, feature kernels, and alignment in a five-hidden-layer linear network.Source: Figure 4 · p. 9, §5.2 · p. 9

6 Feature Learning Dynamics is Preserved at Fixed γ0

At fixed feature-learning strength, the reported CIFAR CNN loss and kernel dynamics are largely independent of width. Higher γ0 is associated with faster training and greater kernel evolution.

6 Feature Learning Dynamics is Preserved at Fixed γ0

  • For CNNs on two-class CIFAR, loss and last-layer kernel dynamics are largely independent of width when N is rescaled with γ to keep γ0 fixed.The trend is reported consistently for channel counts N = 250 and N = 500.Source: §6 · p. 9
  • Higher γ0 corresponds to larger changes in kernel norm, faster alignment to the target, and more rapid training in the tested CNNs.The experiment tracks test loss, classification error, and last-layer kernel alignment on boat-versus-plane CIFAR.Source: §6 · p. 9

7 Discussion

The discussion presents DMFT as a self-consistent account of feature dynamics in infinite networks, spanning lazy and rich regimes. It also notes limits in rigor, asymptotic scope, and computational cost.

7 Discussion

  • DMFT interpolates between static-NTK lazy behavior and rich feature learning, and at γ0 = 1 recovers the stochastic process derived with Tensor Programs.The framework’s saddle-point equations are exactly solvable for deep linear networks and numerically solvable for nonlinear ones.Source: §7 · p. 9
  • For CNNs trained on two CIFAR classes, loss and kernel dynamics are consistent across N = 250 and N = 500 at fixed γ0; higher γ0 brings more active kernel evolution and faster training.Figure 5 reports test loss, classification error, and last-layer feature-kernel alignment.Source: Figure 5 · p. 10
  • The derivation uses heuristic physics techniques and applies in the T, P = O(1) regime; other asymptotic regimes may exhibit phenomena relevant to deep learning.The stated full-batch DMFT scaling is O(P^3T^3), and its compute and memory advantage over a width-N network requires N ≫ PT.Source: §7 · p. 10, Table 1 · p. 10

Checklist

The checklist states that the theory applies in the infinite-width limit and that derivations support the main-text claims. Supplementary figures illustrate kernel and loss dynamics across models, feature-learning strengths, and widths.

Checklist

  • The theory assumes infinite-width architectures, and the paper reports derivations supporting its main-text claims.The checklist also reports error bars for DMFT alignment scores as a function of width; other runs use one wide network based on predicted concentration over initialization.Source: Checklist · p. 16

A Additional Figures

  • In a two-layer ReLU network, self-consistent DFT reproduces preactivation distributions, loss dynamics, and learned kernels, with larger γ0 associated with non-Gaussian activations and stronger kernel block structure.For small γ0, final preactivations are approximately Gaussian; the DMFT distribution estimate is noisy due to finite sampling.Source: Figure 6 · p. 17
  • For the L = 3 tanh network, DMFT predicts loss and final kernels, whose block structure reflects class clustering, while temporal kernel components show nontrivial dynamics.The figure compares a width-2000 network with DMFT solutions.Source: Figure 7 · p. 18
  • In a depth-2 linear network, learning curves shift from exponential convergence at small γ0 to logistic-sigmoidal trajectories at large γ0.The accompanying kernel quantity is its projection onto the task-relevant subspace.Source: Figure 8 · p. 18
  • At γ0 = 1.5, gradient independence underestimates feature-kernel adaptation in the L = 5 deep linear network.The figure also presents loss curves predicted under the gradient-independence ansatz.Source: Figure 9 · p. 19
  • Across widths N ∈{250, 500}, CNN loss and prediction dynamics agree, while deeper-layer feature kernels show stronger finite-size effects and final test accuracy is roughly insensitive to γ0.Higher γ0 is associated with faster convergence in this model.Source: Figure 10 · p. 19

B Algorithmic Implementation

The nonlinear-network solver alternates between sampling fields conditional on kernels and estimating kernels from sampled moments. The experimental details describe the associated weighting scheme and MLP implementation.

B Algorithmic Implementation

  • The nonlinear solver alternates between sampling fields from Gaussian processes conditional on kernels and estimating kernel order parameters from moments of those samples.This sample-and-solve strategy is presented as an efficient numerical approach to the self-consistent equations.Source: B Algorithmic Implementation · pp. 19–20
  • The update speed β controls recency weighting; using β < 1 allows smaller per-iteration sample counts while retaining accurate estimates.The deep-network experiments used β = 0.6, and convergence usually took about 15 steps for the cited depth-4 network.Source: B Algorithmic Implementation · p. 20

C.1 MLP Experiments

  • The MLP experiments used full-batch gradient descent with Gaussian weight initialization, and hidden features were stored during training to compute kernels.The experiments could be reproduced with the provided notebooks.Source: §C.1 · p. 20

C.2 CNN Experiments on CIFAR-10

The CIFAR-10 experiments use a ReLU CNN with global average pooling and a dense readout. Training uses MSE, and test-set embeddings are used to calculate kernel alignment.

C.2 CNN Experiments on CIFAR-10

  • The CNN uses stride-1 ReLU convolutions, global average pooling, and a final dense readout layer.The model is implemented as a JAX parameter pytree.Source: §C.2 · p. 20
  • Training uses MSE loss with base learning rate η0 = 2.0 × 10−4 and batch size 250.The learning rate passed to the optimizer is scaled using γ0.Source: §C.2 · p. 21
  • The experiments calculate alignment between the last-layer test embedding kernel and yy⊤; training uses four NVIDIA GPUs.A width-500, L = 3 network takes roughly one hour to train.Source: §C.2 · p. 21

D Derivation of Self-Consistent Dynamical Field Theory

The appendix derives self-consistent kernel equations through a path-integral field theory and saddle-point analysis. It connects kernel order parameters to stochastic field distributions, NTK dynamics, and predictions, then specializes the framework to varying widths and simpler networks.

D Derivation of Self-Consistent Dynamical Field Theory

  • The derivation uses the MSRDJ path-integral framework and scales feature learning through γ0, with positive γ0 yielding order-one feature evolution.The theory expresses gradients and the NTK in terms of hidden features and recursively evolving fields.Source: D Derivation of Self-Consistent… · p. 21, §D.1 · p. 22

D.2.1 Final L = 1 DMFT equations

  • For a one-hidden-layer network, the kernel order parameters are averages of activation and gradient inner products, and solving them requires only time-diagonal kernels.The initial fields are independent Gaussians in this case.Source: §D.2 · p. 24, §D.2.1 · p. 24

D.3 Path Integral Formulation for Deep Networks

  • For deep networks, the path-integral construction enforces field and kernel definitions, factors the hidden-unit integrals, and yields an action whose stationary points define the infinite-width equations.The fields are treated as causally determined stochastic processes.Source: §D.3 · p. 25, §D.4 · p. 26, §D.5 · p. 27

D.5 Saddle Point Equations

  • The saddle-point analysis obtains self-consistency by averaging observables over the single-site distribution induced by the moment-generating function.The action scales with width, enabling saddle-point integration in the infinite-width limit.Source: §D.5 · p. 27

D.6 Single Site Stochastic Process: Hubbard Trick

  • The Hubbard-trick formulation interprets auxiliary fields as independent Gaussian processes, then determines the activation and gradient fields self-consistently.The final equations rescale order parameters so non-Gaussian corrections to preactivations are of order γ0.Source: §D.6 · pp. 28–29, §D.7 · p. 29

D.7 Final DMFT Equations

  • The final self-consistent kernels determine the NTK dynamics and, consequently, network output predictions.The equations are stated after the rescaling of the auxiliary order parameters.Source: §D.7 · p. 29

D.8 Varying Network Widths and Initialization Scales

  • With layer widths Nℓ = aℓN, the field evolution scales as γ0/√aℓ; wider layers evolve less, and the authors predict more feature learning in early layers of widening CNNs.That CNN prediction is left for future evaluation.Source: §D.8 · p. 31

E Two Layer Networks; F Deep Linear Networks; F.1 Two Layer Linear Network

  • In the two-layer case, constant independent fields permit closed dynamics in H, G, and Δ; the appendix also gives a kernel-update procedure for deep linear networks.The two-layer fields can be described by differential equations, and the stated initial conditions apply for any data Kx and y.Source: E Two Layer Networks · p. 31, F Deep Linear Networks · p. 32, §F.1 · p. 33

F.1.1 Whitened Data in Two Layer Linear

For whitened two-layer linear networks, the scalar dynamics reduce to the task direction and feature evolution accelerates learning as γ0 increases. The supplied passages also describe extensions to deep linear and convolutional networks, multiple outputs, weight decay, and Langevin training.

F.1.1 Whitened Data in Two Layer Linear

  • The whitened two-layer system obeys a conservation law constraining (H_y, y−Δ) to a hyperbola, and its final kernel has a rank-one spike in the target direction.The final-kernel result is also noted in finite-width small-initialization networks and a normative toy model.Source: §F.1.1 · p. 33
  • Increasing γ0 provides strict acceleration in the whitened two-layer learning dynamics.At large γ0, the learning timescale is τ ∼ 1/γ0, but the stated asymptotics use an unstable initial condition and require an alternative condition for realistic learning curves.Source: §F.1.1 · pp. 33–34

F.2 Deep Linear Whitened Data

  • In deep linear networks on whitened data, hidden kernels need only be tracked in the one-dimensional task-relevant subspace along y.The response-function formulation allows arbitrary sample sizes with no sample-size dependence in its computation.Source: §F.2 · p. 34

G Convolutional Networks with Infinite Channels

  • The infinite-channel CNN field theory scales as O(P^3|S|^3T^3), making its solutions computationally intensive.The construction effectively multiplies the number of sample indices by the number of spatial indices.Source: G Convolutional Networks with Infinite… · p. 35

I Multiple Output Channels

J Weight Decay in Deep Homogenous Networks

  • With weight decay, the fixed-point function dynamics yield a representer theorem using the final NTK, and the DMFT is demonstrated for a two-layer ReLU network.For γ0 > 0 the network has a nontrivial long-time fixed point, unlike the γ0 = 0 limit described in the passage.Source: J Weight Decay in Deep Homogenous… · p. 37

K Bayesian/Langevin Trained Mean Field Networks

  • In Langevin dynamics, β controls the injected Gaussian noise scale; β → ∞ recovers usual DMFT, while finite β balances feature updates with noise.At late times, the system equilibrates to its Gibbs distribution.Source: K Bayesian/Langevin Trained Mean Field… · p. 37

K.2 Weak Feature Learning, Long Time Limit

  • In the weak-feature-learning, long-time Langevin limit, the temporal-diagonal equations recover initialization NNGP recursions with weight variance λ^-1.The variance is set by the weight-decay term.Source: §K.2 · p. 40

K.3 Equilibrium Analysis

The appendices examine equilibrium and extensions of the field theory, including momentum and discrete-time training, and derive parameterization constraints for finite-scale prediction and feature evolution. The equilibrium analysis also identifies an unresolved relationship between two orders of limits.

K.3 Equilibrium Analysis

  • The relationship between taking the finite-N Langevin equilibrium limit before N → ∞ and taking the DMFT limits in the opposite order remains open.The finite-N equilibrium may require time extensive in N to reach.Source: §K.3 · p. 40
  • At non-negligible feature-learning strength, the equilibrium last-layer features are non-Gaussian, whereas deep linear network distributions remain Gaussian and close algebraically.Source: §K.3 · p. 41

L Momentum Dynamics

  • The field theory accommodates momentum by changing field definitions; its time constant τ → 0 recovers gradient descent, while τ > 0 accumulates past fields and kernels.Source: L Momentum Dynamics · pp. 41–42

M Discrete Time

  • The discrete-time model loses the NTK as a key player, and its continuous-time limit recovering NTK-governed function evolution remains future work.Source: M Discrete Time · pp. 42–43

N Equivalent Parameterizations

  • The parameterization analysis identifies a one-dimensional family satisfying finite preactivations, finite-time learning, and feature learning, and relates the scheme to µP.The paper states that its stochastic process is equivalent to Yang and Hu’s µP scheme.Source: N Equivalent Parameterizations · p. 43

N.1 Predictions Evolve in ON(1) time

  • When D ∼ ON(N), this simply demands the constraint c = 2aℓ − 1 for every layer ℓ ∈ {0, ..., L}.Source: §N.1 · p. 44

N.2 Fields Are ON(1)

  • The standard scaling a_ℓ = 0, b_ℓ = 1 preserves O(1) feature scale, while O(1) feature-gradient inner products require 2a + b = 1.Source: §N.2 · p. 44

N.3 ON(1) Feature Evolution

  • O(1) feature evolution is necessary and sufficient when d − c − 2a_ℓ + 1/2 = 0.Source: §N.3 · p. 45

N.4 Putting Constraints Together

The appendix constrains parameterizations for feature learning, relates the DMFT process to Tensor Programs, and analyzes perturbative and finite-width kernel corrections. The corrections propagate across layers, while the finite-width analysis leaves their effect on expected test loss unresolved.

N.4 Putting Constraints Together; N.5 ON(1) Learning Rate

  • Requiring O(1) learning rate makes the feature-learning parameterization unique: it is the µP parameterization of Yang and Hu.The O(1) feature-evolution constraints alone allow a one-dimensional family.Source: §N.4 · p. 45, §N.5 · p. 45

N.6 Equivalence of DMFT at γ0 = 1 and Tensor Programs derived Stochastic Process

  • The paper shows that Tensor Programs and DMFT, treated as alternative derivations, give identical descriptions of stochastic processes induced by random initialization and gradient descent (GD) in infinite neural networks.Source: §N.6 · p. 46

O Gradient Independence

  • The gradient independence approximation makes the backward-pass fields independent Gaussian processes, corresponding to Aℓ = Bℓ = 0 in the full DMFT.Source: O Gradient Independence · p. 46

P.1 Small γ0 Expansion

  • The small-γ0 perturbation analysis concerns concentrating kernels, not finite-width kernel fluctuations.The authors distinguish this regime from finite-width perturbative work [27] [28] [35].Source: §P.1 · p. 46

P.2.1 Leading Corrections to Φ1 Kernel is O(γ2; P.3 Forward Pass Induction for Φℓ

  • For nonlinear networks, odd Gaussian moments eliminate linear and cubic terms, so the activation kernel begins at O(γ0^2) and corrections propagate through the forward recursion.The resulting Φℓ expansion is Φℓ,0 + γ0^2Φℓ,2 + O(γ0^4).Source: §P.2.1 · p. 49, §P.3 · p. 50

P.4 Leading Corrections to GL Kernel is O(γ2

  • Odd Gaussian moments also remove the linear and cubic terms in the output-layer gradient-kernel expansion.Source: §P.4 · p. 50

P.6 Form of the Leading Corrections

  • Activation-kernel corrections accumulate from preceding layers, whereas gradient-kernel corrections accumulate from succeeding layers, with Gℓ,2 scaling as O(L + 1 − ℓ).Source: §P.6 · pp. 51–52

P.7 Perturbation Theory in Width N (Finite Size Corrections); P.7.1 DMFT Action Hessian

  • Finite-width kernel-distribution corrections are O(N^-1), approximated by Gaussian order parameters with saddle-point Hessian covariance; whether these fluctuations raise expected test loss remains open.The fluctuations have standard deviation O(N^-1/2).Source: §P.7 · p. 53, §P.7.1 · p. 55
Loading 2205.09653v3…