Asymptotics of representation learning in finite Bayesian neural networks

Abstract

Recent works have suggested that finite Bayesian neural networks may sometimes outperform their infinite cousins because finite networks can flexibly adapt their internal representations. However, our theoretical understanding of how the learned hidden layer representations of finite networks differ from the fixed representations of infinite networks remains incomplete. Perturbative finite-width corrections to the network prior and posterior have been studied, but the asymptotics of learned features have not been fully characterized. Here, we argue that the leading finite-width corrections to the average feature kernels for any Bayesian network with linear readout and Gaussian likelihood have a largely universal form. We illustrate this explicitly for three tractable network architectures: deep linear fully-connected and convolutional networks, and networks with a single nonlinear hidden layer. Our results begin to elucidate how task-relevant learning signals shape the hidden layer representations of wide Bayesian neural networks.

13+28 pages, 4 figures; v3: extensive revision with improved exposition and new section on CNNs, accepted to NeurIPS 2021; v4: minor updates to supplement; v5: post-NeurIPS update, minor typos fixed

References in corpus (12)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• The Principles of Deep Learning Theory
• What Are Bayesian Neural Network Posteriors Really Like?
• Neural Tangents: Fast and Easy Infinite Neural Networks in Python
• Asymptotics of Wide Networks from Feynman Diagrams
• Asymptotics of representation learning in finite Bayesian neural networks
• Contrasting random and learned features in deep Bayesian linear regression
• Bayesian Neural Network Priors Revisited
• Depth induces scale-averaging in overparameterized linear Bayesian neural networks
• Asymptotics of Wide Convolutional Neural Networks
• On the asymptotics of wide networks with polynomial activations
• A self consistent theory of Gaussian Processes captures feature learning effects in finite CNNs

Cited by in corpus (6)

• A statistical mechanics framework for Bayesian deep neural networks beyond the infinite-width limit
• Asymptotics of representation learning in finite Bayesian neural networks
• Depth induces scale-averaging in overparameterized linear Bayesian neural networks
• On neural network kernels and the storage capacity problem
• Bayesian RG Flow in Neural Network Field Theories
• Summary statistics of learning link changing neural representations to behavior