Beyond IID weights: sparse and low-rank deep Neural Networks are also
  Gaussian Processes
The infinitely wide neural network has been proven a useful and manageable
mathematical model that enables the understanding of many phenomena appearing
in deep learning. One example is the convergence of random deep networks to
Gaussian processes that allows a rigorous analysis of the way the choice of
activation function and network weights impacts the training dynamics. In this
paper, we extend the seminal proof of Matthews et al. (2018) to a larger class
of initial weight distributions (which we call PSEUDO-IID), including the
established cases of IID and orthogonal weights, as well as the emerging
low-rank and structured sparse settings celebrated for their computational
speed-up benefits. We show that fully-connected and convolutional networks
initialized with PSEUDO-IID distributions are all effectively equivalent up to
their variance. Using our results, one can identify the Edge-of-Chaos for a
broader class of neural networks and tune them at criticality in order to
enhance their training. Moreover, they enable the posterior distribution of
Bayesian Neural Networks to be tractable across these various initialization
schemes.
Thiziri Nait-Saada, Alireza Naderi, Jared Tanner
# Beyond iid weights: sparse and low-rank deep neural networks are also Gaussian Processes

###### Abstract

The infinitely wide neural network has proven a useful and manageable mathematical model that enables the understanding of many phenomena appearing in deep learning. One example is the convergence of random deep networks to Gaussian Processes that allows a rigorous analysis of the way the choice of activation function and network weights impacts the training dynamics. In this paper, we extend the seminal proof of Matthews et al. (2018) to a larger class of initial weight distributions (which we call Pseudo-iid), including the established cases of iid and orthogonal weights, as well as the emerging low-rank and structured sparse settings celebrated for their computational speed-up benefits. We show that fully connected and convolutional networks initialized with Pseudo-iid distributions are all effectively equivalent up to their variance. Using our results, one can identify the Edge-of-Chaos for a broader class of neural networks and tune them at criticality in order to enhance their training.

## 1 Introduction

Deep neural networks are often studied at random initialization, where in the limit of infinite width, they have been shown to generate intermediate entries which approach Gaussian Processes. Seemingly this was first studied for one-layer networks in Neal (1996) when the weight matrices have identically and independently distributed (iid) Gaussian entries, and became a popular model for deep networks following the seminal results for deep fully connected networks in Lee et al. (2017) and Matthews et al. (2018). Specifically, the latter formulated a proof strategy that quickly became a cornerstone in the field, paving the way for various extensions of the Gaussian Process limit (see Section 1.1). The resulting Gaussian Process has proven to be of practical interest for at least two reasons. First, it emerges as a mathematically motivated choice of prior in the context of Bayesian Neural Networks. Secondly, it helps the analysis of exploding and vanishing gradient phenomena as done in Schoenholz et al. (2017) and Pennington et al. (2018), amongst other network properties.

In this paper, we extend the proof of the Gaussian Process from Matthews et al. (2018) to a larger class of initial weight distributions (which we call Pseudo-iid). Pseudo-iid distributions include the already well-known cases of iid and orthogonal weights. Moreover, we explore other important settings that conform to our more general conditions (e.g. structured sparse and low-rank), yet for which the Gaussian Process limit could not be derived previously due to their violation of the iid assumption.

Why studying low-rank and structured sparse networks at initialization?In recent years, deep learning models have significantly increased in size, while concurrently smaller AI chips have been preferred. The most widely studied approaches for bridging the gap between these two are to reduce the number of parameters by either pruning the network to be sparse or using low-rank factors. The _lottery ticket hypothesis_Frankle and Carbin (2019) has given empirical evidence of the existence of pruned subnetworks (i.e. _winning tickets_) achieving test accuracy that is comparable and even superior to their original network. Unfortunately most pruning methods produce sparsity patterns that are unstructured, thereby limiting potential hardware accelerations (see Hoefler et al. (2021)). Remarkably, Chen et al. (2022) revealed some _winning tickets_ with structured sparsity, that
can be exploited by efficient hardware accelerators on GPUs and FPGAs to reduce greatly their in-memory access, storage, and computational burden (Zhu et al. (2020)). Similarly, low-rank models are similarly being used to improve the efficiency and accuracy of networks Osawa et al. (2017); Yang et al. (2020a,b); Tai et al. (2022); Price and Tanner (2023); Karimi Mahabadi et al. (2021). The emergence of low-rank structures within deep networks during training is also observed where the low-rank factors are associated with learned classes Han et al. (2022) and this phenomenon has motivated thresholding of singular values in Berlyand et al. (2023), which is observed to improve the accuracy similarly to the lottery ticket hypothesis.

If such computationally efficient winning tickets exist, why would one want to train a dense (or full-rank) network at a high computational price only to later reduce it to be sparse (or low-rank)? This question has given rise to single-shot approaches, often referred to as Pruning at Initialization (PaI) such as SNIP Lee et al. (2019), SynFlow Tanaka et al. (2020), GraSP Wang et al. (2020a), NTT Liu and Zenke (2020). Consequently, the present work delves into the properties of such random networks at initialization, within the larger framework of Pseudo-iid regimes.

### Related work

