# Wide stochastic networks:
Gaussian limit and PAC-Bayesian training

Eugenio Clerico111Corresponding author, email<EMAIL_ADDRESS>Department of Statistics, University of Oxford, UK  George Deligiannidis
Department of Statistics, University of Oxford, UK  Arnaud Doucet  Department
of Statistics, University of Oxford, UK

###### Abstract

The limit of infinite width allows for substantial simplifications in the
analytical study of over-parameterized neural networks. With a suitable random
initialization, an extremely large network is well approximated by a Gaussian
process, both before and during training. In the present work, we establish a
similar result for a simple stochastic architecture whose parameters are
random variables. The explicit evaluation of the output distribution allows
for a PAC-Bayesian training procedure that directly optimizes the
generalization bound. For a large but finite-width network, we show
empirically on MNIST that this training approach can outperform standard PAC-
Bayesian methods.

##  1 Introduction

Overparameterized neural networks have become the first-choice algorithm for
several learning tasks. Despite the common idea that overfitting is
unavoidable when the number of trainable parameters exceeds (sometimes by
several orders of magnitudes) the size of the training set, deep architectures
made of layers with thousands of nodes have shown to generalize extremely
well. The complexity bounds from the classical statistical learning theory are
manifestly inadequate tools to study the generalization properties of wide and
deep neural networks [1]. Consequently, the last decades have seen the
flourishing of novel bounds and techniques aiming at explaining the undeniable
success of overparameterized networks.

On the one hand a tremendously large amount of trainable parameters makes a
direct study of the network dynamics extremely hard. On the other hand, the
literature is rich with simple approximations for the limit of infinite width
[2, 3, 4, 5, 6, 7, 8, 9]. For a fully connected multi-layer architecture, the
limit consists in assuming that each layer includes an infinite number of
nodes. However, alternative definitions of “width” makes sense of this limit
for a wide range of architectures [4].
Although unachievable in practice, infinitely wide networks feature many
interesting properties which can help grasp the complex behaviour of large but
finite-size architectures. One most remarkable fact is that infinitely wide
networks behave like Gaussian processes when all the parameters are
independently randomly initialized with the proper scaling. This simple
Gaussian behaviour, which has the major advantage of being fully captured by
the mean and the covariance functions, is valid both before [10, 11] and
during training [12].

In the present work, we focus on establishing a similar result for a simple
stochastic architecture, featuring a single fully connected hidden layer whose
number of nodes is taken to infinity. For a stochastic network, the randomness
is not simply due to the initialization but is intrinsic, as the parameters
themselves are random variables. In our case, we will assume that each
parameter has an independent normal distribution, described by its mean and
its standard deviation. For convenience, we will henceforth refer to the set
of these means and standards deviations as the network hyper-parameters. Under
suitable conditions, the output distribution of our infinitely wide stochastic
network becomes a multivariate Gaussian, whose mean and covariance functions
$M$ and $Q$ can be expressed in terms of the means and the standard deviations
of the random parameters. By correctly tuning the hyper-parameters, we can
thus train the network.

A significant reason behind the interest in studying stochastic networks is
their role in the context of learning with guarantees, where the goal is to
get a bound on the generalization error once the training is completed,
without any held-out test dataset. Despite the numerous sophisticated
techniques from the statistical learning literature, in practice, the gap
between the generalization bounds obtained in this way and the actual error is
still relatively large. Until recently, for networks in the overparameterized
regime, non-vacuous bounds could only be achieved under strong and often
hardly verifiable hypotheses [13]. However, over the last few years, some
promising results, consisting in non-vacuous (but still quite loose) bounds,
have been obtained applying some PAC-Bayesian methods to the training of
stochastic classifiers [14, 15].

The PAC-Bayesian bounds [16, 17, 18, 19] are upper bounds on the expected
generalization error of a stochastic network, holding with high probability
over the choice of the training dataset. The underlying idea is that if the
distribution of the network random parameters does not change much during the
training, then the trained model should not be prone to overfit. The
“distance” between the initial and final network configurations is measured
via the Kullback-Leibler divergence $\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})$
between a data-independent prior distribution $\mathcal{P}$ on the network
stochastic parameters, fixed before the training, and the final (posterior)
distribution $\mathcal{Q}$, obtained with the training. We call a training
method PAC-Bayesian when its goal is to optimize a PAC-Bayesian bound.
However, as long as the distribution of the network output is unknown, this
optimization cannot be tackled directly, and a bounded and differentiable
surrogate of the 0-1 loss needs to be used [14, 15].

Here, we propose to use our Gaussian approximation to train a wide stochastic
network in a way that makes it possible to optimize the PAC-Bayesian bound
directly. We empirically show that wide stochastic networks trained in this
way can achieve better performances than those obtained by training via the
standard PAC-Bayesian methods.

##  2 Stochastic networks

Consider a simple network $\mathbb{R}^{p}\to\mathbb{R}^{q}$ with a single
hidden layer

$F(x)=W^{1}\,\phi(W^{0}x)\,,$

where $W^{1}$ is a $q\times n$ matrix, $W^{0}$ a $n\times p$ matrix, and
$\phi$ the activation function. The network is stochastic, in the sense that
the components of $W^{0}$ and $W^{1}$ are random variables. Note that this is
not only true at initialization, conversely the network is assumed to be
intrinsically stochastic. More precisely, let $(\zeta^{1}_{ij})_{i=1\dots
q}^{j=1\dots n}$ and $(\zeta^{0}_{jk})_{j=1\dots n}^{k=1\dots p}$ be
independent families of iid standard normal random variables. We assume that
the stochastic parameters are distributed as

$W^{1}_{ij}=\tfrac{1}{\sqrt{n}}(\mathfrak{s}^{1}_{ij}\zeta^{1}_{ij}+\mathfrak{m}^{1}_{ij})\,;\qquad
W^{0}_{jk}=\tfrac{1}{\sqrt{p}}(\mathfrak{s}^{0}_{jk}\zeta^{0}_{jk}+\mathfrak{m}^{0}_{jk})\,.$

The means $\mathfrak{m}$’s and the standard deviations $\mathfrak{s}$’s of the
random parameters are what we will call the hyper-parameters of the network.
Given their values at initialization (which might be randomly drawn) they are
deterministic variables.

We are interested in the limit of large $n$, meaning that the network is
assumed to be very large. In the present work we aim at showing that, as
$n\to\infty$, for each fixed input $x$ $F(x)$ converges to a multivariate
normal, whose covariance matrix $Q(x)\in\mathbb{R}^{q}\times\mathbb{R}^{q}$
and mean vector $M(x)\in\mathbb{R}^{q}$ are deterministic functions of the
hyper-parameters $\mathfrak{m}$’s and $\mathfrak{s}$’s. In short, for any
fixed input $x$, we want to show that

$F(x)\to\mathcal{N}(M(x),Q(x))\,.$ (1)

Having an explicit form for the law of the output of the network can be
extremely useful. Indeed, we might be able to evaluate the expected value of
the empirical loss analytically. Moreover, the stochasticity of the network
might turn out to be an advantage for optimization. For example, the 0-1 loss
expected value can have a non-zero gradient and hence allow training the
network with gradient descent, something usually not possible for
deterministic architectures.

###  2.1 Infinite width limit

We start by focusing on the hidden layer $Y^{0}$, whose nodes can be expressed
as

$Y^{0}_{j}(x)=\sum_{k=1}^{p}W^{0}_{jk}x_{k}=\frac{1}{\sqrt{p}}\sum_{k=1}^{p}\mathfrak{s}^{0}_{jk}\zeta^{0}_{jk}x_{k}+\frac{1}{\sqrt{p}}\sum_{k=1}^{p}\mathfrak{m}^{0}_{jk}x_{k}\,$

for any fixed input $x\in\mathbb{R}^{p}$. As the $\zeta^{0}$’s are iid
standard Gaussian random variables, we can write the first sum as

$\frac{1}{\sqrt{p}}\sum_{k=1}^{p}\mathfrak{s}^{0}_{jk}\zeta^{0}_{jk}x_{k}=\left(\frac{1}{p}\sum_{k=1}^{p}(\mathfrak{s}^{0}_{jk}x_{k})^{2}\right)^{1/2}\bar{\zeta}^{0}_{j}\,,$

where $(\bar{\zeta}^{0}_{j})_{j=1\dots n}$ are independent standard normals.
Let us define

$Q^{0}_{jj^{\prime}}(x)=\delta_{jj^{\prime}}\frac{1}{p}\sum_{k=1}^{p}(\mathfrak{s}^{0}_{jk}x_{k})^{2}\,,\qquad\qquad
M^{0}_{j}(x)=\frac{1}{\sqrt{p}}\sum_{k=1}^{p}\mathfrak{m}^{0}_{jk}x_{k}\,,$

so that

$Y^{0}(x)\sim\mathcal{N}(M^{0}(x),Q^{0}(x))$

is a $n$-dimensional multivariate normal with mean vector $M^{0}(x)$ and
covariance matrix $Q^{0}(x)$. As $Q^{0}(x)$ is diagonal, all the components of
$Y^{0}(x)$ are independent, and we can write

$Y^{0}_{j}(x)=\sqrt{Q^{0}_{jj}(x)}\,\bar{\zeta}^{0}_{j}+M^{0}_{j}(x)\,,$ (2)

where the $\bar{\zeta}^{0}_{j}$’s are independent standard Gaussians.

Now define the random variable

$\Phi^{0}_{j}(x)=\phi(Y^{0}_{j}(x))\,.$

Clearly we have $F_{i}(x)=\sum_{j=1}^{n}W^{1}_{ij}\Phi^{0}_{j}(x)$. Expanding
the components of $W^{1}$ we can write

$F_{i}(x)=\frac{1}{\sqrt{n}}\sum_{j=1}^{n}\mathfrak{s}_{ij}^{1}\zeta_{ij}^{1}\Phi^{0}_{j}(x)+\frac{1}{\sqrt{n}}\sum_{j=1}^{n}\mathfrak{m}_{ij}^{1}\Phi_{j}^{0}(x)\,.$

For any fixed input $x$, in the limit $n\to\infty$, we have an infinite sum of
independent random variables, which are not identically distributed. In order
to claim the convergence to a multivariate normal distribution, we will need
to control the variance and some higher moment of these variables, meaning
that the hyper-parameters must have the correct order of magnitude. This is
possible if the network is correctly initialized, and we can prove that the
result is true even during training if the hyper-parameters keep close enough
to their initial values.

Once the Gaussian limit established, we can explicitly evaluate both the mean
vector $M(x)$ and the covariance matrix $Q(x)$ of the output. For the mean, we
have

$M_{i}(x)=\mathbb{E}[F_{i}(x)]=\frac{1}{\sqrt{n}}\sum_{j=1}^{n}\mathfrak{m}_{ij}^{1}\mathbb{E}[\Phi_{j}^{0}(x)]\,.$
(3)

As for $Q(x)$, we have
$Q_{ii^{\prime}}(x)=\mathbb{C}_{ii^{\prime}}[F(x)]=\mathbb{E}[F_{i}(x)F_{i^{\prime}}(x)]-\mathbb{E}[F_{i}(x)]\mathbb{E}[F_{i^{\prime}}(x)]$,
which with a quick manipulation becomes

$Q_{ii^{\prime}}(x)=\delta_{ii^{\prime}}\frac{1}{n}\sum_{j=1}^{n}(\mathfrak{s}^{1}_{ij})^{2}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]+\frac{1}{n}\sum_{j=1}^{n}\mathfrak{m}^{1}_{ij}\mathfrak{m}^{1}_{i^{\prime}j}\mathbb{V}[\Phi^{0}_{j}(x)]\,,$
(4)

where we used the fact that the nodes of the hidden layer are independent.

