Bayesian RG Flow in Neural Network Field Theories

TL;DR

BRG-NNFT provides a unified information-theoretic renormalization framework that links neural network function-space descriptions to dual statistical field theories. Training induces an IR-to-UV flow in the space of SFTs, while information-shell BRG conducts a UV-to-IR coarse-graining, with ERG recovered in momentum-shell cases; analytic results for infinite-width NNGPs and a cos-net special case connect BRG to mass and momentum renormalization in free scalar SFTs. The authors validate the approach analytically and numerically, demonstrating a critical distinguishability scale below which sloppy parameters can be pruned without degrading performance. This framework offers a principled bridge between Bayesian inference and Wilsonian RG for nonlocal NNFTs and suggests avenues for sampling exotic SFTs and informing diffusion-model strategies.

Abstract

The Neural Network Field Theory correspondence (NNFT) is a mapping from neural network (NN) architectures into the space of statistical field theories (SFTs). The Bayesian renormalization group (BRG) is an information-theoretic coarse graining scheme that generalizes the principles of the exact renormalization group (ERG) to arbitrarily parameterized probability distributions, including those of NNs. In BRG, coarse graining is performed in parameter space with respect to an information-theoretic distinguishability scale set by the Fisher information metric. In this paper, we unify NNFT and BRG to form a powerful new framework for exploring the space of NNs and SFTs, which we coin BRG-NNFT. With BRG-NNFT, NN training dynamics can be interpreted as inducing a flow in the space of SFTs from the information-theoretic `IR' $\rightarrow$ `UV'. Conversely, applying an information-shell coarse graining to the trained network's parameters induces a flow in the space of SFTs from the information-theoretic `UV' $\rightarrow$ `IR'. When the information-theoretic cutoff scale coincides with a standard momentum scale, BRG is equivalent to ERG. We demonstrate the BRG-NNFT correspondence on two analytically tractable examples. First, we construct BRG flows for trained, infinite-width NNs, of arbitrary depth, with generic activation functions. As a special case, we then restrict to architectures with a single infinitely-wide layer, scalar outputs, and generalized cos-net activations. In this case, we show that BRG coarse-graining corresponds exactly to the momentum-shell ERG flow of a free scalar SFT. Our analytic results are corroborated by a numerical experiment in which an ensemble of asymptotically wide NNs are trained and subsequently renormalized using an information-shell BRG scheme.

Figures (11)

• Figure 1: Here, $\pi(\theta)$ is the parameter distribution of a NN, $\phi_{\theta}$, which has a dual SFT with action, $S[\phi]$. The BRG scheme coarse grains the original NN parameter distribution to $\pi_{\Lambda}$, indexed by the Fisher information scale, $\Lambda$. Updated NN ensembles have a transformed dual action, $S_{\Lambda}[\phi]$, connected to original action through an information-theoretic BRG flow.
• Figure 2: A commutative diagram which unifies NNFT and BRG. The composition of these maps defines a new proposal for constructing information geometric flows through the space of field theories. We interpret these BRG flows as a suitable replacement for standard field theoretic ERG flows in general contexts where spatially local coarse graining is not sufficient for implementing a meaningful flow.
• Figure 3: NNFT maps a fixed NN architecture, $\phi_{\theta}$, with parameter distribution, $\pi(\theta)$, into a Euclidean field theory with action, $S[\phi]$.
• Figure 4: BRG is an information geometric coarse graining scheme that realizes a family of parameter distributions $\pi_{\Lambda}$ indexed by a distinguishability scale $\Lambda$.
• Figure 5: Applying NNFT to a family of Bayesian renormalized parameter distributions realizes a one parameter family of Euclidean field theories which we interpret as an information theoretically coarse grained flow in the space of NNFTs.