Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On the Neural Feature Ansatz for Deep Neural Networks

Edward Tansley, Estelle Massart, Coralia Cartis

TL;DR

This work develops a depth-aware theoretical foundation for the Neural Feature Ansatz (NFA), linking the first-layer Gram matrix to the AGOP via a depth-dependent exponent α, with α=1/L for L-layer deep linear networks under gradient flow and balanced initialization. It further shows that with weight decay, the NFA persists asymptotically even for unbalanced initializations, while constructing counterexamples demonstrating that nonlinear activations can violate the NFA. Complementary numerical experiments across optimizers, initializations, and architectures validate the theoretical predictions and reveal a robust low-rank structure being captured by the first-layer features when learning multi-index targets. The results highlight both the depth-induced biases shaping feature learning and the limitations of NFA as a universal descriptor of generalization, offering a nuanced view of how depth and regularization influence learned representations. Formally, the core relationships include the proportionality W_1^T W_1 ∝ (A_f)^α with α depending on depth (α=1/L) and A_f = J_f^T J_f, both of which are supported by theory and experiments.

Abstract

Understanding feature learning is an important open question in establishing a mathematical foundation for deep neural networks. The Neural Feature Ansatz (NFA) states that after training, the Gram matrix of the first-layer weights of a deep neural network is proportional to some power $α>0$ of the average gradient outer product (AGOP) of this network with respect to its inputs. Assuming gradient flow dynamics with balanced weight initialization, the NFA was proven to hold throughout training for two-layer linear networks with exponent $α= 1/2$ (Radhakrishnan et al., 2024). We extend this result to networks with $L \geq 2$ layers, showing that the NFA holds with exponent $α= 1/L$, thus demonstrating a depth dependency of the NFA. Furthermore, we prove that for unbalanced initialization, the NFA holds asymptotically through training if weight decay is applied. We also provide counterexamples showing that the NFA does not hold for some network architectures with nonlinear activations, even when these networks fit arbitrarily well the training data. We thoroughly validate our theoretical results through numerical experiments across a variety of optimization algorithms, weight decay rates and initialization schemes.

On the Neural Feature Ansatz for Deep Neural Networks

TL;DR

This work develops a depth-aware theoretical foundation for the Neural Feature Ansatz (NFA), linking the first-layer Gram matrix to the AGOP via a depth-dependent exponent α, with α=1/L for L-layer deep linear networks under gradient flow and balanced initialization. It further shows that with weight decay, the NFA persists asymptotically even for unbalanced initializations, while constructing counterexamples demonstrating that nonlinear activations can violate the NFA. Complementary numerical experiments across optimizers, initializations, and architectures validate the theoretical predictions and reveal a robust low-rank structure being captured by the first-layer features when learning multi-index targets. The results highlight both the depth-induced biases shaping feature learning and the limitations of NFA as a universal descriptor of generalization, offering a nuanced view of how depth and regularization influence learned representations. Formally, the core relationships include the proportionality W_1^T W_1 ∝ (A_f)^α with α depending on depth (α=1/L) and A_f = J_f^T J_f, both of which are supported by theory and experiments.

Abstract

Understanding feature learning is an important open question in establishing a mathematical foundation for deep neural networks. The Neural Feature Ansatz (NFA) states that after training, the Gram matrix of the first-layer weights of a deep neural network is proportional to some power of the average gradient outer product (AGOP) of this network with respect to its inputs. Assuming gradient flow dynamics with balanced weight initialization, the NFA was proven to hold throughout training for two-layer linear networks with exponent (Radhakrishnan et al., 2024). We extend this result to networks with layers, showing that the NFA holds with exponent , thus demonstrating a depth dependency of the NFA. Furthermore, we prove that for unbalanced initialization, the NFA holds asymptotically through training if weight decay is applied. We also provide counterexamples showing that the NFA does not hold for some network architectures with nonlinear activations, even when these networks fit arbitrarily well the training data. We thoroughly validate our theoretical results through numerical experiments across a variety of optimization algorithms, weight decay rates and initialization schemes.
Paper Structure (31 sections, 10 theorems, 43 equations, 8 figures, 3 tables, 1 algorithm)

