Non-Gaussian processes and neural networks at finite widths

I Inception

Gaussian processes model many phenomena in the physical world. A prime example is Brownian motion Brown (1828), modeled as the integral of Gaussian-distributed bumps exerted on a point-like solute Einstein (1905). The theory of elementary particles Weinberg (1995) also becomes a Gaussian process in the free limit where interactions between particles are turned off, and many-body systems as complex as glasses come to be Gaussian in the infinite-dimensional, mean-field, limit Parisi and Zamponi (2010). In the context of machine learning, Neal in Ref. Neal (1996) pointed out that a class of neural networks give rise to Gaussian processes in the infinite-width limit, which can perform exact Bayesian inference from training to test data Williams (1997). They occupy a corner of theoretical playground wherein the karakuri of neural networks is scrutinized Lee et al. (2018); Matthews et al. (2018); Jacot et al. (2018); Chizat et al. (2018); Lee et al. (2019); Geiger et al. (2019).

In reality, Gaussian processes are but mere idealizations. Brownian particles have finite-size structure, elementary particles interact, and many-body systems respond nonlinearly. In order to understand rich phenomena exhibited by these real systems, Gaussian processes rather serve as starting points to be perturbed around. Indeed many edifices in theoretical physics are built upon the successful treatment of non-Gaussianity, with a notable example being renormalization-group flow Kadanoff (1966); Wilson (1971); Weinberg (1996); Goldenfeld (2018). In the quest to elucidate behaviors of real neural networks away from the infinite-width limit, it is thus natural to wonder if the similar treatment of non-Gaussianity yields equally elegant and powerful machinery.

Here we set out on this program, perturbatively treating finite-width corrections to neural networks. Prior distributions of outputs are obtained through progressively integrating out preactivations of neurons layer by layer, yielding non-Gaussian priors. The whole procedure closely resembles renormalization-group flow Goldenfeld (2018); Mehta and Schwab (2014): it bridges probability distributions at different scales through coarse-graining of random variables at microscopic scales; the flow of distributions is traced through running couplings, which in particular capture the degree of non-Gaussianity in these distributions; resulting recursive equations (R1,R2,R3) govern the evolution of these running couplings from lower to higher layers, just as renormalization-group equations do from microscopic to macroscopic scales. Such a recursive approach enables us to treat finite-width corrections to various observables, for networks with arbitrary activation functions.

The rest of the paper is structured as follows. In Section II we review and set up basic concepts. Our master recursive formulae (R1,R2,R3) are derived in Section III, which control the flow of preactivation distributions. After an interlude with concrete examples in Section IV, we extend the Gaussian-process Bayesian inference to non-Gaussian priors in Section V and study inference of neural networks at finite widths. We conclude in Section VI with dreams.

II To infinity and beyond

In this paper we study real finite-width neural networks in the regime where the number of neurons in hidden layers is asymptotically large whereas input and output dimensions are kept constant.

Higher moments are then obtained by Wick’s contractions Wick (1950); Zee (2010). For instance,

For those unfamiliar with Wick’s contractions and connected correlation functions (a.k.a. cumulants), a pedagogical review is provided in Appendix A as our formalism heavily relies on them.

In the infinite-width limit where $n_{1},n_{2},\ldots,n_{L-1}\rightarrow\infty$ (but finite $n_{0}$ and $n_{L}$), it has been argued – with varying degrees of rigor Neal (1996); Lee et al. (2018); Matthews et al. (2018) – that the prior distribution of outputs is governed by the Gaussian process with a kernel

and all the higher moments given by Wick’s contractions. Here, the sample index $\alpha$ labels different inputs in a dataset. There exists a recursive formula that lets us evaluate this kernel for any pair of inputs Lee et al. (2018) [c.f. Equation (R1)]. Importantly, once the values of the kernel are evaluated for all the pairs of $N_{\text{D}}=N_{\text{R}}+N_{\text{E}}$ input data, $\left\{\bm{x}_{\alpha}\right\}_{\alpha=1,\ldots,N_{\text{D}}}$, consisting of $N_{\text{R}}$ training inputs with target outputs and $N_{\text{E}}$ test inputs with unknown targets, we can perform exact Bayesian inference to yield mean outputs as predictions for $N_{\text{E}}$ test data Williams (1997); Williams and Rasmussen (2006) [c.f. Equation (GPM)]. This should be contrasted with stochastic gradient descent (SGD) optimization Robbins and Monro (1951), through which typically a single estimate for the optimal model parameters of the posterior, $\bm{\theta}_{\star}$, is obtained and used to predict outputs for test inputs; Bayesian inference instead marginalizes over all model parameters, performing an ensemble average over the posterior distribution MacKay (1995).