We conclude by stating some rigorous results whose proofs are postponed to the
appendix.

######  Proposition 1.

Fix an input $x$, let $Z(x)\sim\mathcal{N}(M(x),Q(x))$ and denote as
$\mathcal{C}$ the class of measurable convex subsets of $\mathbb{R}^{q}$. Then

$\sup_{C\in\mathcal{C}}|\mathbb{P}(F(x)\in C)-\mathbb{P}(Z(x)\in C)|\leq\kappa
q^{1/4}\frac{B}{\sqrt{n}}\,,$

where $\kappa\leq 4$ is an absolute constant and

$B\leq
q^{1/2}\frac{\sum_{j=1}^{n}\sum_{i=1}^{q}(2|\mathfrak{s}^{1}_{ij}|^{3}+8|\mathfrak{m}^{1}_{ij}|^{3})\mathbb{E}[|\Phi^{0}_{j}(x)|^{3}]}{\sum_{j=1}^{n}\mathbb{E}[\Phi_{j}^{0}(x)^{2}]\min_{i=1\dots
q}(\mathfrak{s}_{ij}^{1})^{2}}\,.$

In particular, if $B=o(\sqrt{n})$ for $n\to\infty$, then $F(x)\to Z(x)$ in
distribution.

As a corollary of the above result, if the stochastic network acts as a
classifier, its performance is related to the one of its Gaussian
approximation.

######  Corollary 1.

Assume that the network deals with a classification problem, where for each
instance $x$ there is a single correct label $y=f(x)\in\\{1\dots q\\}$. Then,
for each fixed input $x\neq 0$, denoting as
$\hat{f}(x)=\operatorname*{argmax}_{i=1\dots q}F_{i}(x)$ and
$\bar{f}(x)=\operatorname*{argmax}_{i=1\dots q}Z_{i}(x)$, we have

$|\mathbb{P}(\hat{f}(x)=f(x))-\mathbb{P}(\bar{f}(x)=f(x))|\leq\kappa
q^{1/4}\frac{B}{\sqrt{n}}\,.$

###  2.2 Initialization and lazy training

With a suitable random initialization of the hyper-parameters and in a lazy
training regime, we show that, as $n\to\infty$, our stochastic network has a
Gaussian limit, in the sense that the quantity $B/\sqrt{n}$ of Proposition 1
vanishes as $n\to\infty$. For simplicity, we will assume that the activation
function $\phi:\mathbb{R}\to\mathbb{R}$ is Lipshitz continuous.

We let all the network hyper-parameters be independently initialized in the
following way:

$\displaystyle\mathfrak{m}_{jk}^{0}\sim\mathcal{N}(0,1)\,;\qquad$
$\displaystyle\mathfrak{m}_{ij}^{1}\sim\mathcal{N}(0,1)\,;$ (5)
$\displaystyle\mathfrak{s}_{jk}^{0}=1\,;\qquad$
$\displaystyle\mathfrak{s}_{ij}^{1}=1\,.$

For convenience we will henceforth denote as $\hat{\mathbb{P}}$ the
probability measure representing the above initialization, while $\mathbb{P}$
denotes the probability measure describing the intrinsic stochasticity of the
network.

######  Proposition 2.

Consider a network initialized according to $\hat{\mathbb{P}}$ as in (5), and
whose activation function $\phi$ is Lipshitz continuous. Then, for any fixed
input $x\neq 0$ and letting $B$ be defined as in Proposition 1, we have
$\frac{B}{\sqrt{n}}\to 0$ as $n\to\infty$, in probability with respect to the
random initialization $\hat{\mathbb{P}}$. More precisely,
$B=O_{\hat{\mathbb{P}}}(1)$ as $n\to\infty$.
In particular, the network tends to a Gaussian limit at initialization, in
distribution wrt the intrinsic stochasticity $\mathbb{P}$ and in probability
wrt $\hat{\mathbb{P}}$.

The next proposition states that the limit will still be valid as long as the
hyper-parameters do not move too much from their initialization (lazy
training).

######  Proposition 3.

Fix a constant $J>0$ independent of $n$ and assume that $\phi$ is Lipshitz.
Consider a network with initial configuration
$(\widetilde{\mathfrak{m}},\widetilde{\mathfrak{s}})$ drawn according to
$\hat{\mathbb{P}}$, as in (5). Now consider a new configuration
$(\mathfrak{m},\mathfrak{s})$ for the network, such that

$\|\mathfrak{m}^{0}-\widetilde{\mathfrak{m}}^{0}\|^{2}_{F,2}+\|\mathfrak{m}^{1}-\widetilde{\mathfrak{m}}^{1}\|_{F,2}^{2}+\|\mathfrak{s}^{0}-\widetilde{\mathfrak{s}}^{0}\|^{2}_{F,2}+\|\mathfrak{s}^{1}-\widetilde{\mathfrak{s}}^{1}\|_{F,2}^{2}\leq
J^{2}\,,$ (6)

where $\|\cdot\|_{F,2}$ denotes the 2-Frobenius norm of a matrix. Let $B$ be
defined as in Proposition 1, referring to the configuration
$(\mathfrak{m},\mathfrak{s})$. For any fixed input $x\neq 0$ we have
$B=O_{\hat{\mathbb{P}}}(1)$ as $n\to\infty$, in probability with respect to
the random initialization $\hat{\mathbb{P}}$.

In the next section, we will see that the lazy training constraint can be
restated in terms of a bound on the Kullback-Leibler divergence between the
initial and final distributions of the stochastic parameters. This fact will
allow us to ensure that the limit is valid when training the network to
optimize a PAC-Bayesian objective.

##  3 PAC-Bayesian framework

###  3.1 PAC-Bayes basics

Consider a standard classification problem, where to each instance
$x\in\mathcal{X}\subseteq\mathbb{R}^{p}$ corresponds a label
$y\in\mathcal{Y}=\\{1\dots q\\}$. The goal is to build an algorithm able to
predict $y$, given $x$. We assume that for each $x$ there is a single correct
label $y=f(x)$ and that the $x$’s are distributed according to some
probability measure $\mathbb{P}_{X}$ on $\mathcal{X}$. To train our algorithm,
we have access to a sample $S=(X_{h})_{h=1\dots m}$, which is correctly
labelled, that is for every $X_{h}\in S$ we know $f(X_{h})$. Each $X_{h}$ is
an independent draw from $\mathbb{P}_{X}$, so that
$\mathbb{P}_{S}=\mathbb{P}_{X}^{\otimes m}$. In order to evaluate the
performance of the network during and after the training, we let $\ell$ be the
0-1 loss, that is $\ell(\hat{y},y)=0$ if $\hat{y}=y$, $1$ otherwise.

We assume our algorithm to be a neural network whose set of parameters we will
denote as $w$. We write as $\hat{f}_{w}(x)$ the prediction of a network with
parameter configuration $w$ for the instance $x$. The empirical loss
$L_{S}(w)$ is defined as the average of the loss of the network on the
training set, $L_{S}(w)=\frac{1}{m}\sum_{x\in S}\ell(\hat{f}_{w}(x),f(x))$.
This quantity can be explicitly evaluated since we have access to the true
labels on the training set. However, in order to judge how the algorithm will
perform on unseen examples, we need to deal with the true loss,
$L_{X}(w)=\mathbb{E}_{X}[\ell(\hat{f}_{w}(X),f(X))]$, which is the expected
value of the loss under $\mathbb{P}_{X}$.

The PAC-Bayesian framework [16, 17, 18, 19] deals with stochastic networks,
architectures whose parameters are random variables. We consider a prior
probability measure $\mathcal{P}$ on the random parameters. $\mathcal{P}$ has
to be chosen independently of the training dataset $S$. After the training,
the parameters will be described by a new probability measure $\mathcal{Q}$
(the so-called posterior), clearly $S$-dependent. The idea is that if
$\mathcal{P}$ and $\mathcal{Q}$ are not too “far” from each other, then the
network will probably generalize well.

Given the probability measure $\mathcal{Q}$ on the parameter space, define the
expected true loss
$L_{X}(\mathcal{Q})=\mathbb{E}_{W\sim\mathcal{Q}}[L_{X}(W)]$ and the expected
empirical loss $L_{S}(\mathcal{Q})=\mathbb{E}_{W\sim\mathcal{Q}}[L_{S}(W)]$.
The PAC-Bayes bounds are upper bounds on $L_{X}(\mathcal{Q})$, which holds
with high probability $\mathbb{P}_{S}$ on the choice of the training sample
$S$. They usually involve the expected empirical error $L_{S}(\mathcal{Q})$,
and a divergence term in the form
$\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})=\mathbb{E}_{\mathcal{Q}}[\log(\mathrm{d}\mathcal{P}/\mathrm{d}\mathcal{Q})]$.
Essentially we will use the following result, from [20].

######  Proposition 4.

Fix a data-independent prior $\mathcal{P}$. With probability higher than
$1-\delta$ on the choice of the training set $S=(X_{h})_{h=1\dots m}$222Here
we assume that the training set $S$ has size $m\geq 8$.

$\mathrm{kl}(L_{X}(\mathcal{Q})\|L_{S}(\mathcal{Q}))\leq\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{m}$

for any posterior $\mathcal{Q}$, where $\mathrm{kl}(u\|v)$ denotes the
$\mathrm{K\\!L}$ divergence between two Bernoulli distributions, with means
$u$ and $v$ respectively.

If we define

$\mathrm{kl}^{-1}(u|c)=\sup\\{v\in[0,1]:\mathrm{kl}(u\|v)\leq c\\}\,,$ (7)

the next corollary follows immediately.

######  Corollary 2.

Fix a data-independent prior $\mathcal{P}$. With probability higher than
$1-\delta$ on the choice of the training set $S=(X_{h})_{h=1\dots m}$2

$L_{X}(\mathcal{Q})\leq\mathrm{kl}^{-1}\left(L_{S}(\mathcal{Q})\bigg{|}\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{m}\right)$
(8)

for any posterior $\mathcal{Q}$.

We can hence imagine the following training algorithm [16]:

  * •

Fix $\delta\in(0,1)$ and a prior $\mathcal{P}$ for the network stochastic
parameters.

  * •

Collect a sample $S$ of $m$ iid datapoints.

  * •

Compute the optimal posterior $\mathcal{Q}$ minimizing (8).

  * •

Implement a stochastic network characterized by the law $\mathcal{Q}$.

Clearly, in general it will not be possible to find an explicit form for the
optimal $\mathcal{Q}$.

###  3.2 PAC-Bayesian training

In practice, we need an explicit form for the prior distribution of the
parameters. Following the approach of [14], we assume that both $\mathcal{P}$
and $\mathcal{Q}$ are multivariate normal distributions with diagonal
covariance matrices. In other words, the random parameters of the network are
independent normal random variables. In this way, a network with $N$
stochastic parameters can be represented with $2N$ deterministic values, their
means and their variances. Let us denote as $\mathfrak{m}$ the $N$-dimensional
vector of means and $\mathfrak{s}$ the one of standard deviations of the
posterior. Let $\widetilde{\mathfrak{m}}$ and $\widetilde{\mathfrak{s}}$ be
the same quantities for the prior. In short we have
$\mathcal{P}=\mathcal{N}(\widetilde{\mathfrak{m}},\operatorname*{diag}(\widetilde{\mathfrak{s}}^{2}))$
and
$\mathcal{Q}=\mathcal{N}(\mathfrak{m},\operatorname*{diag}(\mathfrak{s}^{2}))$.
As we are dealing with multivariate normal distributions, we have a simple
explicit form for the term $\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})$ which
appears in the PAC-Bayesian bound (8),

