A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks

Sanjeev Arora, Nadav Cohen, Noah Golowich, Wei Hu

Introduction

Deep learning builds upon the mysterious ability of gradient-based optimization methods to solve related non-convex problems. Immense efforts are underway to mathematically analyze this phenomenon. The prominent landscape approach focuses on special properties of critical points (i.e. points where the gradient of the objective function vanishes) that will imply convergence to global optimum. Several papers (e.g. Ge et al. (2015); Lee et al. (2016)) have shown that (given certain smoothness properties) it suffices for critical points to meet the following two conditions: (i) no poor local minima — every local minimum is close in its objective value to a global minimum; and (ii) strict saddle property — every critical point that is not a local minimum has at least one negative eigenvalue to its Hessian. While condition (i) does not always hold (cf. Safran and Shamir (2018)), it has been established for various simple settings (e.g. Soudry and Carmon (2016); Kawaguchi (2016)). Condition (ii) on the other hand seems less plausible, and is in fact provably false for models with three or more layers (cf. Kawaguchi (2016)), i.e. for deep networks. It has only been established for problems involving shallow (two layer) models, e.g. matrix factorization (Ge et al. (2016); Du et al. (2018a)). The landscape approach as currently construed thus suffers from inherent limitations in proving convergence to global minimum for deep networks.

A potential path to circumvent this obstacle lies in realizing that landscape properties matter only in the vicinity of trajectories that can be taken by the optimizer, which may be a negligible portion of the overall parameter space. Several papers (e.g. Saxe et al. (2014); Arora et al. (2018)) have taken this trajectory-based approach, primarily in the context of linear neural networks — fully-connected neural networks with linear activation. Linear networks are trivial from a representational perspective, but not so in terms of optimization — they lead to non-convex training problems with multiple minima and saddle points. Through a mix of theory and experiments, Arora et al. (2018) argued that such non-convexities may in fact be beneficial for gradient descent, in the sense that sometimes, adding (redundant) linear layers to a classic linear prediction model can accelerate the optimization. This phenomenon challenges the holistic landscape view, by which convex problems are always preferable to non-convex ones.

Gradient Descent for Deep Linear Neural Networks

We denote by ∥v∥\left\|{\mathbf{v}}\right\| the Euclidean norm of a vector v{\mathbf{v}}, and by ∥A∥F\left\|A\right\|_{F} the Frobenius norm of a matrix AA.

Note that the notation LN(⋅)L^{N}(\cdot) is consistent with that of Equation (2), as a network with depth N=1N=1 precisely reduces to a (directly parameterized) linear model.

We focus on studying the process of training a deep linear neural network by gradient descent, i.e. of tackling the optimization problem in Equation (3) by iteratively applying the following updates:

where η>0\eta>0 is a configurable learning rate. In the case of depth N=1N=1, the training problem in Equation (3) is smooth and strongly convex, thus it is known (cf. Boyd and Vandenberghe (2004)) that with proper choice of η\eta, gradient descent converges to global minimum at a linear rate. In contrast, for any depth greater than 11, Equation (3) comprises a fundamentally non-convex program, and the convergence properties of gradient descent are highly non-trivial. Apart from the case N=2N=2 (shallow network), one cannot hope to prove convergence via landscape arguments, as the strict saddle property is provably violated (see Section 1). We will see in Section 3 that a direct analysis of the trajectories taken by gradient descent can succeed in this arena, providing a guarantee for linear rate convergence to global minimum.

Convergence Analysis

In this section we establish convergence of gradient descent for deep linear neural networks (Equations (4) and (3)) by directly analyzing the trajectories taken by the algorithm. We begin in Subsection 3.1 with a presentation of two concepts central to our analysis: approximate balancedness and deficiency margin. These facilitate our main convergence theorem, delivered in Subsection 3.2. We conclude in Subsection 3.3 by deriving a convergence guarantee that holds with constant probability over a random initialization.

In our context, the notion of approximate balancedness is formally defined as follows:

Note that in the case of -balancedness, i.e. Wj+1⊤Wj+1=WjWj⊤W_{j+1}^{\top}W_{j+1}=W_{j}W_{j}^{\top}, ∀j∈{1,…,N−1}\forall{j}\in\{1,\ldots,N-1\}, all matrices WjW_{j} share the same set of non-zero singular values. Moreover, as shown in the proof of Theorem 1 in Arora et al. (2018), this set is obtained by taking the NN-th root of each non-zero singular value in the end-to-end matrix W1:NW_{1:N}. We will establish approximate versions of these facts for δ\delta-balancedness with δ>0\delta>0, and admit their usage by showing that if the weights of a linear neural network are initialized to be approximately balanced, they will remain that way throughout the iterations of gradient descent. The condition of approximate balancedness at initialization is trivially met in the special case of linear residual networks (d0=⋯=dN=dd_{0}=\cdots=d_{N}=d and W1(0)=⋯=WN(0)=IdW_{1}(0)=\cdots=W_{N}(0)=I_{d}). Moreover, as Claim 2 in Appendix B shows, for a given δ>0\delta>0, the customary initialization via random Gaussian distribution with mean zero leads to approximate balancedness with high probability if the standard deviation is sufficiently small.

The second concept we introduce — deficiency margin — refers to how far a ball around the target is from containing rank-deficient (i.e. low rank) matrices.

The term “deficiency margin” alludes to the fact that if Equation (6) holds, every matrix W′W^{\prime} whose distance from Φ\Phi is no greater than that of WW, has singular values cc-bounded away from zero:

Suppose WW has deficiency margin cc with respect to Φ\Phi. Then, any matrix W′W^{\prime} (of same size as Φ\Phi and WW) for which ∥W′−Φ∥F≤∥W−Φ∥F\left\|W^{\prime}-\Phi\right\|_{F}\leq\left\|W-\Phi\right\|_{F} satisfies σmin(W′)≥c\sigma_{min}(W^{\prime})\geq{c}.

Our proof relies on the inequality σmin(A+B)≥σmin(A)−σmax(B)\sigma_{min}(A+B)\geq\sigma_{min}(A)-\sigma_{max}(B) — see Appendix D.1. ∎

We will show that if the weights W1,…,WNW_{1},\ldots,W_{N} are initialized such that (they are approximately balanced and) the end-to-end matrix W1:NW_{1:N} has deficiency margin c>0c>0 with respect to the target Φ\Phi, convergence of gradient descent to global minimum is guaranteed. In fact, a deficiency margin implies that all critical points in the respective sublevel set (set of points with smaller loss value) are global minima. This however is far from sufficient for proving convergence, as sublevel sets are unbounded, and the loss landscape over them is non-convex and non-smooth. Indeed, we show in Appendix C that deficiency margin alone is not enough to ensure convergence — without approximate balancedness, the lack of smoothness can cause divergence. Moreover, the convergence will outpace a particular rate that gets faster when cc grows larger. This suggests that from a theoretical perspective, it is advantageous to initialize a linear neural network such that the end-to-end matrix has a large deficiency margin with respect to the target. Claim 3 in Appendix B provides information on how likely deficiency margins are in the case of a single output model (scalar regression) subject to customary zero-centered Gaussian initialization. It shows in particular that if the standard deviation of the initialization is sufficiently small, the probability of a deficiency margin being met is close to 0.50.5; on the other hand, for this deficiency margin to have considerable magnitude, a non-negligible standard deviation is required.

Taking into account the need for both approximate balancedness and deficiency margin at initialization, we observe a delicate trade-off under the common setting of Gaussian perturbations around zero: if the standard deviation is small, it is likely that weights be highly balanced and a deficiency margin be met; however overly small standard deviation will render high magnitude for the deficiency margin improbable, and therefore fast convergence is less likely to happen; on the opposite end, large standard deviation jeopardizes both balancedness and deficiency margin, putting the entire convergence at risk. This trade-off is reminiscent of empirical phenomena in deep learning, by which small initialization can bring forth efficient convergence, while if exceedingly small, rate of convergence may plummet (“vanishing gradient problem”), and if made large, divergence becomes inevitable (“exploding gradient problem”). The common resolution of residual connections (He et al., 2016) is analogous in our context to linear residual networks, which ensure perfect balancedness, and allow large deficiency margin if the target is not too far from identity.

2 Main Theorem

Using approximate balancedness (Definition 1) and deficiency margin (Definition 2), we present our main theorem — a guarantee for linear convergence to global minimum:

Assume that gradient descent is initialized such that the end-to-end matrix W1:N(0)W_{1:N}(0) has deficiency margin c>0c>0 with respect to the target Φ\Phi, and the weights W1(0),…,WN(0)W_{1}(0),\ldots,W_{N}(0) are δ\delta-balanced with \delta=c^{2}\big{/}\big{(}256\cdot{N}^{3}\cdot\left\|\Phi\right\|_{F}^{2(N-1)/N}\big{)}. Suppose also that the learning rate η\eta meets:

The assumptions made in Theorem 1 — approximate balancedness and deficiency margin at initialization — are both necessary, in the sense that violating any one of them may lead to convergence failure. We demonstrate this in Appendix C. In the special case of linear residual networks (uniform dimensions and identity initialization), a sufficient condition for the assumptions to be met is that the target matrix have (Frobenius) distance less than 0.50.5 from identity. This strengthens one of the central results in Bartlett et al. (2018) (see Section 5). For a setting of random near-zero initialization, we present in Subsection 3.3 a scheme that, when the output dimension is 11 (scalar regression), ensures assumptions are satisfied (and therefore gradient descent efficiently converges to global minimum) with constant probability. It is an open problem to fully analyze gradient descent under the common initialization scheme of zero-centered Gaussian perturbations applied to each layer independently. We treat this scenario in Appendix B, providing quantitative results concerning the likelihood of each assumption (approximate balancedness or deficiency margin) being met individually. However the question of how likely it is that both assumptions be met simultaneously, and how that depends on the standard deviation of the Gaussian, is left for future work.

