Stochastic Neural Networks with Infinite Width are Deterministic

We first introduce two basic assumptions of the network structure.

Throughout this work, the component functions $g^{1},\ g^{2}$ of a network $f$ are called a block, which can be seen as a generalization of a layer. It is appropriate to call $g^{1}$ the input block and $g^{2}$ the output block. Because the Lipshitz constant of a neural network can be upper bounded by the product of the largest eigenvalue of each weight matrix times the Lipshitz constant of the non-linearity, the assumption that every block is Lipshitz-continuous applies to all existing networks with fixed weights and with Lipshitz-continuous activation functions (such as ReLU, tanh, Swish (Ramachandran et al.,, 2017), Snake (Ziyin et al.,, 2020) etc.).

If we restrict ourselves to feedforward architectures, we can discuss the meaning of an ”increasing width” without much ambiguity. However, in our work, since the definition of blocks (layers) is abstract, it is not immediately clear what it means to ”increase the width of a block.” The following definition makes it clear that one needs to specify a sequence of blocks to define an increasing width.¹¹1Also, note that this definition of ”width” makes it possible to define different ways of ”increasing” the width and is thus more general than the standard procedure of simply increasing the output dimension of the corresponding linear transformation.

Note that if $f=g^{2}\circ g^{1}$, $g^{1}={}^{d_{2},d_{1}}g^{1}$ and $g^{2}={}^{d_{4},d_{3}}g^{2}$, $d_{2}$ must be equal to $d_{3}$; namely, specifying $g^{1}$ constrains the input dimension of the next block. It is appropriate to call $d_{2}$ the width of the block ${}^{d_{2},d_{1}}g$. Since each block is parametrized by its own parameter set, the union of all parameter sets is the parameter set of the neural network $f$: $f=f_{w}$. Since every block comes with the respective indices and equipped with its own parameter set, we omit specifying the indices and the parameter set when unnecessary. The next assumption specifies what it means to have a larger width.

This assumption can be seen as a constraint on the types of block sets $\mathcal{S}(w)$ we can choose. As a concrete example, the following proposition shows that the block set induced by a linear layer with arbitrary input and output dimensions followed by an element-wise non-linearity satisfies our assumption.

Proof. Consider two functions ${}^{d_{1},d_{2}}g_{\{W,b\}}(x)$ and ${}^{d_{1}^{\prime},d_{2}}g_{\{W^{\prime},b^{\prime}\}}(x)$ in $\mathcal{S}(g)$. Let $m$ be an arbitrary mapping from $\{1,...,d_{1}^{\prime}\}\to\{1,...,d_{1}\}.$ It suffices to show that there exist $W^{\prime}$ and $b^{\prime}$ such that $[^{d_{1},d_{2}}g_{\{W,b\}}(x)]_{m(l)}-[^{d_{1}^{\prime},d_{2}}g_{\{W^{\prime},b^{\prime}\}}(x)]_{l}=0$ for all $l$. For a matrix $M$, we use $M_{j:}$ to denote the $j$-th row of $M$. By definition, this condition is equivalent to

which is achieved by setting $b^{\prime}=b$ and $W^{\prime}_{l:}=W_{m(l):}$, where $W_{m(l):}$ is the $m(l)$-th row of $W$. $\square$

Now, we are ready to define a stochastic neural network.

Namely, a stochastic network becomes a proper neural network when averaged over the noise of the stochastic block. To proceed, we make the following assumption about the randomness in the stochastic layer.

This assumption applies to standard stochastic techniques in deep learning, such as dropout or the reparametrization trick used in approximate Bayesian deep learning. Lastly, we assume the following condition for the architecture.

In other words, we assume that the second block $g^{2}$ can always be decomposed as $g^{\prime}\circ M$, such that $M$ is an optimizable linear transformation. This is the only architectural assumption we make. In principle, this can be replaced by weaker assumptions. However, we leave this as an important future work because Assumption 4 is sufficient for the purpose of this work and is general enough for the applications we consider (such as dropout and VAE). We also stress that the condition that $g^{2}$ starts with a linear transformation does not mean that the first actual layer of $g^{2}$ is linear. Instead, $M$ can be followed by an arbitrary Lipshitz activation function as is usual in practice; in our definition, if it exists, this following activation is assumed into the definition of $g^{\prime}$.