II.2 Beyond infinity

collectively dictate the distribution of preactivations. Our aim is to trace the flow of these distributions progressively and cumulatively all the way up to the last layer whereat Bayesian inference is executed. More specifically, we shall inductively and self-consistently show that two-point preactivation correlation functions take the formIn the main text we place tildes on objects that depend only on sample indices $\alpha$’s in order to distinguish them from those that depend both on sample indices $\alpha$’s and neuron indices $i$’s.

and connected four-point preactivation correlation functions

II.3 Related work

Our Schwinger operator approach is orthogonal to the replica approach by Cohen et al. (2019) and, unlike the planar diagrammatic approach by Dyer and Gur-Ari (2019), applies to general activation functions, made possible by accumulating corrections layer by layer rather than dealing with them all at once. See also Antognini (2019). More substantially, in contrast to these previous approaches, we here study finite-width effects on Bayesian inference and find that the renormalization-group picture naturally emerges, with layers playing the role of scales.

III Distributional flow

As auxiliary objects in recursive steps, let us introduce activation correlation functions

Our basic strategy is to establish relations

With the mastery of Wick’s contractions and connected correlation functions, it is simple to derive the following combinatorial hack (Appendix A.4): viewing prior preactivations

for any function $\mathcal{F}$. Here the operators $\mathcal{O}_{S}[\mathbf{z}]$ and $\mathcal{O}_{V}[\mathbf{z}]$ capture $1/n$ corrections due to self-energy and four-point vertex, respectively, and are defined as

where the sample indices are raised by using the inverse core kernel as a metric, meaning

Using the above hack, we can evaluate the activation correlations by straightforward algebra with Wick’s contractions. In particular, as the Gaussian integral is diagonal in the neuron index $i$, we just need to disentangle cases with repeated and unrepeated neuron indices. The solution for this exercise is in Appendix B: it is arguably the most cumbersome algebra in this paper.

III.2 Master recursive flow equations

with the potential $\mathcal{H}[\mathbf{z}]=\mathcal{H}_{0}[\mathbf{z}]+\epsilon\mathcal{H}_{1}[\mathbf{z}]+O(\epsilon^{2})$ where $\epsilon\equiv\frac{1}{n_{L-1}}\ll 1$,

Again, this can be derived through Wick’s contractions. It is important to note that $n_{L}$ is constant and thus $\epsilon\mathcal{H}_{1}[\mathbf{z}]$ can consistently be treated perturbatively.If $n_{L}$ were of order $n\gg 1$, the potential $\mathcal{H}$ would become a large-$n$ vector model, for which we would have to sum the infinite series of bubble diagrams Moshe and Zinn-Justin (2003).

IV Interlude: examples

The recursive relations obtained above can be evaluated numerically Lee et al. (2018) [or sometimes analytically for rectified linear unit (ReLU) activation Cho and Saul (2009)], which is a perfectly adequate approach: at the leading order it involves four-dimensional Gaussian integrals at most. Here, continuing the theme of wearing out Wick’s contractions, we develop an alternative analytic method that works for any polynomial activations Liao and Poggio (2017), providing another perfectly cromulent approach.

For a general polynomial activation of degree $p$, $\sigma(z)=\sum_{k=0}^{p}a_{k}z^{k}$, the nontrivial term in Equation (R1) can be expanded as

Each term can then be evaluated by Wick’s contractions and the same goes for all the terms in Equations (R2) and (R3).The same approach could be adopted for an analytic function but it would in general be difficult to sum the resulting infinite series in a closed form. It could nonetheless be useful in, for example, proving convergence properties. Below and in Appendix C, we illustrate this procedure with simple examples.

and the linearly layer-dependent four-point vertex

It succinctly reproduces the result that can be obtained through planar diagrams in this special setup Dyer and Gur-Ari (2019). Quadratic activation Li et al. (2018) is worked out in Appendix C.1.

IV.2 ReLU with single input

IV.3 Experimental verification: output distributions for a single input

V Bayesian inference