$\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})=\frac{1}{2}\left(\sum_{i}\left(\frac{\mathfrak{s}_{i}}{\widetilde{\mathfrak{s}}_{i}}\right)^{2}-N+\sum_{i}\left(\frac{\mathfrak{m}_{i}-\widetilde{\mathfrak{m}}_{i}}{\widetilde{\mathfrak{s}}_{i}}\right)^{2}+2\sum_{i}\log\frac{\widetilde{\mathfrak{s}}_{i}}{\mathfrak{s}_{i}}\right)\,.$
(9)

In this setting, the most troublesome part of (8) is the term
$L_{S}(\mathcal{Q})$, which in general cannot be computed explicitly. However,
by sampling a from $\mathcal{Q}$, we might get some Monte Carlo (MC) estimate
$\hat{L}_{S}(\mathcal{Q})$ of this term and hence obtain an estimator of the
RHS of (8). Still, given such an estimate, the minimization problem is far
from being solved. A possible approach [14, 15] is to use some standard
gradient-based descent method to minimize the PAC-Bayesian objective.

It is easily checked that (9) is differentiable with respect to $\mathfrak{m}$
and $\mathfrak{s}$ (which are the trainable hyper-parameters of the
posterior). However, a fundamental issue is that $\hat{L}_{S}(\mathcal{Q})$
will have a null gradient almost everywhere, as this is the case for
$L_{S}(w)$ for each realization $w$ used in the estimate. Consequently, we
need to use some surrogate of the 0-1 loss, such as the cross-entropy loss
(CEL). Still, there is one major problem: the unboundedness of CEL. Indeed, a
more general formulation of Proposition 4 holds for losses bounded in $[0,1]$,
but not for the unbounded case. To circumvent this issue [14] uses the fact
that $\text{CEL}/\log 2$ is an upper bound for the 0-1 loss in binary
classification, whilst [15] clamps the output probabilities of the network so
that they are larger than some fixed value.

It must be said that although $\hat{L}_{S}(\mathcal{Q})$ has a null gradient,
this is not the case for $L_{S}(\mathcal{Q})$. Hence, if we know exactly the
output distribution of the stochastic network, we might be able to use the 0-1
loss directly without the use of any surrogate. This is indeed the case in the
Gaussian approximation, as we will see in the next section.

It is worth mentioning that similar considerations hold when using an almost
everywhere constant activation function to train a stochastic network. Indeed,
[21, 22] developed an interesting variant of PAC-Bayesian training for binary
classifiers with the binary activation function ($\phi=\text{sign}$). In this
setting, the simple form of the output of each hidden layer allows for a more
explicit expression of the distribution of the hidden nodes, which permits
overcoming the fact that the binary activation function is non-differentiable.

A further difficulty of PAC-Bayesian training is the fact that the function
$\mathrm{kl}^{-1}$ does not have a simple formula and requires numerical
approximation. For the purpose of gradient-based optimization, it is common to
use explicit upper bounds on this quantity, which have a simpler form. For
instance [14] uses that $\mathrm{kl}^{-1}(u|v)\leq u+\sqrt{v/2}$, whilst [15]
compares this choice with two further upper bounds. However, in the present
work, we propose to use $\mathrm{kl}^{-1}$ directly, getting its gradient by
backpropagating through the Newton method used in its evaluation, which indeed
consists in the iterated differentiable operations. An alternative approach
would be to use that $\mathrm{kl}^{\prime}$ has a simple closed-form
expression and
$(\mathrm{kl}^{-1})^{\prime}=1/(\mathrm{kl}^{\prime}\circ\mathrm{kl}^{-1})$,
where $\mathrm{kl}^{-1}$ can be evaluated via the Newton method. From the
numerical results, it seems that the fact of using $\mathrm{kl}^{-1}$ during
the training is not advantageous when training with CEL, but it is when the
0-1 loss is used.

##  4 PAC-Bayesian training in the Gaussian limit

Instead of doing the standard PAC-Bayesian training with a surrogate loss, we
can train our wide stochastic network by assuming that its Gaussian
approximation is exact. However, once the training over, we will need to
evaluate the final bound without such an assumption.

At the initialization, we know that for a network initialized according to
$\hat{\mathbb{P}}$, as in (5), the Gaussian approximation is asymptotically
true for large $n$. Moreover, training to optimize a PAC-Bayesian bound
ensures that the $\mathrm{K\\!L}$ keeps small during training, independently
of $n$. The following lemma assures that controlling the $\mathrm{K\\!L}$ is
sufficient to claim that the network is in the lazy training regime (6).
Consequently, a wide stochastic network is asymptotically Gaussian throughout
the PAC-Bayesian training.

######  Lemma 1.

Define the multivariate Gaussian distributions
$\mathcal{P}=\mathcal{N}(\widetilde{\mathfrak{m}},\operatorname*{diag}(\widetilde{\mathfrak{s}}^{2}))=N(\widetilde{\mathfrak{m}},\operatorname*{Id})$
and
$\mathcal{Q}=\mathcal{N}(\mathfrak{m},\operatorname*{diag}(\mathfrak{s}^{2}))$
for the parameters of a stochastic network. We have

$\|\mathfrak{m}^{0}-\widetilde{\mathfrak{m}}^{0}\|^{2}_{F,2}+\|\mathfrak{m}^{1}-\widetilde{\mathfrak{m}}^{1}\|_{F,2}^{2}+\|\mathfrak{s}^{0}-\widetilde{\mathfrak{s}}^{0}\|^{2}_{F,2}+\|\mathfrak{s}^{1}-\widetilde{\mathfrak{s}}^{1}\|_{F,2}^{2}\leq
2\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})\,.$

###### Proof.

It suffices to write down explicitly the $\mathrm{K\\!L}$ as in (9). We
conclude noticing that, for all $u>0$, $u^{2}-1-2\log u\geq(u-1)^{2}$. ∎

We start by showing in the next section that it is possible to get a
meaningful gradient from the expected 0-1 loss in the Gaussian limit. Then we
will see that simple activation functions, such as $\phi=\operatorname*{ReLU}$
or $\phi=\sin$, allows the evaluation of the gradients of the output mean and
covariance with respect to the hyper-parameters. Finally, we will discuss how
to obtain a PAC-Bayesian bound after the training.

###  4.1 Training with the 0-1 loss

Consider a Gaussian stochastic network whose output has law (1),
$F(x)\sim\mathcal{N}(M(x),Q(x))$. We assume that we have a classification
problem where for each instance $x\in\mathcal{X}\subseteq\mathbb{R}^{d}$, we
need to establish the correct label $y=f(x)$. There are $q$ possible labels
$i=1\dots q$, and the prediction of the network is given by
$\hat{f}(x)=\operatorname*{argmax}_{i=1\dots q}F_{i}(x)$. We define the 0-1
loss as

$\ell(\hat{y},y)=\begin{cases}0\text{ if }\hat{y}=y\,;\\\ 1\text{
otherwise.}\end{cases}$

It is natural to aim at minimizing $\mathbb{E}[\ell(\hat{f}(x),f(x))]$ (where
the expectation is over the stochastic parameters of the network) since this
quantity is actually equal to the probability of making a mistake in
classifying $x$, $\mathbb{P}(\hat{f}(x)\neq f(x))$. We aim at tackling this
problem by performing gradient descent (GD), assuming that $M(x)$ and $Q(x)$
are differentiable functions of the network trainable hyper-parameters. To do
so, we need to evaluate $\nabla_{M}\mathbb{E}[\ell(\hat{f}(x),y)]$ and
$\nabla_{Q}\mathbb{E}[\ell(\hat{f}(x),y)]$.

Note that $\ell(\hat{f}(x),y)$ has a null gradient almost everywhere for any
random realization of the network. For this reason, a non-stochastic network
cannot be trained directly with the 0-1 loss. However, this is not the case
for a stochastic network, as the gradients of the expected loss are
informative. The reason for that can be intuitively explained by thinking of
what happens if we try to “go down the stairs” with GD. A single realization
of the network will be a point on a horizontal step: there is no way to
understand the right direction in order to go down. However, if we consider
the whole stochastic distribution of the network, it spreads over all the
steps, and it has a global view of the stairs. It is hence not surprising that
the gradient of the expected loss is non-zero.

####  4.1.1 Binary classification ($q=2$)

For binary classification, we can write the expected 0-1 loss as

$\mathbb{E}[\ell(\hat{f}(x),1)]=\mathbb{P}(F_{2}(x)>F_{1}(x))\qquad\text{and}\qquad\mathbb{E}[\ell(F(x),2)]=\mathbb{P}(F_{1}(x)>F_{2}(x))\,.$

This quantities can be computed explicitly. Indeed we have the equality in law

$F_{2}(x)-F_{1}(x)=\sqrt{Q_{11}(x)+Q_{22}(x)-2Q_{12}(x)}\,\zeta+(M_{2}(x)-M_{1}(x))\,,$

where $\zeta\sim\mathcal{N}(0,1)$. So, we find

$\mathbb{E}[\ell(\hat{f}(x),1)]=\mathbb{P}\left(\zeta>\frac{M_{1}(x)-M_{2}(x)}{\sqrt{Q_{11}(x)+Q_{22}(x)-2Q_{12}(x)}}\right)$

and

$\mathbb{E}[\ell(\hat{f}(x),2)]=\mathbb{P}\left(\zeta>\frac{M_{2}(x)-M_{1}(x)}{\sqrt{Q_{11}(x)+Q_{22}(x)-2Q_{12}(x)}}\right)\,.$

Clearly the two quantities above can be written explicitly in terms of the
error function erf, as $\zeta$ is distributed as a standard normal and
$\mathbb{P}(\zeta>u)=\frac{1}{2}(1-\text{erf}\,(u/\sqrt{2}))$. It is then
clear that $\mathbb{E}[\ell(\hat{f}(x),y)]$ is differentiable with respect to
$M$ and $Q$, with non-trivial derivatives.

####  4.1.2 Multiclass classification ($q>2$)

When there are more than two classes, things become more complicated. Now we
can write

$\mathbb{E}[\ell(\hat{f}(x),i^{\star})]=\mathbb{P}\left(F_{i^{\star}}(x)\leq\max_{i\neq
i^{\star}}F_{i}(x)\right)=1-\mathbb{P}\left(F_{i^{\star}}(x)>\max_{i\neq
i^{\star}}F_{i}(x)\right)\,.$

Hence, given a $q$-dimensional Gaussian vector $Y\sim\mathcal{N}(M,Q)$, we
need to find an estimate of $\mathbb{P}(Y_{i^{\star}}>\max_{i\neq
i^{\star}}Y_{i})$.
The most trivial estimator would consist of sampling different realizations of
$Y$ and then give a MC estimate. However, as we are interested in the gradient
of the expected loss, this method will not work. Indeed, the gradient of this
estimate is the sum of the gradients of the 0-1 loss of each sample. As all
these gradients are $0$, we do not obtain anything informative. We have thus
to proceed in a less naive way.

For convenience let us assume that $i^{\star}=q$ (the last label), so that we
are interested in $\mathbb{P}(Y_{q}>\max_{i<q}Y_{i})$. With a Cholesky-like
algorithm we can find a lower triangular matrix $A$ such that

$Y\sim AX+M\,,$

where $X\sim\mathcal{N}(0,\text{Id})$. We have
$Y_{i}=\sum_{i^{\prime}=1}^{q}A_{ii^{\prime}}X_{i^{\prime}}+M_{i}$ and
$A_{iq}=0$ for $i=1\dots(q-1)$, while $A_{qq}>0$. For $i<q$, we can write

$\mathbb{P}(Y_{q}>Y_{i})=\mathbb{P}\left(X_{q}>\sum_{i^{\prime}=1}^{q-1}\frac{A_{ii^{\prime}}-A_{qi^{\prime}}}{A_{qq}}\,X_{i^{\prime}}+\frac{M_{i}-M_{q}}{A_{qq}}\right)\,.$

