On the Initialisation of Wide Low-Rank Feedforward Neural Networks
The edge-of-chaos dynamics of wide randomly initialized low-rank feedforward
networks are analyzed. Formulae for the optimal weight and bias variances are
extended from the full-rank to low-rank setting and are shown to follow from
multiplicative scaling. The principle second order effect, the variance of the
input-output Jacobian, is derived and shown to increase as the rank to width
ratio decreases. These results inform practitioners how to randomly initialize
feedforward networks with a reduced number of learnable parameters while in the
same ambient dimension, allowing reductions in the computational cost and
memory constraints of the associated network.
Thiziri Nait Saada, Jared Tanner
# On the Initialisation of Wide Low-Rank Feedforward Neural Networks

###### Abstract

The edge-of-chaos dynamics of wide randomly initialized low-rank feedforward networks are analyzed. Formulae for the optimal weight and bias variances are extended from the full-rank to low-rank setting and are shown to follow from multiplicative scaling. The principle second order effect, the variance of the input-output Jacobian, is derived and shown to increase as the rank to width ratio decreases. These results inform practitioners how to randomly initialize feedforward networks with a reduced number of learnable parameters while in the same ambient dimension, allowing reductions in the computational cost and memory constraints of the associated network.

Machine Learning, ICML

## 1 Introduction

Neural networks being applied to new settings, limiting transfer learning, are typically initialized with i.i.d. random entries. The edge-of-chaos theory of (Poole et al., 2016) determine the appropriate scaling of the weight matrices and biases so that intermediate layer representations (1) and the median of the input-output Jacobian's spectra (10) are to first order independent of the layer. Without this normalization there is typically an _exponential_ growth in the magnitude of these intermediate representations and gradients as they progress between layers of the network; such a disparity of scale inhibits the early training of the network (Glorot & Bengio, 2010).

For instance, consider an untrained fully connected neural network whose weights and biases are set to be respectively identically and independently distributed with respect to a Gaussian distributions: \(W_{ij}^{(l)}\sim\mathcal{N}(0,\frac{\sigma_{W}^{2}}{N_{l-1}})\), \(b_{i}^{(l)}\sim\mathcal{N}(0,\sigma_{b}^{2})\) with \(N_{l}\) the width at layer \(l\). Starting such a network, with nonlinear activation \(\phi:\mathbb{R}\rightarrow\mathbb{R}\), from an input vector \(z^{0}:=x^{0}\in\mathbb{R}^{N_{0}}\), the data propagation is then given by the following equations,

\[h_{j}^{(l)}=\sum_{k=1}^{N_{l-1}}W_{jk}^{(l)}z_{k}^{(l)}+b_{j}^{(l)},\qquad z_{ k}^{(l)}=\phi(h_{k}^{(l-1)}) \tag{1}\]

where we call \(h^{(l)}\) the preactivation vector at layer \(l\).

It has been shown by (Poole et al., 2016) that the pre-activation vectors \(h^{l}\) have geometric properties of length \(q^{l}:=N_{l}^{-1}\left(h^{l}\right)^{T}h^{l}\) and the pairwise covariance \(q_{12}^{l}:=N_{l}^{-1}\left(h_{1}^{l}\right)^{T}h_{2}^{l}\) of two inputs \(x^{0,1}\) and \(x^{0,2}\) which propagate through the network according to functions of the network entries' variances and nonlinear activation \((\sigma_{b},\sigma_{W},\phi)\). These propagation maps were computed by (Poole et al., 2016) in the limiting setting of infinitely wide networks and either i.i.d. Gaussian entries or scaled randomly drawn orthnormal matrices. Here we extend this setting to their low-rank analogous.

Consider rank \(r_{l}:=\gamma_{l}N_{l}\) weight matrices, \(W^{(l)}\in\mathbb{R}^{N_{l}\times N_{l-1}}\), formed as

\[W_{ij}^{(l)}=\sum_{k=1}^{r_{l}}\alpha_{k,j}^{(l)}(C_{k}^{l})_{i}, \tag{2}\]

