Exact Solutions of a Deep Linear Network

Applications of neural networks have achieved great success in various fields. One central open question is why neural networks, being nonlinear and containing many saddle points and local minima, can sometimes be optimized easily (Choromanska et al., 2015a, ) while becoming difficult and requiring many tricks to train in some other scenarios (Glorot and Bengio,, 2010; Gotmare et al.,, 2018). One established approach is to study the landscape of deep linear nets (Choromanska et al., 2015b, ), which are believed to approximate the landscape of a nonlinear net well. A series of works proved the famous results that for a deep linear net, all local minima are global (Kawaguchi,, 2016; Lu and Kawaguchi,, 2017; Laurent and Brecht,, 2018), which is regarded to have successfully explained why deep neural networks are so easy to train because it implies that initialization in any attractive basin can reach the global minimum without much effort (Kawaguchi,, 2016). However, the theoretical problem of when and why neural networks can be hard to train is understudied.

In this work, we theoretically study a deep linear net with weight decay and stochastic neurons, whose loss function takes the following form in general:

where $\mathbb{E}_{x}$ denotes the expectation over the training set, $U$ and $W^{(i)}$ are the model parameters, $D$ is the depth of the network,¹¹1In this work, we use “depth” to refer to the number of hidden layers. For example, a linear regressor has depth $0$. $\epsilon$ is the noise in the hidden layer (e.g., due to dropout), $d_{i}$ is the width of the $i$-th layer, and $\gamma$ is the strength of the weight decay. Previous works have studied special cases of this loss function. For example, Kawaguchi, (2016) and Lu and Kawaguchi, (2017) study the landscape of $L_{0}$ when $\epsilon$ is a constant (namely, when there is no noise). Mehta et al., (2021) studies $L_{0}$ with (a more complicated type of) weight decay but without stochasticity and proved that all the stationary points are isolated. Another line of works studies $L_{0}$ when the noise is caused by dropout (Mianjy and Arora,, 2019; Cavazza et al.,, 2018). Our setting is more general than the previous works in two respects. First, apart from the mean square error (MSE) loss $L_{0}$, an $L_{2}$ regularization term (weight decay) with arbitrary strength is included; second, the noise $\epsilon$ is arbitrary. Thus, our setting is arguably closer to the actual deep learning practice, where the injection of noises to latent layers is common, and the use of weight decay is virtually ubiquitous (Krogh and Hertz,, 1992; Loshchilov and Hutter,, 2017). One major limitation of our work is that we assume the label $y$ to be $1$-dimensional, and it can be an important future problem to prove whether an exact solution exists or not when $y$ is high-dimensional.

Our foremost contribution is to prove that all the global minimum of an arbitrarily deep and wide linear net takes a simple analytical form. In other words, we identify in closed form the global minima of Eq. (1) up to a single scalar, whose analytical expression does not exist in general. We then show that it has nontrivial properties that can explain many phenomena in deep learning. In particular, the implications of our result include (but are not limited to):

1. 1.

   Weight decay makes the landscape of neural nets more complicated;

   • •

     we show that bad minima²²2Unless otherwise specified, we use the word “bad minimum” to mean a local minimum that is not a global minimum. emerge as weight decay is applied, whereas there is no bad minimum when there is no weight decay. This highlights the need to escape bad local minima in deep learning with weight decay.

2. 2.

   Deeper nets are harder to optimize than shallower ones;

   • •

     we show that a $D\geq 2$ linear net contains a bad minimum at zero, whereas a $D=1$ net does not. This partially explains why deep networks are much harder to optimize than shallower ones in deep learning practice.

3. 3.

   Depending on the task, the common initialization methods (such as the Kaiming init.) can initialize a deep model in the basin of attraction of the bad minimum at zero;

   • •

     common initialization methods initialize the models at a radius of roughly $1/\sqrt{width}$ around the origin; however, we show that the width of the bad minimum is task-dependent and can be larger than the initialization radius for tasks with a small margin ($||\mathbb{E}[xy]||$);