An additional point to make is that Theorem 1 poses a structural limitation on the linear neural network. Namely, it requires the dimension of each hidden layer (di, i=1,…,N−1d_{i},~{}i=1,\ldots,N-1) to be greater than or equal to the minimum between those of the input (d0d_{0}) and output (dNd_{N}). Indeed, in order for the initial end-to-end matrix W1:N(0)W_{1:N}(0) to have deficiency margin c>0c>0, it must (by Claim 1) have full rank, and this is only possible if there is no intermediate dimension did_{i} smaller than min⁡{d0,dN}\min\{d_{0},d_{N}\}. We make no other assumptions on network architecture (depth, input/output/hidden dimensions).

2.2 Proof

The cornerstone upon which Theorem 1 rests is the following lemma, showing non-trivial descent whenever σmin(W1:N)\sigma_{min}(W_{1:N}) is bounded away from zero:

Under the conditions of Theorem 1, we have that for every t=0,1,2,…t=0,1,2,\ldots : Note that the term dL1dW(W1:N(t))\frac{dL^{1}}{dW}(W_{1:N}(t)) below stands for the gradient of L1(⋅)L^{1}(\cdot) — a convex loss over (directly parameterized) linear models (Equation (2)) — at the point W1:N(t)W_{1:N}(t) — the end-to-end matrix of the network at iteration tt. It is therefore (see Equation (5)) non-zero anywhere but at a global minimum.

We prove the lemma here for the idealized setting of perfect initial balancedness (δ=0\delta=0):

and infinitesimally small learning rate (η→0+\eta\to 0^{+}) — gradient flow:

where τ\tau is a continuous time index, and dot symbol (in W˙j(τ)\dot{W}_{j}(\tau)) signifies derivative with respect to time. The complete proof, for the realistic case of approximate balancedness and discrete updates (δ,η>0\delta,\eta>0), is similar but much more involved, and appears in Appendix D.2.

We will see that a stronger version of Equation (10) holds, namely, one without the 1/21/2 factor (which only appears due to discretization).

By (Theorem 11 and Claim 11 in) Arora et al. (2018), the weights W1(τ),…,WN(τ)W_{1}(\tau),\ldots,W_{N}(\tau) remain balanced throughout the entire optimization, and that implies the end-to-end matrix W1:N(τ)W_{1:N}(\tau) moves according to the following differential equation:

where vec(A)vec(A), for an arbitrary matrix AA, stands for vectorization in column-first order, and PW1:N(τ)P_{W_{1:N}(\tau)} is a positive semidefinite matrix whose eigenvalues are all greater than or equal to σmin(W1:N(τ))2(N−1)/N\sigma_{min}(W_{1:N}(\tau))^{2(N-1)/N}. Taking the derivative of L1(W1:N(τ))L^{1}(W_{1:N}(\tau)) with respect to time, we obtain the sought-after Equation (10) (with no 1/21/2 factor):

The first transition here (equality) is an application of the chain rule; the second (equality) plugs in Equation (11); the third (inequality) results from the fact that the eigenvalues of the symmetric matrix PW1:N(τ)P_{W_{1:N}(\tau)} are no smaller than σmin(W1:N(τ))2(N−1)/N\sigma_{min}(W_{1:N}(\tau))^{2(N-1)/N} (recall that ∥⋅∥\left\|\cdot\right\| stands for Euclidean norm); and the last (equality) is trivial — ∥A∥F=∥vec(A)∥\left\|A\right\|_{F}=\left\|vec(A)\right\| for any matrix AA. ∎

With Lemma 1 established, the proof of Theorem 1 readily follows:

Since the coefficients 1−η⋅σmin(W1:N(t))2(N−1)N1-\eta\cdot\sigma_{min}(W_{1:N}(t))^{\frac{2(N-1)}{N}} are necessarily non-negative (otherwise would contradict non-negativity of L1(⋅)L^{1}(\cdot)), we may unroll the inequalities, obtaining:

Now, this in particular means that for every t′=0,1,2,…t^{\prime}=0,1,2,\ldots :

Deficiency margin cc of W1:N(0)W_{1:N}(0) along with Claim 1 thus imply \sigma_{min}\big{(}W_{1:N}(t^{\prime})\big{)}\geq{c}, which when inserted back into Equation (12) yields, for every t=1,2,3,…t=1,2,3,\ldots :

η⋅c2(N−1)N\eta\cdot{c}^{\frac{2(N-1)}{N}} is obviously non-negative, and it is also no greater than 11 (otherwise would contradict non-negativity of L1(⋅)L^{1}(\cdot)). We may therefore incorporate the inequality 1-\eta\cdot{c}^{2(N-1)/N}\leq\exp\big{(}-\eta\cdot{c}^{2(N-1)/N}\big{)} into Equation (13):

from which it follows that L1(W1:N(t))≤ϵL^{1}(W_{1:N}(t))\leq\epsilon if:

3 Balanced Initialization

We define the following procedure, balanced initialization, which assigns weights randomly while ensuring perfect balancedness:

Set WN≃UΣ1/N,WN−1≃Σ1/N,…,W2≃Σ1/N,W1≃Σ1/NV⊤W_{N}\simeq{U}\Sigma^{1/N},W_{N-1}\simeq\Sigma^{1/N},\ldots,W_{2}\simeq\Sigma^{1/N},W_{1}\simeq\Sigma^{1/N}V^{\top}, where the symbol “≃\simeq” stands for equality up to zero-valued padding. These assignments can be accomplished since min⁡{d1,…,dN−1}≥min⁡{d0,dN}\min\{d_{1},\ldots,d_{N-1}\}\geq\min\{d_{0},d_{N}\}. By design W1:N=AW_{1:N}=A and Wj+1⊤Wj+1=WjWj⊤W_{j+1}^{\top}W_{j+1}=W_{j}W_{j}^{\top}, ∀j∈{1,…,N−1}\forall{j}\in\{1,\ldots,N{-}1\} — these properties are actually all we need in Theorem 2, and step (iii) in Procedure 1 can be replaced by any assignment that meets them.

The concept of balanced initialization, together with Theorem 1, leads to a guarantee for linear convergence (applicable to output dimension 11 — scalar regression) that holds with constant probability over the randomness in initialization:

For any constant 0<p<1/20<p<1/2, there are constants d0′,a>0d_{0}^{\prime},a>0 As shown in the proof of the theorem (Appendix D.3), d0′,a>0d_{0}^{\prime},a>0 can take on any pair of values for which: (i) d0′≥20d_{0}^{\prime}\geq 20; and (ii) \big{(}1-2\exp(-d_{0}^{\prime}/16)\big{)}\big{(}3-4F(2/\sqrt{a/2})\big{)}\geq 2p, where F(⋅)F(\cdot) stands for the cumulative distribution function of the standard normal distribution. For example, if p=0.25p=0.25, it suffices to take any d0′≥100,a≥100d_{0}^{\prime}\geq 100,a\geq 100. We note that condition (i) here (d0′≥20d_{0}^{\prime}\geq 20) serves solely for simplification of expressions in the theorem. such that the following holds. Assume dN=1,d0≥d0′d_{N}=1,d_{0}\geq d_{0}^{\prime}, and that the weights W1(0),…,WN(0)W_{1}(0),\ldots,W_{N}(0) are subject to balanced initialization (Procedure 1) such that the entries in W1:N(0)W_{1:N}(0) are independent zero-centered Gaussian perturbations with standard deviation s≤∥Φ∥2/ad02s\leq\|\Phi\|_{2}/\sqrt{ad_{0}^{2}}. Suppose also that we run gradient descent with learning rate \eta\leq(s^{2}d_{0})^{4-2/N}\big{/}\big{(}10^{5}N^{3}\|\Phi\|_{2}^{10-6/N}\big{)}. Then, with probability at least pp over the random initialization, we have that for every ϵ>0\epsilon>0 and:

Experiments

Balanced initialization (Procedure 1) possesses theoretical advantages compared with the customary layer-wise independent scheme — it allowed us to derive a convergence guarantee that holds with constant probability over the randomness of initialization (Theorem 2). In this section we present empirical evidence suggesting that initializing with balancedness may be beneficial in practice as well. For conciseness, some of the details behind our implementation are deferred to Appendix E.

In addition to a three layer network, we also evaluated a deeper, eight layer model (with hidden widths identical to the former — N=8N=8, d1=⋯=d7=32d_{1}=\cdots=d_{7}=32). In particular, using the same experimental protocol as above, we measured convergence time under different choices of standard deviation for the initialization. Figure 1(a) displays the result of this test alongside that of the three layer model. As the figure shows, transitioning from three layers to eight aggravated the instability with respect to initialization — there is now a narrow band of standard deviations that lead to convergence in reasonable time, and outside of this band convergence is extremely slow, to the point where it does not take place within the duration we allowed (10610^{6} iterations). From the perspective of our analysis, a possible explanation for the aggravation is as follows: under layer-wise independent initialization, the magnitude of the end-to-end matrix W1:NW_{1:N} depends on the standard deviation in a manner that is exponential in depth, thus for large depths the range of standard deviations that lead to moderately sized W1:NW_{1:N} (as required for a deficiency margin) is limited, and within this range, there may not be many standard deviations small enough to ensure approximate balancedness. The procedure of balanced initialization (Procedure 1) circumvents these difficulties — it assigns W1:NW_{1:N} directly (no exponential dependence on depth), and distributes its content between the individual weights W1,…,WNW_{1},\ldots,W_{N} in a perfectly balanced fashion. Rerunning the experiment of Figure 1(a) with this initialization replacing the customary layer-wise scheme (using same experimental protocol), we obtained the results shown in Figure 1(b) — both the original three layer network, and the deeper eight layer model, converged quickly under virtually all standard deviations tried.

