📊 statistics

Kernel Renormalization in Bayesian Deep Neural Networks: the Equivalent Wishart Ansatz in the Proportional Regime

This paper introduces an effective approximate approach using an equivalent Wishart Ansatz to predict the generalization performance of Bayesian deep neural networks in the proportional regime, successfully capturing representation learning through renormalized kernels and self-consistent order parameters that align well with empirical sampling experiments.

Original authors: Paolo Baglioni, Christian Keup, Vincenzo Zimbardo, Rosalba Pacelli, Alessandro Vezzani, Raffaella Burioni, Pietro Rotondo

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Paolo Baglioni, Christian Keup, Vincenzo Zimbardo, Rosalba Pacelli, Alessandro Vezzani, Raffaella Burioni, Pietro Rotondo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand how a giant, complex machine (a Deep Neural Network) learns to recognize patterns, like distinguishing between pictures of cats and dogs. Usually, scientists try to understand these machines by pretending they are infinitely large. In this "infinite" world, the machine behaves very predictably, like a simple, smooth curve. This is called the "lazy" regime.

However, real-world machines aren't infinite. They have a specific, finite size. When the number of data points you feed the machine is roughly the same size as the number of neurons inside it (a situation the authors call the "proportional regime"), things get messy. The machine starts to learn in complex, non-linear ways that the "infinite" theories can't explain.

This paper introduces a new way to predict how these finite-sized machines behave without having to run millions of expensive computer simulations. Here is the breakdown of their discovery using simple analogies:

1. The Problem: The "Black Box" of Finite Size

Think of a deep neural network as a multi-story building where each floor processes information.

The Infinite View: If the building were infinitely wide, the information flowing through it would be like water in a perfectly smooth pipe. You could easily predict the output.
The Real View: In a real, finite building, the pipes are narrower. The water (data) creates turbulence, splashes, and eddies. These "finite-width effects" are actually what make deep learning powerful, but they are incredibly hard to calculate mathematically because the interactions between the layers are chaotic.

2. The Solution: The "Equivalent Wishart Ansatz" (EWA)