Let us define the $(q-1)$ dimensional random vector $\tilde{X}$ as

$\tilde{X}=\tilde{A}X+\tilde{M}\,,$

where $\tilde{A}$ is a $(q-1)\times q$ matrix and $\tilde{M}$ is a $(q-1)$
vector, whose elements are given by
$\tilde{A}_{ii^{\prime}}=\frac{A_{ii^{\prime}}-A_{qi^{\prime}}}{A_{qq}}$ and
$\tilde{M}_{i}=\frac{M_{i}-M_{q}}{A_{qq}}$ repectively. With this notation we
have $\mathbb{P}(Y_{q}>Y_{i})=\mathbb{P}(X_{q}>\tilde{X}_{i})$. Now, we have
gained that $X_{q}$ is independent from all the other $X_{i}$’s, and so from
all the $\tilde{X}_{i}$’s. In short,
$(X_{q}|\tilde{X})=X_{q}\sim\mathcal{N}(0,1)$. So we can write

$\mathbb{P}\left(Y_{q}>\max_{i<q}Y_{i}\right)=\mathbb{P}\left(X_{q}>\max_{i<q}\tilde{X}_{i}\right)=\mathbb{E}\left[\mathbb{P}\left(X_{q}>\max_{i<q}\tilde{X}_{i}\bigg{|}\tilde{X}\right)\right]\,.$

Now, if we let $\psi(u)=\tfrac{1}{2}(1-\text{erf}\,(u/\sqrt{2}))$, we get

$\mathbb{P}\left(Y_{q}>\max_{i<q}Y_{i}\right)=\mathbb{E}\left[\psi\left(\max_{i<q}\tilde{X}_{i}\right)\right]\,.$

We can estimate the above expression with MC sampling. Note that it is almost
everywhere differentiable with respect to the components of $M$ and $Q$ (as
the Cholesky transform is differentiable) and the gradient with respect to $M$
and $Q$ is not trivially $0$.

Finally, for a general $i^{\star}\in\\{1\dots q\\}$, we can get
$\mathbb{P}(Y_{i^{\star}}>\max_{i\neq i^{\star}}Y_{i})$ by simply performing a
swap of the two labels $i^{\star}$ and $q$, and then apply the method for
$i^{\star}=q$.

###  4.2 Differentiation of $M$ and $Q$

We have so far established that we can effectively differentiate the expected
loss with respect to $M$ and $Q$. Still, to train the network we will need to
evaluate the gradients with respect to the hyper-parameters $\mathfrak{m}$’s
and $\mathfrak{s}$’s.

Now recall that $\Phi^{0}_{k}(x)$ is in the form $\phi(a\zeta+b)$, with
$a=\sqrt{Q^{0}_{kk}(x)}$, $b=M^{0}_{k}(x)$, and $\zeta\sim\mathcal{N}(0,1)$.
When the activation function $\phi$ is simple enough,
$\mathbb{E}[\phi(a\zeta+b)]$ and $\mathbb{E}[\phi(a\zeta+b)^{2}]$ have
explicit analytic forms as functions of $a$ and $b$. Exploiting this fact, it
is possible to evaluate the derivatives of $M$ and $Q$ with respects to the
$\mathfrak{m}^{0}$’s and $\mathfrak{s}^{0}$’s, needed in order to train the
network with gradient-based methods. This if for instance the case for
$\phi=\operatorname*{ReLU}$ and $\phi=\sin$. We refer to the appendix for more
details.

###  4.3 Final computation of the bound

Once training is complete, we need to do without the Gaussianity assumption to
compute the final bound. The idea is to follow the same approach as [14] and
[15].

Let $W_{1}\dots W_{N}$ be $N$ independent samples of the whole set of network
stochastic parameters, drawn according to $\mathcal{Q}$. We consider a MC
estimate of $L_{S}(\mathcal{Q})$, defined as

$\hat{L}_{S}(\mathcal{Q})=\frac{1}{N}\sum_{h=1}^{N}L_{S}(W_{h})\,.$

Given $\delta^{\prime}\in(0,1)$, with probability at least $1-\delta^{\prime}$
we have, [23],

$L_{S}(\mathcal{Q})\leq\mathrm{kl}^{-1}\left(\hat{L}_{S}(\mathcal{Q})\big{|}\tfrac{1}{N}\log\tfrac{2}{\delta^{\prime}}\right)\,,$
(10)

where $\mathrm{kl}^{-1}$ is defined in (7). As $\mathrm{kl}^{-1}$ is an
increasing function of its first argument, we deduce from Corollary 2 that
with probability higher than $1-\delta-\delta^{\prime}$ we have

$L_{X}(\mathcal{Q})\leq\mathrm{kl}^{-1}\left(\mathrm{kl}^{-1}\left(\hat{L}_{S}(\mathcal{Q})\big{|}\tfrac{1}{N}\log\tfrac{2}{\delta^{\prime}}\right)\bigg{|}\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{m}\right)\,.$
(11)

It is possible to avoid sampling all the $p\times n^{2}\times q$ stochastic
parameters to obtain a realization of the network. Indeed, as all the
parameters are independent Gaussian, we have that each hidden node is
distributed according to (2)

$Y^{0}_{j}(x)=\sqrt{Q^{0}_{jj}(x)}\,\bar{\zeta}^{0}_{j}+M^{0}(x)\,,$

where the $\bar{\zeta}^{0}_{j}$’s are independent standard normals. Hence, for
each of the $n$ hidden nodes we can sample from one single Gaussian (instead
of from $p$ ones).
Moreover, given the hidden layer, the output as well is the sum of independent
Gaussians. We can write

$F_{k}(x)|Y^{0}=\left(\frac{1}{n}\sum_{j=1}^{n}(\mathfrak{s}^{1}_{j}\phi(Y^{0}_{j}))^{2}\right)^{1/2}\bar{\zeta}^{1}_{k}+\frac{1}{n}\sum_{j=1}^{n}\phi(Y^{0}_{j})\mathfrak{m}^{1}_{j}\,,$

where the $\bar{\zeta}^{1}_{k}$’s are iid standard normals. So, once drawn the
hidden layer, for each of the $k$ components of the output we need to draw a
single Gaussian (instead of $n$ ones).
In this way, we can draw according to the exact output law (we are not making
any approximation here) by sampling only $d+n$ random variables. This fact,
which speeds up the MC sampling procedure, can be used even when training with
the standard PAC-Bayesian approach from [14, 15].

As a final remark, we can mention that an alternative way to get an exact
result from the Gaussian approximation would be to use an upper bound such as
the one in Corollary 1 to control the finite-size correction to the expected
empirical loss. However, for networks with $O(10^{3})$ hidden nodes, like the
ones that we used in our experiments, this last approach tends to give looser
bounds compared to the method described above, at least when the number $N$ of
samples used in the MC estimate is of order $O(10^{5})$.

##  5 Experimental results

In this section, we discuss empirical results to validate our theoretical
findings. First, we compare the Gaussian predictions with the output
distribution of the nodes of a wide stochastic network. Then we report the
results obtained by training a stochastic network on MNIST, and on a binary
version of it, with PAC-Bayesian methods like those from [14, 15], and with
our Gaussian method. On both datasets, the Gaussian method led to tighter
final generalization bounds.
The PyTorch code that we wrote to run our experiments is available at the
GitHub repository https://github.com/eclerico/WideStochNet.

For the present work, we decided to keep the experimental setting as simple as
possible. Indeed, we opted for training only the means $\mathfrak{m}$’s,
keeping the standard deviations $\mathfrak{s}$’s fixed to their initial value,
similarly to what is done in [21]. Moreover, coherently with the rest of this
paper, all the networks we used

Figure 1: Distributions of the three output nodes of a wide stochastic network
trained on a toy classification task. In black the theoretical predictions.

had no bias, and the PAC-Bayesian priors were chosen in a completely data-
independent fashion (they coincided with the distribution of the network at
initialization, as suggested by [14, 15]). We leave for future work a more
comprehensive experimental study.

In all the experiments, the training consisted of optimising some PAC-Bayesian
bound via SGD with momentum parameter $0.9$. The PAC parameter $\delta$ was
always chosen equal to $0.025$.

We start by considering a toy dataset. The data points have been sampled from
three multivariate distributions (labelled as $1$, $2$, $3$) in
$\mathbb{R}^{4}$, and then projected on the unit sphere in $\mathbb{R}^{4}$. A
stochastic network with one hidden layer of $n=1200$ nodes has been trained to
predict from which of the three Gaussian clusters each point comes. The
histograms in Figure 1 represent the distributions of network output nodes,
both before and after the training. They have been obtained for a single
example by sampling the random parameter distribution $10^{6}$ times. In black
are the theoretical predictions of the Gaussian profiles in the limit of
infinite width. The agreement with the histograms seems excellent, showing
that the network is essentially Gaussian already for $O(10^{3})$ nodes.

We now focus on a binary version of the MNIST dataset, where the digits from
$0$ to $4$ are relabelled as $0$ and those from $5$ to $9$ as $1$. The
training set is the standard one from MNIST and consists of $m=60000$ images.
We considered a stochastic network with $n=1200$ hidden nodes and ReLU
activation function, initialized as in (5). Note that in our implementation,
the scaling factors $1/\sqrt{p}$ and $1/\sqrt{n}$ are absorbed in the hyper-
parameters, so that during training, we performed the gradient descent on
$\mu^{0}=\mathfrak{m}^{0}/\sqrt{p}$ and $\mu^{1}=\mathfrak{m}^{1}/\sqrt{n}$
(the standard deviations are kept fixed). We trained the network with six
different methods: three of them make use of the Gaussian approximation
introduced in the present work, and the remaining three are taken from [14]
and [15].

Figure 2: Evolution of the expected empirical loss, the $\mathrm{K\\!L}$
penalty, and the Gaussian PAC-Bayesian bound, for different training
procedures of a wide stochastic network on Binary MNIST.

We start by describing the three standard PAC-Bayesian methods. The first one
is inspired by the approach of [14]. The network is trained to optimize an
objective in the form

$L_{S}(\mathcal{Q})+\sqrt{\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{2m}}\,,$
(12)

where $L_{S}(\mathcal{Q})$ is the expectation under $\mathcal{Q}$ of the
empirical cross-entropy loss divided by $\log 2$. In the following we will
refer to this method as “S Standard” (in short “S std”).
The second and third methods are inspired by [15]. We will refer to them as “S
Lambda” (“S lbd”) and “S Quadratic” (“S quad”). “S Lambda” optimizes the
objective

$\frac{L_{S}(\mathcal{Q})}{(1-\lambda/2)}+\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{m\lambda(1-\lambda/2)}\,,$
(13)

where $\lambda\in(0,1)$ is a trainable parameter and $L_{S}(\mathcal{Q})$ the
same as above. As in [15], the training alternates the optimization of the
network hyper-parameters with that of $\lambda$.
The “S Quadratic” optimization objective is

$\left(\sqrt{L_{S}(\mathcal{Q})+\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{2m}}+\sqrt{\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{2m}}\right)^{2}\,,$
(14)

where $L_{S}(\mathcal{Q})$ is again the empirical cross-entropy rescaled by
$\log 2$.
For all three methods described so far, the network is trained by sampling at
each iteration an independent copy of the parameters.

We compared the performance of the three standard methods above with three
Gaussian methods. We call “G Standard” (“G std”) the method described in
Section 4: the PAC-Bayes bound (8) is directly optimized, choosing as
objective

$\mathrm{kl}^{-1}\left(L_{S}(\mathcal{Q})\bigg{|}\frac{\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta}}{m}\right)\,,$
(15)