As a final experiment, we evaluated the effect of balanced initialization in a setting that involves non-linear activation, softmax-cross-entropy loss and stochastic optimization (factors not accounted for by our analysis). For this purpose, we turned to the MNIST tutorial built into TensorFlow (Abadi et al., 2016), https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/tutorials/mnist which comprises a fully-connected neural network with two hidden layers (width 128128 followed by 3232) and ReLU activation (Nair and Hinton, 2010), trained through stochastic gradient descent (over softmax-cross-entropy loss) with batch size 100100, initialized via customary layer-wise independent Gaussian perturbations centered at zero. While keeping the learning rate at its default value 0.010.01, we varied the standard deviation of initialization, and for each value measured the training loss after 1010 epochs. As opposed to the dataset used in our experiments with linear networks, measuring the training loss with MNIST is non-trivial computationally (involves passing through 60K60K examples). Therefore, rather than continuously polling training loss until it reaches a certain threshold, in this experiment we chose to evaluate speed of convergence by measuring the training loss once after a predetermined number of iterations. We then replaced the original (layer-wise independent) initialization with a balanced initialization based on Gaussian perturbations centered at zero (latter was implemented per Procedure 1, disregarding non-linear activation), and repeated the process. The results of this experiment are shown in Figure 1(d). Although our theoretical analysis does not cover non-linear activation, softmax-cross-entropy loss or stochasticity in optimization, its conclusion of balanced initialization leading to improved (faster and more stable) convergence carried over to such setting.

Related Work

Theoretical study of gradient-based optimization in deep learning is a highly active area of research. As discussed in Section 1, a popular approach is to show that the objective landscape admits the properties of no poor local minima and strict saddle, which, by Ge et al. (2015); Lee et al. (2016); Panageas and Piliouras (2017), ensure convergence to global minimum. Many works, both classic (e.g. Baldi and Hornik (1989)) and recent (e.g. Choromanska et al. (2015); Kawaguchi (2016); Hardt and Ma (2016); Soudry and Carmon (2016); Haeffele and Vidal (2017); Nguyen and Hein (2017); Safran and Shamir (2018); Nguyen and Hein (2018); Laurent and Brecht (2018)), have focused on the validity of these properties in different deep learning settings. Nonetheless, to our knowledge, the success of landscape-driven analyses in formally proving convergence to global minimum for a gradient-based algorithm, has thus far been limited to shallow (two layer) models only (e.g. Ge et al. (2016); Du and Lee (2018); Du et al. (2018a)).

An alternative to the landscape approach is a direct analysis of the trajectories taken by the optimizer. Various papers (e.g. Brutzkus and Globerson (2017); Li and Yuan (2017); Zhong et al. (2017); Tian (2017); Brutzkus et al. (2018); Li et al. (2018); Du et al. (2018c; b); Liao et al. (2018)) have recently adopted this strategy, but their analyses only apply to shallow models. In the context of linear neural networks, deep (three or more layer) models have also been treated — cf. Saxe et al. (2014) and Arora et al. (2018), from which we draw certain technical ideas for proving Lemma 1. However these treatments all apply to gradient flow (gradient descent with infinitesimally small learning rate), and thus do not formally address the question of computational efficiency.

Conclusion

For deep linear neural networks, we have rigorously proven convergence of gradient descent to global minima, at a linear rate, provided that the initial weight matrices are approximately balanced and the initial end-to-end matrix has positive deficiency margin. The result applies to networks with arbitrary depth, and any configuration of input/output/hidden dimensions that supports full rank, i.e. in which no hidden layer has dimension smaller than both the input and output.

Our assumptions on initialization — approximate balancedness and deficiency margin — are both necessary, in the sense that violating any one of them may lead to convergence failure, as we demonstrated explicitly. Moreover, for networks with output dimension 11 (scalar regression), we have shown that a balanced initialization, i.e. a random choice of the end-to-end matrix followed by a balanced partition across all layers, leads assumptions to be met, and thus convergence to take place, with constant probability. Rigorously proving efficient convergence with significant probability under customary layer-wise independent initialization remains an open problem. The recent work of Shamir (2018) suggests that this may not be possible, as at least in some settings, the number of iterations required for convergence is exponential in depth with overwhelming probability. This negative result, a theoretical manifestation of the “vanishing gradient problem”, is circumvented by balanced initialization. Through simple experiments we have shown that the latter can lead to favorable convergence in deep learning practice, as it does in theory. Further investigation of balanced initialization, including development of variants for convolutional layers, is regarded as a promising direction for future research.

The analysis in this paper uncovers special properties of the optimization landscape in the vicinity of gradient descent trajectories. We expect similar ideas to prove useful in further study of gradient descent on non-convex objectives, including training losses of deep non-linear neural networks.

This work is supported by NSF, ONR, Simons Foundation, Schmidt Foundation, Mozilla Research, Amazon Research, DARPA and SRC. Nadav Cohen is a member of the Zuckerman Israeli Postdoctoral Scholars Program, and supported by Schmidt Foundation.

References

References

Appendix

By definition, when data is whitened, Λxx\Lambda_{xx} is equal to identity, yielding:

where c:=−12Tr⁡(ΛyxΛyx⊤)+12Tr⁡(Λyy)c:=-\tfrac{1}{2}\operatorname{Tr}(\Lambda_{yx}\Lambda_{yx}^{\top})+\tfrac{1}{2}\operatorname{Tr}(\Lambda_{yy}) does not depend on WW. Hence we arrive at Equation (1).

Appendix B Approximate Balancedness and Deficiency Margin Under Customary Initialization

Two assumptions concerning initialization facilitate our main convergence result (Theorem 1): (i) the initial weights W1(0),…,WN(0)W_{1}(0),\ldots,W_{N}(0) are approximately balanced (see Definition 1); and (ii) the initial end-to-end matrix W1:N(0)W_{1:N}(0) has positive deficiency margin with respect to the target Φ\Phi (see Definition 2). The current appendix studies the likelihood of these assumptions being met under customary initialization of random (layer-wise independent) Gaussian perturbations centered at zero.

For approximate balancedness we have the following claim, which shows that it becomes more and more likely the smaller the standard deviation of initialization is:

In terms of deficiency margin, the claim below treats the case of a single output model (scalar regression), and shows that if the standard deviation of initialization is sufficiently small, with probability close to 0.50.5, a deficiency margin will be met. However, for this deficiency margin to meet a chosen threshold cc, the standard deviation need be sufficiently large.

Appendix C Convergence Failures

In this appendix we show that the assumptions on initialization facilitating our main convergence result (Theorem 1) — approximate balancedness and deficiency margin — are both necessary, by demonstrating cases where violating each of them leads to convergence failure. This accords with widely observed empirical phenomena, by which successful optimization in deep learning crucially depends on careful initialization (cf. Sutskever et al. (2013)).

Claim 4 below shows For simplicity of presentation, the claim treats the case of even depth and uniform dimension across all layers. It can easily be extended to account for arbitrary depth and input/output/hidden dimensions. that if one omits from Theorem 1 the assumption of approximate balancedness at initialization, no choice of learning rate can guarantee convergence:

In terms of deficiency margin, we provide (by adapting Theorem 4 in Bartlett et al. (2018)) a different, somewhat stronger result — there exist settings where initialization violates the assumption of deficiency margin, and despite being perfectly balanced, leads to convergence failure, for any choice of learning rate: This statement becomes trivial if one allows initialization at a suboptimal stationary point, e.g. Wj(0)=0, j=1,…,NW_{j}(0)=0,~{}j=1,\ldots,N. Claim 5 rules out such trivialities by considering only non-stationary initializations.

Appendix D Deferred Proofs

To simplify the presentation we will oftentimes use WW as an alternative (shortened) notation for W1:NW_{1:N} — the end-to-end matrix of a linear neural network. We will also use L(⋅)L(\cdot) as shorthand for L1(⋅)L^{1}(\cdot) — the loss associated with a (directly parameterized) linear model, i.e. L(W):=12∥W−Φ∥F2L(W):=\frac{1}{2}\left\|W-\Phi\right\|_{F}^{2}. Therefore, in the context of gradient descent training a linear neural network, the following expressions all represent the loss at iteration tt:

where we define W1:0(t):=Id0W_{1:0}(t):=I_{d_{0}} and WN+1:N(t):=IdNW_{N+1:N}(t):=I_{d_{N}} for completeness.

where ai,ja_{i,j} is the element in the ii-th row and jj-th column of AA.

Recall that for any matrices AA and BB of compatible sizes σmin(A+B)≥σmin(A)−σmax(B)\sigma_{min}(A+B)\geq\sigma_{min}(A)-\sigma_{max}(B), and that the Frobenius norm of a matrix is always lower bounded by its largest singular value (Horn and Johnson (1990)). Using these facts, we have:

