Gibbs Sampling the Posterior of Neural Networks
In this paper, we study sampling from a posterior derived from a neural
network. We propose a new probabilistic model consisting of adding noise at
every pre- and post-activation in the network, arguing that the resulting
posterior can be sampled using an efficient Gibbs sampler. For small models,
the Gibbs sampler attains similar performances as the state-of-the-art Markov
chain Monte Carlo (MCMC) methods, such as the Hamiltonian Monte Carlo (HMC) or
the Metropolis adjusted Langevin algorithm (MALA), both on real and synthetic
data. By framing our analysis in the teacher-student setting, we introduce a
thermalization criterion that allows us to detect when an algorithm, when run
on data with synthetic labels, fails to sample from the posterior. The
criterion is based on the fact that in the teacher-student setting we can
initialize an algorithm directly at equilibrium.
Giovanni Piccioli, Emanuele Troiani, Lenka Zdeborová
# Gibbs Sampling the Posterior of Neural Networks

###### Abstract

In this paper, we study sampling from a posterior derived from a neural network. We propose a new probabilistic model consisting of adding noise at every pre- and post-activation in the network, arguing that the resulting posterior can be sampled using an efficient Gibbs sampler. The Gibbs sampler attains similar performances as the state-of-the-art Monte Carlo Markov chain methods, such as the Hamiltonian Monte Carlo or the Metropolis adjusted Langevin algorithm, both on real and synthetic data. By framing our analysis in the teacher-student setting, we introduce a thermalization criterion that allows us to detect when an algorithm, when run on data with synthetic labels, fails to sample from the posterior. The criterion is based on the fact that in the teacher-student setting we can initialize an algorithm directly at equilibrium.

## I Introduction

Neural networks are functions parametrized by the so-called weights, mapping inputs to outputs. Neural networks are commonly trained by seeking values of weights that minimize a prescribed loss function. In some contexts, however, we want to sample from an associated probability distribution of the weights. Such sampling is at the basis of Bayesian deep learning [50; 54]. It is used in Bayesian uncertainty estimation [25; 31; 46] or to evaluate Bayes-optimal performance in toy models where the data-generative process is postulated [3]. In this paper, we focus on studying the algorithms and properties of such sampling.

Given training inputs \(X\) in Bayesian learning, one implicitly assumes the labels to be generated according to the stochastic process \(y\sim P(y|X,W)\), where \(W\) are the weight of the network, on which a prior \(P(W|X)\) is placed. At its heart, Bayesian deep learning consists of sampling from the posterior probability of the parameters:

\[P(W|X,y)=\frac{P(y|W,X)P(W|X)}{P(y|X)}, \tag{1}\]

where we simply used Bayes theorem. This sampling problem is, in general, NP-hard [10], with many techniques being developed to sample from (1). In this paper, we look at iterative algorithms that in the large time limit, return samples from the posterior distribution (1). Most available algorithms for this task are based on MCMC methods. We focus on the two following questions:

* **Q1:** Do we have a method to evaluate whether the algorithms have thermalized, i.e., if the samples returned by the MCMC plausibly come from the posterior (1)?
* **Q2:** Which combinations of sampling algorithm and form of the posterior distribution achieve the best performance in terms of ability to thermalize while reaching a low test error?

The first question addresses the long-standing problem of estimating an MCMC's thermalization time, that is, the time at which the MCMC starts sampling well from the posterior. We propose a criterion for thermalization based on the teacher-student setting. The criterion can only be reliably applied to synthetic labels generated by a teacher network. After a comparison with other thermalization heuristics, we argue that the teacher-student criterion is more discriminative, in that it provides a higher lower bound to the thermalization time. The second question explores the interplay between the form of the posterior and the sampling algorithm: since there is more than one way of translating a network architecture into a probabilistic process, we exploit this freedom to introduce a generative process in which noise is added at every pre- and post-activation of the network. We then design a Gibbs sampler tailored to this posterior and compare it to other commonly used MCMCs.
### Related literature