where the scalars \(\left(\alpha_{k,i}^{(l)}\right)_{1\leq i\leq N_{l-1}}\in\mathbb{R}\stackrel{{ \text{iid}}}{{\sim}}\mathcal{N}(0,\frac{\sigma_{\alpha}^{2}}{N_{l-1}})\) and the columns \(C_{1}^{(l)},\dots,C_{r_{l}}^{(l)}\) are drawn jointly as the matrix \(C^{(l)}:=[C_{1}^{(l)},\dots,C_{r_{l}}^{(l)}]\in\mathbb{R}^{N_{l}\times r_{l}}\) from the Grassmannian of rank \(r\) matrices with orthonormal columns having zero mean and variance \(1/N_{l}\). Similarly, consider bias vectors within the same column span as \(W^{(l)}\), given by \(b^{(l)}(C_{1}^{(l)}+\dots+C_{r_{l}}^{(l)})\), where \(b^{(l)}\in\mathbb{R}\sim\mathcal{N}(0,\sigma_{b}^{2})\). It is shown in Appendix A.2 that, in the large width limit, the preactivation vector \(h^{(l)}\) follows a Gaussian distribution over the \(r-\)dimensional column span of \(W^{(l)}\) with a non-diagonal covariance; this differs from the full rank setting in (Poole et al., 2016) where the entries in (1) are independent.

We extend the pre-activation length and correlation maps to this low-rank setting:

\[q^{l} =\gamma_{l}\left(\sigma_{\alpha}^{2}\int_{\mathbb{R}}\phi^{2}( \sqrt{q^{l-1}}z)Dz+\sigma_{b}^{2}\right) \tag{3}\] \[:=\mathcal{V}(q^{(l-1)}|\sigma_{\alpha},\sigma_{b},\gamma_{l}) \tag{4}\]
where \(Dz:=\frac{1}{\sqrt{2\pi}}e^{-\frac{z^{2}}{2}}dz\) is the Gaussian probability measure, and

\[q_{12}^{l} =\gamma_{l}\bigg{(}\sigma_{\alpha}^{2}\int_{\mathbb{R}^{2}}\phi(u_ {1})\phi(u_{2})Dz_{1}Dz_{2}+\sigma_{b}^{2}\bigg{)}, \tag{5}\] \[:=\mathcal{C}(q_{ab}^{l-1},q_{aa}^{l-1},q_{bb}^{l-1}|\sigma_{ \alpha},\sigma_{b},\gamma_{l}) \tag{6}\]