The actual rather restrictive assumption in Assumption 4 is that the function $g^{\prime}$ has a fixed input dimension (like a “bottleneck”). In practice, when one scales up the model, it is often the case that one wants to scale up the width of all other layers simultaneously. For the first block, this is allowed by assumption 2.²²2For example, if $g^{1}$ is a multilayer perceptron, it is easy to check that assumption  2 is satisfied if one increases the intermediate layers of $g^{1}$ simultaneously. We note that this bottleneck assumption is mainly for notational concision. In the appendix A.3, we show that one can also extend the result to the case when the input dimension (and the intermediate dimensions) of $g^{\prime}$ also increases as one increases the width of the stochastic layer.

Notation Summary. To summarize, we require a network $f$ to be decomposed into two blocks: $f=g^{2}\circ g^{1}$ and $g^{1}$ is a stochastic block. Each block is associated with its indices, which specify its input and output dimensions, and a parameter set that we optimize over. For example, we can write a block $g$ as $g={}^{d_{2},d_{1}}g^{i}_{w}$ to specify that $g$ is the $i$-th block in a neural network, is a mapping from $\mathbb{R}^{d_{1}}$ to $\mathbb{R}^{d_{2}}$, and that its parameters are $w$. However, for notational conciseness and better readability, we omit some of the specifications when the context is clear. For the parameter $w$, we abuse the notation a little. Sometimes, we view $w$ as a set and discuss its unions and subsets; for example, let $f_{w}=g^{2}_{w^{2}}\circ g^{1}_{w^{1}}$; then, we say that the parameter set $w$ of $f$ is the union of the parameter set of $g^{1}$ and $g^{2}$: $w=w^{1}\cup w^{2}$. Alternatively, we also view $w$ as a vector in a subset of the real space, so that we can look for the minimizer $w$ in such a space (in expressions such as $\min_{w}L(w)$).

In this work, we restrict our theoretical result to the MSE loss. Consider an arbitrary training set $\{(x_{i},y_{i})\}_{i=1}^{N}$, when training the deterministic network, we want to find

It is convenient to write $[f_{w}(x_{i})-y_{i}]^{2}$ as $L_{i}(w)$.

An overparametrized network can be defined as a network that can achieve zero training loss on such a dataset.

Namely, for a stochastic network, we say that it is overparametrized if its deterministic part is overparametrized. When there is no degeneracy in the data (if $x_{i}=x_{j}$, then $y_{i}=y_{j}$), zero training loss can be achieved for a wide enough neural network, and this definition is essentially equivalent to assuming that there is no data degeneracy and is thus not a strong limitation of our result.

With a stochastic block, the training loss becomes (due to the sampling of the hidden representation)

Note that this loss function can still be reduced to $0$ if $f_{w}(x_{i})=y_{i}$ for all $i$ with probability $1$.

With these definitions at hand, we are ready to state our main result.

Proof Sketch. The full proof is given in Appendix Section A.1. In the proof, we denote the term $L(^{d_{1}}f_{w}(x,\epsilon),y_{i})$ as $L_{i}^{d_{1}}(w)$. Let $w_{*}$ be the global minimizer of ${}\sum_{j}^{N}\mathbb{E}_{\epsilon}\left[L_{j}^{d_{1}}(w)\right]$. Then, for any $w$, by definition of the global minimum,

If $\lim_{d_{1}\to\infty}{}\sum_{j}^{N}\mathbb{E}_{\epsilon}\left[L_{j}^{d_{1}}(w)\right]=0$, we have $\lim_{d_{1}\to\infty}{}\sum_{j}^{N}\mathbb{E}_{\epsilon}\left[L_{j}^{d_{1}}(w_{*})\right]=0$, which implies that $\lim_{d_{1}\to\infty}\mathbb{E}_{\epsilon}\left[L_{j}^{d_{1}}(w_{*})\right]=0$ for all $j$. By bias-variance decomposition of the MSE, this, in turn, implies that $\text{Var}[^{d_{1}}f_{w_{*}}(x_{j})]=0$ for all $j$. Therefore, it is sufficient to construct a sequence of $w$ such that $\lim_{d_{1}\to\infty}{}\mathbb{E}_{\epsilon}\sum_{j}^{N}\left[L_{j}^{d_{1}}(w)\right]=0$. The rest of the proof shows that, with the architecture assumptions we made, such a network can indeed be constructed. In particular, the architectural assumptions allow us to make independent copies of the output of the stochastic block and the linear transformation after it allows us to average over such independent copies to recover the mean with a vanishing variance, which can then be shown to be able to achieve zero loss. $\square$