When running an MCMC one has to wait a certain number of iterations for the algorithm to start sampling from the desired probability measure. We will refer to this burn-in period as the thermalization time or \(T_{\text{therm}}\)[41]. Samples before \(T_{\text{therm}}\) should therefore be discarded. Estimating \(T_{\text{therm}}\) is thus of great practical importance, as it is crucial to know how long the MCMC should be run. More formally, we initialize the MCMC with a \(\delta\)-peaked distribution as its initial condition \(P_{t=0}(x)=\delta_{x,x_{0}}\) where \(x_{0}\in\mathcal{X}\) is a starting state of our liking. Then, the probability from which the MCMC samples evolves according to \(P_{t+1}(x)=\sum_{x^{\prime}\in\mathcal{X}}P_{t}(x^{\prime})P(x_{t+1}=x|x_{t}=x ^{\prime})\)[15], where the kernel \(P(x_{t+1}=x^{\prime}|x_{t}=x)\) contains the transition probabilities of the MCMC. If this kernel satisfies the detailed balance condition with respect to a probability \(\pi(\cdot)\), and the chain is ergodic, then for \(t\to\infty\), \(P_{t}(x)\to\pi(x)\), and thus the MCMC will sample from \(\pi\). In this setting, the thermalization time represents the time after which \(P_{t}\) is sufficiently close to \(\pi\). The actual thermalisation threshold will change depending on the metric we use to measure the distance between \(P_{t}\) and \(\pi\). Most commonly used distances between distributions \(P\) and \(Q\) can be written in the form \(d(P,Q)=\sup_{\varphi\in\mathcal{F}}\big{\{}\sum_{x\in\mathcal{X}}\varphi(x)(P( x)-Q(x))\big{\}}\), with \(\mathcal{F}\) a function space that determines how strict the convergence is. One of the strictest metrics is the total variation distance \(d_{TV}\), obtained by picking \(\mathcal{F}_{TV}=\{\varphi:\varphi\) measurable, \(||\varphi||_{\infty}\leq 1\}\). This distance gives rise to the definition of mixing time as the smallest \(t\) for which \(d_{TV}(P_{t},\pi)\leq 1/4\). Analytical bounds on the mixing time can be obtained using the transition kernel's spectral properties [27]. These results, however, are of difficult practical application, and the convergence metric might be too strict for practical purposes. A looser definition is that of weak convergence, which corresponds to a distance \(d_{w}\) given by \(\mathcal{F}_{w}=\{\varphi\text{ s.t. }||\varphi||_{\infty}+||\varphi||_{Lip}\leq 1\}\)[48]. For \(d_{w}(P_{t},\pi)\) to be small, all Lipschitz bounded functions must have a similar expectation under \(P_{t}\) and \(\pi\). This requirement can still be too strict: for example, in the statistical physics literature [41][57][34] one only includes in \(\mathcal{F}\) functions \(\varphi\) whose output concentrates in the thermodynamic limit, giving rise to \(d_{\text{tatphys}}\). We prefer to leave the definition of \(T_{\text{therm}}\) a bit vague since, from a practical point of view, none of the above convergence metrics can be so far computed efficiently.

Practitioners have instead resorted to a number of heuristics, which provide lower bounds to the thermalization time. These heuristics usually revolve around two ideas. We first have methods involving multiple chains [5; 9; 16; 41]. In different flavours, all these criteria rely on comparing multiple chains with different initializations. Once all the chains have thermalized, samples from different chains should be indistinguishable. Another approach consists of finding functions with known mean under the posterior and verifying whether the empirical mean is also close to its predicted value [9; 13; 19; 20; 58]. The proposed method for detecting thermalization relies instead on the teacher-student framework [57].

Another field we connect with is that of Bayesian learning of neural networks. For an introduction see [18; 24; 32; 50] and references therein. We shall first examine the probabilistic models for Bayesian learning of neural networks and then review the algorithms that are commonly used to sample. In order to specify the posterior (1), one needs to pick the likelihood (or data generating process) \(P(y|X,W)\). The most common model, employed in the great majority of works [23; 39; 47; 52; 54] is \(P(y|X,W)=\frac{1}{Z}\exp\left(-\frac{1}{2\Delta}\sum_{\mu}\ell(y^{\mu},f(X^{\mu},W))\right)\), where \(f(\cdot,W)\) is the neural network function, \(\ell\) is the loss function, \(\mu\) is the sample index, and \(\Delta\) is a temperature parameter. As an alternative, other works have introduced the "stochastic feedforward networks" (SFFN), where noise is added at every layer's pre-activation [37; 42; 56; 45]. Outside of the Bayesian learning of neural networks literature, models where intermediate pre- or post-activations are added as dynamical variables have also been considered in the predictive coding literature [2; 35; 36; 53].

Once a probabilistic model has been chosen, the goal is to obtain samples from the corresponding posterior. A first solution consists of approximating the posterior with a simpler distribution, which is easier to sample. This is the strategy followed by variational inference methods [26; 30; 31; 45; 49; 55]. Although variational inference yields fast algorithms, it is often based on uncontrolled approximations. Another category of approximate methods is that of "altered MCMCs", i.e., Monte Carlo algorithms which have been modified to be faster at the price of not sampling anymore from the posterior [7; 28; 29; 38; 40; 59]. An example of these algorithms is the discretized Langevin dynamics [51]. Restricting the sampling to a subset of the parameters has also been considered in [44] as an alternative training technique.

Finally, we have exact sampling methods: these are iterative algorithms that in the large time limit are guaranteed to return samples from the posterior distribution. Algorithms for exact sampling mostly rely on MCMC methods. The most popular ones are Hamiltonian Monte Carlo (HMC) [12; 39], Metropolis adjusted Langevin algorithm (MALA) [4; 33; 43] and the No U-turn sampler (NUTS)[22]. Within the field of Bayesian learning in neural networks, HMC is the most commonly used algorithm [23; 52; 54]. The proposed Gibbs sampler is inspired to the work of [1], and later [14; 21], that introduced the idea of augmenting the variable space in the context of logistic and multinomial regression.
Teacher-Student Thermalization Criterion