D.2 Proof of Lemma 1

To prove Lemma 1, we will in fact prove a stronger result, Lemma 2 below, which states that for each iteration tt, in addition to (9) being satisfied, certain other properties are also satisfied, namely: (i) the weight matrices W1(t),…,WN(t)W_{1}(t),\ldots,W_{N}(t) are 2δ2\delta-balanced, and (ii) W1(t),…,WN(t)W_{1}(t),\ldots,W_{N}(t) have bounded spectral norms.

For 1≤j≤N−11\leq j\leq N-1, ∥Wj+1⊤(t)Wj+1(t)−Wj(t)Wj⊤(t)∥F≤2δ\|W_{j+1}^{\top}(t)W_{j+1}(t)-W_{j}(t)W_{j}^{\top}(t)\|_{F}\leq 2\delta.

If t≥1t\geq 1, then for 1≤j≤N−11\leq j\leq N-1,

For 1≤j≤N1\leq j\leq N, ∥Wj(t)∥σ≤(4∥Φ∥F)1/N\|W_{j}(t)\|_{\sigma}\leq(4\|\Phi\|_{F})^{1/N}.

First we observe that Lemma 1 is an immediate consequence of Lemma 2.

Notice that condition B(t)\mathcal{B}(t) of Lemma 2 for each t≥1t\geq 1 immediately establishes the conclusion of Lemma 1 at time step t−1t-1. ∎

We next prove some preliminary lemmas which will aid us in the proof of Lemma 2. The first is a matrix inequality that follows from Lidskii’s theorem. For a matrix AA, let Sing⁡(A)\operatorname*{Sing}(A) denote the rectangular diagonal matrix of the same size, whose diagonal elements are the singular values of AA arranged in non-increasing order (starting from the (1,1)(1,1) position).

For any two matrices A,BA,B of the same size, ∥Sing⁡(A)−Sing⁡(B)∥σ≤∥A−B∥σ\|\operatorname*{Sing}(A)-\operatorname*{Sing}(B)\|_{\sigma}\leq\|A-B\|_{\sigma} and ∥Sing⁡(A)−Sing⁡(B)∥F≤∥A−B∥F\|\operatorname*{Sing}(A)-\operatorname*{Sing}(B)\|_{F}\leq\|A-B\|_{F}.

∥D1−OD1O⊤∥F≤2ϵ\|D_{1}-OD_{1}O^{\top}\|_{F}\leq 2\epsilon.

Since D1D_{1} and OD2OTOD_{2}O^{T} are both symmetric positive semi-definite matrices, their singular values are equal to their eigenvalues. Moreover, the singular values of D1D_{1} are simply its diagonal elements and the singular values of OD2OTOD_{2}O^{T} are simply the diagonal elements of D2D_{2}. Thus by Lemma 3 we get that ∥D1−D2∥F≤∥D1−OD2OT∥F≤ϵ\|D_{1}-D_{2}\|_{F}\leq\|D_{1}-OD_{2}O^{T}\|_{F}\leq\epsilon. Since the Frobenius norm is unitarily invariant, ∥D1−D2∥F=∥OD1OT−OD2OT∥F\|D_{1}-D_{2}\|_{F}=\|OD_{1}O^{T}-OD_{2}O^{T}\|_{F}, and by the triangle inequality it follows that

Lemma 5 below states that if W1,…,WNW_{1},\ldots,W_{N} are approximately balanced matrices, i.e. Wj+1⊤Wj+1−WjWj⊤W_{j+1}^{\top}W_{j+1}-W_{j}W_{j}^{\top} has small Frobenius norm for 1≤j≤N−11\leq j\leq N-1, then we can bound the Frobenius distance between W1:j⊤W1:jW_{1:j}^{\top}W_{1:j} and (W1⊤W1)j(W_{1}^{\top}W_{1})^{j} (as well as between Wj:NWj:N⊤W_{j:N}W_{j:N}^{\top} and (WNWN⊤)N−j+1(W_{N}W_{N}^{\top})^{N-j+1}).

and for 1≤j≤N1\leq j\leq N, ∥Wj∥σ≤M\|W_{j}\|_{\sigma}\leq M. Then, for 1≤j≤N1\leq j\leq N,

Moreover, if σmin\sigma_{min} denotes the minimum singular value of W1:NW_{1:N}, σ1,min\sigma_{1,min} denotes the minimum singular value of W1W_{1} and σN,min\sigma_{N,min} denotes the minimum singular value of WNW_{N}, then

Since the Frobenius norm is invariant to orthogonal transformations, we get that

By Lemma 4, we have that ∥Σj+1⊤Σj+1−ΣjΣj⊤∥F≤ν\|\Sigma_{j+1}^{\top}\Sigma_{j+1}-\Sigma_{j}\Sigma_{j}^{\top}\|_{F}\leq\nu and ∥ΣjΣj⊤−Vj+1⊤UjΣjΣj⊤Uj⊤Vj+1∥F≤2ν\|\Sigma_{j}\Sigma_{j}^{\top}-V_{j+1}^{\top}U_{j}\Sigma_{j}\Sigma_{j}^{\top}U_{j}^{\top}V_{j+1}\|_{F}\leq 2\nu. We may rewrite the latter of these two inequalities as

For matrices A,BA,B, we have that ∥AB∥F≤∥A∥σ⋅∥B∥F\|AB\|_{F}\leq\|A\|_{\sigma}\cdot\|B\|_{F}. Therefore, for j+1≤i≤Nj+1\leq i\leq N, we have that

We now argue that ∥(Σi−1Σi−1⊤)k−(Σi⊤Σi)k∥F≤ν⋅kM2(k−1)\|(\Sigma_{i-1}\Sigma_{i-1}^{\top})^{k}-(\Sigma_{i}^{\top}\Sigma_{i})^{k}\|_{F}\leq\nu\cdot kM^{2(k-1)}. Note that ∥Σi−1Σi−1⊤−Σi⊤Σi∥F≤ν\|\Sigma_{i-1}\Sigma_{i-1}^{\top}-\Sigma_{i}^{\top}\Sigma_{i}\|_{F}\leq\nu, verifying the case k=1k=1. To see the general case, since square diagonal matrices commute, we have that

By the triangle inequality, we then have that

By an identical argument (formally, by replacing WjW_{j} with WN−j+1⊤W_{N-j+1}^{\top}), we get that

(19) and (20) verify (17) and (16), respectively, so it only remains to verify (18).

Let us write the eigendecomposition of W1:NW1:N⊤W_{1:N}W_{1:N}^{\top} with an orthogonal eigenbasis as W1:NW1:N⊤=UΣU⊤W_{1:N}W_{1:N}^{\top}=U\Sigma U^{\top}, where Σ\Sigma is diagonal with its (non-negative) elements arranged in non-increasing order and UU is orthogonal. We can write the left hand side of (21) as ∥UΣU⊤−UN(ΣNΣN⊤)NUN⊤∥F=∥Σ−U⊤UN(ΣNΣN⊤)NUN⊤U∥F\|U\Sigma U^{\top}-U_{N}(\Sigma_{N}\Sigma_{N}^{\top})^{N}U_{N}^{\top}\|_{F}=\|\Sigma-U^{\top}U_{N}(\Sigma_{N}\Sigma_{N}^{\top})^{N}U_{N}^{\top}U\|_{F}.

By an identical argument using (20), we get that, in the case that d0≤dNd_{0}\leq d_{N}, if σ1,min\sigma_{1,min} denotes the minimum singular value of Σ1\Sigma_{1}, then

(Notice that we have used the fact that the nonzero eigenvalues of W1:NW1:N⊤W_{1:N}W_{1:N}^{\top} are the same as the nonzero eigenvalues of W1:N⊤W1:NW_{1:N}^{\top}W_{1:N}.) This completes the proof of (18). ∎

Using Lemma 5, we next show in Lemma 6 that if W1,…,WNW_{1},\ldots,W_{N} are approximately balanced, then an upper bound on ∥WN⋯W1∥σ\|W_{N}\cdots W_{1}\|_{\sigma} implies an upper bound on ∥Wj∥σ\|W_{j}\|_{\sigma} for 1≤j≤N1\leq j\leq N.

Suppose ν,C\nu,C are real numbers satisfying C>0C>0 and 0<ν≤C2/N30N20<\nu\leq\frac{C^{2/N}}{30N^{2}}. Moreover suppose that the matrices W1,…,WNW_{1},\ldots,W_{N} satisfy the following:

For 1≤j≤N−11\leq j\leq N-1, ∥Wj+1⊤Wj+1−WjWj⊤∥F≤ν\|W_{j+1}^{\top}W_{j+1}-W_{j}W_{j}^{\top}\|_{F}\leq\nu.

Then for 1≤j≤N1\leq j\leq N, ∥Wj∥σ≤C1/N⋅21/(2N)\|W_{j}\|_{\sigma}\leq C^{1/N}\cdot 2^{1/(2N)}.

For 1≤j≤N1\leq j\leq N, let us write the singular value decomposition of WjW_{j} as Wj=UjΣjVj⊤W_{j}=U_{j}\Sigma_{j}V_{j}^{\top}, where the singular values of WjW_{j} are decreasing along the main diagonal of Σj\Sigma_{j}. By Lemma 4, we have that for 1≤j≤N−11\leq j\leq N-1, ∥Σj+1⊤Σj+1−ΣjΣj⊤∥F≤ν\|\Sigma_{j+1}^{\top}\Sigma_{j+1}-\Sigma_{j}\Sigma_{j}^{\top}\|_{F}\leq\nu, which implies that ∣∥Σj+1⊤Σj+1∥σ−∥ΣjΣj⊤∥σ∣≤ν\left|\|\Sigma_{j+1}^{\top}\Sigma_{j+1}\|_{\sigma}-\|\Sigma_{j}\Sigma_{j}^{\top}\|_{\sigma}\right|\leq\nu.

