Abstract:While the impressive performance of modern neural networks is often attributed to their capacity to efficiently extract task-relevant features from data, the mechanisms underlying this rich feature learning regime remain elusive, with much of our theoretical understanding stemming from the opposing lazy regime. In this work, we derive exact solutions to a minimal model that transitions between lazy and rich learning, precisely elucidating how unbalanced layer-specific initialization variances and learning rates determine the degree of feature learning. Our analysis reveals that they conspire to influence the learning regime through a set of conserved quantities that constrain and modify the geometry of learning trajectories in parameter and function space. We extend our analysis to more complex linear models with multiple neurons, outputs, and layers and to shallow nonlinear networks with piecewise linear activation functions. In linear networks, rapid feature learning only occurs from balanced initializations, where all layers learn at similar speeds. While in nonlinear networks, unbalanced initializations that promote faster learning in earlier layers can accelerate rich learning. Through a series of experiments, we provide evidence that this unbalanced rich regime drives feature learning in deep finite-width networks, promotes interpretability of early layers in CNNs, reduces the sample complexity of learning hierarchical data, and decreases the time to grokking in modular arithmetic. Our theory motivates further exploration of unbalanced initializations to enhance efficient feature learning.

A spring-block theory of feature learning in deep neural networks

A spring-block theory of feature learning in deep neural networks

Low-dimensional Intrinsic Dimension Reveals a Phase Transition in Gradient-Based Learning of Deep Neural Networks

Critical feature learning in deep neural networks

Deep CovDenseSNN: A Hierarchical Event-Driven Dynamic Framework with Spiking Neurons in Noisy Environment

An effective theory of collective deep learning

Learning Curves for Deep Neural Networks: A Gaussian Field Theory Perspective

Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning

Two-Phase Dynamics of Interactions Explains the Starting Point of a DNN Learning Over-Fitted Features

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

A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks

Envisioning Future Deep Learning Theories: Some Basic Concepts and Characteristics

Three Mechanisms of Feature Learning in the Exact Solution of a Latent Variable Model

An Information Theoretic Interpretation to Deep Neural Networks.

Grokking as a First Order Phase Transition in Two Layer Networks

Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks

Unifying Low Dimensional Observations in Deep Learning Through the Deep Linear Unconstrained Feature Model

Imitating Deep Learning Dynamics via Locally Elastic Stochastic Differential Equations

Opening the black box of deep learning

Plastic Learning with Deep Fourier Features

Feature Learning and Generalization in Deep Networks with Orthogonal Weights