In this section we explain how to use the teacher-student setting to build a thermalization test for sampling algorithms. The test gives a lower bound on the thermalization time. We start by stating the main limitation of this approach: the criterion can only be applied to synthetic datasets. In other words, the training labels \(y\) must be generated by a teacher network, using the following procedure.

We first pick arbitrarily the training inputs and organize them into an \(n\times d\) matrix \(X\). Each row \(X^{\mu}\) of the matrix is a different training sample, for a total of \(n\) samples. We then sample the teacher weights \(W_{\star}\) from the prior \(P(W)\). Finally, we generate the noisy training labels as \(y^{\mu}\sim P(y|X^{\mu},W_{\star})\). Our goal is to draw samples from the posterior \(P(W|D)\), where \(D=\{(X^{\mu},y^{\mu})\}_{\mu\in[n]}\) indicates the training set. Suppose we want to have a lower bound on the thermalization time of a MCMC initialized at a particular configuration \(W_{\text{start}}\). The method consists of running two parallel chains \(W_{1}(t)\) and \(W_{2}(t)\). For the first chain, we use an informed initialization, meaning we initialize the chain on the teacher weights, thus setting \(W_{1}(t=0)=W_{\star}\). For second chain we set \(W_{2}(t=0)=W_{\text{start}}\). To determine convergence we consider a test function \(g(W)\). We first run the informed initialization: after some time \(T_{1}\), \(g(W_{1}(t))\) will become stationary. Using samples collected after \(T_{1}\) we compute the expected value of \(g\) (let us call it \(\overline{g}\)). Next, we run the second chain. The lower bound to the thermalization time of \(W_{2}(t)\) is the time where \(g(W_{2}(t))\) becomes stationary and starts oscillating around \(\overline{g}\). In practice, this time is determined by visually inspecting the time series of \(g\) under the two initializations, and observing when the two merge.

At first glance this method does not seem too different from [16] or [5], whose method (described in Appendix A) relies on multiple chains with different initializations. There is however a crucial difference: the informed initialization is already thermalized at \(t=0\). To see this, recall that the pair \(W_{\star},D\) was obtained by first sampling \(W_{\star}\) from \(P(W)\) and then sampling \(D\) from \(P(D|W_{\star})\). This implies that \(W_{\star},D\) is a sample from the joint distribution \(P(W,D)\). Rewriting \(P(W,D)=P(W|D)P(D)\), we see that \(W_{\star}\) is also typical under the posterior distribution. In conclusion, the power of the teacher-student setting lies in the fact that it gives us access to one sample from the posterior, namely \(W_{\star}\). It then becomes easier to check whether a second chain is sampling from the posterior by comparing the value of a test function. In contrast, other methods comparing chains with different initialization have no guarantee that if the two chains "merge" then the MCMC is sampling from the posterior, since it is possible that both chains are trapped together far from equilibrium.

## III The Intermediate Noise Model

In this section, we introduce a new probabilistic model for Bayesian learning of neural networks. We start by reviewing the classical formulation of Bayesian learning of neural networks. Let \(f(x,W)\) be the neural network function, with \(W\) its parameters, and \(x\in\mathbb{R}^{d}\) the input vector. Given a training set \(X\in\mathbb{R}^{n\times d},y\in\mathbb{R}^{n}\) we aim to sample from

\[P(W|X,y)=\frac{1}{P(y|X)}P(W)\exp\left[-\frac{1}{2\Delta}\sum_{\mu=1}^{n}\ell \left(y^{\mu},f(X^{\mu},W)\right)\right], \tag{2}\]

where \(\ell(\cdot,\cdot)\) is the single sample loss function, and \(\Delta\) a temperature parameter. Notice that to derive (2) from (1), we supposed that \(P(W|X)=P(W)\), i.e., \(W\) is independent of \(X\). This is a common and widely adopted assumption in the Bayesian learning literature, and we shall make it in what follows. Most works in the field of Bayesian learning of neural networks attempt to sample from (2). This form of the posterior corresponds to the implicit assumption that the labels were generated by

\[y^{\mu}\sim P_{\text{out}}(y|f(X^{\mu},W)),\text{ with }P_{\text{out}}(y|z) \propto e^{-\frac{1}{2\Delta}\ell(y,z)} \tag{3}\]

where \(W\) are some weights sampled from the prior. We propose an alternative generative model based on the idea of introducing a small Gaussian noise at every pre- and post-activation in the network. The motivation behind this process lies in the fact that we are able to sample the resulting posterior efficiently using a Gibbs sampling scheme. Consider the case where \(f(\cdot,W)\) is a multilayer perceptron with \(L\) layers, without biases and with activation function \(\sigma\left(\cdot\right)\). Hence we have \(f(x,W)=W^{(L)}\sigma\left(W^{(L-1)}\sigma\left(\ldots\sigma(W^{(1)}x)\ldots \right)\right)\). Here \(W^{(\ell)}\in\mathbb{R}^{d_{\ell+1}\times d_{\ell}}\) indicates the weights of layer \(\ell\in[L]\), with \(d_{\ell}\) the width of the layer. We define the pre-activations \(Z^{(\ell)}\in\mathbb{R}^{n\times d_{\ell}}\) and post activations
\(X^{(\ell)}\in\mathbb{R}^{n\times d_{\ell}}\) of layer \(\ell\). Using Bayes theorem and applying the chain rule to the likelihood we obtain

