Beyond NNGP: Large Deviations and Feature Learning in Bayesian Neural Networks
| Authors | Katerina Papagiannouli et al. |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2602.22925 |
| Download | |
| Categories | stat.ML, cs.LG |
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.
Engineering Breakdown
Plain English
This paper studies how Bayesian neural networks behave when they're very wide, moving beyond the Gaussian process (NNGP) approximation that's been standard in the field. The authors use large-deviation theory—a mathematical framework for understanding rare but statistically important fluctuations—to characterize posterior concentration and feature learning in these networks. They show that unlike fixed-kernel NNGP theory, the posterior actually optimizes over both predictors and internal kernel representations jointly, which captures effects like non-Gaussian tails and data-dependent kernel selection that NNGP misses. Experiments on moderately-sized networks demonstrate their framework accurately predicts finite-width behavior, validating the theory against standard approximations.
Core Technical Contribution
The core novelty is replacing the NNGP limit with a large-deviations analysis that provides explicit variational objectives (rate functions) governing posterior concentration in wide Bayesian neural networks. Rather than treating kernels as fixed as in NNGP theory, this framework reveals that the posterior output rate function emerges from joint optimization over both the predictor function and the internal kernel structure—meaning kernels are data-dependent and learned implicitly. This provides a new notion of complexity directly at the functional level and explains feature learning as a consequence of large-deviations asymptotics rather than a limitation of the Gaussian process approximation. The approach yields quantitatively accurate predictions for finite-width networks (not just infinite width), including non-Gaussian posterior tails and kernel deformation that standard theory cannot capture.
How It Works
The method begins with a wide Bayesian neural network and applies large-deviation theory to characterize the rate at which posterior probability concentrates on different predictors as width grows. Instead of the posterior converging to a point (as in NNGP), large-deviations theory describes the exponentially decaying tail probabilities through a rate function that measures the 'cost' of a given predictor under the prior and likelihood. The key insight is that this rate function is obtained by jointly optimizing over: (1) the functional form of the predictor (what the network computes), and (2) the internal kernel representation (how hidden layers interact). Numerically, the framework uses variational optimization to find the rate function, which can then be evaluated to predict posterior behavior and uncertainty quantification. The theory accommodates non-Gaussian posterior deformations and reveals how the network's internal representations adapt to the data, unlike fixed-kernel methods that assume static geometry.
Production Impact
For practitioners, this work provides a more accurate theoretical lens for understanding and tuning Bayesian neural networks at finite widths (say, hundreds to thousands of neurons), rather than relying on asymptotic Gaussian process limits that often break down in practice. If integrated into a production Bayesian inference pipeline, this framework would enable better uncertainty quantification—particularly for detecting when standard posterior approximations are unreliable—and could inform architecture search by revealing which kernel structures emerge as optimal for given datasets. The main production trade-off is computational: the large-deviations rate function requires solving a non-convex variational problem, which is more expensive than forward-passing through an NNGP but potentially cheaper than full MCMC sampling. This makes the approach most valuable for offline model validation, architecture selection, and understanding failure modes in safety-critical systems where accurate uncertainty is essential; it's less suitable for real-time inference where latency is paramount.
Limitations and When Not to Use This
The framework is fundamentally asymptotic—it characterizes behavior as networks grow very wide—so its quantitative predictions degrade for small-to-medium networks (e.g., under 100 hidden units), limiting applicability to tiny edge models. The paper assumes standard Bayesian neural network setups and does not address modern complications like batch normalization, residual connections, or attention mechanisms, which have different limiting behavior. The variational optimization required to compute rate functions is non-convex and may be sensitive to initialization, and the paper does not provide convergence guarantees or algorithmic details for solving these problems at scale. Additionally, while the theory explains feature learning, it remains primarily a descriptive framework—it tells you what to expect from a Bayesian network, but does not directly prescribe how to improve generalization or design better priors.
Research Context
This work extends decades of neural network limit theory (Chizat & Bach, Jacot et al. on NTK/NNGP) by showing that large-deviations theory captures phenomena those Gaussian limits miss, particularly kernel adaptation and non-Gaussian tails. It builds on recent progress in understanding implicit feature learning and the feature/kernel dichotomy in wide networks, positioning large deviations as the right mathematical tool for the finite-width regime where real networks operate. The paper opens a research direction toward non-asymptotic, quantitatively predictive theory for Bayesian neural networks and suggests that variational rate-function optimization could become a new tool for model criticism and architecture search. It connects classical statistical mechanics (large deviations) to modern deep learning, potentially inspiring new methods for uncertainty quantification and robustness analysis in other architectures.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::