Write M=max⁡1≤j≤N∥Wj∥σ=max⁡1≤j≤N∥Σj∥σM=\max_{1\leq j\leq N}\|W_{j}\|_{\sigma}=\max_{1\leq j\leq N}\|\Sigma_{j}\|_{\sigma}. By the above we have that ∥ΣjΣj⊤∥σ≥M2−Nν\|\Sigma_{j}\Sigma_{j}^{\top}\|_{\sigma}\geq M^{2}-N\nu for 1≤j≤N1\leq j\leq N.

Let the singular value decomposition of W1:NW_{1:N} be denoted by W1:N=UΣV⊤W_{1:N}=U\Sigma V^{\top}, so that ∥Σ∥σ≤C\|\Sigma\|_{\sigma}\leq C. Then by (17) of Lemma 5 and Lemma 4 (see also (22), where the same argument was used), we have that

Now recall that ν\nu is chosen so that ν≤C2/N30⋅N2.\nu\leq\frac{C^{2/N}}{30\cdot N^{2}}. Suppose for the purpose of contradiction that there is some jj such that ∥WjWj⊤∥σ>21/NC2/N\|W_{j}W_{j}^{\top}\|_{\sigma}>2^{1/N}C^{2/N}. Then it must be the case that

Next, using (25) and (1+1/x)x≤e(1+1/x)^{x}\leq e for all x>0x>0,

Since ∥(ΣNΣN⊤)N∥σ=∥ΣNΣN⊤∥σN\|(\Sigma_{N}\Sigma_{N}^{\top})^{N}\|_{\sigma}=\|\Sigma_{N}\Sigma_{N}^{\top}\|_{\sigma}^{N}, we get by combining (23) and (26) that

and since 1−e/20>1/(5/4)1-e/20>1/(5/4), it follows that ∥ΣNΣN⊤∥σ<(5/4)1/NC2/N\|\Sigma_{N}\Sigma_{N}^{\top}\|_{\sigma}<(5/4)^{1/N}C^{2/N}, which contradicts (24). It follows that for all 1≤j≤N1\leq j\leq N, ∥WjWj⊤∥σ≤21/NC2/N\|W_{j}W_{j}^{\top}\|_{\sigma}\leq 2^{1/N}C^{2/N}. The conclusion of the lemma then follows from the fact that ∥WjWj⊤∥σ=∥Wj∥σ2\|W_{j}W_{j}^{\top}\|_{\sigma}=\|W_{j}\|_{\sigma}^{2}. ∎

D.2.2 Single-Step Descent

Lemma 7 below states that if certain conditions on W1(t),…,WN(t)W_{1}(t),\ldots,W_{N}(t) are met, the sought-after descent — Equation (9) — will take place at iteration tt. We will later show (by induction) that the required conditions indeed hold for every tt, thus the descent persists throughout optimization. The proof of Lemma 7 is essentially a discrete, single-step analogue of the continuous proof for Lemma 1 (covering the case of gradient flow) given in Section 3.

Assume the conditions of Theorem 1. Moreover, suppose that for some tt, the matrices W1(t),…,WN(t)W_{1}(t),\ldots,W_{N}(t) and the end-to-end matrix W(t):=W1:N(t)W(t):=W_{1:N}(t) satisfy the following properties:

∥Wj(t)∥σ≤(4∥Φ∥F)1/N\|W_{j}(t)\|_{\sigma}\leq(4\|\Phi\|_{F})^{1/N} for 1≤j≤N1\leq j\leq N.

∥W(t)−Φ∥σ≤∥Φ∥F\|W(t)-\Phi\|_{\sigma}\leq\|\Phi\|_{F}.

∥Wj+1⊤(t)Wj+1(t)−Wj(t)Wj⊤(t)∥F≤2δ\|W_{j+1}^{\top}(t)W_{j+1}(t)-W_{j}(t)W_{j}^{\top}(t)\|_{F}\leq 2\delta for 1≤j≤N−11\leq j\leq N-1.

σmin:=σmin(W(t))≥c{\sigma}_{min}:={\sigma}_{min}(W(t))\geq c.

Then, after applying a gradient descent update (4) we have that

For simplicity write M=(4∥Φ∥F)1/NM=(4\|\Phi\|_{F})^{1/N} and B=∥Φ∥FB=\|\Phi\|_{F}. We first claim that

Since c≤σminc\leq\sigma_{min}, for (27) to hold it suffices to have

As the minimum singular value of Φ\Phi must be at least cc, we must have c≤∥Φ∥σc\leq\|\Phi\|_{\sigma}. Since then c∥Φ∥F≤c∥Φ∥σ≤1\frac{c}{\|\Phi\|_{F}}\leq\frac{c}{\|\Phi\|_{\sigma}}\leq 1, it holds that

The second inequality above is trivial, and for the first to hold, since c≤∥Φ∥Fc\leq\|\Phi\|_{F}, it suffices to take

which is guaranteed by the definition of δ\delta in Theorem 1.

Next we continue with the rest of the proof. It follows from (14) thatHere, for matrices A1,…,AKA_{1},\ldots,A_{K} such that AKAK−1⋯A1A_{K}A_{K-1}\cdots A_{1} is defined, we write ∏1j=KAj:=AKAK−1⋯A1\prod_{1}^{j=K}A_{j}:=A_{K}A_{K-1}\cdots A_{1}.

where (⋆)(\star) denotes higher order terms in η\eta. We now bound the Frobenius norm of (⋆)(\star). To do this, note that since L(W)=12∥W−Φ∥F2L(W)=\frac{1}{2}\|W-\Phi\|_{F}^{2}, dLdW(W(t))=W(t)−Φ\frac{dL}{dW}(W(t))=W(t)-\Phi. Then

where the last inequality uses ηMN−2BN≤1/2\eta M^{N-2}BN\leq 1/2, which is a consequence of (27). Next, by Lemma 5 with ν=2δ\nu=2\delta,

Next, by standard properties of tensor product, we have that

Let us write eigenvalue decompositions W1⊤W1=UDU⊤,WNWN⊤=VEV⊤W_{1}^{\top}W_{1}=UDU^{\top},W_{N}W_{N}^{\top}=VEV^{\top}. Then

where the second inequality follows from (D.2.2). Hence the minimum diagonal element of Λ\Lambda is at least (σmin2/(4/3))(N−1)/N≥σmin2(N−1)/N/(4/3)(\sigma_{min}^{2}/(4/3))^{(N-1)/N}\geq\sigma_{min}^{2(N-1)/N}/(4/3).

It follows as a result of the above inequalities that if we write E(t)=vec(W(t+1))−vec(W(t))+η(OΛO⊤)vec(dLdW(W(t)))E(t)=vec(W(t+1))-vec(W(t))+\eta(O\Lambda O^{\top})vec\left(\frac{dL}{dW}(W(t))\right), then

where the first inequality follows since L(W)=12∥W−Φ∥F2L(W)=\frac{1}{2}\|W-\Phi\|_{F}^{2} is 11-smooth as a function of WW. Next, by (29) and (30),

By (27, D.2.2), which bound η,2δ\eta,2\delta, respectively, we have that

D.2.3 Proof of Lemma 2

We use induction on tt, beginning with the base case t=0t=0. Since the weights W1(0),…,WN(0)W_{1}(0),\ldots,W_{N}(0) are δ\delta-balanced, we get that A(0)\mathcal{A}(0) holds automatically. To establish B(0)\mathcal{B}(0), note that since W1:N(0)W_{1:N}(0) has deficiency margin c>0c>0 with respect to Φ\Phi, we must have ∥W1:N(0)−Φ∥F≤σmin(Φ)≤∥Φ∥F\|W_{1:N}(0)-\Phi\|_{F}\leq\sigma_{min}(\Phi)\leq\|\Phi\|_{F}, meaning that L1(W1:N(0))≤12∥Φ∥F2L^{1}(W_{1:N}(0))\leq\frac{1}{2}\|\Phi\|_{F}^{2}.

Finally, by B(0)\mathcal{B}(0), which gives ∥W(0)−Φ∥F≤∥Φ∥F\|W(0)-\Phi\|_{F}\leq\|\Phi\|_{F}, we have that

To show that the above implies C(0)\mathcal{C}(0), we use condition A(0)\mathcal{A}(0) and Lemma 6 with C=2∥Φ∥FC=2\|\Phi\|_{F} and ν=2δ\nu=2\delta. By the definition of δ\delta in Theorem 1 and since c≤∥Φ∥Fc\leq\|\Phi\|_{F}, we have that

as required by Lemma 6. As A(0)\mathcal{A}(0) and (33) verify the preconditions 1. and 2., respectively, of Lemma 6, it follows that for 1≤j≤N1\leq j\leq N, ∥Wj(t)∥σ≤(2∥Φ∥F)1/N⋅21/(2N)<(4∥Φ∥F)1/N\|W_{j}(t)\|_{\sigma}\leq(2\|\Phi\|_{F})^{1/N}\cdot 2^{1/(2N)}<(4\|\Phi\|_{F})^{1/N}, verifying C(0)\mathcal{C}(0) and completing the proof of the base case.