\[P(\{X^{(\ell)}\}_{\ell=2}^{L},\{Z^{(\ell)}\}_{\ell=2}^{L},\{W^{( \ell)}\}_{\ell=1}^{L}|X,y)=\] \[\frac{1}{P(y|X)}P(\{W^{(\ell)}\}_{\ell=1}^{L})P(y,\{X^{(\ell)}\}_{ \ell=2}^{L},\{Z^{(\ell)}\}_{\ell=2}^{L}|\{W^{(\ell)}\}_{\ell=1}^{L},X)=\frac{1 }{P(y|X)}P(\{W^{(\ell)}\}_{\ell=1}^{L})\times\] \[\times\left[\prod_{\ell=2}^{L}P(Z^{(\ell+1)}|X^{(\ell)},W^{(\ell) })P(X^{(\ell)}|Z^{(\ell)})\right]P(Z^{(2)}|X,W^{(1)}) \tag{4}\]

with the constraint \(Z^{(L+1)}=y\). The conditional probabilities are assumed to be

\[P(Z^{(\ell+1)}|X^{(\ell)},W^{(\ell)})=\prod_{\mu=1}^{n}\prod_{ \alpha=1}^{d_{\ell+1}}\mathcal{N}\left(Z_{\alpha}^{(\ell+1)\mu}\Big{|}W_{ \alpha}^{(\ell)T}X^{(\ell)\mu},\Delta_{Z}^{(\ell+1)}\right) \tag{5}\] \[P(X^{(\ell)}|Z^{(\ell)})=\prod_{\mu=1}^{n}\prod_{i=1}^{d_{\ell}} \mathcal{N}\left(X_{\alpha}^{(\ell)\mu}\Big{|}\sigma(Z_{i}^{(\ell)\mu}), \Delta_{X}^{(\ell)}\right), \tag{6}\]

where \(\{\Delta_{Z}^{(\ell)}\}_{\ell=2}^{L+1}\), \(\{\Delta_{X}^{(\ell)}\}_{\ell=2}^{L}\) control the amount of noise added at each pre- and post- activation. This structure of the posterior implicitly assumes that the pre- and post-activations are iteratively generated as

\[Z^{(\ell+1)}=X^{(\ell)}W^{(\ell)T}+\epsilon_{Z}^{(\ell+1)},\quad X^{(\ell+1)}= \sigma(Z^{(\ell+1)})+\epsilon_{X}^{(\ell+1)},\quad\ell\in[L]. \tag{7}\]

\(X^{(1)}=X\in\mathbb{R}^{n\times d}\) are the inputs, \(Z^{(L+1)}=y\in\mathbb{R}^{n}\) represent the labels and \(\epsilon_{Z}^{(\ell)},\epsilon_{X}^{(\ell)}\) are \(n\times d_{\ell}\) matrices of i.i.d. respectively \(\mathcal{N}(0,\Delta_{Z}^{(\ell)})\) and \(\mathcal{N}(0,\Delta_{X}^{(\ell)})\) elements. We will refer to (7) as the intermediate noise generative process.

If we manage to sample from the posterior (4), which has been augmented with the variables \(\{X^{(\ell)}\}_{\ell=2}^{L},\{Z^{(\ell)}\}_{\ell=2}^{L}\), then we can draw samples from \(P(W|X,y)\), just by discarding the additional variables. A drawback of this posterior is that one has to keep in memory all the pre- and post-activations in addition to the weights.

We remark that the intermediate noise generative process admits the classical generative process (3) and the SFNN generative model as special cases. Setting all \(\Delta\)s (and hence all \(\epsilon\)) to zero in (7) except for \(\Delta_{Z}^{(L+1)}\) indeed gives back the classical generative process (3), with \(\ell(y,z)=(y-z)^{2}\) and \(\Delta=\Delta_{Z}^{(\ell+1)}\). Instead, setting \(\Delta_{X}^{(\ell)}=0\) for all \(\ell\), but keeping the noise in the pre-activations gives the SFNN model.

## IV Gibbs sampler for neural networks

