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 with 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.

该研究揭示了神经网络中特征学习的机制及其对学习速度、初始值的敏感性，并提供了使用不平衡初始化来增强高效特征学习的理论依据。