Beyond NNGP: Large Deviations and Feature Learning in Bayesian Neural Networks

TL;DR

This paper explores the behavior of wide Bayesian neural networks beyond Gaussian-process limits, using large deviation theory to understand fluctuations, feature learning, and posterior concentration at the functional level.

Contribution

It introduces a novel variational framework based on rate functions to analyze posterior behavior, incorporating internal kernel optimization unlike traditional fixed-kernel approaches.

Findings

Accurately describes finite-width network behavior

Captures non-Gaussian tail effects

Reveals data-dependent kernel selection

Abstract

We study wide Bayesian neural networks focusing on the rare but statistically dominant fluctuations that govern posterior concentration, beyond Gaussian-process limits. Large-deviation theory provides explicit variational objectives-rate functions-on predictors, providing an emerging notion of complexity and feature learning directly at the functional level. We show that the posterior output rate function is obtained by a joint optimization over predictors and internal kernels, in contrast with fixed-kernel (NNGP) theory. Numerical experiments demonstrate that the resulting predictions accurately describe finite-width behavior for moderately sized networks, capturing non-Gaussian tails, posterior deformation, and data-dependent kernel selection effects.