4. 4.

   Thus, the use of (effective) weight decay is a major cause of various types of collapses in deep learning (for example, see Figure 1).

Organization: In the next section, we discuss the related works. In Section 3, we derive the exact solution for a two-layer net. Section 4 extends the result to an arbitrary depth. In Section 5, we study and discuss the relevance of our results to many commonly encountered problems in deep learning. The last section concludes the work and discusses unresolved open problems. All proofs are delayed to Section B. Moreover, additional theoretical results on the effect of including a bias term is considered in Section D.

Notation. For a matrix $W$, we use $W_{i:}$ to denote the $i$-th row vector of $W$. $||Z||$ denotes the $L_{2}$ norm if $Z$ is a vector and the Frobenius norm if $Z$ is a matrix. The notation $*$ signals an optimized quantity. Additionally, we use the superscript ^∗ and subscript _∗ interchangeably, whichever leads to a simpler expression. For example, $b_{*}^{2}$ and $(b^{*})^{2}$ denote the same quantity, while the former is “simpler."

In many ways, linear networks have been used to help understand nonlinear networks. For example, even at depth $0$, where the linear net is just a linear regressor, linear nets are shown to be relevant for understanding the generalization behavior of modern overparametrized networks (Hastie et al.,, 2019). Saxe et al., (2013) studies the training dynamics of a depth-$1$ network and uses it to understand the dynamics of learning of nonlinear networks. These networks are the same as a linear regression model in terms of expressivity. However, the loss landscape is highly complicated due to the existence of more than one layer, and linear nets are widely believed to approximate the loss landscape of a nonlinear net (Kawaguchi,, 2016; Hardt and Ma,, 2016; Laurent and Brecht,, 2018). In particular, the landscape of linear nets has been studied as early as 1989 in Baldi and Hornik, (1989), which proposed the well-known conjecture that all local minima of a deep linear net are global. This conjecture is first proved in Kawaguchi, (2016), and extended to other loss functions and deeper depths in Lu and Kawaguchi, (2017) and Laurent and Brecht, (2018). Many relevant contemporary deep learning problems can be understood with deep linear models. For example, two-layer linear VAE models are used to understand the cause of the posterior collapse problem (Lucas et al.,, 2019; Wang and Ziyin,, 2022). Deep linear nets are also used to understand the neural collapse problem in contrastive learning (Tian,, 2022). We also provide more empirical evidence in Section 5.

This section finds the global minima of a two-layer linear net. The data point is a $d$-dimensional vector $x\in\mathbb{R}^{d}$ drawn from an arbitrary distribution, and the labels are generated through an arbitrary function $y=y(x)\in\mathbb{R}$. For generality, we let different layers have different strengths of weight decay even though they often take the same value in practice. We want to minimize the following objective:

where $d_{1}$ is the width of the hidden layer and $\epsilon_{i}$ are independent random variables. $\gamma_{w}>0$ and $\gamma_{u}>0$ are the weight decay parameters. Here, we consider a general type of noise with $\mathbb{E}[\epsilon_{i}]=1$ and $\mathbb{E}[\epsilon_{i}\epsilon_{j}]=\delta_{ij}\sigma^{2}+1$ where $\delta_{ij}$ is the Kronecker’s delta, and $\sigma^{2}>0$.³³3While we formally require $\gamma$ and $\sigma$ to nonzero, one can show that the solutions we provided remain global minimizers in the zero limit by applying Theorem 2 from Ziyin and Ueda, (2022). For shorthand, we use the notation $A_{0}:=\mathbb{E}[xx^{T}]$, and the largest and the smallest eigenvalues of $A_{0}$ are denoted as $a_{\rm max}$ and $a_{\rm min}$ respectively. $a_{i}$ denotes the $i$-th eigenvalue of $A_{0}$ viewed in any order. For now, it is sufficient for us to assume that the global minimum of Eq. (2) always exists. We will prove a more general result in Proposition 1, when we deal with multilayer nets.

This section extends our result to multiple layers. We first derive the analytical formula for the global minimum of a general arbitrary-depth model. We then show that the landscape for a deeper network is highly nontrivial.