To the best of our knowledge, the Gaussian Process behaviour in the infinite width regime was first established by Neal (1996) in the case of one-layer fully connected networks when the weights are iid sampled from standard distributions. The result has then been extended in Matthews et al. (2018) to deep fully connected networks, where the depth is fixed, the weights are distributed as iid Gaussians and the widths of the hidden layers grow jointly to infinity. Jointly scaling the network width substantially distinguishes their method of proof from the approach taken in Lee et al. (2017), where the authors considered a sequential limit analysis through layers. That is, analyzing the limiting distribution at one layer when the previous ones have already converged to their limiting distributions, as in Lee et al. (2017), is significantly different from examining it when the previous layers are jointly converging to their limits, as is done in Matthews et al. (2018). The latter approach now serves as a foundational machinery in the infinite width analysis from which numerous subsequent works followed. This paper is one of them.

Since this Gaussian Process behaviour has been established, two main themes of research have further been developed. The first consists of the extension of such results for more general and complex architectures such as convolutional networks with many channels Novak et al. (2020), Garriga-Alonso et al. (2019) or any modern architectures composed of fully connected, convolutional or residual connections, as summarized in Yang (2021), using the Tensor Program terminology. The second line of work concerns the generalization of this Gaussian Process behaviour to other possible weight distributions, such as orthogonal weights in Huang et al. (2021) or, alternatively, any iid weights with finite moments as derived in Hanin (2021). Note that the orthogonal case does not fit into the latter as entries are exchangeable but not independent. The same kind of results for general architectures in the iid setting have been derived in Golikov and Yang (2022) and Yang (2021). Our contribution fits into this line of research, where we relax the independence requirement of the weight matrix entries and instead consider Pseudo-iid distributions made of uncorrelated and exchangeable random variables. This broader class of distributions, Definition 2, enables us to present a unified proof that strictly generalizes the approaches taken so far and encompasses all of them, for two types of architectures, namely, fully connected and convolutional networks.

### Organization of the paper

In Section 2, we focus on fully connected neural networks, formally stating the Pseudo-iid regime and its associated Gaussian Process (GP) in Theorem 1, providing a rigorous proof in Appendix A. Our Pseudo-iid regime unifies the previously studied settings of iid and orthogonal weights, while also allowing for novel initialization schemes conform to the Gaussian Process limit, such as low-rank and structured sparse weights. We expand the definition of Pseudo-iid distributions to convolutional kernels and extend the resulting Gaussian Process limit to convolutional neural networks (CNNs) in Section 2.2. In Section 3, we provide examples of Pseudo-iid distributions in practice, supporting our theoretical results with numerical simulations. Moreover, we allude to two important consequences of the GP limit in 3.3, namely the analysis of stable initialization of deep networks on the so-called Edge-of-Chaos (EoC), as well as the implication of the limiting Gaussian
distribution on Bayesian inference, in our more expansive Pseudo-iid regime. Lastly, in Section 4, we review our main contributions and put forward some further research directions.

## 2 Gaussian Process behaviour in the Pseudo-iid regime

We consider an untrained fully connected neural network with width \(N_{\ell}\) at layer \(\ell\in\{1,\cdots,L+1\}\). Its weights \(W^{(\ell)}\in\mathbb{R}^{N_{\ell}\times N_{\ell-1}}\) and biases \(b^{(\ell)}\in\mathbb{R}^{N_{\ell}}\) at layer \(\ell\) are sampled from centered probability distributions, respectively \(\mu_{W}^{(\ell)}\) and \(\mu_{b}^{(\ell)}\). Starting with such a network, with nonlinear activation \(\phi:\mathbb{R}\to\mathbb{R}\), the propagation of any input data vector \(z^{(0)}\coloneqq x\in\mathcal{X}\subseteq\mathbb{R}^{N_{0}}\) through the network is given by the following equations,

\[h_{i}^{(\ell)}(x)=\sum_{j=1}^{N_{\ell-1}}W_{ij}^{(\ell)}z_{j}^{(\ell-1)}(x)+b_ {i}^{(\ell)},\qquad z_{j}^{(\ell)}(x)=\phi(h_{j}^{(\ell)}(x)), \tag{1}\]

where \(h^{(\ell)}(x)\in\mathbb{R}^{N_{\ell}}\) is referred to as the preactivation vector at layer \(\ell\), or the feature maps.

Pseudo-iid distributions (formally introduced in Def. 2) only need to meet some moment conditions along with an exchangeability requirement that we define below.

**Definition 1** (Exchangeability).: _Let \(X_{1},\cdots,X_{n}\) be scalar, vector or matrix-valued random variables. We say \((X_{i})_{i=1}^{n}\) are exchangeable if their joint distribution is invariant under permutations, i.e. \((X_{1},\cdots,X_{n})\stackrel{{ d}}{{=}}(X_{\sigma(1)},\cdots,X_{ \sigma(n)})\) for all permutations \(\sigma:[n]\to[n]\). A random matrix is called row- (column-) exchangeable if its rows (columns) are exchangeable random vectors, respectively._

A row-exchangeable and column-exchangeable weight matrix \(W\in\mathbb{R}^{m\times n}\) is not in general entrywise exchangeable, which means its distribution is not typically invariant under arbitrary permutations of its entries; particularly, out of \((mn)!\) possible permutations of the entries, \(W\) only needs to be invariant under \(m!n!\) of them -- an exponentially smaller number. We can now define the family of Pseudo-iid distributions.

**Definition 2** (Pseudo-iid).: _Let \(m,n\) be two integers and \(a\in\mathbb{R}^{n}\) be any fixed vector. We will say that the random matrix \(W=(W_{ij})\in\mathbb{R}^{m\times n}\) is in the Pseudo-iid distribution with parameter \(\sigma^{2}\) if_