The authors propose a clever shortcut. Instead of trying to track every single drop of water (every neuron's exact state), they suggest looking at the statistical shape of the turbulence.

The Analogy: Imagine you are trying to describe the weather in a stormy city. Instead of tracking every single raindrop, you realize that the overall pattern of the rain follows a specific, known statistical shape (like a bell curve, but for matrices).
The "Wishart" Magic: The authors discovered that even though the network is non-linear and complex, the "turbulence" (the fluctuations in how the network processes data) behaves mathematically as if it were following a specific, well-understood distribution called the Wishart distribution.
The "Ansatz": This is just a fancy word for a "smart guess." They guessed: "Let's pretend the chaos in every layer of the network follows this specific Wishart pattern."

3. The Result: A Simple Recipe for Complex Behavior

By making this guess, they were able to shrink a massive, impossible-to-solve problem down to a tiny, manageable one.

Before: To understand the network, you needed to solve equations involving millions of variables (one for every connection).
After: The EWA allows you to describe the entire network's behavior using just a few numbers (called "order parameters").
- Think of it like this: Instead of needing a map of every street in a city to predict traffic, you only need to know the average speed on the main highway and the number of cars.
- For a network with $L$ layers, they found that you only need $L$ simple numbers to predict how well the network will learn. These numbers tell you how much the "turbulence" amplifies or dampens the signal as it goes through the building.

4. Testing the Theory

The authors didn't just do math; they tested it against reality.

They built actual neural networks (with about 10 layers and a few hundred neurons) and trained them on real datasets (like MNIST digits and CIFAR-10 images).
They used powerful computer sampling methods (like a high-tech version of rolling dice millions of times) to see what the networks actually did.
The Verdict: Their "smart guess" (EWA) matched the real-world results incredibly well, even for networks with up to 10 layers. It was far more accurate than the old "infinite" theories, which failed to capture the nuances of finite-sized networks.

5. A Surprising Discovery: The "Metastable" Trap

While testing, they found something weird. When the networks got very deep and the data load was high, the computer simulations sometimes got "stuck" in a temporary state.

The Analogy: Imagine a ball rolling down a hill. Usually, it rolls straight to the bottom. But sometimes, it gets stuck in a small dip halfway down. It looks like it has settled, but if you wait long enough (or shake the hill), it will eventually roll out of the dip and reach the true bottom.
The authors found that standard computer simulations often got stuck in these "dips" (metastable states), making it look like the network had stopped learning, when in reality, it just needed more time to find the true solution.

Summary

The paper provides a new "rule of thumb" for understanding deep neural networks that aren't infinitely large. By realizing that the chaos inside these networks follows a predictable statistical pattern (the Wishart distribution), they created a simple mathematical tool that accurately predicts how these networks learn, bridging the gap between simple theory and complex reality.

Title: Kernel Renormalization in Bayesian Deep Neural Networks: the Equivalent Wishart Ansatz in the Proportional Regime

Problem Statement
Deep neural networks (DNNs) are central to modern machine learning, yet their theoretical understanding remains challenging, particularly for multi-layered, nonlinear architectures. While the "lazy" infinite-width limit (Neural Network Gaussian Process or NNGP) provides exact descriptions for wide networks, it fails to capture the finite-width effects responsible for the superior performance of DNNs over kernel methods. Conversely, the "rich" learning regime, where features adapt significantly, often requires high-dimensional order parameters (e.g., $P \times P$ kernel matrices) that are computationally intractable.

The specific gap addressed by this work is the lack of a low-dimensional, effective theory for Bayesian DNNs in the proportional regime, where the number of training samples $P$ and the network width $N$ scale proportionally to infinity ( $P, N \to \infty$ with $\alpha = P/N$ fixed). In this regime, finite-width fluctuations are non-negligible, and existing theories either rely on perturbative corrections (which fail for deep networks) or require solving coupled high-dimensional equations that offer no dimensionality reduction.

Methodology: The Equivalent Wishart Ansatz (EWA)
The authors propose an approximate but effective theory based on the Equivalent Wishart Ansatz (EWA). The core idea is to model the stochastic fluctuations of the hierarchical empirical kernels in a deep nonlinear network as if they were Wishart distributed, similar to the exact behavior found in deep linear networks.

Kernel Hierarchy: The network output prior is described by a hierarchy of empirical kernels $K^{(\ell)}_E$ at each layer $\ell$ . In the proportional limit, the authors assume that the conditional distribution of a kernel at layer $\ell$ , given the kernel at layer $\ell-1$ , behaves as a Wishart distribution:
$K^{(\ell)}_E | K^{(\ell-1)}_E \sim \mathcal{W}_P\left( \frac{\Theta^{L-\ell+1}(K^{(\ell-1)}_E)}{N_\ell}, N_\ell \right)$
where $\Theta$ is the NNGP kernel function and $N_\ell$ is the width.
Dimensionality Reduction: By exploiting the properties of Wishart matrices (specifically the contraction property where quadratic forms follow a $\chi^2$ distribution), the authors reduce the complex integral over the parameter space to a low-dimensional integral over $L$ scalar order parameters $q_\ell$ (or matrices $q_\ell$ for CNNs).
- For Fully Connected (FC) networks, the effective action depends on a scalar product $Q = \prod q_\ell$ .
- For Convolutional Neural Networks (CNNs), a Stacked Equivalent Wishart Ansatz is introduced, leading to a low-dimensional matrix order parameter $Q$ that captures patch-patch correlations.
Effective Action: The partition function is expressed via a large deviation principle, yielding an effective action $S(q)$ that depends only on these order parameters and the dataset. The renormalized kernel is given by $K^{(R)}_Q = Q \Theta^L(C)$ , where $C$ is the input Gram matrix.
Handling Non-Zero Means: The authors demonstrate that for the proportional regime, the distinction between central and non-central Wishart distributions (arising from activation functions with non-zero means like ReLU) is asymptotically negligible. The zero-mean EWA remains a valid effective description.

Key Contributions

Low-Dimensional Theory for Nonlinear DNNs: The paper provides the first low-dimensional approximate description of Bayesian multi-layer perceptrons (MLPs) and CNNs in the proportional regime, capturing non-perturbative finite-width effects without resorting to $P \times P$ self-consistency equations.
Equivalent Wishart Ansatz: It introduces the EWA as a mechanism to capture the dominant stochastic fluctuations of hierarchical empirical kernels, reducing the problem to finding $L$ (or a small matrix) order parameters.
Extension to CNNs: The framework is extended to convolutional architectures, identifying a hierarchical local kernel renormalization mechanism that accounts for finite-width effects in spatially correlated data.
Validation in the Proportional Regime: The theory is tested against extensive numerical experiments (MCMC sampling) on real-world datasets (MNIST, CIFAR-10) and synthetic data, covering depths up to $L \sim 10$ and various activation functions (Erf, ReLU).

Results

Quantitative Agreement: The EWA predictions for training loss and generalization error show "overall very good agreement" with numerical sampling experiments for networks with depths up to $L \sim O(10)$ and various dataset sizes. The theory successfully tracks the behavior of finite-width networks where the large-width NNGP limit fails.
Systematic Deviations: Two types of deviations are observed when both depth $L$ $L$ and load $\alpha$ $α$ are large:
1. Gradual Drift: In Gaussian data experiments at high $\alpha$ , the EWA error grows smoothly with depth, likely due to accumulating errors in the rate function width or neglected directional inhomogeneities.
2. Metastability: For ReLU networks on CIFAR-10 at $L > 5$ and $\alpha \approx 2.5$ , a sharp, abrupt drop in training and test error is observed in MCMC simulations. This phenomenon, interpreted as a metastable transition in the posterior, is not captured by the smooth EWA prediction, suggesting a breakdown of the ansatz in extreme regimes.
Mean-Field vs. Standard Parametrization: The theory successfully predicts generalization performance in both standard (SP) and mean-field ( $\mu$ P) parametrizations. In the rich learning regime ( $\mu$ P), the improvement in generalization is found to be driven primarily by a suppression of predictor variance rather than a strong adaptation of the predictor bias, consistent with the EWA's renormalized kernel picture.

Significance and Claims
The authors claim that the EWA provides a tractable effective theory that bridges the gap between the lazy infinite-width limit and the complex, high-dimensional feature learning regime. By reducing the description of finite-width learning to a few scalar (or low-rank matrix) order parameters, the paper argues that the relative scaling of data size and width ( $\alpha = P/N$ ) is a more fundamental organizing principle for finite-width learning than the traditional "lazy vs. rich" dichotomy.

The work suggests that for moderate depths and loads, the complex adaptation of features in nonlinear DNNs can be effectively captured by a simple global (or local, for CNNs) renormalization of the kernel. This offers a baseline for understanding feature learning, indicating that the large gap between infinite-width predictions and finite-width phenomenology is largely captured by this directionally homogeneous renormalization. The paper concludes that while the EWA is an approximation, its success in predicting empirical learning curves across diverse architectures and datasets establishes it as a powerful tool for analyzing Bayesian deep learning in the proportional regime.