Relevance to nonlinear models. We first caution the readers that the following discussion should be taken with a caveat and is based on the philosophy that deep linear nets can approximate the nonlinear ones. This approximation certainly holds for fully connected models with differentiable activation functions such as tanh or Swish because they are, up to first-order Taylor expansion, a deep linear net around zero, which is the region for which our theory is the most relevant. We empirically demonstrate that close to the origin, the landscape of linear nets can indeed approximate that of nonlinear nets quite well. To compare, we plug in the solution in Theorem 4 to both linear and nonlinear models of the same architecture and compare the loss values at different values of $b$ around $b=0$. For simplicity, we only consider the case $D=1$. The activation functions we consider are ReLU, Tanh, and Swish (Ramachandran et al.,, 2017), a modern and differentiable variant of ReLU. See Fig. 2.

The regressor $x\in\mathbb{R}^{d}$ is sampled as Gaussian random vectors. We consider two methods of generating $y$; the first one (left) is $y=v\cdot x$. The second one (right) is $y=v\cdot\tanh(x)$, where the weight $v\in\mathbb{R}^{d}$ is obtained as a Gaussian random vector. Fig. 2 shows that the landscape consisting of Tanh is always close to the linear landscape. Swish is not as good as Tanh, but the Swish landscape shows a similar tendency to the linear landscape. The ReLU landscape is not so close to the linear landscape either for $b>0$ or $b<0$, but it agrees completely with the linear landscape on the other side, as expected. Besides the quantitative closeness, it is also important to note that all the landscapes agree qualitatively, containing the same number of local minima at similar values of $b$.

Landscape of multi-layer neural networks. The combination of Theorem 3 and Proposition 2 shows that the landscape of a deep neural network can become highly nontrivial when there is a weight decay and when the depth of the model is larger than $2$. This gives an incomplete but meaningful picture of a network’s complicated but interesting landscape beyond two layers (see Figure 1 for an incomplete summary of our results). In particular, even when the nontrivial solution is the global minimum, the trivial solution is still a local minimum that needs to be escaped. Our result suggests the previous understanding that all local minima of a deep linear net are global cannot generalize to many practical settings where deep learning is found to work well. For example, a series of works attribute the existence of bad (non-global) minima to the use of nonlinearities (Kawaguchi,, 2016) or the use of a non-regular (non-differentiable) loss function (Laurent and Brecht,, 2018). Our result, in contrast, shows that the use of a simple weight decay is sufficient to create a bad minimum.⁶⁶6Some previous works do suggest the existence of bad minima when weight decay is present, but no direct proof exists yet. For example, Taghvaei et al., (2017) shows that when the model is approximated by a linear dynamical system, regularization can cause bad local minima. Mehta et al., (2021) shows the existence of bad local minima in deep linear networks with weight decay through numerical simulations. Moreover, the problem with such a minimum is two-fold: (1) (optimization) it is not global and so needs to be “overcome" and (2) (generalization) it is a minimum that has not learned any feature at all because the model constantly outputs zero. To the best of our knowledge, previous to our work, there has not been any proof that a bad minimum can generically exist in a rather arbitrary network without any restriction on the data.⁷⁷7In the case of nonlinear networks without regularization, a few works proved the existence of bad minima. However, the previous results strongly depend on the data and are rather independent of architecture. For example, one major assumption is that the data cannot be perfected and fitted by a linear model (Yun et al.,, 2018; Liu,, 2021; He et al.,, 2020). Some other works explicitly construct data distribution (Safran and Shamir,, 2018; Venturi et al.,, 2019). Our result, in contrast, is independent of the data. Thus, our result offers direct theoretical justification for the widely believed importance of escaping local minima in the field of deep learning (Kleinberg et al.,, 2018; Liu et al.,, 2021; Mori et al.,, 2022). In particular, previous works on escaping local minima often hypothesize landscapes that are of unknown relevance to an actual neural network. With our result, this line of research can now be established with respect to landscapes that are actually deep-learning-relevant.