1. _the matrix is row-exchangeable and column-exchangeable,_
2. _its entries are centered, uncorrelated, with variance_ \(\mathbb{E}(W_{ij}^{2})=\frac{\sigma^{2}}{n}\)_,_
3. \(\mathbb{E}\big{|}\sum_{j=1}^{n}a_{j}W_{ij}\big{|}^{8}=K\|\mathbf{a}\|_{2}^{8} n^{-4}\) _for some constant_ \(K\)_,_
4. _and_ \(\lim_{n\to\infty}\frac{n^{2}}{\sigma^{4}}\mathbb{E}(W_{ia,j}W_{ii,j}W_{i,j^{ \prime}}W_{i,d,j^{\prime}})=\delta_{ia,i_{k}}\delta_{i_{c},i_{d}}\)_, for all_ \(j\neq j^{\prime}\)_._

_When \(W^{(1)}\) has iid Gaussian entries and the other weight matrices \(W^{(\ell)},\,2\leq\ell\leq L+1\), of a neural network (see equation 1) are drawn from a Pseudo-iid distribution, we will say that the network is under the Pseudo-iid regime._

The Pseudo-iid conditions (i)-(iv) are included to serve following purposes: Conditions (i) and (ii) allow for non-iid initializations such as permuted block sparse and low-rank while being sufficient for a CLT-type argument to hold. The variance scaling prevents the preactivations in equation 2 from blowing up as the network's size grows. Condition (iii) requires the moments of the possibly dependent weights to behave like independent ones, while condition (iv) constrains their cross-correlations to vanish fast enough. Together they maintain that in the limit, these variables interact as if they were independent, and this is all our proof needs. The importance of condition (iii) is expanded upon in Appendix G. It is worth noting that because the first layer's weight matrix has only one dimension scaling up with \(n\), the tension between the entries' dependencies and the network
size cannot be resolved, thus we take these weights to be iid 1. Whether some of the conditions are redundant still remains an open question.

Footnote 1: In fact, it is enough for \(W^{(1)}\) to have iid rows. We assume Gaussian iid entries for simplicity, yet this can be further relaxed at the cost of a more complicated formula for the first covariance kernel (see equation 21). Examples of row iid distributions that could be used for \(W^{(1)}\) which are similar to the low-rank and structured ensembles include the product \(W^{(1)}=DPC\) of a diagonal matrix \(D\in\mathbb{R}^{N_{1}[n]\times N_{1}[n]}\) with entries drawn iid, \(P\) a random permutation matrix, and \(C\in\mathbb{R}^{N_{1}[n]\times N_{0}}\) a non-zero low rank or structured sparse matrix.

Throughout this paper, we consider a specific set of activation functions that satisfy the so-called linear envelope property, Definition 3, satisfied by most activation functions used in practice (ReLu, Softmax, Tanh, HTanh, etc.).

**Definition 3**.: _(Linear envelope property) A function \(\phi:\mathbb{R}\to\mathbb{R}\) is said to satisfy the linear envelope property if there exist \(c,M\geq 0\) such that, for any \(x\in\mathbb{R}\),_

\[|\phi(x)|\leq c+M|x|.\]

### The Pseudo-iid regime for fully connected networks

Our proofs of Pseudo-iid networks converging to Gaussian Processes are done in the more sophisticated simultaneous width growth limit as pioneered by Matthews et al. (2018). For a review of the literature on deep networks with sequential vs. simultaneous scaling, see Section 1.1. One way of characterizing such a simultaneous convergence over all layers is to consider that all widths \(N_{\ell}\) are increasing functions of one parameter, let us say \(n\), such that, as \(n\) grows, all layers' widths increase: \(\forall\ell\in\{1,\cdots,L\},N_{\ell}\coloneqq N_{\ell}[n]\). We emphasize this dependence on \(n\) by appending an index \(X[n]\) to the random variables \(X\) when \(n\) is finite and denote by \(X[*]\) its limiting random variable, corresponding to \(n\to\infty\). The input dimension \(N_{0}\) and the final output layer dimension \(N_{L+1}\) are finite thus they do not scale with \(n\). Moreover, the input data are assumed to come from a countably infinite input space \(\mathcal{X}\) (see A.1). Equation equation 1 can thus be rewritten, for any \(x\in\mathcal{X}\), as

\[h_{i}^{(\ell)}(x)[n]=\sum_{j=1}^{N_{\ell-1}[n]}W_{ij}^{(\ell)}z_{j}^{(\ell-1)} (x)[n]+b_{i}^{(\ell)},\qquad z_{j}^{(\ell)}(x)[n]=\phi(h_{j}^{(\ell)}(x)[n]), \tag{2}\]

and the associated Gaussian Process limit is given in Theorem 1, which we can now state.

**Theorem 1** (GP limit for fully connected Pseudo-iid networks).: _Suppose a fully connected neural network as in equation 2 is under the Pseudo-iid regime with parameter \(\sigma_{W}^{2}\) and the activation satisfies the linear envelope property Def. 3. Let \(\mathcal{X}\) be a countably-infinite set of inputs. Then, for every layer \(2\leq\ell\leq L+1\), the sequence of random fields \((i,x)\in[N_{\ell}]\times\mathcal{X}\mapsto h_{i}^{(\ell)}(x)[n]\in\mathbb{R}^ {N_{\ell}}\) converges in distribution to a centered Gaussian Process \((i,x)\in[N_{\ell}]\times\mathcal{X}\mapsto h_{i}^{(\ell)}(x)[*]\in\mathbb{R}^ {N_{\ell}}\), whose covariance function is given by_