The proof of Lemma 2 follows directly from the following inductive claims.

A(t),B(t),C(t)⇒B(t+1)\mathcal{A}(t),\mathcal{B}(t),\mathcal{C}(t)\Rightarrow\mathcal{B}(t+1). To prove this, we use Lemma 7. We verify first that the preconditions hold. First, C(t)\mathcal{C}(t) immediately gives condition 1. of Lemma 7. By B(t)\mathcal{B}(t), we have that ∥W(t)−Φ∥σ≤∥W(t)−Φ∥F≤∥Φ∥F\|W(t)-\Phi\|_{\sigma}\leq\|W(t)-\Phi\|_{F}\leq\|\Phi\|_{F}, giving condition 2. of Lemma 7. A(t)\mathcal{A}(t) immediately gives condition 3. of Lemma 7. Finally, by B(t)\mathcal{B}(t), we have that LN(W1(t),…,WN(t))≤LN(W1(0),…,WN(0))L^{N}(W_{1}(t),\ldots,W_{N}(t))\leq L^{N}(W_{1}(0),\ldots,W_{N}(0)), so σmin(W1:N(t))≥c\sigma_{min}(W_{1:N}(t))\geq c by Claim 1. This verifies condition 4. of Lemma 7. Then Lemma 7 gives that LN(W1(t+1),…,WN(t+1))≤LN(W1(t),…,WN(t))−12σmin(W(t))2(N−1)/Nη∥dLdW(W(t))∥F2L^{N}(W_{1}(t+1),\ldots,W_{N}(t+1))\leq L^{N}(W_{1}(t),\ldots,W_{N}(t))-\frac{1}{2}\sigma_{min}(W(t))^{2(N-1)/N}\eta\left\|\frac{dL}{dW}(W(t))\right\|_{F}^{2}, establishing B(t+1)\mathcal{B}(t+1).

A(0),A′(1),…,A′(t),A(t),B(0),…,B(t),C(t)⇒A(t+1),A′(t+1)\mathcal{A}(0),\mathcal{A}^{\prime}(1),\ldots,\mathcal{A}^{\prime}(t),\mathcal{A}(t),\mathcal{B}(0),\ldots,\mathcal{B}(t),\mathcal{C}(t)\Rightarrow\mathcal{A}(t+1),\mathcal{A}^{\prime}(t+1). To prove this, note that for 1≤j≤N−11\leq j\leq N-1,

By B(0),…,B(t)\mathcal{B}(0),\ldots,\mathcal{B}(t), ∥W1:N(t)−Φ∥F≤∥Φ∥F\|W_{1:N}(t)-\Phi\|_{F}\leq\|\Phi\|_{F}. By the triangle inequality it then follows that ∥W1:N(t)∥σ≤2∥Φ∥F\|W_{1:N}(t)\|_{\sigma}\leq 2\|\Phi\|_{F}. Also A(t)\mathcal{A}(t) gives that for 1≤j≤N−11\leq j\leq N-1, ∥Wj(t)Wj⊤(t)−Wj+1⊤(t)Wj+1(t)∥F≤2δ\|W_{j}(t)W_{j}^{\top}(t)-W_{j+1}^{\top}(t)W_{j+1}(t)\|_{F}\leq 2\delta. By Lemma 6 with C=2∥Φ∥F,ν=2δC=2\|\Phi\|_{F},\nu=2\delta (so that (34) is satisfied),

In the first inequality above, we have also used the fact that for matrices A,BA,B such that ABAB is defined, ∥AB∥F≤∥A∥σ∥B∥F\|AB\|_{F}\leq\|A\|_{\sigma}\|B\|_{F}. (35) gives us A′(t+1)\mathcal{A}^{\prime}(t+1).

We next establish A(t+1)\mathcal{A}(t+1). By B(i)\mathcal{B}(i) for 0≤i≤t0\leq i\leq t, we have that ∥dLdW(W(i))∥F=∥W−Φ∥F≤∥Φ∥F\left\|\frac{dL}{dW}(W(i))\right\|_{F}=\left\|W-\Phi\right\|_{F}\leq\|\Phi\|_{F}. Using A′(i)\mathcal{A}^{\prime}(i) for 0≤i≤t0\leq i\leq t and summing over ii gives

Next, by B(0),…,B(t)\mathcal{B}(0),\ldots,\mathcal{B}(t), we have that L(W(i))≤L(W(0))L(W(i))\leq L(W(0)) for i≤ti\leq t. Since W(0)W(0) has deficiency margin of cc and by Claim 1, it then follows that σmin(W(i))≥c\sigma_{min}(W(i))\geq c for all i≤ti\leq t. Therefore, by summing B(0),…,B(t)\mathcal{B}(0),\ldots,\mathcal{B}(t),

where (2) follows from the definition of η\eta in (7), and the last equality follows from definition of δ\delta in Theorem 1. By (36), it follows that

A(t),B(t)⇒C(t)\mathcal{A}(t),\mathcal{B}(t)\Rightarrow\mathcal{C}(t). We apply Lemma 6 with ν=2δ\nu=2\delta and C=2∥Φ∥FC=2\|\Phi\|_{F}. First, the triangle inequality and B(t)\mathcal{B}(t) give

verifying precondition 2. of Lemma 6. A(t)\mathcal{A}(t) verifies condition 1. of Lemma 6, so for 1≤j≤N1\leq j\leq N, ∥Wj(t)∥σ≤(4∥Φ∥F)1/N\|W_{j}(t)\|_{\sigma}\leq(4\|\Phi\|_{F})^{1/N}, giving C(t)\mathcal{C}(t).

The proof of Lemma 2 then follows by induction on tt. ∎

D.3 Proof of Theorem 2

Theorem 2 is proven by combining Lemma 8 below, which implies that the balanced initialization is likely to lead to an end-to-end matrix W1:N(0)W_{1:N}(0) with sufficiently large deficiency margin, with Theorem 1, which establishes convergence.

Then, with probability at least (1−ϵ)⋅3−4F(2/b1)2(1-\epsilon)\cdot\frac{3-4F(2/\sqrt{b_{1}})}{2}, VV will have deficiency margin ∥Φ∥2/(b2d)\|\Phi\|_{2}/(b_{2}d) with respect to Φ\Phi.

The proof of Lemma 8 is postponed to Appendix D.5, where Lemma 8 will be restated as Lemma 16.

By Lemma 9 with k=d0/16k=d_{0}/16, we have that

We next use Lemma 8, with b1=∥Φ∥22/(2s2d02),b2=2∥Φ∥22/(s2d02)b_{1}=\|\Phi\|_{2}^{2}/(2s^{2}d_{0}^{2}),b_{2}=2\|\Phi\|_{2}^{2}/(s^{2}d_{0}^{2}); note that since a≥2a\geq 2, b1≥1b_{1}\geq 1, as required by the lemma. Lemma 8 then implies that with probability at least

W1:N(0)W_{1:N}(0) will have deficiency margin s2d0/2∥Φ∥2{s^{2}d_{0}}/{2\|\Phi\|_{2}} with respect to Φ\Phi. By the definition of balanced initialization (Procedure 1) W1(0),…,WN(0)W_{1}(0),\ldots,W_{N}(0) are -balanced. Since 24⋅6144<1052^{4}\cdot 6144<10^{5}, our assumption on η\eta gives

so that Equation (7) holds with c=s2d02∥Φ∥2c=\frac{s^{2}d_{0}}{2\|\Phi\|_{2}}. The conditions of Theorem 1 thus hold with probability at least that given in Equation (38). In such a constant probability event, by Theorem 1 (and the fact that a positive deficiency margin implies L1(W1:N(0))≤12∥Φ∥22L^{1}(W_{1:N}(0))\leq\frac{1}{2}\|\Phi\|_{2}^{2}), if we choose

then L1(W1:N(t0))≤18∥Φ∥22L^{1}(W_{1:N}(t_{0}))\leq\frac{1}{8}\|\Phi\|_{2}^{2}, meaning that ∥W1:N(t0)−Φ∥2≤12∥Φ∥2=∥Φ∥2−12σmin(Φ)\|W_{1:N}(t_{0})-\Phi\|_{2}\leq\frac{1}{2}\|\Phi\|_{2}=\|\Phi\|_{2}-\frac{1}{2}\sigma_{min}(\Phi). Moreover, by condition A(t0)\mathcal{A}(t_{0}) of Lemma 2 and the definition of δ\delta in Theorem 1, we have, for 1≤j≤N−11\leq j\leq N-1,

We now apply Theorem 1 again, verifying its conditions again, this time with the initialization (W1(t0),…,WN(t0))(W_{1}(t_{0}),\ldots,W_{N}(t_{0})). First note that the end-to-end matrix W1:N(t0)W_{1:N}(t_{0}) has deficiency margin c=∥Φ∥2/2c=\|\Phi\|_{2}/2 as shown above. The learning rate η\eta, by Equation (39), satisfies Equation (7) with c=∥Φ∥2/2c=\|\Phi\|_{2}/2. Finally, since

for d0≥2d_{0}\geq 2, by Equation (41), the matrices W1(t0),…,WN(t0)W_{1}(t_{0}),\ldots,W_{N}(t_{0}) are δ\delta-balanced with δ=∥Φ∥2/N(1/2)2256N3\delta=\frac{\|\Phi\|^{2/N}(1/2)^{2}}{256N^{3}}. Iteration t0t_{0} thus satisfies the conditions of Theorem 1 with deficiency margin ∥Φ∥2/2\|\Phi\|_{2}/2, meaning that for