Gibbs sampling [17] is an MCMC algorithm that updates each variable in sequence by sampling it from its conditional distribution. For a probability measure with three variables \(P(\theta_{1},\theta_{2},\theta_{3})\), one step of Gibbs sampling can be described as follows. Starting from the configuration \(\theta_{1}(t),\theta_{2}(t),\theta_{3}(t)\), we first draw \(\theta_{1}(t+1)\sim P(\theta_{1}|\theta_{2}(t),\theta_{3}(t))\), then we draw \(\theta_{2}(t+1)\sim P(\theta_{2}|\theta_{1}(t+1),\theta_{3}(t))\) and finally \(\theta_{3}(t+1)\sim P(\theta_{3}|\theta_{1}(t+1),\theta_{2}(t+1))\). Repeating this procedure one can prove [6] that, in the limit of many iterations (\(t\gg 1\)) and provided that the chain is ergodic, the samples \((\theta_{1}(t),\theta_{2}(t),\theta_{3}(t))\) will come from \(P(\theta_{1},\theta_{2},\theta_{3})\). We now present a Gibbs sampler for the intermediate noise posterior (4), with Gaussian prior. More specifically the prior on \(W^{(\ell)}\) is i.i.d. \(\mathcal{N}(0,1/\lambda_{W}^{(\ell)})\) over the weights' coordinates. The full derivation of the algorithm is reported in Appendix C, here we sketch the main steps. To define the sampler we need to compute the distributions of each of \(X^{(\ell)},Z^{(\ell)},W^{(\ell)}\) conditioned on all other variables (here indicated by "All"). For \(X^{(\ell)}\) the conditional distribution factorizes over samples \(\mu\in[n]\), leading to

\[P(X^{(\ell)\mu}|\text{All})=P(X^{(\ell)\mu}|Z^{(\ell)\mu},W^{(\ell)},Z^{(\ell+1 )\mu})=\mathcal{N}(X^{(\ell)\mu}|(m^{(X_{\ell})})^{\mu},\Sigma^{(X_{\ell})}). \tag{8}\]

This is a multivariate Gaussian with covariance \(\Sigma^{(X_{\ell})}=\left(\frac{1}{\Delta_{Z}^{(\ell+1)}}W^{(\ell)T}W^{(\ell)} +\frac{1}{\Delta_{X}^{(\ell)}}\mathbb{I}_{d_{\ell}}\right)^{-1}\), and mean \((m^{(X_{\ell})})^{\mu}=\Sigma^{(X_{\ell})}\left(\frac{1}{\Delta_{X}^{(\ell)}} \sigma(Z^{(\ell)\mu})+\frac{1}{\Delta_{Z}^{(\ell+1)}}W^{(\ell)T}Z^{(\ell+1)\mu }\right).\)

Considering \(W^{(\ell)}\), we exploit that the conditional factorizes over the rows \(\alpha\in[d_{\ell+1}]\).

\[P(W_{\alpha}^{(\ell)}|\text{All})=P(W_{\alpha}^{(\ell)}|X^{(\ell)},Z_{\alpha}^{( \ell+1)})=\mathcal{N}(W_{\alpha}^{(\ell)}|(m_{W}^{(\ell)})_{\alpha},\Sigma_{W}^{ (\ell)}), \tag{9}\]

with \(\Sigma_{W}^{(\ell)}=\left(\frac{1}{\Delta_{Z}^{(\ell+1)}}X^{(\ell)T}X^{(\ell)}+ \lambda_{W}^{(\ell)}\mathbb{I}_{d_{\ell}}\right)^{-1}\), and \((m_{W}^{(\ell)})_{\alpha}=\frac{1}{\Delta_{Z}^{(\ell+1)}}\Sigma_{W}^{(\ell)T}Z _{\alpha}^{(\ell+1)}\).
For \(Z^{(\ell+1)}\) the conditional factorizes both over samples and over coordinates. We have

\[P(Z_{\alpha}^{(\ell+1)\mu}|\text{All})=P(Z_{\alpha}^{(\ell+1)\mu}|X _{\alpha}^{(\ell+1)\mu},W_{\alpha}^{(\ell)},X^{(\ell)\mu})\propto \tag{10}\] \[\exp\left[-\frac{1}{2\Delta_{Z}^{(\ell+1)\mu}}\left(Z_{\alpha}^{( \ell+1)\mu}-W_{\alpha}^{(\ell)T}X^{(\ell)\mu}\right)^{2}-\frac{1}{2\Delta_{X}^ {(\ell+1)}}\left(\sigma(Z_{\alpha}^{(\ell+1)\mu})-X_{\alpha}^{(\ell+1)\mu} \right)^{2}\right].\]

Notice that the conditional distributions of \(W_{\alpha}^{(\ell)}\) and \(X^{(\ell)\mu}\) are multivariate Gaussians and can be easily sampled. Instead \(Z_{\alpha}^{(\ell)\mu}\) is a one-dimensional random variable with non Gaussian distribution. Appendix E provides recipes for sampling it for sign, ReLU and absolute value activations.