\[\mathbb{E}\Big{[}h_{i}^{(\ell)}(x)[*]\cdot h_{j}^{(\ell)}(x^{\prime})[*]\Big{]} =\delta_{i,j}K^{(\ell)}(x,x^{\prime}), \tag{3}\]

_where_

\[K^{(\ell)}(x,x^{\prime})=\begin{cases}\sigma_{b}^{2}+\sigma_{W}^{2}\mathbb{E }_{(u,v)\sim\mathcal{N}(\mathbf{0},K^{(\ell-1)}(x,x^{\prime}))}[\phi(u)\phi(v) ],&\ell\geq 2,\\ \sigma_{b}^{2}+\frac{\sigma_{W}^{2}}{N_{0}}\langle x,x^{\prime}\rangle,&\ell=1.\end{cases} \tag{4}\]

### The Pseudo-iid regime for Convolution Neural Networks

We consider a CNN with \(C_{\ell}\) number of channels at layer \(\ell\in\{1,\cdots,L+1\}\) and two-dimensional convolutional filters \(\mathbf{U}_{i,j}^{(\ell)}\in\mathbb{R}^{k\times k}\) mapping the input channel \(j\in\{1,\cdots,C_{\ell-1}\}\) to the output channel \(i\in\{1,\cdots,C_{\ell}\}\). The input signal \(\mathbf{X}\) (also two-dimensional) has \(C_{0}\) channels \(\mathbf{x}_{1},\cdots,\mathbf{x}_{C_{0}}\)
and its propagation through the network is given by

\[\mathbf{h}_{i}^{(\ell)}(\mathbf{X})[n]=\begin{cases}b_{i}^{(1)}\mathbf{1}+\sum_{j= 1}^{C_{0}}\mathbf{U}_{i,j}^{(1)}\star\mathbf{x}_{j},&\ell=1,\\ b_{i}^{(\ell)}\mathbf{1}+\sum_{j=1}^{C_{\ell-1}}\mathbf{U}_{i,j}^{(\ell)}\star \mathbf{z}_{j}^{(\ell-1)}(\mathbf{X})[n],&\ell\geq 2,\end{cases} \tag{5}\]

In equation 5, the symbol \(\mathbf{1}\) should be understood as having the same size as the convolution output, the non-linearity \(\phi(\cdot)\) is applied entrywise, and we emphasize on the simultaneous scaling with \(n\) through appending \([n]\) to the feature maps \(\mathbf{h}_{i}^{(\ell)}(\mathbf{X})[n]\). We denote spatial (multi-) indices by boldface Greek letters \(\mathbf{\mu}\), \(\mathbf{\nu}\), etc., that are ordered pairs of integers taking values in the range of the size of the array. For example, if \(\mathbf{X}\) is an RGB (\(C_{0}=3\)) image of \(H\times D\) pixels, \(i=2\), and \(\mathbf{\mu}=(\alpha,\beta)\), then \(X_{i,\mathbf{\mu}}\) returns the green intensity of the \((\alpha,\beta)\) pixel. Moreover, we define \([\mathbf{\mu}]\) to be the patch centered at the pixel \(\mathbf{\nu}\) covered by the filter, e.g. if \(\mathbf{\mu}=(\alpha,\beta)\) and the filter covers \(k\times k=(2k_{0}+1)\times(2k_{0}+1)\) pixels, then \(\llbracket\mathbf{\mu}\rrbracket=\{(\alpha^{\prime},\beta^{\prime})\mid\alpha-k_{ 0}\leq\alpha^{\prime}\leq\alpha+k_{0},\ \beta-k_{0}\leq\beta^{\prime}\leq\beta+k_{0}\}\), with the usual convention of zero-padding for the out-of-range indices. Sufficient conditions for Pseudo-iid CNNs to converge to a Gaussian Process in the simultaneous scaling limit are given in Definition 4.

**Definition 4** (Pseudo-iid for CNNs).: _Consider a CNN with random filters and biases \(\{\mathbf{U}_{i,j}^{(\ell)}\}\) and \(\{b_{i}^{(\ell)}\}\) as in equation 5. It is said to be in the Pseudo-iid regime with parameter \(\sigma^{2}\) if \(\mathbf{U}^{(1)}\) has iid\(\mathcal{N}(0,\frac{\sigma^{2}}{C_{0}})\) entries and, for \(2\leq\ell\leq L+1\), such that for any fixed vector \(\mathbf{a}:=(a_{j,\mathbf{\nu}})_{j,\mathbf{\nu}}\),_