Let us take off from the terminal point of Section III: we have obtained the recursive equations (R0-R3) for the Gaussian-process kernel and the leading finite-width corrections and codified them in the weakly non-Gaussian prior distributions $p[\mathbf{z}]$ (D0-D2) of outputs

dictated by the potential $\mathcal{H}[\mathbf{z}]=\mathcal{H}_{0}[\mathbf{z}]+\epsilon\mathcal{H}_{1}[\mathbf{z}]+O(\epsilon^{2})$ with $\epsilon\equiv\frac{1}{n_{L-1}}\ll 1$. Examples in Section IV illustrate that finite-width corrections stay perturbative typically when $\frac{\text{depth}}{\text{width}}\ll 1$. Let us now divide $N_{\text{D}}$ inputs into $N_{\text{R}}$ training and $N_{\text{E}}$ test inputs as

and the training inputs come with target outputs

We shall develop a procedure to infer outputs for test inputs a lá Bayes, perturbatively extending the textbook Williams and Rasmussen (2006). For field theorists, our calculation is just a background-field calculation Weinberg (1996) in disguise.

Taking the liberty of notations, we let the number of input-data arguments dictate the summation over sample indices $\alpha$ inside the potential $\mathcal{H}$, and denote the joint probabilities

Given the training targets $\bm{y}_{\text{R}}$, the posterior distribution of test outputs are given by Bayes’ rule:

The leading Gaussian-process contributions can be segregated out through the textbook manipulation Williams and Rasmussen (2006) [c.f. Appendix D]: denoting the full Gaussian-process kernel in the last layer as

and the Gaussian-process posterior mean prediction as

and defining a fluctuation $\left(\textnormal{z}_{\text{E}}\right)_{i;\dot{\gamma}}\equiv\left(\textnormal{y}^{\text{GP}}_{\text{E}}\right)_{i;\dot{\gamma}}+\left(\delta\textnormal{z}_{\text{E}}\right)_{i;\dot{\gamma}}$ and a matrix $\widetilde{K}_{\Delta}\equiv\widetilde{K}_{\text{EE}}-\widetilde{K}_{\text{ER}}\widetilde{K}_{\text{RR}}^{-1}\widetilde{K}_{\text{RE}}$,

For any function $\mathcal{F}$, its expectation over the Bayesian posterior (Bayes) then turns into

where the deviation kernel $\left\langle\left(\delta\textnormal{z}_{\text{E}}\right)_{i_{1};\dot{\gamma}_{1}}\left(\delta\textnormal{z}_{\text{E}}\right)_{i_{2};\dot{\gamma}_{2}}\right\rangle_{K_{\Delta}}\equiv\delta_{i_{1}i_{2}}\left(\widetilde{K}_{\Delta}\right)_{\dot{\gamma}_{1}\dot{\gamma}_{2}}$ and the normalization factor

In particular the mean posterior output is given by

Stringing together $\overline{\phi}_{i;\alpha}\equiv[\left(y_{\text{R}}\right)_{i;\bar{\beta}},\left(\textnormal{y}^{\text{GP}}_{\text{E}}\right)_{i;\dot{\gamma}}]$, recalling Equation (D2) for $\mathcal{H}_{1}$, and using Wick’s contractions for one last time, the mean prediction becomes

VI Dreams

In this paper, we have developed the perturbative formalism that captures the flow of preactivation distributions from lower to higher layers. The resemblance between our recursive equations and renormalization-group flow equations in high-energy and statistical physics is highly appealing. It would be exciting to investigate the structure of fixed points away from the Gaussian asymptopia Schoenholz et al. (2016) and fully realize the dream articulated in Ref. Mehta and Schwab (2014) – the audacious hypothesis that neural networks wash away microscopic irrelevancies and extract relevant features – beyond their limited example of a mapping between two antiquated techniques.

In addition we have developed the perturbative Bayesian inference scheme universally applicable whenever prior distributions are weakly non-Gaussian, and have applied it to the specific cases of neural networks at finite widths. In light of possible finite-width regularization effects, it would be prudent to revisit the empirical comparison between SGD optimization and Bayesian inference at finite widths Lee et al. (2018); Novak et al. (2019), especially for convolutional neural networks.