where $L_{S}(\mathcal{Q})$ is the expected empirical 0-1 loss in the Gaussian
limit, which for binary classification can be evaluated directly without the
need of any sampling.
We then trained the network with the methods “G Lambda” (“G lbd”) and “G
Quadratic” (“G quad”), which optimize objectives equivalent to (13) and (14),
where $L_{S}(\mathcal{Q})$ is the expected 0-1 loss (in the Gaussian
approximation).

All of the above methods have been used to train the same stochastic network,
initialized as in (5). We tried two different learning rate (LR) schedules,
the first consisting of $10000$ epochs with LR $\eta=10^{-5}$ and the second
of $100$ epochs with $\eta=10^{-2}$, followed by $1000$ epochs with
$\eta=10^{-3}$ and $5000$ epochs with $\eta=10^{-4}$. For all the methods but
“S std”, the multi-LR schedule achieved better results than the single-LR one.

For the multi-LR schedule, Figure 2 illustrates the evolution during the
training of different quantities, evaluated as if the Gaussian approximation
was exact. The first row shows the expected bound on the true error, given by
(15) where $L_{S}(\mathcal{Q})$ is always the expected empirical 0-1 loss (in
the Gaussian approximation). The second row depicts the values of
$L_{S}(\mathcal{Q})$, whilst the plots in the last row are the
$\mathrm{K\\!L}$ penalties
$(\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta})/m$.

Table 2 report the final results of the experiment. All the reported values
refer to the multi-LR schedule, except for “S std”, for which we obtained a
better final bound with the single-LR schedule. The column “Bound” indicates
to the final PAC-Bayesian bound (11). For the upper bound (10) on the
empirical error, we used $N=150000$ independent realizations of the net and
$\delta^{\prime}=0.01$. As $\delta=0.025$, we have that

Table 1: Binary MNIST Method | Bound | G Bound | G Loss | Penalty
---|---|---|---|---
G std | .1772 | .1741 | .0673 | .0494
G lbd | .1822 | .1790 | .0891 | .0324
G quad | .1826 | .1793 | .0897 | .0322
S std | .1921 | .1889 | .0716 | .0559
S lbd | .1854 | .1822 | .0519 | .0751
S quad | .1853 | .1821 | .0517 | .0752

Table 2: MNIST Method | Bound | G Bound | G Loss | Penalty
---|---|---|---|---
G std | .2807 | .2773 | .1114 | .0820
S std | .4158 | .4120 | .3265 | .0155
S lbd | .3736 | .3699 | .2717 | .0216
S quad | .3735 | .3698 | .2715 | .0217

the final generalization bounds hold with probability higher than
$1-\delta-\delta^{\prime}=0.965$ on the random selection of the training set.
The two central columns of the table refer to quantities evaluated within the
Gaussian approximation: “G Bound” is the bound given by (8) and “G Loss” is
the expected 0-1 loss. The last column “G Penalty” is the $\mathrm{K\\!L}$
penalty
$(\mathrm{K\\!L}(\mathcal{Q}\|\mathcal{P})+\log\frac{2\sqrt{m}}{\delta})/m$.
All three Gaussian methods gave a tighter final bound than their standard
counterparts, the best result being achieved by the “G Standard“ method.

Note that the result obtained by “S std” is slightly worse than the one from
[14], where for a similar network of $1200$ hidden nodes a final bound of
$.179$ was obtained, whilst our result is $.1921$. However, our setting is
more simplistic: our network has no bias, the standard deviations are not
trained, and there is no choice of the optimal prior among different
initializations.

Finally, we report the results of a similar experiment run on the MNIST
dataset (with the original $10$ labels). The network is essentially the same
one used for the Binary MNIST dataset, with $1200$ hidden nodes and ReLU
activation function. The main difference is that now we have $10$ output
nodes.
We considered four different training methods. “S Standard” (“S std”), “S
Lambda” (“S lbd”), and “S Quadratic” (“S quad”) correspond to the training
objectives (12), (13), and (14), but with a different $L_{S}(\mathcal{Q})$. As
we were not dealing anymore with binary classification, we used a bounded
version of the cross-entropy loss, as in [15]. To achieve this, we fixed
$p_{0}=10^{-5}$ and constrained the probabilities in the definition of the
cross-entropy to be greater or equal than $p_{0}$, see [15] for more details.
In this way, the loss is bounded by $\log(1/p_{0})$, and by rescaling it of
the same factor we can get a loss bounded in $[0,1]$. $L_{S}(\mathcal{Q})$ is
the expected value under $\mathcal{Q}$ of this bounded cross-entropy, averaged
on the training set. As we previously did for the experiment on the binary
MNIST, during the training we estimated $L_{S}(\mathcal{Q})$ by sampling once
per iteration all the hyper-parameters of the network.
The method “G Standard” (“G std”) uses the Gaussian approximation. We trained
the network with the optimization objective (15) where $L_{S}(\mathcal{Q})$ is
the expected empirical 0-1 loss. As previously discussed, for multiclass
classification we do not have a simple expression for the 0-1 loss, but we can
use the method described in Section 4.1.2. Per each iteration, the loss has
been evaluated by an MC estimate of $10^{4}$ independent drawings.

On the MNIST dataset, all the training procedures consisted of $10000$ epochs
with learning rate $\eta=10^{-5}$. Again the training dataset has size
$m=60000$ and the PAC parameter $\delta$ is equal to $0.025$.

Table 2 reports the results of the experiment on MNIST, where for the estimate
of the final bound we used $N=150000$ and $\delta^{\prime}=0.01$. Overall, the
Gaussian method obtained a tighter final bound, with almost a $.1$-gap with
those achieved by the other procedures. It is worth noticing that the
penalties of the standard methods are much lower than the respective losses,
something which was not the case in Table 2. We conjecture that this behaviour
is due to the different rescaling of the cross-entropy loss. Indeed, during
the training on MNIST, we divided the loss by $\log 10^{5}$, while it was
divided by $\log 2$ on Binary MNIST. Hence, in the MNIST experiment, we expect
the $\mathrm{K\\!L}$ penalty to weigh much more during the training. On the
other hand, this is not the case for the Gaussian method, as the loss does not
require any rescaling.

##  6 Conclusions

In the present work, we derived a Gaussian limit for a simple 1-layer
stochastic architecture, and we pointed out how this result can be used in
practice for the PAC-Bayesian training of wide networks. We first proved the
validity of the limit rigorously at initialization and in a lazy training
regime, then showed empirically that the proposed training method can
outperform some standard PAC-Bayesian training procedures.

For the sake of clarity, we focused on an elementary network, but we believe
that the results presented here can be extended in several ways. For instance,
it is straightforward to see that similar conclusions can be derived for an
architecture with biases. A perhaps more interesting direction worth exploring
is the case of multi-layer architectures, a topic on which we are currently
working. It is not hard to derive a recurrence relation for each layer’s mean
and covariance matrix, but the proof of convergence to a multivariate Gaussian
becomes more tricky as the nodes of every hidden layer, apart from the first,
are not independent. However, a lazy-training constraint equivalent to (6)
seems to force the hyper-parameters of all the central layers to be frozen at
their initial values, so that in the limit of infinite width the fact of
having several layers may actually be detrimental for the PAC-Bayesian
training. We leave a more detailed discussion of these observations as future
work.

Finally, we believe that a more comprehensive empirical test is worth being
performed. For instance, it might be interesting to consider the case of data-
dependent priors.

Acknowledgments: This work has been partially supported by the UK Engineering
and Physical Sciences Research Council (EPSRC) through the grants EP/R513295/1
(DTP scheme) and CoSInES EP/R034710/1.

## References

  * [1] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals.  Understanding deep learning (still) requires rethinking generalization.  Commun. ACM, 64(3), 2021.
  * [2] R. M. Neal.  Bayesian learning for neural networks.  Springer Science & Business Media, 118, 1995.
  * [3] S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein.  Deep information propagation.  ICLR, 2017.
  * [4] G. Yang.  Tensor programs I: Wide feedforward or recurrent neural networks of any architecture are gaussian processes.  NeurIPS, 2019.
  * [5] S. Hayou, A. Doucet, and J. Rousseau.  On the impact of the activation function on deep neural networks training.  ICML, 2019.
  * [6] J. Lee, L. Xiao, S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington.  Wide neural networks of any depth evolve as linear models under gradient descent.  NeurIPS, 2019.
  * [7] S. Hayou, E. Clerico, B. He, G. Deligiannidis, A. Doucet, and J. Rousseau.  Stable resnet.  AISTATS, 2021.
  * [8] J. Sirignano and K. Spiliopoulos.  Mean field analysis of neural networks: A law of large numbers.  SIAM Journal on Applied Mathematics, 80(2), 2020.
  * [9] V. De Bortoli, A. Durmus, X. Fontaine, and U. Simsekli.  Quantitative propagation of chaos for SGD in wide neural networks.  NeurIPS, 2020.
  * [10] A. G. Matthews, J. Hron, M. Rowland, R.E. Turner, and Z. Ghahramani.  Gaussian process behaviour in wide deep neural networks.  ICLR, 2018.
  * [11] J. Lee, Y. Bahri, R. Novak, S. Schoenholz, J. Pennington, and J. Sohl-Dickstein.  Deep neural networks as Gaussian processes.  ICLR, 2018.
  * [12] A. Jacot, F. Gabriel, and C. Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  NeurIPS, 2018.
  * [13] Z. Allen-Zhu, Y. Li, and Y. Liang.  Learning and generalization in overparameterized neural networks, going beyond two layers.  NeurIPS, 2019.
  * [14] G. K. Dziugaite and D. M. Roy.  Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.  UAI, 2017.
  * [15] M. Pérez-Ortiz, O. Rivasplata, J. Shawe-Taylor, and C. Szepesvári.  Tighter risk certificates for neural networks.  arXiv:2007.12911, 2020.
  * [16] D. A. McAllester.  Some PAC-Bayesian theorems.  COLT, 1998.
  * [17] D. A. McAllester.  PAC-Bayesian model averaging.  COLT, 1999.
  * [18] B. Guedj.  A primer on PAC-Bayesian learning.  Proceedings of the Second congress of the French Mathematical Society, 2019.
  * [19] O. Catoni.  PAC-Bayesian supervised classification: The thermodynamics of statistical learning.  IMS Lecture Notes Monograph Series, 2007.
  * [20] A. Maurer.  A note on the PAC Bayesian theorem.  arXiv:0411099, 2004.
  * [21] G. Letarte, P. Germain, B. Guedj, and F. Laviolette.  Dichotomize and generalize: PAC-Bayesian binary activated deep neural networks.  NeurIPS, 2019.
  * [22] P. Germain, A. Lacasse, F. Laviolette, and M. Marchand.  PAC-Bayesian learning of linear classifiers.  ICML, 2009.
  * [23] J. Langford and R. Caruana.  (Not) bounding the true error.  NeurIPS, 2002.
  * [24] V. Bentkus.  A Lyapunov-type bound in Rd.  Theory of Probability & Its Applications, 49(2), 2005.

##  Appendix A Appendix

Throughout this section we will use several notations for the norms of vectors
and matrices. For $\gamma>0$ and a vector $v$,
$\|v\|_{\gamma}=(\sum_{i}|v_{i}|^{\gamma})^{1/\gamma}$. If $A$ is a matrix
$\|A\|_{F,\gamma}=(\sum_{ij}|A_{ij}|^{\gamma})^{1/\gamma}$ and
$\|A\|_{\gamma}=\sup_{v:\|v\|_{\gamma}=1}\|Av\|_{\gamma}$.
Moreover we will use the notation $\Gamma=O_{\hat{\mathbb{P}}}(n^{\gamma})$
meaning that $\limsup_{n\to\infty}\frac{|\Gamma|}{n^{\gamma}}<\infty$ in
probability wrt $\hat{\mathbb{P}}$, and
$\Gamma=\Omega_{\hat{\mathbb{P}}}(n^{\gamma})$ for
$\limsup_{n\to\infty}\frac{|\Gamma|}{n^{\gamma}}>0$ in probability wrt
$\hat{\mathbb{P}}$.