Recall that this entire analysis holds only with the probability given in Equation (38). As lim⁡d→∞(1−2exp⁡(−d/16))=1\lim_{d\rightarrow\infty}(1-2\exp(-d/16))=1 and lim⁡a→∞(3−4F(22/a))/2=1/2\lim_{a\rightarrow\infty}(3-4F(2\sqrt{2/a}))/2=1/2, for any 0<p<1/20<p<1/2, there exist a,d0′>0a,d_{0}^{\prime}>0 such that for d0≥d0′d_{0}\geq d_{0}^{\prime}, the probability given in Equation (38) is at least pp. This completes the proof.

In the context of the above proof, we remark that the expressions 1−2exp⁡(−d0/16)1-2\exp(-d_{0}/16) and (3−4F(22/a))/2(3-4F(2\sqrt{2/a}))/2 converge to their limits of 11 and 1/21/2, respectively, as d0,a→∞d_{0},a\rightarrow\infty quite quickly. For instance, to obtain a probability of greater than 0.250.25 of the initialization conditions being met, we may take d0≥100,a≥100d_{0}\geq 100,a\geq 100.

D.4 Proof of Claim 2

We first consider the probability of δ\delta-balancedness holding between any two layers:

Note that for 1≤i,j≤d1\leq i,j\leq d, let XijX_{ij} be the random variable (ATA−BBT)ij(A^{T}A-BB^{T})_{ij}, so that

Then (43) follows from Markov’s inequality. ∎

Now the proof of Claim 2 follows from a simple union bound:

By (43) of Lemma 10, for each 1≤j≤N−11\leq j\leq N-1, k≥1k\geq 1,

and the claim follows with δ=ks210dmax3\delta=ks^{2}\sqrt{10d_{max}^{3}}. ∎

D.5 Proof of Claim 3

To establish Claim 3, we will use the following low-degree anti-concentration result of Carbery and Wright (2001) (see also Lovett (2010); Meka et al. (2016)):

There is an absolute constant C0C_{{0}} such that the following holds. Suppose that h{h} is a multilinear polynomial of KK variables X1,…,XKX_{1},\ldots,X_{K} and of degree NN. Suppose that X1,…,XKX_{1},\ldots,X_{K} are i.i.d. Gaussian. Then, for any ϵ>0\epsilon>0:

The below lemma characterizes the norm of the end-to-end matrix W1:NW_{1:N} following zero-centered Gaussian initialization:

Let f(W1,…,WN)=∥W1:N∥22f(W_{1},\ldots,W_{N})=\|W_{1:N}\|_{2}^{2}, so that ff is a polynomial of degree 2N2N in the entries of W1,…,WNW_{1},\ldots,W_{N}. Notice that

Next, by Lemma 11, there is an absolute constant C0>0C_{0}>0 such that for any ϵ>0\epsilon>0, and any 1≤i0≤d01\leq i_{0}\leq d_{0},

Since f2≥gi02f^{2}\geq g_{i_{0}}^{2} for each i0i_{0}, it follows that

Next, given 0<C2<10<C_{2}<1, choose ϵ=(1−C2)/(2C0N)\epsilon=(1-C_{2})/(2C_{0}N), and B1=2/(1−C2)B_{1}=2/(1-C_{2}). Then by (44) and (45) and a union bound, we have that

The result of the lemma then follows by taking C1=max⁡{21−C2,2C01−C2}C_{1}=\max\left\{\frac{2}{1-C_{2}},\frac{2C_{0}}{1-C_{2}}\right\}. ∎

For d≥20d\geq 20, choose any 0≤h≤10\leq h\leq 1. Then, the area of a (d,1)(d,1)-hyperspherical cap of height hh is at least

In Chudnov (1986), it is shown that the area of a (d,1)(d,1)-hyperspherical cap of height hh is given by 1−Cd−2(h)/Cd−2(0)2\frac{1-C_{d-2}(h)/C_{d-2}(0)}{2}, where

Next, by the inequality 1−t2≥exp⁡(−2t2)1-t^{2}\geq\exp(-2t^{2}) for 0≤t≤1/20\leq t\leq 1/2,

where the last inequality follows from the standard estimate F(x)≥1−exp⁡(−x2/2)F(x)\geq 1-\exp(-x^{2}/2) for x≥1x\geq 1. Also, since 1−t2≤exp⁡(−t2)1-t^{2}\leq\exp(-t^{2}) for all tt,

Therefore, for d≥20d\geq 20, by (46) and (47),

where the second inequality has used 1/(1−y)≤1+2y1/(1-y)\leq 1+2y for all 0<y<1/20<y<1/2 (and where y=2exp⁡((−(d−3)/4))<2exp⁡(−17/4)<1/2y=2\exp((-(d-3)/4))<2\exp(-17/4)<1/2), and the final inequality uses 1+4exp⁡(−(d−3)/4)≤21+4\exp(-(d-3)/4)\leq\sqrt{2} for d≥20d\geq 20. The above chain of inequalities gives us the desired result. ∎

By rescaling, we may assume without loss of generality that ∥Φ∥2=1\|\Phi\|_{2}=1, so that r=1/adr=1/\sqrt{ad}. Let D\mathcal{D} denote the intersection of Bd(r)B^{d}(r) with the open dd-ball of radius 1−1/(ad)1-1/(ad) centered at Φ\Phi. Let C⊂Bd(r)\mathcal{C}\subset B^{d}(r) denote the (d,r)(d,r)-hyperspherical cap of height r\cdot\big{(}1-2/(\sqrt{ad})\big{)}=r-2/(ad) whose base is orthogonal to the line between 0\mathbf{0} and Φ\Phi (see Figure 2). Note that σd,r(∂D∩Sd(r)){\sigma}^{d,r}(\partial\mathcal{D}\cap S^{d}(r)), the Haar measure of the portion of ∂D\partial\mathcal{D} intersecting Sd(r)S^{d}(r), gives the probability that VV belongs to the boundary of D\mathcal{D}. By Lemma 14 above (along with rescaling arguments), since d≥20d\geq 20, σd,r(∂C∩Sd(r))≥12⋅(3−4F(2/a))\sigma^{d,r}(\partial\mathcal{C}\cap S^{d}(r))\geq\frac{1}{2}\cdot(3-4F(2/\sqrt{a})), and therefore V∈∂CV\in{\partial}\mathcal{C} with at least this probability.

The calculation of ∥P−Q∥2\left\|P-Q\right\|_{2} is simply an application of the law of cosines: letting θ\theta be the angle determining the intersection of ∂Bd(r)\partial B^{d}(r) and T\mathcal{T} (see Figure 2), note that

Using that C⊆D\mathcal{C}\subseteq\mathcal{D}, we continue with the proof. Notice the fact that C⊆D\mathcal{C}\subseteq\mathcal{D} is equivalent to ∂C∩Sd(r)⊆∂D∩Sd(r)\partial\mathcal{C}\cap S^{d}(r)\subseteq\partial\mathcal{D}\cap S^{d}(r), by the structure of C\mathcal{C} and D\mathcal{D}. Since the probability that VV lands in ∂C\partial\mathcal{C} is at least 3−4F(2/a)2\frac{3-4F(2/\sqrt{a})}{2}, this lower bound applies to VV landing in ∂D\partial\mathcal{D} as well. Since all V∈∂DV\in\partial\mathcal{D} have distance at most 1−1/(ad)1-1/(ad) from Φ\Phi, and since σmin(Φ)=∥Φ∥2=1\sigma_{min}(\Phi)=\|\Phi\|_{2}=1, it follows that for any V∈∂DV\in\partial\mathcal{D}, ∥V−Φ∥2≤σmin(Φ)−1/(ad)\|V-\Phi\|_{2}\leq\sigma_{min}(\Phi)-1/(ad). Therefore, with probability of at least 3−4F(2/a)2\frac{3-4F(2/\sqrt{a})}{2}, VV has deficiency margin ∥Φ∥2/(ad)\|\Phi\|_{2}/(ad) with respect to Φ\Phi. ∎

Then, with probability at least (1−ϵ)⋅3−4F(2/b1)2(1-\epsilon)\cdot\frac{3-4F(2/\sqrt{b_{1}})}{2}, VV will have deficiency margin ∥Φ∥2/(b2d)\|\Phi\|_{2}/(b_{2}d) with respect to Φ\Phi.

By rescaling we may assume that ∥Φ∥2=1\|\Phi\|_{2}=1 without loss of generality. Then the deficiency margin of VV is equal to 1−∥V−Φ∥21-\|V-\Phi\|_{2}. μ\mu has a well-defined density, so we can set μ^\hat{\mu} to be the probability density function of ∥V∥2\|V\|_{2}. Since μ\mu is rotation-invariant, we can integrate over spherical coordinates, giving

where the first inequlaity used Lemma 15 and the fact that the distribution of VV conditioned on ∥V∥2=r\|V\|_{2}=r is uniform on Sd(r)S^{d}(r). ∎

Then Lemma 16, with d=d0d=d_{0}, b1=ab_{1}=a and b2=ad02C1⋅(C1N)2Nb_{2}=ad_{0}^{2}C_{1}\cdot(C_{1}N)^{2N}, implies that with probability at least 9991000⋅3−4F(2/a)2\frac{999}{1000}\cdot\frac{3-4F(2/\sqrt{a})}{2}, WW has deficiency margin ∥Φ∥2/(ad03C12N+1N2N)\|\Phi\|_{2}/(ad_{0}^{3}C_{1}^{2N+1}N^{2N}) with respect to Φ\Phi. But a≥105a\geq 10^{5} implies that this probability is at least 0.490.49, and from (48),

