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

Daniel Kunin, Allan Raventos, Clémentine Carla Juliette Dominé, Feng Chen, David Klindt, Andrew M Saxe, Surya Ganguli

Keywords: feature learning, rich regime, lazy regime, exact solutions, conserved quantities, balanced initialization, neural tangent kernel, grokking

TL;DR: We derive exact solutions to a minimal model that transitions between lazy and rich learning, precisely elucidating how unbalanced initialization variances and learning rates determine the degree of feature learning in a finite-width network.

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. 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 conspire to influence the degree of feature learning through a set of conserved quantities that constrain and modify the geometry of learning trajectories. We extend our analysis to more complex linear models and to shallow nonlinear networks with piecewise linear activation functions. In linear networks, rapid feature learning only occurs with balanced initializations, 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.

Student Paper: Yes

Submission Number: 59