```
Input: training inputs \(X\), training labels \(y\), noise variances \(\{\Delta_{Z}^{(\ell)}\}_{\ell=2}^{L+1}\), \(\{\Delta_{X}^{(\ell)}\}_{\ell=2}^{L}\), prior inverse variances \(\{\lambda_{W}^{(\ell)}\}_{\ell=1}^{L}\), initial condition \(\{X^{(\ell)}\}_{\ell=2}^{L}\),\(\{W^{(\ell)}\}_{\ell=1}^{L}\), \(\{Z^{(\ell)}\}_{\ell=2}^{L}\), length of the simulation \(t_{\text{max}}\) Output: a sequence \(S\) of samples \(X^{(1)}\gets X\) \(Z^{(L+1)}\gets y\) \(S\leftarrow\left[\{\{W^{(0)}\}_{\ell=1}^{L}\), \(\{X^{(0)}\}_{\ell=2}^{L}\), \(\{Z^{(\ell)}\}_{\ell=2}^{L}\)\right]\) for\(t=1,\ldots,t_{\text{max}}\)do\(W^{(1)}\sim P(W^{(1)}|X,Z^{(2)})\)\(\triangleright\) See eq. (9) for\(\ell=2,\ldots,t_{\text{do}}\)\(X^{(\ell)}\sim P(W^{(\ell)}|Z^{(\ell)},W^{(\ell)},Z^{(\ell+1)})\)\(\triangleright\) See eq. (8) \(W^{(\ell)}\sim P(W^{(\ell)}|X^{(\ell)},Z^{(\ell+1)})\)\(\triangleright\) See eq. (9) \(Z^{(\ell)}\sim P(Z^{(\ell)}|X^{(\ell-1)},W^{(\ell-1)},X^{(\ell)})\)\(\triangleright\) See eq. (10) endfor \(S\).append \(\left(\{W^{(\ell)}\}_{\ell=1}^{L}\), \(\{X^{(\ell)}\}_{\ell=2}^{L}\), \(\{Z^{(\ell)}\}_{\ell=2}^{L}\right)\) endfor
```

**Algorithm 1** Gibbs sampler for MLP