We want to prove a rigorous result of convergence to the Gaussian limit for a
wide stochastic network. We will essentially make use of the next result, due
to [24].

######  Theorem 1.

Let $X_{1},\dots,X_{n}$ be independent random vectors taking values in
$\mathbb{R}^{q}$, such that $\mathbb{E}[X_{j}]=0$ for all $j$. Let
$Y=\frac{1}{\sqrt{n}}\sum_{j=1}^{n}X_{j}$ and assume that the covariance
matrix $\mathbb{C}[Y]$ is non singular. Let
$Z\sim\mathcal{N}(0,\mathbb{C}[Y])$. Denoted as
$\frac{1}{\sqrt{\mathbb{C}}[Y]}$ the inverse of the positive square root of
the matrix $\mathbb{C}[Y]$, let
$B_{j}=\mathbb{E}[\|\frac{1}{\sqrt{\mathbb{C}}[Y]}\,X_{j}\|_{2}^{3}]$ and
$B=\frac{1}{n}\sum_{j=1}^{n}B_{j}$. Let $\mathcal{C}$ denote the class of all
convex subsets of $\mathbb{R}^{p}$.
There exists an absolute positive constant $\kappa<4$ such that

$\sup_{C\in\mathcal{C}}|\mathbb{P}(Y\in C)-\mathbb{P}(Z\in C)|\leq\kappa
q^{1/4}\frac{B}{\sqrt{n}}\,.$ (16)

Our goal is to prove a Gaussian limit as $n\to\infty$ for $F(x)$, whose
components are given by

$F_{i}(x)=\frac{1}{\sqrt{n}}\sum_{j=1}^{n}\mathfrak{s}_{ij}^{1}\zeta_{ij}^{1}\Phi^{0}_{j}(x)+\frac{1}{\sqrt{n}}\sum_{j=1}^{n}\mathfrak{m}_{ij}^{1}\Phi_{j}^{0}(x)\,.$

Let us denote by $X_{j}$ the $q$-dimensional vector $X_{j}=(X_{1j}\dots
X_{qj})$, with

$X_{ij}=\mathfrak{s}_{ij}^{1}\zeta_{ij}^{1}\Phi^{0}_{j}(x)+\mathfrak{m}_{ij}^{1}(\Phi_{j}^{0}(x)-\mathbb{E}[\Phi_{j}^{0}(x)])\,.$

Since all the $\zeta^{1}_{ij}$’s and the $\Phi^{0}_{j}$’s are independent, the
$X_{j}$’s constitute a family of $n$ centred independent $q$-dimensional
random vectors.

Clearly we have $F(x)=\mathbb{E}[F(x)]+\frac{1}{\sqrt{n}}\sum_{j=1}^{n}X_{j}$.
Let us define $Y=\frac{1}{\sqrt{n}}\sum_{j=1}^{n}X_{j}$. Note that, for all
$x$, the covariance matrix $\mathbb{C}[Y]$ is given by $Q(x)$, no matter if
$Y$ is Gaussian or not. Using the same notations of Theorem 1, assuming that
$\mathbb{C}[Y]$ is non-singular, we have

$\sup_{C\in\mathcal{C}}|\mathbb{P}(Y\in C)-\mathbb{P}(Z\in C)|\leq\kappa
q^{1/4}\frac{B}{\sqrt{n}}\,.$

To prove that the $Y$ behaves as a Gaussian for large $n$, we will show that
$B=O(1)$ for large $n$. We can easily upperbound each $B_{j}$ as

$B_{j}\leq\frac{1}{\lambda[Y]^{3/2}}\mathbb{E}[\|X_{j}\|_{2}^{3}]\,,$

where $\lambda[Y]>0$ is the smallest eigenvalue of $\mathbb{C}[Y]$.

Define $G=\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\|X_{j}\|_{2}^{3}]$ and
$\Lambda=\frac{1}{n}\sum_{j=1}^{n}\lambda[X_{j}]$. Clearly, as the $X_{j}$’s
are independent, we have
$\mathbb{C}[Y]=\frac{1}{n}\sum_{j=1}^{n}\mathbb{C}[X_{j}]$. In particular, we
can easily find that $\lambda[Y]\geq\Lambda$.

We summarize what we have so far in the next lemma.

######  Lemma 2.

With the notations introduced above, assuming that $\mathbb{C}[Y]$ is not
singular,

$B\leq\frac{1}{n}\sum_{j=1}^{n}\frac{1}{\Lambda^{3/2}}\mathbb{E}[\|X_{j}\|_{2}^{3}]=\frac{G}{\Lambda^{3/2}}\,.$

Now, from Hölder’s inequality we have
$\|X_{j}\|_{2}\leq\|\mathbf{1}_{q}\|_{6}\|X_{j}\|_{3}=q^{1/6}\|X_{j}\|_{3}$,
and so

$\|X_{j}\|_{2}^{3}\leq
q^{1/2}\|X_{j}\|_{3}^{3}=q^{1/2}\sum_{i=1}^{q}|X_{ij}|^{3}\,.$ (17)

Then, with some simple algebraic manipulation, and applying Jensen’s
inequality, we obtain

$\mathbb{E}[|X_{ij}|^{3}]\leq(|\mathfrak{s}_{ij}^{1}|^{3}\mathbb{E}[|\zeta_{ij}^{1}|^{3}]+8|\mathfrak{m}_{ij}^{1}|^{3})\mathbb{E}[|\Phi_{j}^{0}(x)|^{3}]\,.$

For convenience, we introduce the following notations

$H_{ij}=(2|\mathfrak{s}_{ij}^{1}|^{3}+8|\mathfrak{m}_{ij}^{1}|^{3})\mathbb{E}[|\Phi_{j}^{0}(x)|^{3}]\,;\qquad
H_{j}=\sum_{i=1}^{q}H_{ij}\,;\qquad H=\frac{1}{n}\sum_{j=1}^{n}H_{j}\,,$

so that we have $G\leq q^{1/2}H$, since
$\mathbb{E}[|\zeta|^{3}]=2\sqrt{2/\pi}<2$ for $\zeta\sim\mathcal{N}(0,1)$.

On the other hand, we can find a lowerbound for $\Lambda$ as well. Indeed,
first we can notice that

$\mathbb{C}_{ii^{\prime}}[X_{j}]=\delta_{ii^{\prime}}(\mathfrak{s}^{1}_{ij})^{2}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]+\mathfrak{m}^{1}_{ij}\mathfrak{m}^{1}_{i^{\prime}j}\mathbb{V}[\Phi^{0}_{j}(x)]\,.$

The first term is a diagonal matrix, while the second one is non-negative
definite. Hence we can write

$\lambda[X_{j}]\geq\mathbb{E}[\Phi^{0}_{j}(x)^{2}]\min_{i=1\dots
q}(\mathfrak{s}^{1}_{ij})^{2}\,.$

Defining

$\Theta=\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]\min_{i=1\dots
q}(\mathfrak{s}^{1}_{ij})^{2}\leq\Lambda$ (18)

we have the following corollary of Lemma 2.

######  Corollary 3.

With the same notations as above, if $\mathbb{C}[Y]$ is non singular, we have

$B\leq q^{1/2}\frac{H}{\Theta^{3/2}}\,.$

Note that both $H$ and $\Theta$ can be evaluated explicitly, given the
parameters of the networks, as long as $\phi$ allows for an explicit
evaluation of $\mathbb{E}[|\Phi^{0}(x)|^{\gamma}]$, for $\gamma=1,2,3$. This
means that we can give an exact upper bound to the finite-size error of the
predicted 0-1 loss, for any configuration of the network.

We can now prove Proposition 1 and Corollary 1 from the main text.

######  Proposition 1.

Fix an input $x$, let $Z(x)\sim\mathcal{N}(M(x),Q(x))$ and denote as
$\mathcal{C}$ the class of measurable convex subsets of $\mathbb{R}^{q}$. Then

$\sup_{C\in\mathcal{C}}|\mathbb{P}(F(x)\in C)-\mathbb{P}(Z(x)\in C)|\leq\kappa
q^{1/4}\frac{B}{\sqrt{n}}\,,$

where $\kappa\leq 4$ is an absolute constant and

$B\leq
q^{1/2}\frac{\sum_{j=1}^{n}\sum_{i=1}^{q}(2|\mathfrak{s}^{1}_{ij}|^{3}+8|\mathfrak{m}^{1}_{ij}|^{3})\mathbb{E}[|\Phi^{0}_{j}(x)|^{3}]}{\sum_{j=1}^{n}\mathbb{E}[\Phi_{j}^{0}(x)^{2}]\min_{i=1\dots
q}(\mathfrak{s}_{ij}^{1})^{2}}\,.$

In particular, if $B=o(\sqrt{n})$ for $n\to\infty$, then $F(x)\to Z(x)$ in
distribution.

###### Proof.

The result is a straight consequence of Theorem 1 and Corollary 3. Note
moreover that $\mathcal{C}[Y]$ is non singular as long as all the
$\mathfrak{s}$’s are non-zero, so as long as the bound in the statement is
finite. ∎

######  Corollary 1.

Assume that the network deals with a classification problem, where for each
instance $x$ there is a single correct label $y=f(x)\in\\{1\dots q\\}$. Then,
for each fixed input $x\neq 0$, denoting as
$\hat{f}(x)=\operatorname*{argmax}_{i=1\dots q}F_{i}(x)$ and
$\bar{f}(x)=\operatorname*{argmax}_{i=1\dots q}Z_{i}(x)$, we have

$|\mathbb{P}(\hat{f}(x)=f(x))-\mathbb{P}(\bar{f}(x)=f(x))|\leq\kappa
q^{1/4}\frac{B}{\sqrt{n}}\,.$

###### Proof.

The result follows from Proposition 1 as, for each $k=\\{1\dots q\\}$, the set
$\\{z\in\mathbb{R}^{q}:z_{k}>\max_{i\neq k}z_{k}\\}$ is convex. ∎

In the next section we show how, with a suitable random initialization, we can
assure that the network has an almost Gaussian behaviour. Successively, we
will show that this behaviour is preserved during training, as long as the
hyper-parameters do not move too much from their initial values.

###  A.1 Initialization

We consider the random initialization:

$\displaystyle\mathfrak{m}_{jk}^{0}\sim\mathcal{N}(0,1)\,;\qquad$
$\displaystyle\mathfrak{m}_{ij}^{1}\sim\mathcal{N}(0,1)\,;$ (5)
$\displaystyle\mathfrak{s}_{jk}^{0}=1\,;\qquad$
$\displaystyle\mathfrak{s}_{ij}^{1}=1\,.$

As now we have two sources of randomness, the initialization and the intrinsic
stochasticity of the network, to avoid confusion we will denote as
$\hat{\mathbb{E}}$, $\hat{\mathbb{P}}$, $\hat{\mathbb{V}}$ the expectations,
probabilities and variances with respect to the initialization, whilst
$\mathbb{E}$, $\mathbb{P}$ and $\mathbb{V}$ refer to the network intrinsic
stochasticity.

######  Lemma 3.

Define $H$ and $\Theta$ as in the previous section and assume that $\phi$ is
Lipshitz continuous. Assume that the hyper-parameters are initialized
according to $\hat{\mathbb{P}}$, as in (5). Then, for any fixed $x\neq 0$, as
$n\to\infty$ $H\to h>0$ and $\Theta\to\theta>0$ in probability, with respect
to the random initialization, where both $h$ and $\theta$ are finite.