At a high level, one might wonder why the optimized model achieves zero variance. Our results suggest that the form of the loss function may be crucial. The MSE loss can be decomposed into a bias term and a variance term:

Minimizing the MSE loss involves both minimizing the bias term and the variance term, and the key step of our proof involves showing that a neural network with a sufficient width can reduce the variance to zero. We thus conjecture that convergence to a zero-variance model can be generally true for a broad class of loss functions. For example, one possible candidate for this function class is the set of convex loss functions, which favor a mean solution more than a solution with variance (by Jensen’s inequality), and a neural network is then encouraged to converge to such solutions so as to minimize the variance. However, identifying this class of loss functions is beyond the scope of the present work, and we leave it as an important future step. Lastly, we also stress that the main results are not a trivial consequence of the MSE being convex. When the model is linear, it is rather straightforward to show that the variance reduces to zero in the large-width limit because taking the expectation of the model output is equivalent to taking the expectation of the latent noise. However, this is not trivial to prove for a genuine neural network because the net $f$ is, in general, a nonlinear and nonconvex function of the latent noise $\epsilon$.

Since the noise of dropout is independent, one can immediately apply the main theorem and obtain the following corollary.

Our result thus formally proves the intuitive hypothesis in the original dropout paper (Srivastava et al.,, 2014) that applying dropout to training has the effect of encouraging an averaging effect in the latent space.

We now extend our result to the case where a soft prior constraint exists in the loss function. The model setting of this section is equivalent to a special type of Bayesian latent variable model (e.g., a VAE) whose prior strength decreases towards zero as the width of the model is increased.⁴⁴4In the context of VAE, many previous works have noticed that the VAE has problems with its prediction variance (Dai and Wipf,, 2019; Wang and Ziyin,, 2022). However, none of the previous works discusses the behavior of a VAE when its width tends to infinity. The result in this section is more general and involved than Theorem 1 because the soft prior constraint matches the latent variable to the prior distribution in addition to the MSE loss. The existence of the prior term regularizes the model and prevents a perfect fitting of the original MSE loss, and the main message of this section is to show that even in the presence of such a (weak) regularization term, a vanishing variance can still emerge, but relatively slowly.

While it is natural that a latent variable model $f=g^{2}\circ g^{1}$ can be decomposed into two blocks, where $g^{2}$ is the decoder, and $g^{1}$ is an encoder, we require one additional condition. Namely, the stochastic block ends with a linear transformation layer (which is satisfied by a standard VAE encoder (Kingma and Welling,, 2013), for example).

Note that we explicitly require the weight matrix $W_{2}$ to come with a bias. For the other linear transformations, we have omitted the bias term for notational simplicity. One can check that this definition is consistent with Definition 2. Namely, a net with an encoder block is indeed a type of stochastic net. When the training loss only consists of the MSE, it should follow as an immediate corollary of Theorem 1 that the variance converges to zero. However, the prior term complicates the problem. The following definition specifies the type of the prior loss under our consideration.

We have abstracted away the actual details of the definitions of the prior loss. For our purpose, it is sufficient to say that the equation $[W_{1}h(x_{i})]_{j}=0$ means that the loss function encourages the posterior to have a zero mean and $[W_{2}h(x)+b]_{j}=1$ encourages a unit variance. As an example, one can check that the standard ELBO loss for VAE satisfies this definition. With this architecture, we prove a similar result. The proof is given in Section A.2.

One might wonder whether a vanishing prior strength makes this setting trivial. The proof shows that it is far from trivial and that even a prior with vanishing strength can have a very strong influence on how fast the variance decays towards zero. The proof suggests that if the prior strength decays as $1/d_{1}^{\alpha}$, the variance should decay roughly as $d_{1}^{-\frac{\min(1,\alpha)}{2}}$. Namely, the smaller the $\alpha$, the slower the rate of convergence towards $0$. Additionally, the fastest exponent at which the variance can decay is $-0.5$, significantly smaller than the case for Theorem 1, where the proof suggests an exponent of $-1$. This means that even a vanishing regularization strength can have a strong impact on the prediction variance of the model. Quantitatively, as $\alpha$ approaches zero, the variance can decay arbitrarily slowly. Qualitatively, our result implies that having a regularization term or not qualitatively changes the nature of the global minima of the stochastic model. Also, we note that our problem setting is formally equivalent to a generalized⁵⁵5It is generalized because the MSE is not necessarily a reconstruction loss. $\beta$-VAE (Higgins et al.,, 2016) with $\beta\sim 1/width^{\alpha}$. While it is rarely the case for practitioners to scale $\beta$ towards zero, we stress that the main value of Theorem 2 is the theoretical insight it offers on the qualitative distinction between having a prior term and having not.