Previous works also argue that having a deeper depth does not create a bad minimum (Lu and Kawaguchi,, 2017). While this remains true, its generality and applicability to practical settings now also seem low. Our result shows that as long as weight decay is used, and as long as $D\geq 2$, there is indeed a bad local minimum at $0$. In contrast, there is no bad minimum at $0$ for a depth-$2$ network: the point $b=0$ is either a saddle or the global minimum.⁸⁸8Of course, in practice, the model trained with SGD can still converge to the trivial solution even if it is a saddle point (Ziyin et al.,, 2021) because minibatch SGD is, in general, not a good estimator of the local minima. Having a deeper depth thus alters the qualitative nature of the landscape, and our results agree better with the common observation that a deeper network is harder, if not impossible, to optimize.

We note that our result can also be relevant for more modern architectures such as the ResNet. Using ResNet, one needs to change the dimension of the hidden layer after every bottleneck, and a learnable linear transformation is applied here. Thus, the “effective depth” of a ResNet would be roughly between the number of its bottlenecks and its total number of blocks. For example, a ResNet18 applied to CIFAR10 often has five bottlenecks and 18 layers in total. We thus expect it to have qualitatively similar behavior to a deep linear net with a depth in between. See Figure 1. The experimental details are given in Section A.

Learnability of a neural network. Now we analyze the solution when $D$ tends to infinity. We first note that the existence condition bound in (15) becomes exponentially harder to satisfy as $D$ becomes large:

When this bound is not satisfied, the given neural network cannot learn the data. Recall that for a two-layer net, the existence condition is nothing but $||\mathbb{E}[xy]||^{2}>\gamma^{2}$, independent of the depth, width, or stochasticity in the model. For a deeper network, however, every factor comes into play, and the architecture of the model has a strong (and dominant) influence on the condition. In particular, a factor that increases polynomially in the model width and exponentially in the model depth appears.

A practical implication is that the use of weight decay may be too strong for deep networks. If one increases the depth or width of the model, one should also roughly decrease $\gamma$ according to Eq. (20).

Insufficiency of the existing initialization schemes. We have shown that $0$ is often a bad local minimum for deep learning. Our result further implies that escaping this local minimum can be highly practically relevant because standard initialization schemes are trapped in this local minimum for tasks where the signal $\mathbb{E}[xy]$ is weak. See Inequality (16): any nontrivial global minimum is lower-bounded by a factor proportional to $(\gamma/||\mathbb{E}[xy]||^{1/(D-1)})/d_{0}$, which can be seen as an approximation of the radius of the local minimum at the origin. In comparison, standard deep learning initialization schemes such as Kaiming init. initialize at a radius roughly $1/\sqrt{d_{0}}$. Thus, for tasks $\mathbb{E}[xy]\ll\gamma/\sqrt{d_{0}}$, these initialization methods are likely to initialize the model in the basin of attraction of the trivial regime, which can cause a serious failure in learning. To demonstrate, we perform a numerical simulation shown in the right panel of Figure 3, where we train $D=2$ nonlinear networks with width $32$ with SGD on tasks with varying $||\mathbb{E}[xy]||$. For sufficiently small $||\mathbb{E}[xy]||$, the model clearly is stuck at the origin.⁹⁹9There are many natural problems where the signal is extremely weak. One well-known example is the problem of future price prediction in finance, where the fundamental theorem of finance forbids a large $||\mathbb{E}[xy]||$ (Fama,, 1970). In contrast, linear regression is never stuck at the origin. Our result thus suggests that it may be desirable to devise initialization methods that are functions of the data distribution.

Prediction variance of stochastic nets. A major extension of the standard neural networks is to make them stochastic, namely, to make the output a random function of the input. In a broad sense, stochastic neural networks include neural networks trained with dropout (Srivastava et al.,, 2014; Gal and Ghahramani,, 2016), Bayesian networks (Mackay,, 1992), variational autoencoders (VAE) (Kingma and Welling,, 2013), and generative adversarial networks (Goodfellow et al.,, 2014). Stochastic networks are thus of both practical and theoretical importance to study. Our result can also be used for studying the theoretical properties of stochastic neural networks. Here, we present a simple application of our general solution to analyze the properties of a stochastic net. The following theorem summarizes our technical results.