1. _the convolutional kernel_ \(\mathbf{U}^{(\ell)}\in\mathbb{R}^{C_{\ell}\times C_{\ell-1}\times k\times k}\) _is row-exchangeable and column-exchangeable, that is its distribution is invariant under permutations of the first and second indices,_
2. _the filters' entries are centered, uncorrelated, with variance_ \(\mathbb{E}[(\mathbf{U}_{i,j,\mathbf{\mu}}^{(\ell)})^{2}]=\sigma^{2}/C_{\ell-1}\)_,_
3. \(\mathbb{E}\big{|}\sum_{j=1}^{C_{\ell-1}}\sum_{\mathbf{\nu}}a_{j,\mathbf{\nu}} \mathbf{U}_{i,j,\mathbf{\nu}}^{(\ell)}\big{|}^{8}=K\|\mathbf{a}\|_{2}^{8}(C_{\ell -1})^{-4}\) _for some constant_ \(K\)_,_
4. _and_ \(\lim_{n\to\infty}\frac{C_{\ell-1}[n]^{2}}{\sigma^{4}}\mathbb{E}\big{(} \mathbf{U}_{i_{\alpha},j,\mathbf{\mu}_{a}}^{(\ell)}\mathbf{U}_{i_{b},j,\mathbf{\mu}_{ b}}^{(\ell)}\mathbf{U}_{i_{c},j^{\prime},\mathbf{\mu}_{a}}^{(\ell)}\mathbf{U}_{i _{d},j^{\prime},\mathbf{\mu}_{d}}^{(\ell)}\big{)}=\delta_{i_{a},i_{b}}\delta_{i_{ c},i_{d}}\delta_{\mathbf{\mu}_{a},\mathbf{\mu}_{b}}\delta_{\mathbf{\mu}_{c},\mathbf{\mu}_{d}}\) _for all_ \(j\neq j^{\prime}\)_._

Similar to the fully connected case, the filters' entries can now exhibit some interdependencies as long as their cross-correlations vanish at a fast enough rate dictated by conditions (iii)-(iv).

**Theorem 2** (GP limit for CNN Pseudo-iid networks).: _Suppose a CNN as in equation 5 is under the Pseudo-iid regime with parameter \(\sigma_{W}^{2}\) and the activation satisfies the linear envelope property Def. 3. Let \(\mathcal{X}\) be a countably-infinite set of inputs and \(\mathbf{\mu}\in\mathcal{I}\) denote a spatial (multi-) index. Then, for every layer \(1\leq\ell\leq L+1\), the sequence of random fields \((i,\mathbf{X},\mathbf{\mu})\in[C_{\ell}]\times\mathcal{X}\times\mathcal{I}\mapsto h_ {i,\mathbf{\mu}}^{(\ell)}(\mathbf{X})[n]\) converges in distribution to a centered Gaussian Process \((i,\mathbf{X},\mathbf{\mu})\in[C_{\ell}]\times\mathcal{X}\times\mathcal{I}\mapsto h _{i,\mathbf{\mu}}^{(\ell)}(\mathbf{X})[*]\), whose covariance function is given by_

\[\mathbb{E}\Big{[}h_{i,\mathbf{\mu}}^{(\ell)}(\mathbf{X})[*]\cdot h_{j,\mathbf{\mu}^{ \prime}}^{(\ell)}(\mathbf{X}^{\prime})[*]\Big{]}=\delta_{i,j}\Big{(}\sigma_{b }^{2}+\sigma_{W}^{2}\sum_{\mathbf{\nu}\in[\mathbf{\mu}]\cap[\mathbf{\mu}^{\prime}]}K_{\bm {\nu}}^{(\ell)}(\mathbf{X},\mathbf{X}^{\prime})\Big{)}, \tag{6}\]

_where_

\[K_{\mathbf{\nu}}^{(\ell)}(\mathbf{X},\mathbf{X}^{\prime})=\begin{cases}\mathbb{E}_{ (u,v)\sim\mathcal{N}(\mathbf{0},K_{\mathbf{\nu}}^{(\ell-1)}(\mathbf{X},\mathbf{X}^ {\prime}))}[\phi(u)\phi(v)],&\ell\geq 2,\\ \frac{1}{C_{0}}\sum_{i=1}^{C_{0}}X_{i,\mathbf{\nu}}X_{i,\mathbf{\nu}}^{\prime},&\ell=1. \end{cases} \tag{7}\]

These equations resemble those derived in the fully connected case, apart from the introduction of additional terms accounting for the pixels in some patch \(\mathbf{\mu}\), that vanish when the filter kernel size is reduced to \(k=1\). Our proof of Theorem 2 can be found in Appendix B.
Pseudo-iid in practice

In this section, we illustrate some important Pseudo-iid distributions through examples. We further show that the Gaussian behaviour can serve, among others, two purposes of practical significance (see section 3.3). First, we elucidate the analysis of signal propagation in neural networks. Secondly, this Gaussian Process limit simplifies the initially complicated task of choosing priors for Bayesian Neural Networks.

### Examples of Pseudo-iid distributions

