Technical Summary: Bayesian Inference with Shaped Deep Non-linear MLPs

Problem Statement

Deep learning theory aims to characterize how neural networks make predictions when model parameters (width $N$, depth $L$) and dataset size ($P$) are simultaneously large. A fundamental challenge arises because the limits of $N, L, P \to \infty$ do not necessarily commute. Existing frameworks often rely on the infinite-width limit at fixed depth (yielding Neural Tangent Kernels or Neural Network Gaussian Processes) or specific linear network regimes. However, these approaches often fail to capture the role of depth in non-linear networks or the interplay between feature learning and kernel methods in proportional scaling regimes where $L \propto N$ and $P \propto N$.

The specific problem addressed is the rigorous characterization of Bayesian inference in deep, non-linear Multi-Layer Perceptrons (MLPs) under a joint scaling limit where the number of training samples ($P$), input dimension ($N_0$), hidden layer width ($N$), and number of hidden layers ($L$) all diverge. The authors seek to determine:

1. How the predictive posterior behaves in this regime.
2. Under what conditions depth provides a benefit (increases Bayesian model evidence).
3. Whether and how non-linear networks can be approximated by data-dependent kernel methods.

Methodology

1. Shaped Activation Functions

To ensure non-degenerate behavior as $L, N \to \infty$, the authors employ shaped activation functions $\phi_s$. For a base non-linearity $\phi_{base}$, the activation is scaled as $\phi_s(x) = s \phi_{base}(x/s)$. This mechanism ensures that the cumulative effect of non-linearities across many layers remains controllable and comparable to linear effects. The shaping factor $s$ is chosen such that the ratio $t = L/N$ uniquely controls the initialization distribution.

2. Neural Covariance SDE (NSDE)

The core technical tool is the Neural Covariance Stochastic Differential Equation (NSDE). The authors analyze the evolution of the conjugate kernel (Gram matrix of hidden features) $\Phi^{(\ell)}$ across layers.

• In the limit $N, L \to \infty$ with $L/N \to t$, the discrete evolution of the kernel converges to a diffusion process $\Phi_t$.
• The process is governed by:
  $$d\Phi_t = \alpha_{act} \cdot b(\Phi_t) dt + \Phi_t^{1/2} dB_t \Phi_t^{1/2}$$
  where $b(\Phi_t)$ is a drift term dependent on the base non-linearity, and $B_t$ is a Brownian motion matrix.
• To handle large $P$, a time-change $\tau = P t$ is introduced, defining a "depth-time" $\tau^* = LP/N$. The dynamics are re-scaled to balance the diffusion and drift terms.

3. Bayesian Partition Function via SDE

The authors reformulate the Bayesian partition function $Z_\beta(x_0, \kappa)$ (which encodes the predictive posterior) as an expectation over the prior distribution of the conjugate kernel.

• Using a Gaussian integral identity, the partition function is expressed as:
  $$Z_\beta(x_0, \kappa) \propto \int_{\mathbb{R}^P} \exp\left(-\frac{1}{2\beta}\|p\|^2 + ip^\top Y\right) \mathbb{E}_{prior}\left[\exp\left(-\frac{1}{2}v^\top \Phi v\right)\right] dp$$
  where $v = [p^\top, \kappa]^\top$ and $\Phi$ is the conjugate kernel at the penultimate layer.
• By Theorem 1, in the joint limit, the expectation over the MLP prior converges to the expectation over the solution of the NSDE. This reduces the problem of Bayesian inference in deep MLPs to solving the SDE for $\Phi_\tau$ and computing the resulting expectation.

4. Perturbative Analysis

Since solving the SDE for arbitrary $\tau^*$ is intractable, the authors perform a perturbative expansion for small effective depth $\tau = LP/N$. They compute the first-order derivative of the log-partition function with respect to $\tau$ at $\tau=0$. This allows them to separate the contributions of the linear network component (diffusion-only) and the non-linear kernel component (drift-only).

Key Contributions and Results

1. Rigorous Unifying Framework

The paper provides the first mathematically rigorous framework for analyzing Bayesian inference in non-linear deep MLPs under the proportional limit $LP/N \in \Theta(1)$. It bridges the gap between infinite-width kernel methods and deep linear networks, covering both smooth and ReLU activation functions at arbitrary temperatures ($\beta$).

2. Recovery and Generalization of Prior Results

• Deep Linear Networks: The framework recovers exact results for deep linear networks (where the partition function involves Meijer-G integrals) as a special case where the drift term vanishes.
• Infinite-Width Limit: It recovers the result that in the limit $N \gg LP$, Bayesian inference reduces to a data-agnostic kernel method with a deterministic feature map derived from the Neural ODE.
• Weakly Non-linear Networks: The results generalize perturbative findings from physics literature (e.g., [HZ24]) regarding weakly non-linear networks, but with full mathematical rigor and without relying on unproven simplifications.

3. Perturbative Expansion of Bayesian Evidence

The authors derive a simple criterion for when increasing depth (increasing $\tau = LP/N$) improves the Bayesian model evidence (marginal likelihood).

• Takeaway 2.1: The first-order change in the log-evidence is the sum of a linear network term and a non-linear kernel term.
• Takeaway 3.1 & 3.2: For specific data distributions (normalized, approximately orthogonal inputs), the authors derive explicit inequalities determining when depth is beneficial.
  • For smooth activations, depth is beneficial if the data statistics and activation derivatives satisfy a specific inequality involving the signal-to-noise ratio and the curvature of the activation function.
  • For ReLU activations, depth is beneficial if and only if the variance of the labels exceeds a threshold related to the asymmetry of the ReLU ($c_+ - c_-$).
• Phase Transition: The results suggest a phase transition where depth is only beneficial if the non-linearity is not "too strong" relative to the data structure.

4. Equivalence to Data-Dependent Kernel Methods

• Takeaway 2.2: At first order in $\tau$, the Bayesian predictive posterior for test data is equivalent to a kernel method, but with a data-dependent feature map.
• Unlike the infinite-width limit where the feature map is fixed, the effective feature map here includes a correction term proportional to the training data labels and the orthogonal component of the test input. This provides a theoretical justification for "feature learning" in the Bayesian sense: the network learns a feature map that is slightly adapted to the training set, distinct from the prior feature map.

Significance and Claims

The authors claim that their work offers a mathematically rigorous alternative to physics-inspired derivations that often rely on heuristic assumptions. By utilizing the NSDE and perturbative expansions, they:

1. Clarify the role of depth: They demonstrate that depth is not universally beneficial; its utility depends on the interplay between the data generating process, the activation function, and the temperature.
2. Unify Kernel and Feature Learning: They show that even in the deep, non-linear regime, the predictive posterior can be viewed as a kernel method, but with a feature map that is dynamically adjusted by the training data (to first order).
3. Provide Analytical Tools: The framework allows for the computation of Bayesian statistics (evidence, posterior mean, variance) in regimes previously considered intractable, specifically for non-linear networks with finite but large width and depth.

The paper explicitly states that it does not claim to solve the general case for arbitrary $\tau$ or to provide experimental validation on real-world datasets. Instead, it establishes a theoretical foundation for understanding the mechanisms of Bayesian inference in deep networks through the lens of stochastic differential equations and proportional scaling limits.