Putting all ingredients together, we obtain the Gibbs sampling algorithm, whose pseudocode is reported in Algorithm 1. The main advantages of Gibbs sampling lie in the fact that it has no hyperparameters to tune and, moreover, it is a rejection-free sampling method. In the case of MCMCs, hyperparameters are defined to be all parameters that can be changed without affecting the probability measure that the MCMC asymptotically samples. The Gibbs sampler can also be parallelized across layers: a parallelized version of Algorithm 1 is presented in Appendix D. Finally, one can also extend this algorithm to more complex architectures: Appendices F and G contain respectively the update equations for biases and convolutional networks. We release an implementation of the Gibbs sampler at [https://github.com/SPOC-group/gibbs-sampler-neural-networks](https://github.com/SPOC-group/gibbs-sampler-neural-networks).

## V Numerical results

In this section we present numerical experiments to support our claims. We publish the code to reproduce these experiments at [https://github.com/SPOC-group/numerics-gibbs-sampling-neural-nets](https://github.com/SPOC-group/numerics-gibbs-sampling-neural-nets).

### Teacher student convergence method

In section II we proposed a thermalization criterion based on having access to an already thermalized initialization. Here we show that it is more discriminative than other commonly used heuristics. We first briefly describe these heuristics.

* **Stationarity**. Thermalization implies stationarity since once the MCMC has thermalized, it samples from a fixed probability measure. Therefore any observable, plotted as a function of time should oscillate around a constant value. The converse (stationarity implies thermalization) is not true. Nevertheless observing when a function becomes stationary gives a lower bound on \(T_{\text{therm}}\).
* **Score method**[13]. Given a probability measure \(P(x)\), we exploit the fact that \(\mathbb{E}_{x\sim P}\left[\frac{\partial\log P(x)}{\partial x}\right]=\int \frac{\partial P(x)}{\partial x}dx=0\). We then monitor the function \(\frac{\partial\log P(x)}{\partial x}\) along the dynamics. The time at which it starts fluctuating around zero is another lower bound to \(T_{\text{therm}}\).
* \(\hat{R}\) **statistic [16]**. Two (or more) MCMCs are run in parallel starting from different initializations. The within-chain variance is compared to the total variance, obtained by merging samples from both chains. Call the ratio of these variances \(\hat{R}\) (a precise definition of which is given in Appendix A). If the MCMC has thermalized, the samples from the two chains should be indistinguishable, thus \(\hat{R}\) will be close to 1. The time at which \(\hat{R}\) gets close to 1 provides yet another lower bound to the thermalization time.

We compare these methods in the case of a one hidden layer neural network, identical for the teacher and the student, with input dimension \(d_{1}=50\), \(d_{2}=10\) hidden units and a scalar output. This corresponds to the function

\[f(x,W)=b^{(2)}+W^{(2)}\sigma\left(W^{(1)}x+b^{(1)}\right), \tag{11}\]

where \(\sigma(x)=\max(0,x)\) and \(W\) indicates the collection of all parameters: \(W^{(1)}\in\mathbb{R}^{d_{2}\times d_{1}}\) and \(W^{(2)}\in\mathbb{R}^{1\times d_{2}}\). We specify the prior by setting \(\lambda_{W}^{(1)}=\lambda_{b}^{(1)}=d_{1},\lambda_{W}^{(2)}=\lambda_{b}^{(2) }=d_{2}\), the prior on the bias \(b^{(\ell)}\) is \(\mathcal{N}(0,1/\lambda_{b}^{(\ell)})\) i.i.d. over the coordinates of the bias vector. Let \(n=2084\) be the size of the training set. We pick \(n\) to be four times the number of parameters in the network anticipating that the training set contains enough information to learn the teacher. We start by generating the matrix of training inputs \(X\in\mathbb{R}^{n\times d_{1}}\) with i.i.d. standard Gaussian entries, then we sample the teacher's weights \(W_{*}\) from the Gaussian prior. For concreteness we set \(\Delta_{Z}^{(2)},\Delta_{Z}^{(2)},\Delta_{Z}^{(3)}\) to the same value \(\Delta\) and set \(\Delta=10^{-4}\). To generate the training labels \(y\), we feed \(X\), the teacher's weights \(W_{*}\) and \(\Delta\) into the generative process (7), adapted to also add the biases. For the test set, we first sample \(X_{\text{test}}\), with i.i.d. standard Gaussian entries. Both the test labels and the test predictions are generated in a noiseless way (i.e., just passing the inputs through the network). In this way, the test error takes the following form: test \(\text{MSE}=\frac{1}{\eta_{\text{test}}}\sum_{\mu=1}^{n_{\text{test}}}\left(f( X_{\text{test}}^{\mu},W_{*})-f(X_{\text{test}}^{\mu},W)\right)^{2}.\) The full details about this experiment set are in Appendix B. We run the Gibbs sampler on the intermediate noise posterior starting from three different initializations: informed, zero and random. Respectively the student's variables are initialized to the teacher's counterparts, to zero, or are sampled from the prior. In this particular setting, the Gibbs sampler initialized at zero manages to thermalize, while the random initializations fail to do so. Two independent random initializations are shown, in order to be able to use the multiple chains method.

Figure 1 illustrates a representative result of these experiments. In the left panel, we aim to find the highest lower bound to the thermalization time of the zero-initialized chain. Looking at the score method we plot \(U=\Delta\frac{1}{d_{1}d_{2}}\sum_{i=1}^{d_{1}}\sum_{\alpha=1}^{d_{2}}\frac{ \partial\log P}{\partial W_{i}^{(1)}}\), where \(P\) indicates the posterior distribution; this is the score rescaled by \(\Delta\) and averaged

Figure 1: Comparison of different thermalization measures. In the legend, next to each method we write between parentheses the initialization (or pair of initializations) the method is applied to. The circles on the \(x\) axis represent the thermalization times estimated by each method. **Left:** We compare the predictions for the thermalization time of the zero-initialized MCMC. The red \(y\) scale on the right refers uniquely to the lines in red. All the other quantities should be read on the black \(y\) scale. **Right:** We compare the predictions for the thermalization time of two chains initialized independently at random. The pink \(y\) scale refers uniquely to the pink line. All other quantities should be read on the black logarithmic scale. The randomly initialized runs fail to thermalize and their test MSEs get stuck on a plateau. However, \(\hat{R}\), whose time series on the plateau is stationary and close to 1, fails to detect this lack of thermalization.
over the first layer weights. In the zero-initialized chain, \(U\) starts oscillating around zero already at \(t=20\). Then we consider the \(\hat{R}\) statistics computed on the outputs of the two chains with zero and informed initializations. The criterion estimates that the zero-initialized chain has thermalized after \(t=6\times 10^{4}\), when \(\hat{R}\) approaches 1 and becomes stationary. Next, we consider the teacher-student method, with the test MSE as the test function (\(g\) in our previous discussion). According to this method, the MCMC thermalizes after the test MSE time series of the informed and zero-initialized chains merge, which happens around \(t=10^{5}\). Finally, the stationarity criterion, when applied to the test MSE or to \(\hat{R}\) gives a similar estimate for the thermalization time. The \(x\)-axis of the left plot provides a summary of this phenomenology, by placing a circle at the thermalization time estimated by each method. In summary, the teacher-student method is the most conservative, but the \(\hat{R}\) statistics-based method is also reasonable here.

The right panel of Fig. 1 then shows a representative situation where thermalization is not reached yet the \(\hat{R}\) statistics-based method would indicate it is. In the right panel, two randomly initialized chains, denoted by _random 1_ and _random 2_ are considered. Neither of these chains actually thermalizes, in fact looking at the test MSE time series we see that both chains get stuck on the same plateau around MSE=\(10^{-3}\) and are unable to reach the MSE of the informed initialization. However, as soon as both chains reach the plateau, \(\hat{R}\) quickly drops to a value close to the order of 1 and thereafter becomes stationary, mistakenly signalling thermalization. This example exposes the problem at the heart of the multiple-chain method: the method can be fooled if the chains find themselves close to each other but far from equilibrium. Similarly, since the chains become stationary after they hit the plateau, the stationarity criterion would incorrectly predict that they have thermalized. To conclude, we have shown an example where common thermalization heuristics fail to recognize that the MCMC has not thermalized; instead, the teacher-student method detects the lack of thermalization.

### Gibbs sampler

In this section, we show that the combination of intermediate noise posterior and Gibbs sampler is effective in sampling from the posterior by comparing it to HMC, run both on the classical and intermediate noise posteriors, and to MALA, run on the classical posterior. We provide the pseudocode for these algorithms in Appendix H. For the first set of experiments, we use the same network architecture as in the previous section. The teacher weights \(W_{*}\), as well as \(X,X_{\text{test}}\) are also sampled in the same way. The intermediate noise and the classical generative process prescribe different ways of generating the labels. However, to perform a fair comparison, we use the same dataset for all MCMCs and posteriors; thus we generate the training set in a noiseless way, i.e., setting \(y^{\mu}=f(X^{\mu},W_{\star})\). We generate 72 datasets according to this procedure, each time using independently sampled inputs and teacher's weights. The consequence of generating datasets in a noiseless way is that the noise level used to generate the data is different from the one in the MCMC, implying that the informed initialization will not exactly be a sample from the posterior. However, the noise is small enough that we did not observe any noticeable difference in the functioning of the teacher-student criterion.

First, we aim to characterize how often each algorithm thermalizes, when started from an uninformed initialization. Uninformed means that the network's initialization is agnostic to the teacher's weights. For several values of \(\Delta\), and for all the 72 datasets, we run the four algorithms (Gibbs, classical HMC, intermediate HMC, classical MALA) starting from informed and uninformed initializations. More information about the initializations and hyperparameters of these experiments is contained in Appendix I.

The left panel of figure 2 depicts the proportion of the 72 datasets in which the uninformed initialization thermalizes within 5:30h of simulation. The \(x-\)axis is the equilibrium test MSE, i.e., the average test MSE reached by the informed initialization once it becomes stationary. When \(\Delta\), and thus the test MSE, decreases, the proportion of thermalized runs drops for all algorithms, with the Gibbs sampler attaining the highest proportion, in most of the range. In the right panel, we plot the dynamics of the test error under each algorithm for a run where they all thermalize. For the same \(\Delta\)s of this plot (respectively \(\Delta=10^{-3},4.64\times 10^{-4}\) for the classical and intermediate noise posterior), we compute the average thermalization time among the runs that thermalize. Classical HMC, MALA, Gibbs, and intermediate HMC take, respectively on average around \(130,2700,3200,12500\) seconds to thermalize. This shows that the classical HMC, when it thermalizes, is the fastest method, while MALA and Gibbs occupy the second and third position, with similar times. However classical HMC thermalizes about 20% less often than the Gibbs sampler. Therefore in cases where it is essential to reach equilibrium, the Gibbs sampler represents the best choice.

We now move from the abstract setting of Gaussian data to more realistic inputs and architectures. As an architecture we use a one-hidden layer MLP with 12 hidden units and ReLU activations, and a simple convolutional network (CNN) with a convolutional layer, followed by average pooling, ReLU activations, and a fully connected layer. See Appendix J for a description of both models and of the experimental details. In this setting, we resort to the stationarity criterion to check for thermalization, since the teacher-student method is inapplicable. We compare the Gibbs sampler with HMC and MALA both run on the classical posterior, picking MNIST as dataset. Figure 3 shows the test error as a
function of time for the two architectures. We choose the algorithms \(\Delta\)s such that they all reach a comparable test error at stationarity. We then compare the time it takes each algorithm to reach this error. The results of the experiments are depicted in figure 3. For the MLP all algorithms take approximately the same time to become stationary, around \(500s\). In the CNN case, after HMC and MALA reach stationarity in \(100s\), compared to \(800s\) for Gibbs. We however note that for HMC and MALA to achieve these performances we had to carry out an extensive optimization over hyperparameters, thus the speed is overall comparable.

Figure 3: Gibbs on the intermediate noise posterior and HMC, MALA both on the classical posterior, compared on MNIST. **Left:** MLP with one hidden layer with 12 hidden units. **Right:** CNN network.

Figure 2: Thermalization experiments on synthetic data. **Left:** Proportion of the 72 runs that thermalize plotted against the equilibrium test MSE. **Right:** Example of the dynamics of the test MSE in a particular run where all four algorithms thermalize. In order to get a similar equilibrium test MSE in the classical and intermediate noise posteriors, we pick respectively \(\Delta=10^{-3}\) and \(\Delta=4.64\times 10^{-4}\). The transparent lines represent the informed initializations.
Conclusion

In this work, we introduced the intermediate noise posterior, a probabilistic model for Bayesian learning of neural networks, along with a novel Gibbs sampler to sample from this posterior. We compared the Gibbs sampler to MALA and HMC, varying also the form of the posterior. We found that HMC and MALA both on the classical posterior and Gibbs, on the intermediate noise posterior, each have their own merits and can be considered effective in sampling the high dimensional posteriors arising from Bayesian learning of neural networks. Gibbs compares favourably to the other algorithms in terms of the ability to thermalize, moreover, no hyperparameter tuning is required, it can be applied to non-differentiable posteriors, and can be parallelized across layers.

We further proposed the teacher-student thermalization criterion: a method to obtain stringent lower bounds on the thermalization time of an MCMC, within a synthetic data setting. We first provided a simple theoretical argument to justify the method and subsequently compared it to other thermalization heuristics, finding that the teacher-student criterion consistently gives the highest lower bound to \(T_{\mathrm{therm}}\).

## VII Acknowledgment

We thank Lucas Clarte for introducing us to the blocked Gibbs sampler, and Christian Keup for the useful discussions on predictive coding and stochastic neural networks.