Interestingly, the scaling of prediction variance in asymptotic $\sigma^{2}$ is different for different widths. The third result shows that the prediction variance decreases exponentially fast in $D$. In particular, this result answers a question recently proposed in Ziyin et al., (2022): does a stochastic net trained on MSE have a prediction variance that scales towards $0$? We improve on their result in the case of a deep linear net by (a) showing that the $d_{0}^{-1}$ is tight in general, independent of the depth or other factors of the model, and (b) proving a bound showing that the variance also scales towards zero as depth increases, which is a novel result of our work. Our result also offers an important insight into the cause of the vanishing prediction variance. Previous works (Alemi et al.,, 2018) often attribute the cause to the fact that a wide neural network is too expressive. However, our result implies that this is not always the case because a linear network with limited expressivity can also have a vanishing variance as the model tends to an infinite width.

Collapses in deep learning. Lastly, we comment briefly on the apparent similarity between different types of collapses that occur in deep learning. For neural collapse, our result agrees with the recent works that identify weight decay as a main cause (Rangamani and Banburski-Fahey,, 2022). For Bayesian deep learning, Wang and Ziyin, (2022) identified the cause of the posterior collapse in a two-layer VAE structure to be that the regularization of the mean of the latent variable $z$ is too strong. More recently, the origin and its stability have also been discussed as the dimensional collapse in self-supervised learning (Ziyin et al.,, 2023). Although appearing in different contexts of deep learning, the three types of collapses share the same phenomenology that the model converges to a “collapsed" regime where the learned representation becomes low-rank or constant, which agrees with the behavior of the trivial regime we identified. We refer the readers to Ziyin and Ueda, (2022) for a study of how the second-order phase transition framework of statistical physics can offer a possible unified explanation of these phenomena.

In this work, we derived the exact solution of a deep linear net with arbitrary depth and width and with stochasticity. Our work sheds light on the highly complicated landscape of a deep neural network. Compared to the previous works that mostly focus on the qualitative understanding of the linear net, our result offers a more precise quantitative understanding of deep linear nets. Quantitative understanding is one major benefit of knowing the exact solution, whose usefulness we have also demonstrated with the various implications. The results, although derived for linear models, are also empirically shown to be relevant for networks with nonlinear activations. Lastly, our results strengthen the line of thought that analytical approaches to deep linear models can be used to understand deep neural networks, and it is the sincere hope of the authors to attract more attention to this promising field.

Ziyin is financially supported by the GSS scholarship of the University of Tokyo and the JSPS fellowship. Li is financially supported by CNRS. X. Meng is supported by JST CREST Grant Number JPMJCR1912, Japan.

1. 1.

   For all authors…

   1. (a)

      Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]

   2. (b)

      Did you describe the limitations of your work? [Yes]

   3. (c)

      Did you discuss any potential negative societal impacts of your work? [N/A]

   4. (d)

      Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]

2. 2.

   If you are including theoretical results…

   1. (a)

      Did you state the full set of assumptions of all theoretical results? [Yes]

   2. (b)

      Did you include complete proofs of all theoretical results? [Yes] See Appendix.

3. 3.

   If you ran experiments…

   1. (a)

      Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [No] The experiments are only for demonstration and are straightforward to reproduce following the theory.

   2. (b)

      Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]

   3. (c)

      Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No] The fluctuations are visually negligible.

   4. (d)

      Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] They are done on a single 3080Ti GPU.

4. 4.

   If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…

   1. (a)

      If your work uses existing assets, did you cite the creators? [N/A]

   2. (b)

      Did you mention the license of the assets? [N/A]

   3. (c)

      Did you include any new assets either in the supplemental material or as a URL? [N/A]

   4. (d)

      Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]

   5. (e)

      Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]

5. 5.

   If you used crowdsourcing or conducted research with human subjects…

   1. (a)

      Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]

   2. (b)

      Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]

   3. (c)

      Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]