Bayesian Inference with Deep Weakly Nonlinear Networks
Boris Hanin, Alexander Zlokapa
TL;DR
The paper analyzes Bayesian inference with deep fully connected networks that use the shaped activation $\phi(t)=t+\frac{\psi}{3L}t^3$, in a regime where the training set size $P$, input dimension $N_0$, layer widths $N$, and depth $L$ are all large with $P<N_0$. It shows that, at leading order in width, posterior inference reduces to a kernel method with a $\psi$-dependent feature map $x_\psi$, and that first-order $1/N$ corrections introduce data-dependent, cubic feature-learning corrections controlled by the emergent ratio $LP/N$ (an effective posterior depth). The analysis is built on a novel combinatorial model for prior moments expressed via random graphs and a self-loop process, enabling perturbative expansions of the partition function $Z_\beta(x;\tau)$ and the predictive posterior. In the zero-temperature limit, the authors show conditions under which depth enhances model evidence and generalization for certain data spectra (e.g., power laws) and discuss benign overfitting in deep linear vs nonlinear networks, with depth increasingly beneficial when $\alpha<2$ and data are well-aligned. Overall, the work connects kernel methods with data-dependent learning in deep nonlinear networks and provides a controlled framework to quantify how depth and nonlinearity influence Bayesian inference in high-dimensional, overparameterized regimes.
Abstract
We show at a physics level of rigor that Bayesian inference with a fully connected neural network and a shaped nonlinearity of the form $φ(t) = t + ψt^3/L$ is (perturbatively) solvable in the regime where the number of training datapoints $P$ , the input dimension $N_0$, the network layer widths $N$, and the network depth $L$ are simultaneously large. Our results hold with weak assumptions on the data; the main constraint is that $P < N_0$. We provide techniques to compute the model evidence and posterior to arbitrary order in $1/N$ and at arbitrary temperature. We report the following results from the first-order computation: 1. When the width $N$ is much larger than the depth $L$ and training set size $P$, neural network Bayesian inference coincides with Bayesian inference using a kernel. The value of $ψ$ determines the curvature of a sphere, hyperbola, or plane into which the training data is implicitly embedded under the feature map. 2. When $LP/N$ is a small constant, neural network Bayesian inference departs from the kernel regime. At zero temperature, neural network Bayesian inference is equivalent to Bayesian inference using a data-dependent kernel, and $LP/N$ serves as an effective depth that controls the extent of feature learning. 3. In the restricted case of deep linear networks ($ψ=0$) and noisy data, we show a simple data model for which evidence and generalization error are optimal at zero temperature. As $LP/N$ increases, both evidence and generalization further improve, demonstrating the benefit of depth in benign overfitting.