We elaborate here on some typical initialization schemes and their belonging to the Pseudo-iid class. To the best of our knowledge, rigorous proofs of the Gaussian Process convergence in the literature are restricted to the iid cases only (see Section 1.1). An important non-iid case that also results in a Gaussian Process is the random orthogonal initialization, partially derived in Huang et al. (2021) for fully connected networks. Nonetheless, we regret the authors remained elusive on the treatment of the first layer, an aspect which we have since then enhanced. Our proposed Pseudo-iid regime encompasses both iid and orthogonal cases (see Appendix D) but also allows for a broader class of weight distributions such as random low-rank or structured sparse matrices, for which the Gaussian Process limit has not been established before, despite their practical significance (see 1.

Low-rank weights.Low-rank structures are widely recognized for speeding up matrix multiplications and can be used to reduce memory requirements, see 1. Whilst such structures inevitably impose dependencies between the weight matrix entries \(A\in\mathbb{R}^{m\times n}\), thus breaking the iid assumption, Nait Saada and Tanner (2023) introduced a low-rank framework that falls within our Pseudo-iid regime. Let \(C\coloneqq[C_{1},\cdots,C_{r}]\in\mathbb{R}^{m\times r}\) be a uniformly drawn orthonormal basis for a random \(r\)-dimensional subspace. Suppose \(P=(P_{ij})\in\mathbb{R}^{r\times n}\) has iid entries \(P_{ij}\overset{\text{iid}}{\sim}\mathcal{D}\). Setting \(A\coloneqq CP\), the columns of \(A\) are spanned by those of \(C\). The row and column exchangeability of \(A\) follows immediately from that of \(C\) and \(P\), and the moment conditions (iii) and (iv) are controlled by the choice of distribution \(\mathcal{D}\). Direct computation of the four-cross product that appears in condition (iv) gives us

\[\mathbb{E}(A_{i_{a},1}A_{i_{b},1}A_{i_{c},2}A_{i_{d},2})=s^{2}\sum_{1\leq k,k^ {\prime}\leq r}\mathbb{E}\Big{[}C_{i_{a},k}C_{i_{b},k}C_{i_{c},k^{\prime}}C_{ i_{d},k^{\prime}}\Big{]},\]

where \(s\coloneqq\mathbb{E}(P_{1,1}^{2})\). Using the expression in Lemma 3 of Huang et al. (2021) we can calculate the above expectation and deduce condition (iv) when \(r\) is linearly proportional to \(m\).

Structured sparse weights.Block-wise pruned networks have recently been under extensive study for their efficient hardware implementation Dao et al. (2022b). Once the sparsifying mask is fixed, we may apply random row and column permutations on the weight matrices without compromising the accuracy or the computational benefit. Let \(A=(A_{ij})\in\mathbb{R}^{m\times n}\) have iid entries \(A_{ij}\overset{\text{iid}}{\sim}\mathcal{D}\) and \(B\in\mathbb{R}^{m\times n}\) be the binary block-sparse mask. Let \(\tilde{A}=P_{m}(A\odot B)P_{n}\), where \(P_{m}\) and \(P_{n}\) are random permutation matrices of size \(m\times m\) and \(n\times n\) respectively, and \(\odot\) represents entrywise multiplication. Then, by construction, \(\tilde{A}\) is row- and column-exchangeable and, for suitable choices of underlying distribution \(\mathcal{D}\), it satisfies the moment conditions of Definition 2. Appendix H contains some illustrations for reference.

Orthogonal CNN filters.Unlike the fully connected case, it is not obvious how to define the orthogonality of a convolutional layer and, once defined, how to randomly generate such layers for initialization. Xiao et al. (2018) defines an orthogonal convolutional kernel \(\mathbf{U}\in\mathbb{R}^{c_{out}\times c_{in}\times k\times k}\) made of \(c_{out}\) filters of size \(k\times k\) through the energy preserving property \(\|\mathbf{U}+\mathbf{\times}\|_{2}=\|\mathbf{X}\|_{2}\), for any signal \(\mathbf{X}\) with \(c_{in}\) input channels. Wang et al. (2020) requires the matricized version of the kernel to be orthogonal, while Qi et al. (2020) gives a more stringent definition imposing isometry, i.e.

\[\sum_{i=1}^{c_{out}}\mathbf{U}_{i,j}\star\mathbf{U}_{i,j^{\prime}}=\begin{cases} \mathbf{\delta},&j=j^{\prime},\\ 0,&\text{otherwise},\end{cases}\]

where \(\mathbf{\delta}\) is 1 at \((0,0)\), and \(0\) elsewhere. Another definition in Huang et al. (2021) calls for orthogonality of "spatial" slices \(\mathbf{U}_{\mathbf{\mu}}\in\mathbb{R}^{c_{out}\times c_{in}}\), for all positions \(\mathbf{\mu}\).
We take a different approach than Wang et al. (2020b) for matricizing the tensor convolution operator, setting the stride to 1 and padding to 0: reshape the kernel \(\mathbf{U}\) into a matrix \(\mathbf{\tilde{U}}\in\mathbb{R}^{c_{out}\times k^{2}c_{in}}\) and unfold the signal \(\mathbf{X}\) into \(\mathbf{\tilde{X}}\in\mathbb{R}^{k^{2}c_{in}\times d}\), where \(d\) is the number of patches depending on the sizes of the signal and the filter. This allows \(\mathbf{\tilde{U}}\) to be an arbitrary unstructured matrix rather than the doubly block-Toeplitz matrix in Wang et al. (2020b), as shown in Figure 1. Matricizing the tensor convolution operator imposes the structure on the signal \(\mathbf{\tilde{X}}\) rather than the filter \(\mathbf{\tilde{U}}\). Orthogonal (i.e. energy-preserving) kernels \(\mathbf{\tilde{U}}\) can then be drawn uniformly random with orthogonal columns, such that

\[\mathbf{\tilde{U}}^{\top}\mathbf{\tilde{U}}=\frac{1}{k^{2}}I \tag{8}\]

and then reshaped into the original tensor kernel \(\mathbf{U}\). Note that this construction is only possible when \(\mathbf{\tilde{U}}\) is a tall matrix with trivial null space, that is when \(c_{out}\geq k^{2}c_{in}\), otherwise the transpose might be considered. We emphasize that equation 8 is a sufficient (and not necessary) condition for \(\mathbf{U}\) to be energy-preserving, since \(\mathbf{\tilde{X}}\), by construction, belongs to a very specific structure set \(T\subseteq\mathbb{R}^{k^{2}c_{in}\times d}\), and, therefore, \(\mathbf{\tilde{U}}\) only needs to preserve norm on \(T\) (and not everywhere). Therefore, we do not claim the generated orthogonal convolutional kernel \(\mathbf{U}\) is "uniformly distributed" over the set of all such kernels.

Now let us check that what we defined to be orthogonal filters in the previous section verify the conditions of Definition 4. Each filter \(\mathbf{U}_{i,j}\) is flattened as \(\mathbf{\tilde{U}}_{i,j}\in\mathbb{R}^{1\times k^{2}}\) and forms part of a row of \(\mathbf{\tilde{U}}\) as shown below:

\[\mathbf{\tilde{U}}=\left[\begin{array}{cccc}\mathbf{\tilde{U}}_{1,1}& \mathbf{\tilde{U}}_{1,2}&\cdots&\mathbf{\tilde{U}}_{1,\text{e}in}\\ \mathbf{\tilde{U}}_{2,1}&\mathbf{\tilde{U}}_{2,2}&\cdots&\mathbf{\tilde{U}}_ {2,\text{e}in}\\ \vdots&\vdots&\ddots&\vdots\\ \mathbf{\tilde{U}}_{\text{e}out,1}&\mathbf{\tilde{U}}_{\text{e}out,2}&\cdots& \mathbf{\tilde{U}}_{\text{e}out,\text{e}in}\end{array}\right]. \tag{9}\]

Applying permutations on the indices \(i\) and \(j\) translates to permuting rows and "column blocks" of the orthogonal matrix \(\mathbf{\tilde{U}}\), which does not affect the joint distribution of its entries. Hence, the kernel's distribution is unaffected, and therefore \(\mathbf{U}\) is row- and column-exchangeable. The moment conditions are both straightforward to check as \(\mathbf{U}_{i,j,\boldsymbol{\mu}}=\mathbf{\tilde{U}}_{i,(j-1)k^{2}+\mu},\;\mu \in\{1,\cdots,k^{2}\}\), where \(\mu\) is the counting number of the pixel \(\boldsymbol{\mu}\). To check condition (iv), note that \(\mathbb{E}\big{(}\mathbf{U}_{i_{a},1,\boldsymbol{\mu}_{a}}^{(l)}\mathbf{U}_{i_ {b},1,\boldsymbol{\mu}_{b}}^{(l)}\mathbf{U}_{i_{c},2,\boldsymbol{\mu}_{c}}^{( l)}\mathbf{U}_{i_{d},2,\boldsymbol{\mu}_{d}}^{(l)}\big{)}=\mathbb{E}\big{(} \mathbf{\tilde{U}}_{i_{a},\mu_{a}}\mathbf{\tilde{U}}_{i_{b},\mu_{b}}\mathbf{ \tilde{U}}_{i_{c},k^{2}+\mu_{c}}\mathbf{\tilde{U}}_{i_{d},k^{2}+\mu_{d}}\big{)}\), that is a four-cross product of the entries of an orthogonal matrix, whose expectation is explicitly known to be \(\frac{C_{t}+1}{(C_{t}-1)C_{t}(C_{t}+2)}\delta_{i_{a},i_{b}}\delta_{i_{c},i_{d} }\delta_{\mu_{a},\mu_{b}}\delta_{\mu_{c},\mu_{d}}\)(Huang et al., 2021, Lemma 3).

Figure 1: There are various ways to compute convolutions \(\mathbf{U}\star\mathbf{X}\) between a tensor filter \(\mathbf{U}\) and a 2D signal \(\mathbf{X}\) (in the middle) from matrix multiplications. We illustrate the approach taken in Garriga-Alonso et al. (2019) on the left, where the reshaping procedure is applied to the filter, whilst the method we followed, shown on the right, consists of reshaping the signal instead in order to define special structures on the CNN filters such as orthogonality, sparsity and low-rank.
Simulations of the Gaussian Processes in Theorem 1 for fully connected networks with Pseudo-iid weights

Theorem 1 establishes that random fully connected Pseudo-iid networks converge to Gaussian Processes in the infinite width limit. Here we conduct numerical simulations which validate this for modest dimensions of width \(N_{\ell}=n\) for \(n=3\), \(30\), and \(300\). Fig. 2 shows histograms of the empirical distribution of the preactivations when the weights are drawn from various Pseudo-iid distributions. 2 Even at \(n=30\) there is excellent agreement of the histograms with the limiting variance predicted by our Theorem 1. These histograms in Fig. 2 are further quantified with Q-Q plots in Appendix E.

Footnote 2: Experiments conducted for Fig. 2-5 used a fully connected network with activation \(\phi(x)=\tanh(x)\), weight variance \(\sigma_{w}=2\) and without bias. Dropout used probability \(1/2\) of setting an entry to zero, low-rank used rank \(\lceil n/2\rceil\), and block-sparsity used randomly permuted block-diagonal matrices with block-size \(\lceil n/5\rceil\). The code to reproduce all these figures can be found at [https://shorturl.at/gNQO0](https://shorturl.at/gNQO0).

Fig. 3 illustrates the rate with which two independent inputs \(x_{a}\) and \(x_{b}\) generate feature maps corresponding to uncorrelated Gaussian Processes, when passing through a network initialized with Pseudo-iid weights. The experiment is done by examining the joint empirical distribution of \(h_{i}^{(\ell)}(x_{a})[n]\) and \(h_{i}^{(\ell)}(x_{b})[n]\) at the same neuron index \(i\) and layer \(\ell\) through the same network. The level curves depict the theoretical joint Gaussian distribution given by the covariance kernel in Theorem 1. Interestingly, we observe the fastest convergence to the Gaussian Process in the orthogonal initialization. The other Pseudo-iid distributions considered exhibit good agreement at \(n=30\) which improves at \(n=300\). Some more experiments for fully connected networks and CNNs initialized with orthogonal filters can be found in respectively, Appendix E and F.

### Implications of the Gaussian Process Limit

Bayesian Neural Network and Gaussian Process.As opposed to frequentist approaches, Bayesian inference aims at finding a predictive distribution that serves to quantify the model's predictions uncertainty and prevent overfitting. Starting from a prior, this posterior distribution is updated once the data is observed. Considering Bayesian neural networks, these priors are to be defined on the weights and biases if adopting a parametric perspective, or, alternatively, directly on the input-output function represented by the model. Note how one can jump from one approach to the other given

Figure 2: For different instances of the Pseudo-iid regime, in the limit, the preactivation given in the first neuron at the fifth layer tends to a Gausssian whose moments are given by Theorem 1. The experiments were conducted \(10000\) times on a random 7-layer deep fully connected network with input data sampled from \(\mathbb{S}^{8}\).
that random initialization of the parameters of a neural network induces a distribution on the input-output function \(h^{(L+1)}(x)\) of the network. Nonetheless, finding a meaningful prior over the space of functions is in general a delicate task.

In the large width limit, when the weights are drawn Pseudo-iid in a \(L\)-layer deep network, Theorem 1 secures that the function represented by the network is _exactly_ a Gaussian Process with known covariance kernel \(K^{(L+1)}\), which only depends on the variances \(\sigma_{b},\sigma_{W}\) and the activation function. Therefore, the exact posterior distribution can be computed, yielding the Neural Network Gaussian Process (NNGP) equivalence described in Novak et al. (2020). Since the covariance kernel obtained in our less restrictive setting recovers the one in Novak et al. (2020), Matthews et al. (2018) and Garriga-Alonso et al. (2019), their experiments still hold in our case and we refer to these works for practical considerations on how to compute the posterior distribution. Specifically, NNGP yields better accuracy compared to finite point-estimate neural networks Lee et al. (2017) and approximates finite Bayesian networks very well (Matthews et al., 2018, Section 5).

Edge of Chaos (EoC).Random initialization of deep neural networks plays an important role in their trainability by regulating how quantities like the variance of preactivations and the pairwise correlation between input signals propagate through layers. An inappropriate initialization causes the network to collapse to a constant function or become a chaotic mapping oversensitive to input perturbations. The Edge of Chaos (EoC) initialization strategy rectifies these issues, additionally causing the network to have gradients of consistent magnitude across layers. Calculation of the EoC requires integration with respect to the distribution of preactivations in intermediary layers, which, in general, is an intractable task. However, the Gaussian Process limit simplifies the EoC analysis as carried out in Poole et al. (2016) and Xiao et al. (2018). Moreover, under the same distributional assumption, Schoenholz et al. (2017); Pennington et al. (2018) show that initialization on the EoC achieves _dynamical isometry_, making backpropagation of errors stable. Our main contributions of Theorems 1 and 2 allow similar calculations to be made for Pseudo-iid networks such as the examples of low-rank, structured sparse, and orthogonal CNN as shown in Section 3.1. The EoC for low-rank networks has been calculated in Nait Saada and Tanner (2023) under the assumption of Theorems 1 and 2 which was at that time unproven. Interestingly, in terms of signal propagation, structured sparse or low-rank initializations are equivalent to their dense and full-rank counterparts up to a rescaling of the variance by a fractional factor of sparsity or rank.

Figure 3: The empirical joint distribution of the preactivations generated by two distinct inputs flowing through the network. The large width limiting distribution as defined in Theorem 1 is included as level curves. The input data \(x_{a},x_{b}\) are drawn iid from \(\mathbb{S}^{9}\) and \(10000\) experiments were conducted on a \(7\)-layer fully connected network. The horizontal and vertical axes in each subplot are respectively \(h_{1}^{(5)}(x_{a})\) and \(h_{1}^{(5)}(x_{b})\).
Conclusion

We proved in this paper a new Gaussian Process limit for deep neural networks initialized with possibly inter-dependent entries. Examples include orthogonal, low-rank and structured sparse networks which are particularly of interest due to their efficient implementation and their empirically observed enhanced accuracy. Our result makes possible exact Bayesian inference as well as tractable Edge of Chaos analysis, for a broader class of either fully connected or convolutional networks. We expect the present work paves the way for a better understanding of the training dynamics of the emerging deep neural networks with parsimonious representations.

## Acknowledgments

We thank Juba Nait Saada for insightful discussion and feedback on the manuscript, as well as Alex Cliffe for valuable comments and help in drawing Figure 1. Thiziri Nait Saada is supported by the UK Engineering and Physical Sciences Research Council (EPSRC) through the grant EP/W523781/1. Jared Tanner is supported by the Hong Kong Innovation and Technology Commission (InnoHK Project CIMDA).