###### Proof.

First notice that, fixed an input $x\neq 0$ and fixed $n$, all the
$\Phi^{0}_{j}$’s are iid, with respect to $\hat{\mathbb{P}}$, as all the
$\mathfrak{m}^{0}$’s and the $\mathfrak{s}_{0}$’s are. As a consequence all
the $H_{j}$’s are iid with respect to $\hat{\mathbb{P}}$ (note that they have
different distribution for different $n$ as the law of the $\Phi_{j}^{0}$’s
depends on $n$). Now, thanks to the fact that $\phi$ is Lipshitz continuous,
we have that $\limsup_{n\to\infty}\hat{\mathbb{V}}[H_{j}]<\infty$. Hence, by a
standard application of the CLT for triangular arrays, we get that

$H-\hat{\mathbb{E}}[H]=\frac{1}{n}\sum_{j=1}^{n}(H_{j}-\hat{\mathbb{E}}[H_{j}])\to
0$

in distribution, and hence in probability, as $0$ is a constant. It is quickly
verified that the limit $h=\lim_{n\to\infty}\hat{\mathbb{E}}[H]$ exists,
finite and positive.
The proof for $\Theta$ is analogous. ∎

Now we can easily prove Proposition 2.

######  Proposition 2.

Consider a network initialized according to $\hat{\mathbb{P}}$ as in (5), and
whose activation function $\phi$ is Lipshitz continuous. Then, for any fixed
input $x\neq 0$ and letting $B$ be defined as in Proposition 1, we have
$\frac{B}{\sqrt{n}}\to 0$ as $n\to\infty$, in probability with respect to the
random initialization $\hat{\mathbb{P}}$. More precisely,
$B=O_{\hat{\mathbb{P}}}(1)$ as $n\to\infty$.
In particular, the network tends to a Gaussian limit at initialization, in
distribution wrt the intrinsic stochasticity $\mathbb{P}$ and in probability
wrt $\hat{\mathbb{P}}$.

###### Proof.

It is a straight consequence of Lemma 3. ∎

###  A.2 Lazy training

We have established that the Gaussian limit holds at initialization. In the
present section we will see that, as far as the hyper-parameters of the
network do not move too much from their initial values, the limit keeps its
validity.

######  Proposition 3.

Fix a constant $J>0$ independent of $n$ and assume that $\phi$ is Lipshitz.
Consider a network with initial configuration
$(\widetilde{\mathfrak{m}},\widetilde{\mathfrak{s}})$ drawn according to
$\hat{\mathbb{P}}$, as in (5). Now consider a new configuration
$(\mathfrak{m},\mathfrak{s})$ for the network such that

$\|\mathfrak{m}^{0}-\widetilde{\mathfrak{m}}^{0}\|^{2}_{F,2}+\|\mathfrak{m}^{1}-\widetilde{\mathfrak{m}}^{1}\|_{F,2}^{2}+\|\mathfrak{s}^{0}-\widetilde{\mathfrak{s}}^{0}\|^{2}_{F,2}+\|\mathfrak{s}^{1}-\widetilde{\mathfrak{s}}^{1}\|_{F,2}^{2}\leq
J^{2}\,,$

where $\|\cdot\|_{F,2}$ denotes the 2-Frobenius norm of a matrix. Let $B$ be
defined as in Proposition 1, referring to the configuration
$(\mathfrak{m},\mathfrak{s})$. For any fixed input $x\neq 0$ we have
$B=O_{\hat{\mathbb{P}}}(1)$ as $n\to\infty$, in probability with respect to
the random initialization $\hat{\mathbb{P}}$.

###### Proof.

For convenience we will write with a tilde all the quantities relative to the
network at initialization. We denote with a delta the difference between the
final and the initial values of these quantities. For instance
$\Theta=\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]\min_{i=1\dots
q}(\mathfrak{s}^{1}_{ij})^{2}$,
$\widetilde{\Theta}=\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\widetilde{\Phi}^{0}_{j}(x)^{2}]\min_{i=1\dots
q}(\widetilde{\mathfrak{s}}^{1}_{ij})^{2}$, and
$\Delta\Theta=\Theta-\widetilde{\Theta}$.

We will show that for $n\to\infty$, $\Theta=\Omega_{\hat{\mathbb{P}}}(1)$ and
$G=O_{\hat{\mathbb{P}}}(1)$, so that we can conclude using that
$\Lambda\geq\Theta$ and Lemma 2.

Fix an input $x$. First, we need a bound on
$\|\Delta\Phi^{0}(x)\|_{2}=\|\Phi^{0}(x)-\widetilde{\Phi}^{0}\|_{2}$. We have
that $\Phi^{0}(x)=\phi(Y^{0}(x))$. Hence, letting $L$ be the Lipshitz constant
of $\phi$, we have $\|\Delta\Phi^{0}(x)\|_{2}\leq L\|\Delta Y^{0}(x)\|_{2}$.
Now, as $\Delta
Y^{0}_{j}(x)=\frac{1}{\sqrt{p}}\sum_{k=1}^{p}\Delta\mathfrak{m}^{0}_{jk}x_{k}+\frac{1}{\sqrt{p}}\sum_{k=1}^{p}\Delta\mathfrak{s}^{0}_{jk}\zeta^{0}_{jk}x_{k}$,
we have

$\|\Delta\Phi^{0}(x)\|_{2}\leq\frac{L}{\sqrt{p}}(\|\Delta\mathfrak{m}^{0}\|_{2}+\|\Delta\mathfrak{s}^{0}\odot\zeta^{0}\|_{2})\|x\|_{2}\,,$

where $\odot$ denotes the Hadamard product.

Notice that we have

$\mathbb{E}[\|\Delta\mathfrak{s}^{0}\odot\zeta^{0}\|^{2}_{2}]\leq\mathbb{E}[\|\Delta\mathfrak{s}^{0}\odot\zeta^{0}\|_{F,2}^{2}]=\sum_{j=1}^{n}\sum_{k=1}^{p}(\Delta\mathfrak{s}^{0}_{jk})^{2}\mathbb{E}[(\zeta^{0}_{jk})^{2}]=\|\Delta\mathfrak{s}^{0}\|_{F,2}^{2}\leq
J^{2}\,,$

where, as usual, the expectation $\mathbb{E}$ is the one with respect to the
intrinsic stochasticity of the network, due to the $\zeta$’s. We can define a
constant $C\geq 0$, independent of $n$, such that

$\mathbb{E}[\|\Delta\Phi^{0}(x)\|_{2}^{2}]\leq\frac{4L^{2}J^{2}\|x\|_{2}^{2}}{p}=C^{2}$

as $\|\Delta\mathfrak{m}^{0}\|_{2}\leq\|\Delta\mathfrak{m}^{0}\|_{F,2}\leq J$.

Now, recalling the definition of $\Theta$ and using that
$\mathfrak{s}^{1}=1+\Delta\mathfrak{s}^{1}$, we have

$\Theta=\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]\min_{i=1\dots
q}(\mathfrak{s}^{1}_{ij})^{2}\geq\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}](1-2\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|)\,.$

We will show that $\Theta-\widetilde{\Theta}\to 0$.

First notice that

$\frac{1}{n}\mathbb{E}[\|\Phi^{0}(x)\|_{2}^{2}]\geq\frac{1}{n}\mathbb{E}[\|\widetilde{\Phi}^{0}(x)\|_{2}^{2}]-\frac{2}{n}\mathbb{E}[|\widetilde{\Phi}^{0}(x)\cdot\Delta\Phi^{0}(x)|]+\frac{1}{n}\mathbb{E}[\|\Delta\Phi^{0}_{j}(x)\|_{2}^{2}]\,.$

We know that
$\frac{1}{n}\mathbb{E}[\|\widetilde{\Phi}^{0}(x)\|_{2}^{2}]=\widetilde{\Theta}$
by definition. On the other hand we have

$\frac{2}{n}\mathbb{E}[|\widetilde{\Phi}^{0}(x)\cdot\Delta\Phi^{0}(x)|]\leq\frac{2}{n}\mathbb{E}[\|\widetilde{\Phi}^{0}(x)\|_{2}\|\Delta\Phi^{0}(x)\|_{2}]\leq\frac{2C}{\sqrt{n}}\left(\frac{1}{n}\mathbb{E}[\|\widetilde{\Phi}^{0}(x)\|_{2}^{2}]\right)^{1/2}=\frac{2C\widetilde{\Theta}^{1/2}}{\sqrt{n}}=O_{\hat{\mathbb{P}}}(1/\sqrt{n})\,.$

Since $\frac{1}{n}\mathbb{E}[\|\Delta\Phi^{0}_{j}(x)\|_{2}^{2}]\leq C^{2}/n$,
we have that
$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]-\widetilde{\Theta}\to
0$, in probability with respect to the random initialization
$\hat{\mathbb{P}}$.

We still need to show that
$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|\to 0$. Again we can decompose the term in
$\Phi^{0}$ and we have

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|\leq\frac{1}{n}\sum_{j=1}^{n}(\mathbb{E}[\widetilde{\Phi}^{0}_{j}(x)^{2}]+2\mathbb{E}[|\widetilde{\Phi}^{0}_{j}(x)\Delta\Phi^{0}_{j}(x)|]+\mathbb{E}[\Delta\Phi^{0}_{j}(x)^{2}])\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|\,.$

Clearly, for every $j$ we have $\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|\leq J$, and so we can write

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\Phi^{0}_{j}(x)^{2}]\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|\leq\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\widetilde{\Phi}^{0}_{j}(x)^{2}]\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|+\frac{2J}{n}(\mathbb{E}[|\widetilde{\Phi}^{0}(x)\cdot\Delta\Phi^{0}(x)|]+\mathbb{E}[\|\Delta\Phi^{0}(x)\|_{2}^{2}])\,.$

We have already shown that
$\frac{1}{n}(2\mathbb{E}[|\widetilde{\Phi}^{0}(x)\cdot\Delta\Phi^{0}(x)|]+\mathbb{E}[\|\Delta\Phi^{0}(x)\|_{2}^{2}])=O_{\hat{\mathbb{P}}}(1/\sqrt{n})$.
As for the other term, we have

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\widetilde{\Phi}^{0}_{j}(x)^{2}]\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|\leq\frac{1}{\sqrt{n}}\left(\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\widetilde{\Phi}^{0}_{j}(x)^{4}]\right)^{1/2}\left(\sum_{j=1}^{n}\min_{i=1\dots
q}(\Delta\mathfrak{s}^{1}_{ij})^{2}\right)^{1/2}\,.$

Using an argument analogous to the one in the proof of Proposition 2, we have
that $\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\widetilde{\Phi}^{0}_{j}(x)^{4}]$
has a finite limit (in probability wrt $\hat{\mathbb{P}}$). On the other hand,
we have

$\sum_{j=1}^{n}\min_{i=1\dots
q}(\Delta\mathfrak{s}^{1}_{ij})^{2}\leq\sum_{j=1}^{n}\sum_{i=1}^{q}(\Delta\mathfrak{s}^{1}_{ij})^{2}=\|\Delta\mathfrak{s}^{1}\|_{F,2}^{2}\leq
J^{2}\,.$

We have thus obtained that
$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\widetilde{\Phi}^{0}_{j}(x)^{2}]\min_{i=1\dots
q}|\Delta\mathfrak{s}^{1}_{ij}|=O_{\hat{\mathbb{P}}}(1/\sqrt{n})$ and so we
can conlude that $\Theta-\widetilde{\Theta}\to 0$ and so
$\Theta=\Omega_{\hat{\mathbb{P}}}(1)$, in probability wrt the random
initialization $\hat{\mathbb{P}}$.