Next recall the assumption in the hypothesis that s≥C1N(c⋅∥Φ∥2/(d1⋯dN−1))1/2Ns\geq{C}_{1}N(c\cdot\left\|\Phi\right\|_{2}/(d_{1}\cdots{d}_{N-1}))^{1/2N}. Then the deficiency margin in (49) is at least

D.6 Proof of Claim 4

The target matrices Φ\Phi that will be used to prove the claim satisfy σmin(Φ)=1\sigma_{min}(\Phi)=1. We may assume without loss of generality that c≥3/4c\geq 3/4, the reason being that if a matrix has deficiency margin cc with respect to Φ\Phi and c′<cc^{\prime}<c, it certainly has deficiency margin c′c^{\prime} with respect to Φ\Phi.

We first consider the case d=1d=1, so that the target and all matrices are simply real numbers; we will make a slight abuse of notation in identifying 1×11\times 1 matrices with their unique entries. We set Φ=1\Phi=1. For all choices of η\eta, we will set the initializations W1(0),…,WN(0)W_{1}(0),\ldots,W_{N}(0) so that W1:N(0)=cW_{1:N}(0)=c. Then

so that W1:N(0)=cW_{1:N}(0)=c. Then since LN(W1,…,WN)=12(1−WN⋯W1)2L^{N}(W_{1},\ldots,W_{N})=\frac{1}{2}(1-W_{N}\cdots W_{1})^{2}, the gradient descent updates are given by

where we view W1(t),…,WN(t)W_{1}(t),\ldots,W_{N}(t) as real numbers. This gives

Since 3/4≤c<13/4\leq c<1 and −η(c−1)c(N−1)/NA≥0-\eta(c-1)c^{(N-1)/N}A\geq 0, we have that A/2≤3A/4≤Wj(1)A/2\leq 3A/4\leq W_{j}(1) for 1≤j≤N/21\leq j\leq N/2. Next, since 1−c1−c1/N≤N\frac{1-c}{1-c^{1/N}}\leq N for 0≤c<10\leq c<1, we have that A2≥ηN≥η(1−c)1−c1/NA^{2}\geq\eta N\geq\frac{\eta(1-c)}{1-c^{1/N}}, which implies that A2≥c1/NA2+η(1−c)A^{2}\geq c^{1/N}A^{2}+\eta(1-c), or c1/NA+η(1−c)A≤Ac^{1/N}A+\frac{\eta(1-c)}{A}\leq A. Thus Wj(1)≤AW_{j}(1)\leq A for N/2<j≤NN/2<j\leq N. Similarly, using the same bound 3/4≤c<13/4\leq c<1 and the fact that η(1−c)c(N−1)/NA≥2\eta(1-c)c^{(N-1)/N}A\geq 2 we get 316ηA≤Wj(1)≤ηA\frac{3}{16}\eta A\leq W_{j}(1)\leq\eta A for N/2<j≤NN/2<j\leq N. In particular, for all 1≤j≤N1\leq j\leq N, we have that min⁡{η,1}10A≤Wj(1)≤max⁡{η,1}A\frac{\min\{\eta,1\}}{10}A\leq W_{j}(1)\leq\max\{\eta,1\}A.

We prove the following lemma by induction:

For each t≥1t\geq 1, the real numbers W1(t),…,WN(t)W_{1}(t),\ldots,W_{N}(t) all have the same sign and this sign alternates for each integer tt. Moreover, there are real numbers 2≤B(t)<C(t)2\leq B(t)<C(t) for t≥1t\geq 1 such that for 1≤j≤N1\leq j\leq N, B(t)≤∣Wj(t)∣≤C(t)B(t)\leq|W_{j}(t)|\leq C(t) and ηB(t)2N−1≥20C(t)\eta B(t)^{2N-1}\geq 20C(t).

First we claim that we may take B(1)=min⁡{η,1}10AB(1)=\frac{\min\{\eta,1\}}{10}A and C(1)=max⁡{η,1}AC(1)=\max\{\eta,1\}A. We have shown above that B(1)≤Wj(1)≤C(1)B(1)\leq W_{j}(1)\leq C(1) for all jj. Next we establish that ηB(1)2N−1≥20C(1)\eta B(1)^{2N-1}\geq 20C(1). If η≤1\eta\leq 1, then

where the inequality follows from A≥(20⋅102N−1η2N)1/(2N−2)A\geq\left(\frac{20\cdot 10^{2N-1}}{\eta^{2N}}\right)^{1/(2N-2)} by definition of AA. If η≥1\eta\geq 1, then

where the inequality follows from A≥2000≥(20⋅102N−1)1/(2N−2)A\geq 2000\geq\left(20\cdot 10^{2N-1}\right)^{1/(2N-2)} by definition of AA.

Now, suppose the statement of Lemma 17 holds for some tt. Suppose first that Wj(t)W_{j}(t) are all positive for 1≤j≤N1\leq j\leq N. Then for all jj, as B(t)≥2B(t)\geq 2, and ηB(t)2N−1≥20C(t)\eta B(t)^{2N-1}\geq 20C(t),

which establishes that Wj(t+1)W_{j}(t+1) is negative for all jj. Moreover,

Now set B(t+1)=9C(t)B(t+1)=9C(t) and C(t+1)=ηC(t)2N−1C(t+1)=\eta C(t)^{2N-1}. Since N≥2N\geq 2, we have that

The case that all Wj(t)W_{j}(t) are negative for 1≤j≤N1\leq j\leq N is nearly identical, with the same values for B(t+1),C(t+1)B(t+1),C(t+1) in terms of B(t),C(t)B(t),C(t), except all Wj(t+1)W_{j}(t+1) will be positive. This establishes the inductive step and completes the proof of Lemma 17. ∎

By Lemma 17, we have that for all t≥1t\geq 1, LN(W1(t),…,WN(t))=12(W1:N(t)−1)2≥12(2N−1)2>0L^{N}(W_{1}(t),\ldots,W_{N}(t))=\frac{1}{2}(W_{1:N}(t)-1)^{2}\geq\frac{1}{2}(2^{N}-1)^{2}>0, thus completing the proof of Claim 4 for the case where all dimensions are equal to 1.

For the general case where d0=d1=⋯=dN=dd_{0}=d_{1}=\cdots=d_{N}=d for some d≥1d\geq 1, we set Φ=Id\Phi=I_{d}, and given c,ηc,\eta, we set Wj(0)W_{j}(0) to be the d×dd\times d diagonal matrix where all diagonal entries except the first one are equal to 1, and where the first diagonal entry is given by Equation (51), where AA is given by Equation (50). It is easily verified that all entries of Wj(t)W_{j}(t), 1≤j≤N1\leq j\leq N, except for the first diagonal element of each matrix, will remain constant for all t≥0t\geq 0, and that the first diagonal elements evolve exactly as in the 1-dimensional case presented above. Therefore the loss in the dd-dimensional case is equal to the loss in the 1-dimensional case, which is always greater than some positive constant. ∎

D.7 Proof of Claim 5

where the last inequality follows since Φ\Phi has a negative eigenvalue. To analyze gradient descent we use the following result, which was established in Bartlett et al. (2018):

If W1(0),…,WN(0)W_{1}(0),\ldots,W_{N}(0) are all initialized to identity, Φ\Phi is symmetric, Φ=UDU⊤\Phi=UDU^{\top} is a diagonalization of Φ\Phi, and gradient descent is performed with any learning rate, then for each t≥0t\geq 0 there is a diagonal matrix D^(t)\hat{D}(t) such that Wj(t)=UD^(t)U⊤W_{j}(t)=U\hat{D}(t)U^{\top} for each 1≤j≤N1\leq j\leq N.

By Lemma 18, for any choice of learning rate η\eta, the end-to-end matrix at time tt is given by W1:N(t)=UD^(t)NU⊤W_{1:N}(t)=U\hat{D}(t)^{N}U^{\top}. As long as some diagonal element of DD is negative, say equal to −λ<0-\lambda<0, then

Appendix E Implementation Details

Below we provide implementation details omitted from our experimental report (Section 4).

The platform used for running the experiments is PyTorch (Paszke et al., 2017). For compliance with our analysis, we applied PCA whitening to the numeric regression dataset from UCI Machine Learning Repository. That is, all instances in the dataset were preprocessed by an affine operator that ensured zero mean and identity covariance matrix. Subsequently, we rescaled labels such that the uncentered cross-covariance matrix Λyx\Lambda_{yx} (see Section 2) has unit Frobenius norm (this has no effect on optimization other than calibrating learning rate and standard deviation of initialization to their conventional ranges). With the training objective taking the form of Equation (1), we then computed cc — the global optimum — in accordance with the formula derived in Appendix A. In our experiments with linear neural networks, balanced initialization was implemented with the assignment written in step (iii) of Procedure 1. In the non-linear network experiment, we added, for each j∈{1,…,N−1}j\in\{1,\ldots,N-1\}, a random orthogonal matrix to the right of WjW_{j}, and its transpose to the left of Wj+1W_{j+1} — this assignment maintains the properties required from balanced initialization (see Footnote 7). During all experiments, whenever we applied grid search over learning rate, values between 10−410^{-4} and 11 (in regular logarithmic intervals) were tried.