with \(u_{1}=\sqrt{q_{11}^{l-1}}z_{1},u_{2}=\sqrt{q_{22}^{l-1}}(c_{12}^{l-1}z_{1}+ \sqrt{1-(c_{12}^{l-1})^{2}}z_{2}\) and

\[c_{12}^{l}=q_{12}^{l}(q_{11}^{l}q_{22}^{l})^{-\frac{1}{2}}. \tag{7}\]

Equations (3) and (5) are derived in Appendix A.3 and Appendix A.4 respectively. These equations exactly recover the equations by (Poole et al., 2016) when \(\gamma_{l}=1\), and show that by appropriately rescaling \(\sigma_{2}^{2}\) and \(\sigma_{b}^{2}\) by \(\gamma_{r}\) the low-rank maps remain consistent with the full rank setting.

These two mappings (3) and (5) are functions of the network entries variances, the rank at each layer \(\gamma_{l}\) and the nonlinear activation \((\sigma_{b},\sigma_{W},\gamma_{l},\phi)\) which determine the existence of eventual stable fixed points of \(q^{l}\) and \(q_{ab}^{l}\) as well as the dynamics they follow through the network.

The dominant quantity determining the dynamics of the network is

\[\chi_{\gamma}:=\gamma\sigma_{\alpha}^{2}\int_{\mathbb{R}}\bigg{(}\phi^{\prime }(\sqrt{q^{*}}z)\bigg{)}^{2}Dz \tag{8}\]

which is equal to two fundamental quantities. First, \(\chi_{\gamma}\) is equal to the gradient of the correlation function (7) evaluated at correlation \(c_{12}^{l}=1\),

\[\chi_{\gamma}=\frac{\partial c_{12}^{l}}{\partial c_{12}^{l-1}}|_{c_{12}^{l-1 }=1} \tag{9}\]

A detailed derivation of the equivalence of (8) and (9) is given in Appendix A.5. When there exists a fixed point \(q^{*}\) such that \(\mathcal{V}(q^{*})=q^{*}\), and \(\chi<1\), then inputs with small initial correlation converge to correlation 1 at an exponential rate; this phase is referred to as _ordered_. Alternatively, when \(\chi>1\) the fixed point \(c^{*}=1\) becomes unstable, meaning that an input and its arbitrarily small perturbation have correlation \(q_{ab}^{l}\) decreasing with layers; this is referred to as the _chaotic_ phase due to all nearby points on a data manifold diverging as they progress through the network. In the ordered phase, the output function of the network is constant whereas in the chaotic phase it is non-smooth everywhere.

In both cases (\(\chi>1\) or \(\chi<1\)), in (Schoenholz et al., 2016), the mappings \(\mathcal{V}\) and \(\frac{1}{q^{*}}\mathcal{C}\) are shown to converge exponentially fast to their fixed point, when they exist. Therefore, the data geometry is quickly lost as it is propagated through layers. The boundary between these phases, where \(\chi=1\), is referred to as the edge-of-chaos and determines the scaling of \((\sigma_{w},\sigma_{b},\gamma_{l})\), as functions of nonlinear activation \(\phi(\cdot)\), which ensures a sub-exponential asymptotic behaviour of these maps towards their fixed point and thus a deeper data propagation along layers which facilitates early training of the network.

Second, the quantity \(\chi_{\gamma}\) in (8) is equal to the median singular value of the the matrix \(D^{(l)}W^{(l)}\) where \(D^{(l)}\) is the diagonal matrix with at layer \(l\) with entries \(D_{ii}^{(l)}=\phi^{\prime}(h_{i}^{l})\); for details see Appendix A.9. Defining the Jacobian matrix \(J\in\mathbb{R}^{N_{L}\times N_{0}}\) of the input-output map as

\[J:=\frac{\partial z^{L}}{\partial z^{0}}=\prod_{l=1}^{L}D^{(l)}W^{(l)}, \tag{10}\]

we see that the average singular value of \(J\) is equal to \(\chi_{\gamma}^{L}\). If \(\chi_{\gamma}=1\) the average singular value of \(J\) is fixed at 1 throughout the network, while if \(\chi_{\gamma}\) is greater than or less than 1 the average singular value deviates from 1 at an exponential rate. Further note that the growth of a perturbation from a layer to the following one is given by the average squared singular value of \(D^{(l)}W^{(l)}\).

### Main contributions

This manuscript extends the edge-of-chaos analysis of random feed-forward networks to the setting of low-rank matrices, following the work of (Poole et al., 2016). This work is motivated by the recent challenges faced to store in memory the constantly growing number of parameters used to train large Deep Learning models, see (Price and Tanner, 2022) and references therein.

As shown in equations (3), (6), and (8), despite the dependence between entries in the low-rank weight matrices (2), that the edge-of-chaos curve defined by \(\chi_{\gamma}=1\) can be retained by scaling the weight and bias variances \(\sigma_{w}^{2}\) and \(\sigma_{b}^{2}\) respectively by the ratio of the weight matrix rank \(r_{l}\) to layer width \(\gamma_{l}:=r_{l}/N_{l}\), see Figure 1 and contrast with Figure 10. That is, a simple re-scaling retains the dominant first order dynamics of a feedforward network when the weight matrices are initialized to be low-rank.

In Section 2 we show that additional first order dynamics are similarly modified through a multiplicative scaling by the rank to width factor \(\gamma_{l}=r_{l}/N_{l}\). In particular, we demonstrate the role of \(\gamma_{l}\) on the length and correlation depth scale as well as the training gradient vectors.

However, in Section 3 we show that important second order properties of the dynamics, specifically the variance of the singular values of the input-output Jacobian given in (10), is modified by the reduced rank in a way that cannot be overcome with simple re-scaling. This result alerts practitioners
to anticipated greater variability in training low-rank weight matrices and suggests that methods to reduce the variance of the spectrum may be increasingly important in this setting, see (Murray et al., 2021).

The manuscript then concludes with numerical experiments in Section 4 which demonstrate that empirical measurements on the Jacobian are consistent with the established formula and a brief summary and future work in Section 5.

## 2 Network dynamics and data propagation

The parameter \(\chi_{\gamma}\) further controls the length and correlation depth scaling as well as the relative magnitude of training gradients computed via back-propagation for the sum-of-squares loss function.

### Depth scales a functions of \(\chi_{\gamma}\)

The role of \(\chi_{\gamma}\) on the achievable depth scale was pioneered by (Schoenholz et al., 2016) for full-rank feedforward networks. In this subsection we extend their results to the low-rank setting with the suitably adapted spectral mean \(\chi_{\gamma}\) given in (8).

#### 2.1.1 Length depth scale

Assuming there exists a fixed point \(q^{*}\) such that \(\mathcal{V}(q^{*})=q^{*}\), then the dynamics of \(\mathcal{V}(q)\) can be linearized around \(q^{*}\) to obtain stability conditions and a rate of decay which determine how deeply data can propagate through the network before converging towards the fixed point. Following the computations done in (Schoenholz et al., 2016), setting a perturbation around the fixed point \(q^{*}+\epsilon_{l}\), then around the fixed point, \(\epsilon_{l}\) evolves as \(e^{-\frac{l}{\xi_{l}\cdot\gamma}}\), when \(\gamma_{l}=^{\prime}gamma\) is fixed along layers and we define the following quantities,

\[\xi_{q,\gamma}^{-1}:=-\log\bigg{(}\chi_{\gamma}+\gamma\sigma_{\alpha}^{2}\int Dz \phi^{\prime\prime}(\sqrt{q^{*}}z)\phi(\sqrt{q^{*}}z)\bigg{)}.\]

Details are given in Appendix A.6. Given that \(\gamma\in(0,1]\), we can see the convergence gets faster towards the fixed point when increasing \(\gamma\). Note that when \(\gamma\sigma_{\alpha}^{2}=\sigma_{W}^{2}\), we recover the results of a full-rank feedforward neural network in (Schoenholz et al., 2016).

#### 2.1.2 Correlation depth scale

Similarly, we compute the dynamical evolution of the correlation map around its fixed point by considering a perturbation \(\epsilon_{l}\) and we obtain that (see Appendix A.7), when all the ranks are set to be proportional to the width with the same coefficient of proportionality \(\gamma_{l}=\gamma\) at any layer \(l\), then the perturbation vanishes exponentially fast \(\epsilon_{l}=\mathcal{O}(e^{-\frac{l}{\xi_{l}\cdot\gamma}})\) where

\[\xi_{c,\gamma}^{-1}:=-\log\big{(}\chi_{\gamma}\big{)}.\]

We recover that the correlation depth scale diverges to \(+\infty\) when \(\chi_{\gamma}\to 1\), yielding again the key role of this quantity, even in the low-rank case. As \(\gamma\in(0,1]\), we can see the convergence gets faster towards the fixed point when increasing \(\gamma\), which highlights the tension between low-rank and the depth to which data can propagate along layers. Note again we recover previous results from (Schoenholz et al., 2016) after appropriate scaling of the variance.

### Layerwise scaling of the training gradient for the sum-of-squares loss function

As already shown in previous works ((Schoenholz et al., 2016), and (Poole et al., 2016)), there exists an direct link between the capacity for a network to propagate data through layers of a network in the forward pass and to backpropagate gradients of any given error function \(E\). In this section, we extend the results known for full-rank feedforward neural networks with infinite width to the low-rank case, with rank \(r_{l}=\gamma_{l}N_{l}\) evolving proportionally to the width.

Figure 1: Edge of Chaos curve of a low-rank neural network where the rank is proportional to the width by a factor \(\gamma\) and nonlinear activation \(\phi(x)=tanh(x)\). The plot is generated with \(\gamma=\frac{1}{4}\), where the axis are rescaled by \(\gamma\).
The derivative of the training error follows by the chain rule,

\[\frac{\partial E}{\partial h_{i}^{(l)}}:=\delta_{i}^{l}=\big{(}\sum_ {k=1}^{N_{l+1}}\delta_{k}^{l+1}W_{ki}^{(l+1)}\big{)}\phi^{\prime}(h_{i}^{(l)}),\] \[\frac{\partial E}{\partial W_{ij}^{(l)}}=\delta_{i}^{l}\phi(h_{j} ^{(l-1)}),\] \[\frac{\partial E}{\partial\alpha_{ij}^{(l)}}=\bigg{(}\sum_{m=1}^ {r_{l}}\delta_{m}^{l}(C_{i}^{l})_{m}\bigg{)}\phi(h_{j}^{(l-1)}).\]

Consider the propagation of the gradients \(\frac{\partial E}{\partial\alpha_{ij}^{(l)}}\) of the error with respect to our trainable parameters \(\alpha^{(l)}:=\big{(}\alpha_{i,j}^{(l)}\big{)}_{i,j}\), which are initalized \(\Big{(}\alpha_{k,i}^{(l)}\big{)}_{1\leq i\leq N_{l-1}}\in\mathbb{R}\stackrel{{ \text{iid}}}{{\sim}}\mathcal{N}(0,\frac{\sigma_{a}^{2}}{N_{l-1}})\). The length of this gradient along layers \(||\nabla_{\alpha^{(l)}}E||_{2}^{2}\) is proportional to \(\tilde{q}^{l}:=\mathbb{E}((\delta_{1}^{l})^{2})\) (see Appendix A.8 for proofs). In our analysis of the variance of the training error we treat the backpropagated weights as independent from the forwarded weights, which wile not strictly true is commonly done due to its efficacy in aiding computations which reflect the observed backward dynamics of the network, see (Pennington & Bahri, 2017). Considering an input vector \(x^{0,a}\), and \(\tilde{q}_{aa}^{l}:=\tilde{q}^{l}(x^{0,a})\),

\[\tilde{q}_{aa}^{l}=\tilde{q}_{aa}^{l+1}\frac{N_{l+1}}{N_{l}}\chi_ {l+1},\]

see Appendix A.8.

With constant width along layers \(\frac{N_{l+1}}{N_{l}}\approx 1\), then the sequence is asymptotically exponential and \(\tilde{q}_{aa}^{l}=\tilde{q}_{aa}^{L}\prod_{k=l+1}^{L}\chi_{\gamma_{k}}\), or, if the proportional coefficient of the rank \(\gamma_{l}=\gamma\) is constant along layers, \(\tilde{q}_{aa}^{l}=\mathcal{O}(e^{\frac{l}{\Delta_{i}\gamma}})\), where

\[\xi_{\Delta,\gamma}^{-1}:=-\log(\chi_{\gamma})\]

The same critical point is observed in the low-rank setting \(\gamma<1\) as in previous works (Schoenholz et al., 2016) given by \(\chi_{\gamma}=1\) :

* When \(\chi_{\gamma}>1\), then \(||\nabla_{\alpha^{(l)}}E||_{2}^{2}\) grows exponentially after \(|\xi_{\nabla,\gamma}|\) layers. This is the chaotic phase with the network is being exponentially-sensitive to perturbations.
* When \(\chi_{\gamma}<1\), then \(||\nabla_{\alpha^{(l)}}E||_{2}^{2}\) vanishes at an exponential rate after \(\xi_{\nabla,\gamma}\) layers. This is the ordered phase with the network is being insensitive to perturbations.
* When \(\chi_{\gamma}=1\), then \(||\nabla_{\alpha^{(l)}}E||_{2}^{2}\) remains of the same scale across even after an infinite number of layers which is referred to as the edge-of-chaos.

## 3 Dynamical isometry

Using tools from Random Matrix Theory, (Pennington et al., 2018) provides a method to compute the moments of the spectral distribution of the Jacobian, revealing secondary information beyond the mean of the spectra. We review the most essential equations to derive the variance of the Jacobian'spectrum here but we refer the reader to (Tao, 2012) for more details on the random matrix transforms.

### Review of the computation of the variance of the Jacobian

In this section, we review a set of definitions of random matrix transforms that allow the calculation of the spectra of the product of matrices in terms of their individual spectra. Let \(X\) be a random matrix with spectral density \(\rho_{X}\)

\[\rho_{X}(\lambda):=\langle\frac{1}{N}\sum_{i=1}^{N}\delta(\lambda- \lambda_{i})\rangle_{X},\]

where \(\langle.\rangle\) is the average with respect to the distribution of the random matrix \(X\), and \(\delta\) is the usual dirac distribution.

For a probability density \(\rho_{X}\) and \(z\in\mathbb{C}\setminus\mathbb{R}\), the Stieltjes Transform \(G_{X}\) and its inverse are given by

\[G_{X}(z): =\int\frac{\rho_{X}(t)}{t-z}dt,\] \[\rho_{X}(\lambda) =-\pi^{-1}\lim_{\epsilon\to 0^{+}}\text{Im}\big{(}G_{X}( \lambda+\epsilon i)\big{)}.\]

The moment generating function is \(M_{X}(z):=zG_{x}(z)-1=\sum\limits_{k=1}^{\infty}m_{k}z^{-k}\) and the \(\mathcal{S}_{X}\) Transform is defined as \(\mathcal{S}_{X}(z):=\frac{1+z}{xM_{X}^{-1}(z)}\). The interest of using the \(\mathcal{S}\) Transform here is that it has the following multiplicative property, which in our case is desirable as the Jacobian is a product of random matrices: if \(X\) and \(Y\) are freely independent, then \(\mathcal{S}_{XY}=\mathcal{S}_{X}\mathcal{S}_{Y}\).

In (Pennington et al., 2018), the authors start with establishing \(\mathcal{S}_{JJ^{T}}=S_{D^{2}}^{L}S_{W^{T}W}^{L}\), assuming the input vector is chosen such that \(q^{l}\approx q^{*}\) so that distribution of \(D^{2}\) is independent of \(l\) and we already had the weights identically distributed along layers. The strategy here to compute the spectral density of \(\rho_{JJ^{T}}\) (and thus the density of the singular values of the Jacobian \(J\)) starts with computing the \(\mathcal{S}\) Transforms of \(W^{T}W\) and \(D^{2}\) from their spectral density, determined by respectively, the way of sampling the weights at initialisation and the choice of the activation function in the network. Note that in this study we focus only on two possible distributions for the low-rank weights matrix - either scaled Gaussian weights or scaled orthogonal matrices, that are defined more precisely in the next sections. Once that \(\mathcal{S}_{JJ^{T}}\) is obtained by multiplying \(S_{W^{T}W}\) and \(S_{D^{2}}\), rather than inverting it back to find \(\rho_{JJ^{T}}\), the authors show
there is a way to shortcut these steps and obtain directly the moments of \(\rho_{JJ^{T}}\) based on the following set of equations. Defining

\[m_{k} :=\int\lambda^{k}\rho_{JJ^{T}}(\lambda)d\lambda\] \[S_{W^{T}W}(z) :=\gamma^{-1}\sigma_{\alpha}^{-2}\big{(}1+\sum_{k=1}^{\infty}s_{k }z^{k}\big{)}\] \[\mu_{k} =\int Dz\big{(}\phi^{\prime}(\sqrt{q^{*}}z)\big{)}^{2k}\]

then as derived in (Pennington et al., 2018), the first two moments of the spectrum of the Jacobian are

\[m_{1} =(\gamma\sigma_{\alpha}^{2}\mu_{1})^{L}\] \[m_{2} =(\gamma\sigma_{\alpha}^{2}\mu_{1})^{2L}L\left(\frac{\mu_{2}}{ \mu_{1}^{2}}+\frac{1}{L}-1-s_{1}\right).\]

The first moment \(m_{1}\) recovers the previous statement that the average squared singular value is equal to \(m_{1}=\chi_{\gamma}^{L}\) and the edge-of-chaos given by \(\chi_{\gamma}=\gamma\sigma_{\alpha}^{2}\mu_{1}=1\) is consistent with previous results as the gradient either vanishes or grows exponentially along with the median of the Jacobian's spectra. Moreover, the variance of the spectrum of \(JJ^{T}\) about its mean \(\chi_{\gamma}=1\) can now be computed

\[\sigma_{JJ^{T}}^{2}:=m_{2}-m_{1}^{2}=L\left(\frac{\mu_{2}}{\mu_{1}^{2}}-1-s_{1 }\right). \tag{11}\]

The variance \(\sigma_{JJ^{T}}^{2}\) grows linearly with depth as in the full-rank setting, recovering the full-rank result when \(\gamma=1\). As in the edge-of-chaos axes scaling in Figure 1, \(\gamma\sigma_{\alpha}^{2}\) plays the same role as \(\sigma_{W}^{2}\). Note that \(\frac{\mu_{2}}{\mu_{1}^{2}}\geq 1\) and consequently \(\sigma_{JJ^{T}}^{2}\) as given in (11) is only independent of depth \(L\) if \(s_{1}=0\) which is only achieved here in the case of full-rank, i.e. \(\gamma=1\) orthogonal matrices.

### Low-Rank Orthogonal weights

Consider a weight matrix whose \(r\) first columns are orthonormal columns sampled from a normal distribution, and the rest is 0, such that \(W^{T}W=\left(\begin{array}{c|c}\sigma_{\alpha}^{2}\mathbb{I}_{r}&0\\ \hline 0&\mathbb{O}_{N-r}\end{array}\right)\). Therefore the spectral distribution of \(\sigma_{\alpha}^{-2}W^{T}W\) is trivially given by

\[\rho_{\sigma_{\alpha}^{-2}W^{T}W}(z)=\gamma\delta(z-1)+(1-\gamma)\delta(z),\]

from which the \(\mathcal{S}\) Transform is computed, see Appendix A.11, to obtain \(s_{1}=-(\gamma^{-1}-1)\). When \(\gamma=1\), the known result in the full-rank orthogonal case is retrieved.

### Low-Rank Gaussian weights

With weights at any layer \(l\) given by 2, the matrix can be rewritten as the product \(W^{l}=C^{l}A^{l}\), where \(C^{l}\in\mathbb{R}^{N_{l}\times r_{l}}\) with \(C^{l}_{ij}=(C^{l}_{j})_{i}\), and \(A^{l}\in\mathbb{R}^{r_{l}\times N_{l-1}}\) with \(A^{l}_{ij}\stackrel{{ iid}}{{\sim}}\mathcal{N}(0,\frac{\sigma_{ \alpha}^{2}}{N_{l-1}})\). As \({C^{l}}^{T}C^{l}=\mathbb{I}_{r_{l}}\) by construction, then \({W^{(l)}}^{T}W^{(l)}={A^{l}}^{T}{C^{l}}^{T}C^{l}A^{l}={A^{l}}^{T}A^{l}\) which is a Wishart matrix, whose spectral density is known and given by the Marcenko Pastur distribution (Marcenko and Pastur, 1967) where some mass is added at 0 since the matrix \(A^{l}\) is not full-rank and contains some 0 eigenvalues. Recall that \(r_{l}=\gamma N_{l}\).

\[\rho_{A^{T}A}(\lambda)=(1-\gamma)_{+}\delta(\lambda)+\gamma\frac{\sqrt{( \lambda^{+}-\lambda)(\lambda-\lambda^{-})}}{2\pi\lambda\sigma_{\alpha}^{2}} \mathds{1}_{[\lambda^{-},\lambda^{+}]}(\lambda),\]

where \(x_{+}=\text{max}(0,x),\lambda^{-}=(1-\frac{1}{\gamma})^{2}\) and \(\lambda^{+}=(1+\frac{1}{\gamma})^{2}\). The \(\mathcal{S}\) Transform \(S_{W^{T}W}\) can be computed (see Appendix A.10) and expanded around 0, which gives \(s_{1}=-\frac{1}{\gamma}\). Note that when \(\gamma=1\), one recovers the result given in (Pennington et al., 2018).

The \(\mathcal{S}\) Transforms and first moments in both orthogonal and Gaussian cases are summarized in Table 1.

## 4 Numerical experiments

In this section, we give empirical evidence in agreement with the theoretical results established above. Its interest is two-fold:

* The variance of the spectrum of the Jacobian does indeed still grow with depth even in the low-rank setting as emphasized in Figure 2. Moreover, at a fixed depth, the rank to width ratio plays a key role in how the spectrum of the Jacobian spreads out around its mean value, which is 1 when the network is initialised on the edge-of-chaos.
* Figure 3 shows that the advantage that Scaled Orthogonal Weights have over Scaled Gaussian Weights in Feedforward networks presented in (Pennington et al., 2018) is lost for low-rank matrices. Indeed, from (11), one can see that in both situations, it is not possible to adjust either the activation function nor \(q^{*}\) through a careful choice of variances for the weights and the biases, unless \(\gamma=1\) and \(W^{(l)}\) is a scaled orthonormal matrix.

In Figure 2 and Figure 3, the variance of the spectrum of the Jacobian is computed in the low-rank Gaussian and Or

\begin{table}
\begin{tabular}{l c c} \hline Random Matrix W & \(S_{W^{T}W^{(2)}}\) & \(s_{1}\) \\ \hline LR Scaled Orthogonal & \(\gamma^{-1}\sigma_{\alpha}^{-2}\frac{1+z}{1+\gamma-1z}\) & \(1^{-\frac{1}{\gamma}}\) \\ LR Scaled Gaussian & \(\gamma^{-1}\sigma_{\alpha}^{-2}\frac{1+z}{1+z(1+\gamma^{-1})+\gamma-1z^{2}}\) & \(-\frac{1}{\gamma}\) \\ \hline \end{tabular}
\end{table}
Table 1: Transforms of weights. LR stands for Low-Rank.
thogonal cases when the activation function is chosen to be the identity. Although such a choice of activation function completely destroys the network's expressivity power, it is a simple example of situations in the full-rank case where Gaussian distributed weight matrices lead to ill conditioned Jacobians as depth increases. This still holds in the low-rank setting as shown in the plot since the variance \(\sigma_{JJ^{T}}^{2}>0\). Simulations are performed on a \(1000\)- layer wide feedforward network, initialised and fed with a random input, whose length is set to be equal to \(q^{*}\) so that the network would already be at its equilibrium state without passing by a transient phase.

The source code can be found at shorturl.at/syLP9.

## 5 Summary and further work

Herein the edge-of-chaos theory of (Poole et al., 2016) and (Schoenholz et al., 2016) has been extended from the setting of full-rank weight matrices to the low-rank setting. Suitable scaling by the ratio of the rank to width factor \(\gamma_{l}:=r_{l}/N_{l}\) recovers the phenomenon driven by the mean of the Jacobian's spectra which defines the edge-of-chaos. Moreover, the variance of the Jacobian's spectra is shown to be strictly increasing with decreasing \(\gamma_{l}\) which suggests greater variability in the initial training of low-rank feedforward networks.

The edge-of-chaos initialisation scheme has been successfully generalised to a large set of different settings, including changes of architectures as CNNs (Xiao et al., 2018), LSTMs and GRUs (Gilboa et al., 2019), RNNs (Chen et al., 2018), ResNets (Yang and Schoenholz, 2017) and to extra features like dropout (Schoenholz et al., 2016), (Huang et al., 2019) or batch normalisation (Yang et al., 2019) and pruning (Hayou et al., 2020). It has been improved with changes of activation functions (Hayou et al., 2019), (Murray et al., 2021) to enable the data to propagate even deeper through the network. As a future work, each of these settings could be extended to the setting of low-rank weight matrices.

Figure 3: Evolution of the variance of the spectrum of \(JJ^{T}\) with respect to \(\gamma\) where \(\gamma\) is the proportionality coefficient giving the rank of the weights matrices at layer \(l\), whose width is \(N_{l}\), \(r_{l}=\gamma N_{l}\). Points are obtained empirically and averaged over 3 simulations when the lines are derived from the theory, see (11). Confidence intervals of 1 standard deviation around each mean point are shown. The weights are chosen to be low-rank Scaled Orthogonal and the activation function is linear \(\phi:x\mapsto x\). The same seed is used to initialise the weight matrices for each simulation and \(q^{*}\) is set to 0.5. The \(y-\)axis is shown in log scale.

Figure 2: Evolution of the variance of the spectrum of \(JJ^{T}\) with respect to \(\gamma\) where \(\gamma\) is the proportionality coefficient giving the rank of the weights matrices at layer \(l\), whose width is \(N_{l}\), \(r_{l}=\gamma N_{l}\). Points are obtained empirically and averaged over 5 simulations when the lines are derived from the theory, see (11). Confidence intervals of 1 standard deviation around each mean point are shown. The weights are chosen to be low-rank Scaled Gaussian and the activation function is linear \(\phi:x\mapsto x\). The same seed is used to initialise the weight matrices for each simulation and \(q^{*}\) is set to 0.5. The \(y-\)axis is shown in log scale.
## Acknowledgments

TNS is financially supported by the Engineering and Physical Sciences Research Council (EPSRC). JT is supported by the Hong Kong Innovation and Technology Commission (InnoHK Project CIMDA) and thanks UCLA Department of Mathematics for kindly hosting him during the completion of this manuscript.