Finally, given surging interests in SGD dynamics within the large-width regime Jacot et al. (2018); Chizat et al. (2018); Lee et al. (2019); Cohen et al. (2019); Dyer and Gur-Ari (2019), it would be natural to adapt our formalism for investigating corrections to neural tangent kernels, and even aspire to capture a transition from lazy-learning to feature-learning regimes.

Acknowledgments

The author thanks Yasaman Bahri for the discussion that seeded the idea for this project, Boris L. Hanin for persistently preaching about Gaussian processes, and David J. Schwab for permission to call his example limited with our friendship intact. The author also thanks Ethan S. Dyer, Mario Geiger, Guy Gur-Ari, Eric T. Mintun, Stephen H. Shenker, and Lexing Ying for substantially useful discussions, and Daniel A. Roberts for the quality control of all the jokes and more.

References

Appendix A Wick’s tricks

Here is all you need to know in order to follow the calculations in the paper. In the main text, Wick’s contractions are used both for trivially integrating out biases and weights as straightforward applications of Appendix A.1 and for nontrivially integrating out preactivations, with concepts of cumulants reviewed in Appendix A.2 and A.3, culminating in the hack derived in Appendix A.4. The random variables are generically indexed by $\mu=1,\ldots,N$ throughout this Appendix: when applying formulae for biases, $\mu=i$; for weights $\mu=(i,j)$; for full preactivations $\mu=(i,\alpha)$; for single-neuron preactivations $\mu=\alpha$.

For Gaussian-distributed variables $\mathbf{z}=\left\{\textnormal{z}_{\mu}\right\}_{\mu=1,\ldots,N}$ with a kernel $K_{\mu\mu^{\prime}}$, moments

For any odd $m$ such moments identically vanish. For even $m$, Isserlis-Wick’s theorem states that

where the sum is over all the possible pairings of $m$ variables, $(k_{1},k_{2}),\ldots,(k_{m-1},k_{m})$. In general, there are $(m-1)!!=(m-1)\cdot(m-3)\cdots 1$ such pairings. For a proof, see for example Zee (2010). In order to understand and use the theorem, it is instructive to look at a few examples:

A.2 Connected correlations

Given general (not necessarily Gaussian) random variables, connected correlation functions are defined inductively through

where the sum is over all the possible subdivisions of $m$ variables into $s>1$ clusters of sizes $(\nu_{1},\ldots,\nu_{s})$ as $(k^{}_{1},\ldots,k^{}_{\nu_{1}}),\ldots,(k^{[s]}_{1},\ldots,k^{[s]}_{\nu_{s}})$. In order to understand the definition, it is again instructive to look at a few examples. Assuming that all the odd moments vanish,

If these examples do not suffice, here is yet another example to chew on:

We emphasize that these are just renderings of the definition (A.2). The power of this definition will be illustrated in the next two subsections.

A.3 Hierarchical clustering

We often encounter situations with the hierarchy

where $\epsilon\ll 1$ is a small perturbative parameter and here again odd moments are assumed to vanish. Often comes with the hierarchical structure is the asymptotic limit $\epsilon\rightarrow 0$ where

with the Gaussian kernel $K_{\mu_{1}\mu_{2}}$ at zero $\epsilon$ and the leading self-energy correction $S_{\mu_{1}\mu_{2}}$. Let us also denote the leading four-point vertex

A.4 Combinatorial hack

With the review of connected correlation functions passed us, first note that

where in the last equality Wick’s theorem was used backward.

Below, let us use the inverse kernel $\left(K^{-1}\right)^{\mu_{1}\mu_{2}}$ as a metric to raise indices:

Then, in order to simplify the second set of terms in Equation (CLUSTER’) involving self-energy, note that

where the symmetry $\mu_{1}\leftrightarrow\mu_{2}$ of $S_{\mu_{1}\mu_{2}}$ was used. Hence, defining

The similar algebraic exercise renders the other term in Equation (CLUSTER’) to be

In summary, for any function $\mathcal{F}[\mathbf{z}]$ of random variables $\textnormal{z}_{\mu}$

The operators in Equations (OS’) and (OV’) then become

i.e., the operators in Equations (OS) and (OV) in the main text.

Appendix B Full condensed proof

Studiously disentangling cases with different numbers of repetitions in neuron indices $(j_{1},\ldots,j_{k})$, we notice that at order $O\left(\frac{1}{n}\right)$, terms without repetition or with only one repetition contribute, finding