Now, we will show that $G=O_{\hat{\mathbb{P}}}(1)$. We have

$G=\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\|X_{j}\|_{2}^{3}]\leq\frac{4}{n}\sum_{j=1}^{n}\mathbb{E}[\|\widetilde{X}_{j}\|_{2}^{3}]+\frac{4}{n}\sum_{j=1}^{n}\mathbb{E}[\|\Delta
X_{j}\|_{2}^{3}]\,.$

Let us write $X_{j}=U_{j}+V_{j}$, with
$U_{j}=(\zeta^{1}_{ij}\Phi^{0}_{j}(x))_{i=1\dots q}$ and
$V_{j}=(\mathfrak{m}^{1}_{ij}(\Phi^{0}_{j}(x)-\mathbb{E}[\Phi^{0}_{j}(x)]))_{i=1\dots
q}$. Then $\|\Delta X_{j}\|_{2}^{3}\leq 4(\|\Delta U_{j}\|_{2}^{3}+\|\Delta
V_{j}\|_{2}^{3})$.

First, denoting as $\zeta^{1}_{j}$ and $\Delta\mathfrak{s}^{1}_{j}$ the
vectors $(\zeta^{1}_{ij})_{i=1\dots q}$ and
$(\Delta\mathfrak{s}^{1}_{ij})_{i=1\dots q}$, we can write

$\Delta
U_{j}=\Delta\Phi^{0}_{j}(x)\zeta^{1}_{j}+\widetilde{\Phi}^{0}_{j}(x)\Delta\mathfrak{s}^{1}_{j}\odot\zeta^{1}_{j}+\Delta\Phi^{0}_{j}(x)\Delta\mathfrak{s}^{1}_{j}\odot\zeta^{1}_{j}\,,$

where $\odot$ represents the Hadamart product. Since $\Phi^{0}$ and
$\zeta^{1}$ are independent and using that
$\mathbb{E}[|\zeta|^{3}]=2\sqrt{2/\pi}<2$ for $\zeta\sim\mathcal{N}(0,1)$, we
have

$\mathbb{E}[\|\Delta U_{j}\|_{2}^{3}]\leq
54(q^{3/2}\mathbb{E}[|\Delta\Phi^{0}_{j}(x)|^{3}]+\mathbb{E}[|\widetilde{\Phi}^{0}_{j}(x)|^{3}]\mathbb{E}[\|\Delta\mathfrak{s}^{1}_{j}\|_{2}^{3}]+\mathbb{E}[|\Delta\Phi^{0}_{j}(x)|^{3}]\mathbb{E}[\|\Delta\mathfrak{s}^{1}_{j}\|_{2}^{3}])\,.$

Using that $\|\Delta\Phi^{0}(x)\|_{3}^{3}\leq\|\Delta\Phi^{0}(x)\|_{2}^{3}\leq
C^{3}$, we have that

$\frac{1}{n}\sum_{j=1}^{n}q^{3/2}\mathbb{E}[|\Delta\Phi^{0}_{j}(x)|^{3}]\leq\frac{q^{3/2}C^{3}}{n}=O_{\hat{\mathbb{P}}}(1/n)\,.$

Then we can notice that
$\|\Delta\mathfrak{s}^{1}_{j}\|_{2}\leq\|\mathbf{1}_{q}\|_{3}\|\Delta\mathfrak{s}^{1}_{j}\|_{6}=q^{1/3}\|\Delta\mathfrak{s}^{1}_{j}\|_{6}$
by Hölder’s inequality. Hence

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\widetilde{\Phi}^{0}_{j}(x)|^{3}]\mathbb{E}[\|\Delta\mathfrak{s}^{1}_{j}\|_{2}^{3}]\leq\frac{q\|\Delta\mathfrak{s}^{1}\|_{F,6}^{3}}{\sqrt{n}}\left(\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\widetilde{\Phi}^{0}_{j}(x)|^{6}]\right)^{1/2}\leq\frac{qJ^{3}}{\sqrt{n}}\left(\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\widetilde{\Phi}^{0}_{j}(x)|^{6}]\right)^{1/2}=O_{\hat{\mathbb{P}}}(1/\sqrt{n})\,.$

where the last equality comes from the usual argument that
$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\widetilde{\Phi}^{0}_{j}(x)|^{6}]$ has a
finite limit in probability (with respect to the random initialization).

Finally we can notice that $|\Phi^{0}_{j}(x)|\leq\|\Phi^{0}(x)\|_{2}\leq C$
for all $j$, so that

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\Delta\Phi^{0}_{j}(x)|^{3}]\mathbb{E}[\|\Delta\mathfrak{s}^{1}_{j}\|_{2}^{3}]\leq\frac{q^{1/2}}{n}C^{3}\|\Delta\mathfrak{s}^{1}\|_{F,3}^{3}\leq\frac{q^{1/2}C^{3}J^{3}}{n}=O_{\hat{\mathbb{P}}}(1/n)\,,$

where we used that
$\|\Delta\mathfrak{s}^{1}_{j}\|_{2}\leq\|\mathbf{1}_{q}\|_{6}\|\Delta\mathfrak{s}^{1}_{j}\|_{3}=q^{1/6}\|\Delta\mathfrak{s}^{1}_{j}\|_{3}$
by Hölder’s inequality. We can hence conclude that

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\|\Delta
U_{j}\|_{2}^{3}]=O_{\hat{\mathbb{P}}}(1/\sqrt{n})\,.$

Now we need to bound $\|\Delta V_{j}\|_{2}$. Letting
$\mathfrak{m}^{1}_{j}=(\mathfrak{m}^{1}_{ij})_{i=1\dots q}$ and
$\delta\Phi^{0}_{j}(x)=\Phi^{0}_{j}(x)-\mathbb{E}[\Phi^{0}_{j}(x)]$, it can be
easily shown that

$\|\Delta
V_{j}\|_{2}\leq|\delta\widetilde{\Phi}^{0}_{j}(x)|\|\Delta\mathfrak{m}^{1}_{j}\|_{2}+|\Delta\delta\Phi^{0}_{j}(x)|\|\widetilde{\mathfrak{m}}^{1}_{j}\|_{2}+\|\Delta\mathfrak{m}^{1}_{j}\|_{2}|\Delta\delta\Phi^{0}_{j}(x)|\,.$

So we have

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\|V_{j}\|_{2}^{3}]\leq\frac{27}{n}\sum_{j=1}^{n}\mathbb{E}[|\delta\widetilde{\Phi}^{0}_{j}(x)|^{3}]\|\Delta\mathfrak{m}^{1}_{j}\|_{2}^{3}+\frac{27}{n}\sum_{j=1}^{n}\mathbb{E}[|\Delta\delta\Phi^{0}_{j}(x)|^{3}]\|\widetilde{\mathfrak{m}}^{1}_{j}\|_{2}^{3}+\frac{27}{n}\sum_{j=1}^{n}\mathbb{E}[|\Delta\delta\Phi^{0}_{j}(x)|^{3}]\|\Delta\mathfrak{m}^{1}_{j}\|_{2}^{3}\,.$

Starting from the first term, we have that

$\sum_{j=1}^{n}\mathbb{E}[|\delta\widetilde{\Phi}^{0}_{j}(x)|^{3}]\|\Delta\mathfrak{m}^{1}_{j}\|_{2}^{3}\leq\left(\sum_{j=1}^{n}\mathbb{E}[|\delta\widetilde{\Phi}^{0}_{j}(x)|^{6}]\right)^{1/2}\left(\sum_{j=1}^{n}\|\Delta\mathfrak{m}^{1}_{j}\|_{2}^{6}\right)^{1/2}\,.$

From Hölder’s inequality we have
$\|\Delta\mathfrak{m}^{1}_{j}\|_{2}\leq\|\mathbf{1}_{q}\|_{3}\|\Delta\mathfrak{m}_{j}\|_{6}=q^{1/3}\|\Delta\mathfrak{m}_{j}\|_{6}$
and hence

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\delta\widetilde{\Phi}^{0}_{j}(x)|^{3}]\|\Delta\mathfrak{m}^{1}_{j}\|_{2}^{3}\leq\frac{q\|\Delta\mathfrak{m}^{1}\|_{F,6}^{3}}{\sqrt{n}}\left(\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\delta\widetilde{\Phi}^{0}_{j}(x)|^{6}]\right)^{1/2}\leq\frac{qJ^{3}}{\sqrt{n}}\left(\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\delta\widetilde{\Phi}^{0}_{j}(x)|^{6}]\right)^{1/2}=O_{\hat{\mathbb{P}}}1/\sqrt{n})\,.$

As
$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\delta\widetilde{\Phi}^{0}_{j}(x)|^{6}]$
tends in probability (wrt the random initialization) to a finite limit, we
have that the term above is of order $O_{\hat{\mathbb{P}}}1/\sqrt{n})$.

Proceeding analogously, and noting that the $L$-Lipshitzianity of $\phi$
implies that $\mathbb{E}[|\Delta\delta\Phi^{0}_{j}(x)|^{3}]\leq
8L^{3}\mathbb{E}[|\Delta Y^{0}_{j}(x)|^{3}]$, we get

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\Delta\delta\Phi^{0}_{j}(x)|^{3}]\|\widetilde{\mathfrak{m}}^{1}_{j}\|_{2}^{3}\leq\frac{8C^{3}}{\sqrt{n}}\left(\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\|\widetilde{\mathfrak{m}}^{1}_{j}\|_{2}^{6}]\right)^{1/2}=O_{\hat{\mathbb{P}}}1/\sqrt{n})\,,$

and again we can use the fact that
$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\|\widetilde{\mathfrak{m}}^{1}_{j}\|_{2}^{6}]$
converges in probability (wrt the random initialization) to a finite quantity
to show that the above term is of order $O_{\hat{\mathbb{P}}}1/\sqrt{n})$.

Finally, in a similar way we get

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[|\Delta\delta\Phi^{0}_{j}(x)|^{3}]\|\Delta\mathfrak{m}^{1}_{j}\|_{2}^{3}\leq\frac{8qJ^{3}C^{3}}{n}\,.$

We can hence conclude that

$\frac{1}{n}\sum_{j=1}^{n}\mathbb{E}[\|V_{j}\|_{2}^{3}]=O_{\hat{\mathbb{P}}}1/\sqrt{n})$

and so that, as $n\to\infty$, $G\geq O_{\hat{\mathbb{P}}}1)$ in probability
with respect to the random initialization. This ends the proof. ∎

###  A.3 Expected values for $\operatorname*{ReLU}$ and $\sin$ activations

Let $a>0$, $b\in\mathbb{R}$, $\zeta\sim\mathcal{N}(0,1)$. The following
formulae are easily verified by direct calculations:

$\displaystyle\mathbb{E}[\sin(a\zeta+b)]=e^{-a^{2}/2}\sin b\,;$
$\displaystyle\mathbb{E}[\sin(a\zeta+b)^{2}]=\frac{1}{2}(1-e^{2a^{2}}\cos(2b))\,;$
$\displaystyle\mathbb{E}[\operatorname*{ReLU}(a\zeta+b)]=\frac{ae^{-b^{2}/(2a^{2})}}{\sqrt{2\pi}}+\frac{b}{2}\left(1+\operatorname*{erf}\frac{b}{a\sqrt{2}}\right)\,;$
$\displaystyle\mathbb{E}[\operatorname*{ReLU}(a\zeta+b)^{2}]=\frac{abe^{-b^{2}/(2a^{2})}}{\sqrt{2\pi}}+\frac{1}{2}(a^{2}+b^{2})\left(1+\operatorname*{erf}\frac{b}{a\sqrt{2}}\right)\,.$