Critical feature learning in deep neural networks

arXiv (Cornell University)

Abstract

A key property of neural networks driving their success is their ability to learn features from data. Understanding feature learning from a theoretical viewpoint is an emerging field with many open questions. In this work we capture finite-width effects with a systematic theory of network kernels in deep non-linear neural networks. We show that the Bayesian prior of the network can be written in closed form as a superposition of Gaussian processes, whose kernels are distributed with a variance that depends inversely on the network width N . A large deviation approach, which is exact in the proportional limit for the number of data points $P = αN \rightarrow \infty$, yields a pair of forward-backward equations for the maximum a posteriori kernels in all layers at once. We study their solutions perturbatively to demonstrate how the backward propagation across layers aligns kernels with the target. An alternative field-theoretic formulation shows that kernel adaptation of the Bayesian posterior at finite-width results from fluctuations in the prior: larger fluctuations correspond to a more flexible network prior and thus enable stronger adaptation to data. We thus find a bridge between the classical edge-of-chaos NNGP theory and feature learning, exposing an intricate interplay between criticality, response functions, and feature scale.

Authors 5

1. Javed Lindner
2. David Dahmen

   Forschungszentrum Jülich · Jülich Aachen Research Alliance

   Affiliation as printed

   Institute for Advanced Simulation (IAS- 6) , Computational and Systems Neuroscience , Jülich Research Centre , Jülich , Germany

3. Zohar Ringel

   Hebrew University of Jerusalem

   Affiliation as printed

   The Racah In- stitute of Physics , The Hebrew University of Jerusalem , Jerusalem , Israel

4. M. Krämer
5. Moritz Helias Aachen Department of Physics

   RWTH Aachen University · Forschungszentrum Jülich · Jülich Aachen Research Alliance

   Affiliation as printed

   Department of Physics , RWTH Aachen University , Aachen , Germany

   Institute for Advanced Simulation (IAS- 6) , Computational and Systems Neuroscience , Jülich Research Centre , Jülich , Germany