As special cases, we obtain expressions advertised in the main text to be contained in this Appendix:

Nowhere in our derivation had we assumed anything about the form of activation functions. The only potential exceptions to our formalism are exponentially growing activation functions – which we never see in practice – that would make the Gaussian integrals unintegrable.

Appendix C Bestiary of concrete examples

Let us take multilayer perceptrons with quadratic activation, $\sigma(z)=z^{2}$, and study the distributions of preactivations in the second layer as another illustration of our technology. From the master recursion relations (R1-R3) with the initial condition (R0), Wick’s contractions yield

where $\widetilde{K}^{(1)}_{\alpha_{1}\alpha_{2}}=C_{b}^{(1)}+C_{W}^{(1)}\cdot\left(\frac{\bm{x}_{\alpha_{1}}\cdot\bm{x}_{\alpha_{2}}}{n_{0}}\right)$. These expressions are used in the main text for the experimental study of finite-width corrections on Bayesian inference.

C.2 Details for single-input cases

For monomial activations, $\sigma(z)=z^{p}$, such as in deep linear networks Saxe et al. (2013) and quadratic activations Li et al. (2018),

In particular the four-point vertex solution is given by

C.2.2 ReLU with single input

C.3 More experiments on output distributions

Appendix D Finite-width corrections on Bayesian inference

In order to massage Equation (NGPM) into an actionable form, first playing with the metric inversions and defining $\overline{\phi}_{i}^{\ \alpha}\equiv\sum_{\alpha^{\prime}}(\widetilde{K}^{-1})^{\alpha\alpha^{\prime}}\overline{\phi}_{i;\alpha^{\prime}}$, the mean prediction becomes

This expression simplifies drastically through the identity

which can be checked explicitly, recalling $\widetilde{K}_{\Delta}\equiv\widetilde{K}_{\text{EE}}-\widetilde{K}_{\text{ER}}\widetilde{K}_{\text{RR}}^{-1}\widetilde{K}_{\text{RE}}$. Incidentally, this identity can also be used to prove Equation (GP$\Delta$). Now equipped with this identity, recalling $\overline{\phi}_{i;\alpha}\equiv[\left(y_{\text{R}}\right)_{i;\bar{\beta}},\left(\textnormal{y}^{\text{GP}}_{\text{E}}\right)_{i;\dot{\gamma}}]$, we notice that $\overline{\phi}_{i}^{\ \bar{\beta}}=\sum_{\bar{\beta}^{\prime}}\left(\widetilde{K}_{\text{RR}}^{-1}\right)^{\bar{\beta}\bar{\beta}^{\prime}}\left(y_{\text{R}}\right)_{i;\bar{\beta}^{\prime}}$ and $\overline{\phi}_{i}^{\ \dot{\gamma}_{1}}=0$. Similarly

and other components [i.e. with one or both of training components $(\bar{\beta}_{2},\bar{\beta}_{3})$ replaced by test components $\dot{\gamma}$] vanish. Equation (S46) thus simplifies to

Finally, denoting the matrix inside the parenthesis to be

and noticing $\sum_{\dot{\gamma}_{1}}\left(\widetilde{K}_{\Delta}\right)_{\dot{\gamma}\dot{\gamma}_{1}}\left(\widetilde{K}^{-1}\right)^{\dot{\gamma}_{1}\bar{\beta}_{0}}=-\left(\widetilde{K}_{\text{ER}}\widetilde{K}^{-1}_{\text{RR}}\right)_{\dot{\gamma}}^{\ \ \ \bar{\beta}_{0}}$ and $\sum_{\dot{\gamma}_{1}}\left(\widetilde{K}_{\Delta}\right)_{\dot{\gamma}\dot{\gamma}_{1}}\left(\widetilde{K}^{-1}\right)^{\dot{\gamma}_{1}\dot{\gamma}_{0}}=\delta_{\dot{\gamma}}^{\ \dot{\gamma}_{0}}$,

is the mean prediction. Equations (NGPM’) and (D) with $\overline{\phi}_{i}^{\ \bar{\beta}}=\sum_{\bar{\beta}^{\prime}}\left(\widetilde{K}_{\text{RR}}^{-1}\right)^{\bar{\beta}\bar{\beta}^{\prime}}\left(y_{\text{R}}\right)_{i;\bar{\beta}^{\prime}}$ are actionable, i.e., easy to program.