In this experiment, we let the target function be $y=\sin(x)+\eta$ for $x$ uniformly sampled from the domain $[-3,3]$. The target is corrupted by a weak noise $\eta\sim\mathcal{N}(0,0.1)$. The model is a feedforward network with tanh activation with the number of neurons $1\to d\to 1$ where $d\in\{10,50,500\}$, and a dropout with dropout probability $p=0.1$ is applied in the hidden layer. See Figure 1.

We now systematically explore the prediction variance of a model trained with dropout. As an extension of the previous experiment, we let both the input and the target be vectors such that $x,\ y\in\mathbb{R}^{d}$. The target function is $y_{i}=x_{i}^{3}+\eta_{i}$ for $i\in\{1,...,d\}$, where $x_{i},\eta_{i}\sim\mathcal{N}(\mathbf{0},1)$ and the noise $\eta\in\mathbb{R}^{d}$ is also a vector. We let $d=20$ and sample $1000$ data point pairs for training. The model is a three-layer MLP with ReLU activation functions, with the number of neurons $20\to d_{h}\to 20$, where $d_{h}$ is the width of the hidden layer. Dropout is applied to the post-activation values of the hidden layer. In the experiments, we set the dropout probability $p$ to be 0.1, and we independently sample outputs $3000$ times to estimate the prediction variance. Training proceeds with Adam for $4500$ steps, with an initial learning rate of $0.01$, and the learning rate is decreased by a factor of 10 every $1500$ step. See Figure 2(a) for a log-log plot of width vs. prediction variance. We see that the prediction variance of the model on the training set decreases towards zero as the width increases, as our theory predicts. For completeness, we also plot the prediction variance for an independently sampled test set for completeness. We see that, for this task, the prediction variance of the test points agree well with that of the training set. A linear regression on the slope of the tail of the width-variance curve shows that the variances decrease roughly as $d^{-0.7}$, close to what our proof suggests ($d^{-1}$); we hypothesize that the exponent is slightly smaller than $1$ because the training is stopped at a finite time and the model has not fully reached the global minimum.⁶⁶6Moreover, with a finite-learning rate, SGD is a biased estimator of a minimum (Ziyin et al., 2021a, )..

In this section, we conduct experiment on a $\beta-$VAE with latent Gaussian noise. The input data $x\in\mathbb{R}^{d}$ is sampled from a standard Gaussian distribution with $d=20$. We generate $100$ data points for training. The VAE employs a standard encoder-decoder architecture with ReLU nonlinearity. The encoder is a two-layer feedforward network with neurons $20\to 32\to 2\times d_{h}$. The decoder is also a two-layer feedforward network with architecture $d_{h}\to 32\to 20$. Note that our theory requires the prior term $\ell_{\rm vae}$ not to increase with the width; we, therefore, choose $\beta=0.1/d_{h}$. The training objective is the minus standard Evidence Lower Bound (ELBO) composed of reconstruction error and the KL divergence between the parameterized variable and standard Gaussian. We independently sample outputs $100$ times to estimate the prediction variance for estimating the variance. The results in Fig. 2(b) show that the variances of both training and test set decrease as the width increases and follow the same pattern.

Weight decay often has the Bayesian interpretation of imposing a normal distribution prior over the variables, and sometimes it is believed to prevent the model from making a deterministic prediction (Gal and Ghahramani,, 2016). We, therefore, also perform the same experiments as in the previous two subsections, with a weight decay strength of $\lambda=5e-4$. See Appendix Sec. B.1 for the results. We notice that the experimental result is similar and that applying weight decay is not sufficient to prevent the model from reaching a vanishing variance. Conventionally, we expect the prediction of a dropout net to be of order $\lambda$ and not directly dependent on the width (Gal and Ghahramani,, 2016); however, this experiment suggests that width may wield a stronger influence on the prediction variance than the weight decay strength. Since a model cannot reach the actual global minimum when weight decay is applied, this experiment is beyond the applicability of our theory; therefore, this result suggests a conjecture that even in a local minimum, it can still be highly likely for stochastic models to reach a solution with vanishing variance, and proving or disproving this conjecture can be an important future task.