This paper contains 31 sections, 10 theorems, 43 equations, 8 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Notation:
  3. Preliminaries
  4. Low-rank functions.
  5. Average Gradient Outer Product (AGOP).
  6. The Neural Feature Ansatz (NFA).
  7. The NFA for deep linear networks
  8. Removing the balancedness assumption
  9. The NFA for nonlinear networks
  10. What about wider or deeper networks?
  11. The NFA and generalization.
  12. Numerical Experiments
  13. Data generation.
  14. Validating our theoretical results
  15. Architectures considered.
  16. ...and 16 more sections

Key Result

Lemma 3.1

For time $t \geq 0$, let us define $f_{t}({\bm{x}}) = {\bm{W}}_{L,t}{\bm{W}}_{L-1,t}\cdots {\bm{W}}_{1,t}{\bm{x}}$ for ${\bm{x}} \in \mathbb{R}^{d}$. Suppose that ${\bm{W}}_{1,t},{\bm{W}}_{2,t},\ldots,{\bm{W}}_{L,t}$ follow the gradient flow dynamics given by eq:gradient_flow, then for any $t\geq0$

Figures (8)

  • Figure 1: The NFA holds after training even with unbalanced initialization. (a) ${{\bm{W}}_{1}^{\top}{\bm{W}}_{1}}$ before training. (b) $({{\bm{A}}_{f}})^{1/L}$ before training. (c) ${{\bm{W}}_{1}^{\top}{\bm{W}}_{1}}$ after training. (d) $({{\bm{A}}_{f}})^{1/L}$ after training, with $L=5$ linear layers. Alignment as measured by cosine similarity (\ref{['def:cosine_sim']}): before training = 0.915; after training = 1.000. We plot the same experiment as in \ref{['fig:cosine_similarity_though_time']} (a) which is described in \ref{['sec:numerics']}.
  • Figure 2: Illustration of the impact of initialization and weight decay ($\lambda$) on the alignment between ${{\bm{W}}_{1}^{\top}{\bm{W}}_{1}}$ and $({{\bm{A}}_{f}})^{1/L}$ at the end of training (SGD). A learning rate of $\eta = 10^{-4}$ was used.
  • Figure 3: Illustration of the impact of the optimization algorithm on the NFA, with weight decay ($\lambda = 10^{-2}$). When momentum is used, it is weighted by a parameter $\beta = 0.9$. The learning rates for SGD and GD are $\eta = 10^{-4}$ and $\eta = 10^{-3}$, respectively. Here we plot the first $60,000$ epochs before the learning rate decrease.
  • Figure 4: Singular values decay for ${{\bm{W}}_{1}^{\top}{\bm{W}}_{1}}$ vs ${{\bm{A}}_{f^{*}}}$. Normalized Singular Values are computed by dividing by $\sigma_1$ for the respective matrix
  • Figure 5: Illustration of the impact of the weight decay ($\lambda$) on the alignment between ${{\bm{W}}_{1}^{\top}{\bm{W}}_{1}}$ and $({{\bm{A}}_{f}})^{1/L}$ at the end of training (Adam). A learning rate of $\eta = 10^{-4}$ was used.
  • ...and 3 more figures

Theorems & Definitions (20)

  • Definition 2.1: cartis_learning_2024
  • Definition 3.1: Balanced matrices
  • Lemma 3.1: arora_optimization_2018
  • Theorem 3.1: NFA for deep linear networks
  • proof
  • Lemma 3.2: kobayashi_weight_2024
  • Theorem 3.2: NFA for deep linear networks, without balanced initialization
  • Corollary 3.1
  • Example 1
  • Definition 5.1
  • ...and 10 more