Convergence of Generative Deep Linear Networks Trained with Bures-Wasserstein Loss
Authors: Pierre Bréchet
Katerina Papagiannouli
Jing An
Guido Montufar
Venue: Submitted to ICLR 2023 (2023)
Keywords: deep linear network
low-rank approximation
Bures-Wasserstein distance
optimal transport
implicit generative model
critical points

Abstract: We consider a deep matrix factorization model of covariance matrices trained with the Bures-Wasserstein distance. While recent works have made important advances in the study of the optimization problem for overparametrized low-rank matrix approximation, much emphasis has been placed on discriminative settings and the square loss. In contrast, our model considers another interesting type of loss and connects with the generative setting. We characterize the critical points and minimizers of the Bures-Wasserstein distance over the space of rank-bounded matrices. For low-rank matrices the Hessian of this loss can blow up, which creates challenges to analyze convergence of optimizaton methods. We establish convergence results for gradient flow using a smooth perturbative version of the loss and convergence results for finite step size gradient descent under certain assumptions on the initial weights

Under review as a conference paper at ICLR 2023
CRITICAL POINTS AND CONVERGENCE ANALYSIS OF GENERATIVE DEEP LINEAR NETWORKS TRAINED WITH WAISESSASSNR LOSS
Anonymous authors Paper under double-bling brind review
ABSTRACT
We consider a deep matrix factorization model covariance matrices trained with the Sures-Wasserstein distance. While recent works have made important ad- vances in the study of the optimization problem for overparametrized low-rank matrix approximation, much emphasis has been placed discriminative settings and the square loss. In contrast, our model considers another interesting type loss and connects with the generative setting. We characterize the critical points and minimizers of the Bures- Wasssersenn distance over the space of rank- moo-boddddd matrices. For low-rank matrices the Hessian of this loss can blow up, which cre- ates challenges to analyze convergence of optimizaton methods. We establish convergence results for gradient flow using a smooth perturbative version of the loss and convergence results for finite step size gradient descent under certain as- sumptions on the initial weights.
INTRODUCTION
We investigate generative deep linear networks and their optimization using the distance. More precisely, we consider the problem of approximating a target Gaussian distribution with a deep linear neural network generator of Gaussian distributions by minimizing the Bures- Wasserstein distance. This problem of interest in two important ways. First, it pertains the opti- mization of deep linear networks for a type of loss that is qualitatively different from the well-studied and very particular square loss. Second, it can be regarded as a simplified but instructive instance of the parameter optimization problem in generative networks, specifically Wasserstein generative adversarial networks, which are currently not as well understood as discriminative networks.
The optimization landscapes and the properties of parameter optimization procedures for neural networks are among the most puzzling and actively studied topies theoretical deep learning (see, e.g. Mei et al.] 2018, Liu et al. 2022). Deep linear networks, i.c., neural networks having the
identity as activation function, serve as a simplified model for such investigations (Baldi & Hornik
1989 Kawaguchi 2016 Trager et al. 2020 Kohn et al., 2022 Bah al.) 2021). ) The study of
linear networks has guided the development of several useful notions and intuitions in the theoretical
analysis of neural networks, from the absence of bad local minima to the role of parametrization and overparametrization in gradient optimization (Arora et al. 2018 2011a b). Many previous works have focused on discriminative or autoregressive settings and have emphasized the square loss. Although the square loss is indeed popular choice in regression tasks, it interacts in a very special way with the particular geometry of linear networks (Trager et al. (020). The behavior of linear networks optimized with different losses has also been considered in several works (Laurent)
Brecht 2018 Lu & Kawaguchi] 2017 Trager al. 2020 but less well understood.
The Bures- Wasserstein distance was introduced by Bures 1969) to study Hermitian operators in quantum information, particularly density matrices. induces a metric on the space of positive semi-definite matrices. The Bures-Wassersseer distance corresponds to the 2-Wasserstein distance between two centered Gaussian distributions (Bhatia et al. 2019). Wasserstein distances enjoy sev- eral properties, e.g. they remain well defined between disjomtly supported measures and have duality formulations that allow for practical implementations (Villani 2003), that make them good candi- dates and indeed popular choices of loss for learning generative models, with well-known case

---

Under review as a conference paper at ICLR 2023
being the Wasserstein Generative Adversarial Networks (GANs) (Arjovsky al. 201 7). While the 1-Wasserstein distance has been most commonly used in this context, the Bures- Wasserstein distance has also attracted much interest, e.g. in the works of Muzellee & Cuturi 2018); Chewi let al. 2020); Mallasto et al. (2022), and has also appeared in the context of linear quadratic Wasser- stein generative adversarial networks (Feizi al. 2020). A 2-Wasserstein GAN is a minimum 2- 2-Wasserstein distance estimator expressed in Kantorovich duality (see details in Appendix B). This model can serve as an attractive platform to develop the theory particularly when the inner prob- lem can be solved in closed-form. Such formula is available when comparing pairs of Gaussian distributions. In the case of centered Gaussians this corresponds precisely the distance. Strikingly, even in this simple case, the optimization properties of the corresponding prob- lem are not well understood; which we aim to address in the present work.
1.1 CONTRIBUTIONS
We establish series of results on the optimization of deep linear networks trained with the Bures- Wasserstein loss, which we can summarize as follows.
We obtain an analogue of the Eckart- Fckart-Young-Mint Young-Mirsky theorem characterizing the critical points and minimizers of the Wasssssessr distance over matrices of a given rank (Theorem 4.2).
To circumvent the non- non-smooth behaviour of the warseesnssr loss when the matrices drop rank, we introduce a smooth perturbative version (Definition 5 and Lemma 3.3) and characterize its critical points and minimizers over rank-constrained matrices (Theorem and link them to the critical points on the parameter space (Proposition 4.5)
For the smooth Bures-Wasseree Wasserstenn loss, in Theorem 5.6 we show exponential convergence of the gradient flow assuming balanced initial weights (Denniiion 2.1 and uniform margin deficiency condition (Definition 5.2).
For the Bures- Wassersein loss and its smooth version, in Theorem 5.7 we show convergence of gradient descent provided the step size is small enough and assuming balanced initial weights.
1.2 RELATED WORKS
Low rank matrix approximation The function space of linear network corresponds m matrices of rank at most d, the lowest width of the network. Hence optimization in function space closely related the problem of approximating given data matrix by low-rank matrix. When the approximation error is measured in Frobenius norm, Eckart& & Young (1936) characterized the optimal bounded-rank approximation of given matrix in terms its singular value decomposition. Mirsky (1960) obtained the same charactcrization for the more general case of unitary invariant matrix norms, which include the Euclidean operator norm and the Schatten-p norms. There are generalizations to certain weighted norms (Ruben & Zamir 1979. Dutta & Li, (017). However, for general norms the problem is known be difficult (Song et al. 2017 Gillis & Vavasis 2018 Gillis & Shitov 2019).
Loss landscape of deep linear networks For the square loss, the optimization landscape of lin- ear networks has been studied in numerous works. The pioneering work of Baldi & Hornik (1989) showed, focusing on the two-layer case, that there is single minimum (up to a trivial parametriza- tion symmetry) and all other critical points are saddle points. Kawaguchi (206)) obtained corre- sponding results for deep linear nets and showed the existence of bad saddles (with no negative Hessian eigenvalues) for networks with more than three layers. Chulhee et al. (2018) found sets of parameters such that any critical point in this set is global minimum and any critical point outside saddle. Variations include the study of critical points for different types of architectures, such as deep lincar residual networks (Hardt & Ma 2017) and deep linear convolutional networks (Kohn cct al. 2022).
For losses different from the square loss there are also several results. Laurent & Brecht (2018) showed that deep linear nets with no bottlenecks have no local minima that are not global for ar- bitrary convex differentiable loss. Lu & Kawaguchi (2017) showed that if the loss is such that any local minimizer in parameter space can be perturbed to an equally good minimizer with full-rank factor matrices, then all local minima in parameter space are local minima in function space. Trager
2

---

Under review as a conference paper at ICLR 2023
et al. 20200) found that for linear networks with arbitrarily rank-constrained function space, only for the square loss does one have the nonexistence of non-global local minima. However, for ar- bitrary convex losses, non-global local minima when they exist are always pure, meaning that they correspond to local minima in function space.
Optimization dynamics of deep linear networks Saxe et al. (2014) studied the learning dynam- ies of deep linear networks under different classes of initial conditions. Arora et al. (2019b) obtained a closed-form expression for the parametrization along time in deep linear network for the square loss. Notably, the authors found that solutions with a lower rank are preferred as the depth of the network increases. [Arora et al. (2018) derive several invariances of the flow and compare the dy- namics in parameter and function space. For the square loss |rroa et al. (2019a) proved linear convergence for linear networks with no bottlenecks, approximately balanced initial weights, and initial loss smaller than for any rank-deficient solution. A detailed analysis of the dynamics in the shallow case with square loss was conducted by Tarmoun et al. (202T); Min et al. (202T) including symmetric factorization. The role on inbalance was remarked in those works. For the deep case, also focusing on the square loss, |Bah et al. (202T) showed the gradient flow converges to a critical point and a global minimizer on the manifold of fixed rank matrices of some rank. More recently, Nguegnang et al. (2021) extended this analysis to obtain corresponding results for gradient descent.
Bures- Wasserstein distance Chewi et al. (020)) studied the convergence of gradient descent al- gorithms for the Bures- Wassseseessr barycenter, proving linear rates of convergence for the gradient descent. In contrast to our work, they consider a Polyak-Lojasiewicz inequality derived from the optimal transport theory to circumvent the non geodesical convexity of the barycenter. In the same vein, Muzellec & Cuturi 2018) exploit properties of optimal transport theory to optimize the dis- tance between two elliptical distributions. To avoid rank deficiency, they perturbed the diagonal elements of the covariance matrix by a small parameter. Feizi et al. (2020) characterized the optimal solution of 2- Wasseesttein GAN with rank-k linear generator as the k-PCA solution. We will obtain an analogue result in our settings, along with a description of critical points.
1.3 NOTATIONS
... n). We adopt the following notations throughout the paper. For any N, denote [n] {1, 2.... Let S(n) be the spaces of real symmetric matrices of size n. We denote by S+ (resp. S (n) the space of real symmetric positive semi-definite (resp. definite) matrices of size n. Given k & n, the set of rank k positive semi-definite matrices is denoted by (k; We use Mr (resp. Max) to denote the of matrices of rank exactly k (resp. of rank at most k), with the format of the matrix understood from context. The scalar product between two matrices A, B E Rnxm is (A, B) AA" B. and its associated (Frobenius) norm noted I1- The identity matrix of size n will be denoted I,, or 7 when is clear. For (Fréchet) differentiable function f: X -> Y, we denote its differential at x E X in the direction v by d (x) Finally, Crit((f) gives the set of critical points of f. i.c. the set of points at which the differential of is 0.
2 LINEAR NETWORKS AND THEIR GRADIENT DYNAMICS
We consider a linear network with do inputs and N layers of widths di,..., dd. which is a model of linear functions of the form
I WN N.-... ... WiS,
parametrized by the weight matrices W, E RR d; 1, E [N]. We will denote the tuple of weight matrices by W = (W1.... WN) and the space of all such tuples by O This the parameter space of our model. To slightly simplify the notation we will also denote the input and output dimensions by m = dg and n = dy. respectively, and write W = WN W1 for the end-to-end matrix. For any 1 < i < j < N. we will also write Wf W,- ... W. for the matrix product of layer up to j- We note that the represented function is linear in the network input x, but the parametrization is not linear in the parameters W. We denote the parametrization map by
p: θ Rdx xdo W (W1, (W,...., WN) Wxa WN WT.
3

---

Under review as a conference paper at ICLR 2023
The function space of the network is the set of linear functions it can represent, corresponds to the set of possible end-to-end matrices, which are the n x m matrices of rank at most d' := minfdg..... ). ). When d min dx J. the function space is a vector space, but otherwise, when there is a bottleneck so that d < do. dy, it is a non-convex subset of Rm determined by polynomial constraints, namely the vanishing of the (d+11 1) (d 1) minors.
Next we collect a few results on the gradient dynamics of linear networks for general differentiable losses, which have been established in previous works even when in some cases the focus was on the on the square loss (Kawaguchi] 2016 Bah et al., 202T; Chitour et al. 2022 Arora et al., 2018). In the interest of conciseness, here we only provide the main takeaways and defer a more detailed discussion to Appendix(C For the remainder of this section let L': R Ruxm - I be any differentiable
loss and let L be defined through the parametrization μ as L (W) L2 u(W). For any such loss, the gradient flow t + (t) is defined by
(t VL (W()) df
dw, Vw, I. (W) Ws(()). (GF)
Vj E [N],
dt
This governs the evolution of the parameters during gradient minimization of the loss. Further, we observe that the gradient of L is given by
Vw, L (W Wn) = W ... W (W)WT ...JI for j E NJ. (I) As it turns our, the gradient flow dynamics preserves the difference of the Gramians of subsequent layer weight matrices, which are thus invariants of the gradient flow,
d Jे
d di (W)()),(()T)
The important notion of balancedness for the weights of linear networks was first introduced by Fukumizu (1998) in the shallow case and generalized to the deep case Du et al. (018). This is useful in particular as a way of removing the redundancy of the parametrization when investigating the dynamics in function space and has been considered in numerous works.
Definition 2.1 (Balanced weights). The weights WI,..., W are said to be balanced if, for all j E [N- 1], W,W, WJ+ W
Assuming balanced initial weights, if the flow of each W, is defined and bounded, then the rank of the end-to-end matrix remains constant during training (Bah et al. 2021] Proposition 4.4). More- over, the products W W
W N:1 and W N:1 W can be written in concise manner; namely, we have WN W..1 = (Wx WW)WT) W N and W WN = W.) N which simplifies many computations. N:1
Remark 2.2. In order to relax the balanced initial weights assumption, some works also consider approximate balancedness (Arora et al.. 2019a). which requires only that there exists 8 > 0 such that for IN 1]. We will use exactly balanced initialization in our proofs, but they would also go through by invoking approximate balancedness. Another alternative initialization has been proposed by Yun et al. (2021). We defer such extensions to future work favoring here the discussion of the Bures- Wasserstein loss.
[[,, W,W,J, w
3 WASSERSTEN GENERATIVE LINEAR NETWORK
3.1 THE WAISEESSSS LOSS
(BW) distance defined on the space of positive semi-definite matrices (or covariance spuce) (n). Here we collect some of the key properties and discuss the gradient, and refer the reader to Bhatia et al. (0019) for further details on this interesting distance. Definition 3.1 (Bures- Wassersesnn distance). Given two symmetric positive semidefinite matri- ces Eg, Σ E S+ (n), the squared Mures-Wasserstein distance between En and Σ is defined as
The
(E, Eo) = tr + - /200///1////1
(Krohhnin et al. 2021 Lemma A.3) shows that the square root is differentiable on the set of positive definite matrices and as consequence we can differentiate the BW distance at Eg, E E S++ (n).
4

---

Under review as a conference paper at ICLR 2023
However, the mapping Σ + 132 Eu) is not differentiable over all of R"*" Indeed, letting rQr' be a spectral decomposition /// ////2>1/// 2 1/2
(2 can be written as
B (2,
/ + | - 2t
(3)
Due to the square root on Q, the map Σ + (E, Eo) is not differentiable when the rank of
2. ////////// i.e. the number of positive eigenvalues of Q. changes. More specifically, while one can compute the gradient over the set of matrices of rank & for any given k. the norm of the gradient blows up if the matrix changes rank. The gradient of 1 restricted to the set of full-rank matrices is given inAppendix
3.2 LINEAR WASSERSTEIN GAN
The distance defined in (2) corresponds to the 2- 2-Wasserstein distance between two zero-centered Gaussians and can be used loss for training models of Gaussian distributions and in particular generative linear networks. Recall that zero-centered Gaussian distributions are completely specified by their covariances. Given a bias-free linear network and latent Gaussian probability measure (0, m) linear network generator computes a push-forward of the latent distribution, which is again a Gaussian distribution. If Z ~ (O, (m) and = W WZ, then
X ~ N(, WW =:
Given a target distribution un = (0, Σo) or simply the covariance matrix Eg (which may be sample covariance matrix), one can select W by minimizing (" -oo) = WM vo) so that the network approximates the distribution (0, Eg). We will denote the map that takes the end-to- end matrix W to the covariance matrix WW by x: R=xm R"*n, W + WWT.
Loss in covariance, function, and parameter spaces We consider the following losses, which differ only on the choice of the search variable, taking function space or parameter space view.
First, we denote the loss over covariance matrices E E (+ as L: Σ (E, Eo)
Secondly, given W WWT (+)) we define the loss in function space, i.e., over end-to-end matrices W E Rux L' W L TWW). This is given by
WWE E Raxm
/WW ('(W)= tr (WWT- + E. 2 WW w /322222)
(4)
This loss not convex on RTxm as can be seen the one-dimensional case.
Lastly, for tuple of weight matrices W (W1,... WN), we compose L with the parametriza-
tion map ju: W + WN.I, to define the loss parameter space as LN W + Lo u(W), for Wh O. Observe that this is, again, non-convex loss.
Thus, for W E L,, = L' (a(()) L(T ((TT)((W)))))))))))()) = (TTTo En). While the gradient flow GH is defined on the parameters W. the covariance space perspective is useful since it leads to a convex objective function, even if this may subject to non-convex constraints. One of our goals will be to translate propertics between L, L', and L.
Smooth perturbative loss
As mentioned before, the Bures- Wasserstein loss non-smooth ma- trices with vanishing singular values. In turn, the usual analysis tools to prove uniqueness and convergence of the gradient flow do not apply for this loss. Therefore, we introduce smooth per- turbative version. Consider the perturbation map fr: Σ + Σ + TIa. where T > 0 plays the role of regularization strength. Then the perturbative loss on the covariance space can be expressed as ī,, L Pr, and the perturbative loss on function space as L I. æ. More explicitly, we let
LW = tr (WWT + +TI, + Eg 2(= (WWT IIm)
(5)
This function is smooth enough to apply usual convergence arguments for the gradient flow. Like- wise, I.N = Fup y well-defined and smooth O.
5

---

Under review as a conference paper at ICLR 2023
Remark 3.2. The perturbative loss 5 is differentiable. Many results from Brom/Bah Bah et al. [2021 can be carried over for the differentiable Bures- Wasserstein loss. For example, the uniform boundedness of the end-to-end matrix holds, [|W & ( (W))) + tr Σn. Similar observations may apply for the case of L' in the case that the matrix WWT remains positive definite throughout training, in which case the gradient flow remains well defined and the loss is monotonically decreasing. We expand on this in Appendix
The next lemma, proved in Appendix B.4 provides a quantitative bound for the difference between the original and the perturbative loss. For this lemma, we use the fact that the rank is constant. To compare the two losses, we fix the parameters to be the same. Recall that Σ, WW +TI.
Lemma 3.3. Let W E __xm >>0 O, and let M, = WWT T +5I. Assume that rank(WW W =
rank Eg L' (W) is given by 4 and L(E,) is given by 5 Then
[(E,) - L' (W)
1/1
We observe that the upper bound 66 is tight, since it goes to zero as goes to zero.
4 CRITICAL POINTS
In this section, we characterize the critical points of the Bures- Wasserstein loss restricted to matrices of given rank. The proofs of results in this section are given in Appendix D
For E N, denote by Mr the manifold rank-k matrices, i.e.
M₁ ( E " rank W =k].
Similarly, denote MMc. the set of matrices of rank at most k. The format x of the matrices to be inferred from the context. The manifold MM viewd an embedded submanifold of the linear space (I ()), with codimension (n (m k) (Boumal 2022 $2.6; Uschmajew Vandereycken 2020 $9.2.2). Given function f: Ra + R. its restriction on MM is denoted by flM, M₁ W f(W). Even if function is not differentiable over all of R= its restriction on Ma may be differentiable.
Definition 4.1. Let M be a smooth manifold. Let f: R=xm restriction on MM is differentiable. A point W .M is said to be a critical point for J'lM if the differential of f'lm at W is the zero function, dV)
be any function such that its
4.1 CRITICAL POINTS OF OVER M₁
Given matrix A R and set To C [n] of indices, denote by Az. E Rank the sub- matrix of A consisting of the columns with index in Ja- If the matrix A is diagonal, we denote by and columns with index in Jk- The
Ag. R the diagonal submatrix which extracts the following result characterizes the critical points of the loss in function space. It can be regarded as type of Eckart-Youee Youuart-YYu---booe Young-Mirsky result for the case of the Bures-Wassersie Wasserstein loss.
Theorem 4.2 (Critical points of L'). Assume Eo has distinct, positive eigenvalues. Let Eo = QAQ be a spectral decomposition of Eg (so Ω E U()).. Let E [min (n, m]]. A matrix W MM is critical point of L' [M if and only if W -1, 1/2 for some Ts S. [n] and V Rmxk with VTV Ig. In particular, the number of critical points is (2). The minimum attained for Ts = [k]. In particular, inf L' (W) min,, L' (W) and minm, (W) min (WW.
The proof relies on an expression of the gradient V (see Lemma D.3 and evaluating its zeros. The value of the loss evaluated at these critical points allows rank them and identify those that minimal.
Remark 4.3. Interestingly, the critical points and the minimizer of L' characterized the above result agree with those of the square loss (Eckart & Young 1936 Mirsky| 1960). Nonetheless, we observe that (2) is only defined for positive semidefinite matrices. Hence the notion of unitary invariance considered by Mirsky (1960) here only makes sense for left and right multiplication by

---

Under review as a conference paper at ICLR 2023
the same matrix. Moreover, while we can establish unitary invariance for variational extension of the distance, this still is not a norm in the sense that there no function B: R=*x + RR such that Eo) B Σg), and hence it does not fall into the framework of Mirsky (1960). We offer more details about this in Appendix
4.2 CRITICAL POINTS OF THE PERTURBATIVE LOSS
For the critical points of the perturbative loss (W) we obtain the following results.
Theorem 4.4 (Critical points of L.). Let Eg QAΩ be spectral decomposition of Eg. A point W E Mr is critical point for L, if and only if there exists a subset Js C [n]] and semi- orthogonal matrix V E B (i.e., so that V I) such that W 22 (Ag. =I&) =Ix) V2VT. The (unique) minimum over Mag is attained for Js =
Note that the above characterization of the critical points imposes an upper bound on т. In other words, for given W to be critical point, one must have that t X for all j E Tk. because the cigenvalues of ÅΛ. TIy are positive.
In order to link the critical points in parameter space to the critical points in the function space, we appeal to the correspondence drawn in Trager et al. (2020) Propositions 6 and 7). For the Bures- Wasserstein loss, this allows to conclude the following.
Proposition 4.5. Assume a full-rank targer Eg, with distinct, decreasing eigenvalues, and spectral decomposition No = QAN Let T E Am). If W E Crit(LE),, then E = (P((W)) is critical point of the loss Ly (k.n). where k rank Moreover, if d. then W local minimizer for the loss LA if and only E ((((W*)) is local minimizer, and therefore the global minimizer, of the loss L ls, (d.a). In this case, 2;; - + TIn is the T-best d-rank approximation of the target the covariance space, the sense that Ei; (Aw 7
ST. 7)
Proposition 4.5 ensures that, under the assumption that the solution of the gradient flow is (local) minimizer and has the highest possible rank d given the network architecture, the solution in the covariance space is the best d-rank approximation of the target in the sense of the 7-smoothed distance.
Remark 4.6. Under the balancedness assumption, one can show that the rank of the end-to-end matrix does not drop during training (Bah al. 202T Proposition 4.4), and that one escape almost surely the strict saddle points (Bah et al. [021] Theorem 6.3). If the initialization of the network has rank d, the matrices W(1), L > maintain rank throughout training. There can be issues in the limit, since Ma is not closed. Proving that the limit point also belongs to Ma is an interesting open problem that leave for future work.
5 CONVERGENCE ANALYSIS
The Bures- Wasserstein distance can be viewed through the lens of the Procrustes metric (Dryden
et al. 2009 Masarotto et al. 2019). In fact, it can be obtained by the following minimization
problem.
Lemma 5.1 (Bhatia et al. 2019 Theorem 1). For Σ, Eo E SS+ (n).
(o)
min (eU(m)
|2/ U1?
where (n) denotes the orthogonal matrix group. Moreover, the minimizer Ü occurs in the polar decomposition of 1//221/22
We emphasize that in the above description the Bures Wasserstein distance, the minimizer U de- pends on W, so that 8 fundamentally differs from squared Frobenius norm. Moreover, the square root on Eg can lead to singularities when differentiating the loss. Nonetheless, the expression 77) can be used avoid such singularities, by leveraging the following deficiency margin concept.
7

---

Under review as a conference paper at ICLR 2023
Definition 5.2 (Modified deficiency margin). Given target matrix Eg E R and positive constant c O, we say that Eg E Ru** has a modified deficiency margin c with respect to Ep if
12 2 S Omin (2)
min VEU(n)
(8)
With slight abuse of denomination, we will say that W has a uniform deficiency margin WWT has one. This deficiency margin idea can be traced back to Arora et al. (2019a). Note that we can write VWWT- // and this square root can be realized by Cholesky decomposition. Notice that if we initialize close to the target then the above bound ( holds trivially. In fact, if the initial condition (0) satisfies the uniform deficient margin, then we have that the least singular value of (K)) remains bounded below by for all times k 2 0, for the gradient flow or gradient descent with decreasing L.N:
Lemma 5.3. Suppose (O) has a modified deficiency margin c with respect to Ey. Then
( (VW(L)W(H)T)) cc for z O.
muin
The proof of this and all results in this section are provided in Appendix We note that, while the modified margin deficiency assumption is sufficient for Lemma 5.3 to hold, by no means necessary. We will assume that the modified margin deficiency assumption holds for the simplicity of exposition, but the gradient flow analysis in the next paragraph only requires the less restrictive Lemma 5.3 to hold.
Convergence of the gradient flow for the smooth loss Because we cannot exclude the possibility that the rank of WWT drops along the gradient flow of the BW loss, we consider the smooth perturbation way to avoid singularities. We consider the gradient flow (GF) for the perturbative loss 5 The gradient of (5) is given by
(L- 2 - (2 (WWT +TI,) - 1/2 2 W))
On the other hand, we may denote Σ, == WWT + TIa as regularized covariance matrix, and express the I,- loss terms of the optimal transport plan between Σ, and En (Kroshnin et al. We have
2021.
i(Σ-)= tr (E, + En 2 = I[ 1)
(10)
= tr (T - 1) tr (T - 1),
where T 1/2 (5) د, 1/) Σ 1/2 (-) E 1/,,,,/" 1/2
1/2 ///2///////
The perturbation TIn ensures strict convexity as shown in the following result.
Lemma 5.4. The function E, i(E,) strictly convex on (4).
Proof. First we observe that the function f(X) = tr 1/2 is strictly concave on S++(m): for details we refer the reader to Bhatia et al. (2019. Theorem 7). As a result, the function
Σ + i(Ex)) = tr Eo + tr E, - 2tr (I)) / (/1/22,511/1)
", convex on (4). Then Σ, + ir (EX) (E//25, 2 1/2 is an injective linear map since Σ, is positive definite matrix. This means that L is strictly convex.
What's more, the loss L. is strongly-convex on S++ as stated the next lemma.
Lemma 5.5. The Hessian of the loss L satisfies V L(2,) K Ln for any D, S4+ with
/Aiin where Co = =E())( + E
K
203
8

---

Under review as a conference paper at ICLR 2023
Let denote the minimizer of the perturbative loss L(E,) by Σ. Let Ai (0) - Ex be the distance of the initialization from the optimal solution. Equipped with the strict convexity by where
Lemma 5.4, we are ready to show that the gradient flow has convergence rate K' is the constant from the Hessian bound given by Lemma E.5 and Ke,N is constant which depends on the modified margin deficiency and the depth of the linear neural network. Recall that
E, = WN W + TIn, so we prove convergence of gradient flow for the parametrization W.
Theorem 5.6 (Convergence of gradient flow). Assume both balancedness (Definition 2.1 and the modified uniform deficiency margin conditions (Definition verges as
Then the gradient flow GF con-
21AN
((E+()) - < SNc
KI A
(11)
where K Armin
20% is the strong convexity parameter from Lemma 5.5 with Co
2(L(E,()) + tr(E)))
Convergence of gradient descent for the BW loss Assuming that the initial end-to-end matrix W have a uniform deficiency margin, we can establish the following convergence result for gradient descent with finite step sizes, which is valid for the perturbed loss and also for the non-perturbed
original loss. Given an initial value W (0), we consider the gradient descent iteration
w(k 11) = (k) v/VL" (W)),,, k= 0,
(GD)
where 17 > 0 is the learning rate or step size and the gradient is given by ①
Theorem 5.7 (Convergence of gradient descent). Assume that the initial values (.) 1 SiSNN
are balanced and are so that (00) = WN)(0) has uniform deficiency margin c. If the learning rate 7) > 0 satisfies
2²
Nc
1) S min
8M (W)())
2A
4Nc
2
where A
2 GeN N2 (33-3//N 33/N 1/2 (Eo) + 8N(N
1)M 4//N (MA/ + (2///11))) and
M 2 ( / then, for any
0, one can achieve e loss by the gradient
descent GD at iteration
L' L'()()))
k Ne
log
Remark 5.8. Our Theorems
5.6 and 5.7 show that the depth of the network can accelerate the
convergence of the gradient algorithms.
CONCLUSION
In this work, studied the training of generative linear neural networks using the Bures- Wasserstein distance. We characterized the critical points and minimizers of this loss in function space or over the set of matrices of fixed rank k. We introduced smooth approximation of the BW loss obtained by regularizing the covariance matrix and also characterized its critical points in function space. Furthermore, under the assumption of balanced initial weights satisfying uniform deficiency margin condition, we established a convergence guarantee to the global minimizer for the gradient flow of the perturbative loss with exponential rate of convergence. Finally, we also con- sidered the finite-step size gradient descent optimization and established a linear convergence result for both the original and the perturbed loss, provided the step size small enough depending on the uniform margin deficiency condition. We collect our results in Table I These results contribute to the ongoing efforts to better characterize the optimization problems that arise in learning with deep neural networks beyond the commonly considered discriminative settings with the square loss.
In future work, it would be interesting to refine our characterization of critical points of the Bures- Wasserstein loss in parameter space. Moreover, the uniform margin deficiency condition that invoked in order to establish our convergence results constrains the parametrization to be full rank. Relaxing this assumption is interesting endeavor pursue.

---

Under review as a conference paper at ICLR 2023
REFERENCES
P. A. Absil, R. Mahony, and B. Andrews. Convergence of the Iterates of Descent Methods for Analytic Cost Funetions. SIAM Journal on Optimization, 16(2):531-544 2005.
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein Generative Adversarial Networks. In Pro- ceedings of the 34th International Conference on Machine Learning, pp. 214-223. PMLR, 2017.
S. Arora, N. Cohen, and E. Hazan. On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization. Proceedings of the 35th International Conference on Machine Learning, pp. 244-253, 2018.
Arora, N. Cohen, N. Golowich, and W. Hu. A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks. In International Conference on Learning Representations, 2019a.
S. Arora, N. Cohen, W. Hu, and Y. Luo. Implicit Regularization in Deep Matrix Factorization. In Advances in Neural Information Processing Systems, volume 32, 2019b.
B. Bah, H. Rauhut, U. Terstiege, and M. Westdickenberg. Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers. Information and Inference: A Journal of the IMA, 1(():3077333, 2021.
Baldi and K. Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural Networks, 2(1):53-58, 1989.
R. Bhatia, Jain, and Lim. On the Bures- Warres-Wasserttiin distance between positive definite matri- ces. Expositiones Mathematicae, 37(2):165-191, 2019.
N. Boumal. An introduction to optimization on smooth manifolds. To appear with Cambridge net University Press, 2022. htttt::/www.niooleee. htt.laassommmaneee/oooto.
D. Bures. An extension of Kakutani's theorem on infinite product measures to the tensor product of semifinite w alggebaaa... Transactions of the American Mathematical Society, 135:199-212, 1969.
Chewi, Maunu, Rigollet, and A. Stromme. Gradient descent algorithms for Bures- Wasserstein barycenters. In Conference on Learning Theory, pp. 1276-1304, 2020.
Y. Chitour, Z. Liao, and R. Couillet. A Geometric Approach of Gradient Descent Algorithms in Linear Neural Networks. Mathematical Control and Related Fields, 2022.
Y. Chulhee, S. Suvrit, and Ali. Global optimality conditions for deep neural networks. In Interna- tional Conference on Learning Representations, 2018. URL https::/ net/
forum? B= CG
D. C Dowson and B. V Landau. The Fréchet distance between multivariate normal distributions. Journal of Multivariate Analysis, 12(3):450-455 4550004555111 1982.
I. L Dryden, A. Koloydenko, and D. Zhou. Non-Euclidean statistics for covariance matrices with applications to diffusion tensor imaging. The Annals of Applied Statistics, 102 2009.
S. S.S. Du, W. Hu, J. D Lee. Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced. In Advances in Neural Information Processing Systems, volume 31, 2018.
A. Dutta and X. Li. On problem of weighted low-rank approximation of matrices. SIAM Journal on Matrix Analysis and Applications, 38(2):530-553, 2017.
C. Eckart and G. Young. The approximation of one matrix by another of lower rank. Psychometrika, 133):211-218, 1936.
S. Feizi. F. Farnia. T. Ginart, and D. Tse. Understanding GANs in the LQG setting: Formulation, generalization and stability. JEEE Journal on Selected Areas in Information Theory, 1(1):304- 311, 2020.
10

---

Under review as a conference paper at ICLR 2023
K. Fukumizu. Effect of batch learning in multilayer neural networks. In In Proceedings of ICONIP 98, 1998.
N. Gillis and Y. Shitov. Low-rank matrix approximation in the infinity norm. Linear Algebra and its Applications, 581:367-382, 2019.
N. Gillis and S. A. Vavasis. On the Complexity of Robust PCA and 21-Norm Low-Rank Matrix Approximation. Mathematics of Operations Research, 4374):1072-1084, 2018.
M. Hardt and T. Ma. Identity matters in deep learning. In International Conference on Learning Representations, 2017. URL /ttp:e/e/opprrreeee..sttteoooees
L. Kantorovitch. On the translocation of masses. Management Science, 5(1):1-4, 1958. URL http://wws //ww: jstor.......ooo orggstaalee/62636
K. Kawaguchi. Deep learning without poor local minima. In Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
URL httpss./pproceednns..nneet neuripss http:///w...rrrrrooooooooe////// cctppaperr/211///lllll
K. Kohn, T. Merkh, G. Montúfar, and M. Trager. Geometry of linear convolutional networks. SIAM Journal on Applied Algebra and Geometry, 6(3):368-4000) 2022.
A. Kroshnin, V. Spokoiny, and A. Suvorikova. Statistical Inference for Burss-Wasserstein Barycen- ters. The Annals of Applied Probability, 31)3), 2021.
Laurent and Brecht. Deep Linear Networks with Arbitrary Loss: All Local Minima Are Global. In Proceedings of the 35th International Conference Machine Learning, pp. 2902-2907. PMLR, 2018. URL
html
E. Levin, J. Kileel, and N. Boumal. The effect of smooth parametrizations on nonconvex optimiza- tion landscapes, 2022. URL
C. Liu, L. Zhu, and M. Belkin. Loss landscapes and optimization in aver-parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85-116, 2022.
H. Lu and K. Kawaguchi. Depth Creates No Bad Local Minima.
2017.
J. R. Magnus and H. Neudecker. Matrix Differential Calculus with Applications in Statistics and Econometrics. Wiley Series in Probability and Statistics. Wiley, third edition edition, 2019.
A. Mallasto, A. Gerolin, and Q Minh. Entropy-regularized 2- Wasserstein distance between Gaus- sian measures. Information Geometry, 5(1):289-33))) 2022.
V. Masarotto, V. Panaretos, and Y. Zemel. Procrustes metrics on Covariance operators and optimal transportation of Gaussian processes. Sankhya A, 81(1):172-213, 2019.
E. Massart and P.-A. Absil. Quotient Geometry with Simple Geodesics for the Manifold of Fixed- Rank Seative-Seemdeffiine Matrices. SIAM Journal on Matrix Analysis and Applications, 41, 2020. BtRLhttss: //tttps://epuss s am. org/ /t1/10.1111010/// 1tt.111111/12222221
12///M9M12 BM1231389
S. Mei, A. Montanari, and P. Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 1 115(33::E7665 2018.
H. Min, S. Tarmoun, R. Vidal, and E. Mallada. On the Explicit Role of Initialization the Con- vergence and Implicit Bias of Overparametrized Linear Networks. In Proceedings of the 38th International Conference on Machine Learning, pp. 7760-7768, 2021.
I. Mirsky. Symmetric Gauge Functions and Unitary Invariant Norms. The Quarterly Journal of Mathematics, 11(1):50-59, 1960.
B. Muzellec and M. Cuturi. Generalizing Point Embeddings using the Wasserstein Space of Ellipti- cal Distributions. In Advances in Neural Information Processing Systems, 2018.

---

Under review as a conference paper at ICLR 2023
G. M. Nguegnang. H. Rauhut, and U. Terstiege. Convergence of gradient descent for learning linear neural networks, 2021. URL hetps: awwiv wttg.////aaaba....a.b. .20000
O. Pele and M. Werman. Fast and robust earth mover's distances. In 2009 IEEE 12th International Conference on Computer Vision, pp. 460-467, 2009.
G. Ruben and Zamir. Lower rank approximation of matrices by least squares with any choice of weights. Technometrics, 21(4):489 498, 1979. URL http::/wwwtt jstor. org stable/ 1268288
A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In 2nd International Conference Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL
http:/ www.ii.vvv................... org/ abs 1112.1111.00 6120
Z. Song, D. Woodruff, and P. Zhong. Low Rank Approximation with Entrywise LI-Norm Error. In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, STOC 2017, pp. 688-701, New York, NY, USA, 2017. Association for Computing Machinery.
Tarmoun, G. Franca, B. Haeffele, and Vidal. Understanding the Dynamics of Gradient Flow in Overparameterized Linear models. In Proceedings of the 38th International Conference on Machine Learning, pp. 10153-10161, PMLR, 2021. URL http:pps///rrrreeeggg mlr. press/ v ww111/t/a0 htt.ttarmuuuo.. 21 html
M. Trager, K. Kohn, and J. Bruna. Pure and spurious critical points: geometric study of linear networks. In International Conference on Learning Representations, 2020. URL https: openreview. Letp::/eeeooomeeeeoollocc.c rkgolcVVVV
A. Uschmajew and Vandereycken. Geometric Methods on Low-Rank Matrix and Tensor Mani- folds. In Handbook of Variational Methods for Nonlinear Geometric Data, pp. 261-313. Springer International Publishing, 2020.
C. Villani. Topics in Optimal Transportation. Graduate studies in mathematics. Amer- ican Mathematical Society, 2003. URL. https://books: .ooog le. com/books?
idyFAwAAQRAJ
C. Villani. Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften.
Springer Berlin Heidelberg. 2008. URL https://books www.le..oooooooooooooooooooooo.. wwt..co.../.a.cccoe ?id=
hV86557_55t
C. Yun, S. Krishnan, and H. Mobahi. A unifying view on implicit bias in training linear neural networks. In International Conference on Learning Representations, 2021. URL https:// openreview net www.ooororrroaooo.ooo.oo. 41MOKH
12

---

Under review as a conference paper at ICLR 2023
Loss
Parametrization
Critical points
Initialization
Convergence
L WN.1
A PZ7 Ar. 24 -TI) VT
1/2
L,
WN1
Lr E,
Eg
Balanced, MDM Balanced, MDM
GF: Exponential GD: Ollog(1/e))
WN.1 T
ΩAV
Table I: Summary of the results. The target is assumed full rank, with distinct eigenvalues, and spec- tral decomposition Eg QAQT. The end-to-end matrix is WN-1. and the regularized covariance is Σ, V E Rm** is any semi-orthogonal matrix, and Tk C [n]] is an index set. Balanced stands for balanced weights (Definition 2.1) MDM stands for modified deficiency margin (Definition 5.2)
APPENDIX
The appendix is organized follows.
Appendix presents a table summarizing the different geometrical and convergence results.
Appendix discusses the background on the Bures-Wasserstiin Loss and related Optimal Trans- port topics.
Appendix C presents some general properties of a linear neural network and classical results on convergence in the parameter space.
Appendix D gathers the proof's that were omitted in Section
Finally, Appendix presents the proofs from Section
A
SUMMARY OF THE RESULTS
Table presents summary of the different results obtained in this paper. Note that, even the different losses are expressed in the function or the covariance spaces, the gradient flow and gradient descent algorithm are performed on the parameter space O.
B BACKGROUND ON THE BURES- WASSERSTEIN DISTANCE
B. DEFINITION OF W
The Bures- waeess-sssrsttein distance has natural connection with the 2- Wassersstii distance on metric space. the case of zero-centered Gaussian measures, those distances are identical. We present here more general definition of the 2- Wasseerstei distance, which enjoys desirable properties.
Given a metric space (X, II- the 2-Wasserstein distance is a well-established metric on the space of quadratically integrable probability measures (2))..
Definition B.1 (2- Wassersstein distance). Given two quadratically integrable measures (v, 2n) e (P2(X))) the 2- Wasserstein distance is defined as the following minimization problem
inf fIle dm
(2, rETT(E,)))
(12)
where 11(1, mo) is the set of fixed marginals distributions: 11(1, vo) = {x E x
V, Ø2 vo), with m, the marginal of along the ith variable.
It is known that this distance metrizes the weak convergence on the space P2, see e.g. (Villani 2008
Theorem 6.9), and can therefore be leveraged when designing system that relies on comparing probability distributions such as GAN. On the other hand, the computational burden of such a loss can quickly become prohibitive (Pele & Werman, 2009). In a very few cases, efficient computations be done for the loss 12 This constrasts with usual WGAN, which was first introduced by (Arjovsyy et al. (2017). where the loss is computed using a neural network, based on the dual expression of the (1- (a-sessrnnnnein distance.
13

---

Under review as a conference paper at ICLR 2023
Indeed, between two Gaussian measures, the 2- Wasserstein distance has closed-fome expression or a closed-form expression for the discriminator so that adversarial training not needed. We will consider two centered Gaussian distributions, which described by their covariance matrices. In the case of centered Gaussian distributions, the 2-Wasserstein distance reduces to the Bures- Wasserstein distance between the covariance matrices Eo and Σ (Dowson & Landau 1982
Lemma B.2. If (m, Σ) and vn (mo, Eo), then
(V v)) = m - m 12 B (E,
It is well known (see Kantorovitch (1958) or Villani 2003 Theorem 1.3) or (Villan 2008 Theorem 10)) that the squared Wasserstein distance has the following dual expression, also known as the Kantorovich duality:
I f(x)
W vo)=
[(x) duo(s)) + g(x)) ((x)))))))
sup (vo))
(13)
((a, f(x) + g(v) < - -11?
where L is the set of the integrable functions with respect measure v. Therefore, the dual variables f' and are required to be integrable with respect to the source and target measures, and to fulfil the cost inequality.
Remark B.3. In the context of GANs it is common consider the I-Wasserstein distance with cost given by the distance | all. which has dual expression, referred to as the Kantorovich- Rubinstein formula (Villani 2008 $6.2) that allows for more tractable computation in practice with for instance only one dual variable. Nonetheless, there is no closed-form solution known when c(z, y)
| yll-
B.2 FURTHER PROPERTIES OF THE BW LOSS
In this section, we provide further background on the Bures- Wasserstein distance. First, we show that, except in some particular cases (Lemma B.44. the Bures- Wassss-esserttein distance between two covariance matrices is not translation invariant (Lemma (.55.
Lemma B.4. the case that Eg and E commute, the Bures Wasserstein distance reduces to the Hellinger distance:
(E,,Eo) 1-/12-1100000 -1//2/10//1116 113.
EAΣ Σnn
Proof. It simply follows from the fact that, if Σ and Eg commute, so do 1//2 and 2/ so that x ///">//// (5/// and
tr tr 32 -2(2://22111)) tr -103///2///3300//// - /T)
= =|18//2 -
as claimed.
From this, one remarks that the problem of minimizing the BW distance between two covariance matrices that commute does fall under the framework of the Eckart Young-Mirsky theorem if the
optimization variable is 1//2 = (WW Nonetheless, in the case where Σ and En do not commute, do have such correspondence, as in general, the BW distance is not translation invariant, neither when considered as function of (E, Σo) when considered
as it reduces to a problem cast with the square loss.
function of (E1/2 =/12)
Lemma B.5 (BW is not translation invariant). There exist positive semidefinite matrices (E, En) (4) SS (n) and translation T E S. (n)) such that (E+T, Eo+T) + (E, Eo). The same statement also holds for the loss L defined on the matrix square roots, L(E1/2 5-0 1 (F, (D))))
14

---

Under review as a conference paper at ICLR 2023
Proof. For the first part of the statement, taking
Σ ( .). Eo =
=6=(62)) 2) T= (
t 0)
0 ;), 1> 0,
then + + T) 1² 20) = ( + vI+ (v2 1)", which is non-zero.
For the second part of the statement, if
2///2 1 2)
y = 21 T (0 9).
one computes
D/ //// / =I =| + ||
1/ - 2 tr (2) 2 / 1/2
2
6
1/2
=12 2tr
6
20
and
C(E)/2 + T, T.
1//2
+ +T) = =1-1// + +T|1 + |=///
+ TI
-2 (( +T) + +T)(""²+7)(E) + +T)) 1/2
1/2
20
30
=28 2tr
30
90
which gives the difference C( + T, ///² + +T) L( /2 ~ 0.121229 + O.
Lemma B.5 therefore implies that, in the general case, one cannot express the Bures- Wasserstein their square roots)
distance (either the covariance loss would be translation invariant). This hinders a direct application of the Sckart-Young-MMe theorem, where the problem is cast minx ||A X|| ll with fixed A for some unitary invariant norm. Even if this is not possible in general, a similar expression exists, relying on the following variational formulation.
of a difference (otherwise, the
Lemma B.6. The Bures-Wassere Wasserstein distance between two covariance matrices En and & on S coincides with the variational formulation Z
min (EU(A)
1/1/2- ---//' 1/- ir ED + 10--222 2
(14)
Proof. We write
11/2 -UE///117 (= u2 (x!
12
min (EU(n)
Let /1/32/1 0
= V/ be the singular value decomposition for V1, 1/2 unitary, and R = [2//2n//2)
E4/250
/ Therefore U /2/25//²) ir (ViUT R2/22 is max-
imized when V.U V/2 is the identity. Thus, we get 14 and the proof complete.
Lemma B.7, For any symmetric matrix C' E S. for any matrices (A, B) E (RRxX))² 2 one has
tr (CACBT) = tr (CBA)).
(15)
Proof.
tr (CABT) = tr (CT AA =r (BA C) tr (CBAT).
15

---

Under review as a conference paper at ICLR 2023
B.3 GRADIENT OF THE BURES- WASSERSTEIN LOSS
We give here the gradient of the squared-Bures- -al---Buuuus-aasssttteiiii distance between two full-rank covari- matrices.
Lemma B.8 (Gradient of B for full-rank matrices). Suppose Σ, Eu E S++(n). Then the gradient of B is given by
/ / Vx, = 2:1/2 (2:0 2 - 1/222/
(16)
The proof of this Lemma is given in Appendix B.4 The right hand side of (66 is the optimal
transport plan between two centered Gaussian distributions (Bhatia et al. 2019, Muzellec & Cuturi
2018 eq. 7), whose Fréchet differentiability has been explored by Kroshnin al. (2021 Lemma
A.2). This is formulation that we use in the computation of upper bounds for the Hessian in
Appendix E.
B.4 PROOFS OF SECTION 3
In this section, we provide the proof' of Section
Proof of Lemma B.8 Recall the BW distance is given by B (E, = trE + Ir Er
2 tr 11/25511 //
The gradient of the BW is given by
Vr (E, Eo) = - 2 tr (2) /n/ 1/2
(17)
Since E, En are positive definite, we differentiate the f(E, En) = tr / /3)
1/2 with respect
to Σ
Vr, = (22 1/2)
20 x 1/2
///2 1/2 -1
O5. (xx (xي )
81+18
/22 =/// |2!
1/2
( 1 1/2 I
@I+10
1 1211/2 & 21/2 /لا0
-1/2
1/2
( ΣΣ! /22
- 1/2
(18)
Substituting the above expression to
we get
-1/2 2
V E En) = 1 x2/2
(19)
Proof of Lemma 3.3 given by
First note that the difference between the original loss and the perturbative loss
(L(E,) - ((E,)-L((W)) rn 2tr
( ((x)
5/2/2 (1)//2,,>01)))
1/2
12/ WS 22/2) 1/2)
1/2)
(20)
1/2
-(x: WW T 1/2)
Tn + 2tr
Let the singular value decomposition of / ////////2 /2 QA where
+T
A, =
X +
16

---

Under review as a conference paper at ICLR 2023
\ > X2 ...> X, and r rank(WWT]] Similarly, we get the singular value decomposition of (20 WW T /2//////22 UA1/2VT where
A1
We know that tr -0 = tr (At
- (A + T)1/2 since the Frobenius norm
unitary invariant. Likewise get tr / WW D observe that the eigenvalues are distinct and in descending order. This means that we can upper bound the eigenvalues as,
tr(A//2) Ei= X!? Next
Σ + <
E1!'
+ 1
2=1
Therefore, we get back to Lemma
3.3 and get that
(iCE,) - L'(W) rn + 1/2
C GENERAL RESULTS FOR LINEAR NETWORKS
This section deals with general properties of linear networks and their convergence in parameter space. We first recall well-known results that hold for any differentiable loss L' and its parametriza- tion LN L μ..
Lemma C.1 (Gradient flow, Bah et al. 2021 Lemma 2.1). For any differentiable loss L and parametrization LN L' μ., such that n I....... WN) WN- W1, one has
1. For all j [N].
Vw, L (WI, Wn) WJ.
WWW. WJ-1-
(21)
Assume each of the Wt(t) satisfies the flow
GF.
Then, the product WNA1 = WN
W1 satisfies
dW())
N
WN ... WT 5+1 WW)T ... WT W Wi. dt Σ
(22)
For all j [N], and all > O, we have that
d
d (W,(tW;T((t)
(23)
dt
dt
4. (O), (0) are balanced, then for all t 0, (1) = W,WOW (t), and
aw(I) R(t) = dt
N
+ Σ (W(t)WT(()) (WWW)((T(WWW) = 0.
(24)
The BW loss satisfies the Łojasiewicz inequality. Indeed, the following equality can be computed. Lemma C.2. For any W E Rdd xd., for the loss L defined in (4) we have
IV = (W).
(25)
Proof. This equality can obtained directly by computation. Since
(W) 2W -22/ //W /
(26)
17

---

Under review as a conference paper at ICLR 2023
have
Ww)
=4r (WW WW T - /2 W) (WT - / (E/ WW
- /
= =4ir(WWT) - Air (WW (20 WW -//2 -
4 (26/
/ + troo)
(27)
Note that the middle two terms above are the same, and they can be further simplified as
tr (WW'
- 1.1/ 22 WW / -1/32////)) (2 (20
1/2
WW
20
/
(E
1///2 WWT 2 =
(28)
Combining all the terms together, we get the equality 25)
In the case general, twice differentiable loss L and the parametrization L LL p. one can express the second-orde differential structures of the loss.
Lemma C.3 (Second-orde differential). Let (U,V) E e be two parameters, U
(U, UN). V (Vi,..., VN). The second-order differential of the loss LN at W
(W WN) i
L
(U, WT ...... VJ WWWW ... W.L
i=1 1ボ
N N
(Wi-a vU.(U T
& (WNo+1) T. (W) (W,-1n) @ (Wwg+1) vec(Vy),
(29)
where (W) e R²-xn" is the matrix such that, Y(U.V) (R"*n)) (W)(()) voc(U) L (W) vec(V).
Corollary C.4 (Hessian of the Loss). The Hessian LN, (2) can be represented as a dh x db
matrix. It is a block matrix, the blocks corresponding different layers. Each block V W. w, I," (W) has dimension d.d;-1 d,dj-1, and corresponds to the differential d L (W [T i U where T.=(0. 10....,0,U,. U...O.....0) 0). The block diagonals elements are
Vir (W)) (W,1a & (WN++1))) (oa+))))-''L((WW))) (WL1aa) & (WNa+1), the off-diagonal terms are
(30)
V (W)) = (Wi-1a @ (WNiH1) V2L (WW-11) @ WN
+ +[ W.VT. T Wy. WH+1) & (W ... W. Kdjdz-1
(31)
where Kyg the pg-commutation matrix (for X RPxq Kpg vec X vec X
The invariance property the gradient flow (FF) (Lemma C is key in numerous analyses.
Another useful property of the gradient flow (GF) its convergence, under mild assumption on the
loss L' critical point of N Namely, if the trajectory + W((t) remains bounded for all t > 0, and is an analytic function (i.e. locally given by a power series), then GF converges to critical point of I,N, i.e., point 0* so that V N (0*) 0. This is stated in the next theorem.
18

---

Under review as a conference paper at ICLR 2023
Theorem C.5 (Gradient flow converges to critical point of ) Let L' be analytic, such that the trajectory t + p(((t)) remains bounded under the gradient flow evolution 0 = - [ o μ (0). Then, the flows of Wt(t) given by (GF) and of (t) given by (2) are defined and bounded for all t 3
and
WN) converges to a critical point of L
o pu as t + oo.
This result relies on the Lojasiewicz' argument for the convergence of gradient flows Absil et al. 2005 Bah et al. (2021) show how to bound each of the different {W;) once the end-to-end product W N:1 is bounded. The boundedness of |[W] can be showed depending on the loss that is considered. For example, it holds for the regularized loss L,. In Appendix C we collect further general results for linear networks.
In the case of the perturbative loss introduced training. Since the loss I.₃ is analytic, one immediately gets the following result.
on can bound the norm of W throughout
We give a simple test to show the boundedness a trajectory under (GF) This is allowed by the decrease of the loss along training.
Lemma C.6. Assume that, for a given loss I,", there exists there exists an increasing function such that, for any t > 0, [t)lI & ((I' bounded.
Then, the trajectory + (()) under (GF) is
Proof. Under gradient flow, for any t 0, L' (W)(t) < L' (W(()). Indeed, writing the chain rule and the gradient flow 22).
Lw Σ Dw, L (W). (
3
EIV (W),..,,wx)IE < O.
Therefore, for any t > 0, (W(())) & (W(O))).
Now, let f: R if for any t 0, [[((t) & f(L ((L((W(t))) then [[(()]l & ((L(WW(t)) & f(L (W()) is bounded.
+ R be an increasing function, so that f(L (W(t)) (W))) Therefore,
The assumption of Lemma C.6 is satisfied for couple of losses, including the square loss (Bah et al. 2 and the I. loss, as shown in Lemma C.8 It allows to consider losses that "grow with the weights", so that the end-to-ent matrix is bounded when the loss converges to zero
We now show the boundedness of the weights when considering the Bures Wasserstein loss
Lemma C.7 (Boundedness for the BW loss Ż). The loss L(E) can be lower-bounded by the quan- tity? 1/2 ir E ir Eo.
Proof. By definition of dual expression of the Wasserstein distance 13 i(E) = W up) =
sup) f f(x) dug + f flH-Ii (u) deo(y), with va (0, M), vo No0,, Eo) and fll 12 the 11-11- transorrasfoomm of f defined as Vy E IRd, fll-l* (v) inf || y 2 f(x).
With f: + 111 the trla-frrasfomm of f is ful-1P; y + -1lv|12 and we get
1 i(Σ) (M vo) 2 ||a|| dx) / |1ul duo(y) 2 E -r-tre En. (vo,
(32)
as claimed.
Lemma C.8 (Boundedness for the loss L.). The norm the end-to-end matrix W is upper-bounded when using the loss L, defined in (5).
Proof. With pf (Σ)= = E. + TIp =: E,,, the loss I.f satisfies
(W) [((p,(T(W))
1 tr tr Eo
-r tr - tr Eg
n
(33)
(W) 2tr En nt > |
(34)
19

---

Under review as a conference paper at ICLR 2023
Therefore, there exists increasing function f such that |[]]| < (W)) Since the loss de- creases under gradient flow, one has
|W(OI) (2T)(W(0)) + 2tr En nT;
(35)
and the boundedness of t + (t) is shown.
Corollary C.9. For the Bures- Wasserstein loss, if WWW T is positive definite, and loss is differen- tiable, and the norm of the weight throughout the training uniformly bounded
|| S (W - Lo.
(36)
by using similar arguments as in Lemma C.8
Lemma C.10. The gradient flow (GF) on the perturbative loss of L in the limit.
converges to a critical point 0
This property of the gradient flow is necessary in order prove the convergence of the training to a minimizer of L.P. At first glance, the critical points of L do not correspond in general to critical points of L since the parametrization μ also comes into play. This led Trager et al. (2020)
to distinguish between the pure and spurious critical points; i.e., the points that are shared between I,N and I,=, and those that are exclusive to T.N
D PROOFS OF SECTION
In this section, we provide the proofs for the critical points of L' (,M, and LM
D.1 CRITICAL POINTS OF L M
First, the loss L' is expressed on the manifolds Mg (Lemma D.2g where it is differentiable (Lemma D.3). Then, necessary conditions (Lemma D.5) the critical points can be expressed, leading to the proof of Theorem 4.2 The second part of Theorem is then proven by evaluating the loss at the critical points found, and ranking them.
Recall Definition 4.1 Computing the differential of the restriction I [M, will allow to characterize the different critical points.
Definition D.1 (Gradient), Given an embedded manifold M and a function with a smooth restriction flm. the gradient of flM
M is the (unique) element of the tangent space T&M such that, for all v E T,M, ( = (z)),
We begin by expressing the loss L' l.M with the Singular Value Decomposition (SVD) of E//2 W.
Lemma D.2. Let USVT = / be a thin SVD of Iy / so that U E Raxk V E Rmxk UTU VTV Ik, S diag(si,..., sk) ) E BRxk where rank 22/2 11/2W- rank W. The loss I,2 from (4) on Mr can be expressed as
I' = | 1/ 2t
(37)
Proof. If USVT = 2 W is a thin SVD of / then (2) W (2il/2)) T 1/2 Therefore, the expression of the loss L given by (4) can be written as
USUT.
L' (W) tr(WW Eo - 2t (USUT) |W]| + ||>/ 113 -2r S
as claimed.
We then give the gradient L [Ma-
Lemma D.3 (Gradient of L. (m,) Let (n, m) N2, and let l min m). The loss I Mm
[M, (as
given (37) is twice continuously differentiable on Mr. With W e My and USV = E 1/2 W a thin SVD of 2/22 1/2 W. its gradient is
V Me (W) 2W 2VV
(38)
20

---

Under review as a conference paper at ICLR 2023
In order to derive this expression, the differential of the singular values is required. But first, note on the differential notation used throughout the derivations.
Notation (Differential). The differential of a function f can be written using different formalisms. Explicitly, [[( is the differential of f at X in the direction H. Sometimes, with Y = (X)), the shorthand notation dY preferred, where the same symbol used for both the variable and the function. In this case, it is assumed that the direction H is small perturbation dXX around X. For instance, ( HX XH with the full notation.
= f(X) = XX then dY = dxx + XxxX which would be written
Lemma D,4 (Differential of the SVD). Let k & min (n, m) and let X E Ma be matrix with rank X k. Let USV X be thin SVD of X, with U Baak S Ta x V E Rmmn diagonal and U VTV lg. Then, the differential d8 is
o 1 (U dXV),
where A @ B denotes the Hadamard product between A and B.
Proof. Let USV
the decomposition as given in the lemma statement. The differential
rules ensure that
dx dUSV SUVT +USAVT.
This implies that
T dXV = dUsV + VU + +UTUSaTTy
= UT LUS + +
ds = UT dxv -U -US-SdvVy
Since UTU T, dU + U 0, and A == UT dU B = VT dV is also antisymmetric. The matrices A and B being antisymmetric, their diagonals are null; hence so are the diagonals of AS and SB, i.e. T ④ (AS) I @ (SB) = Since S constrained to be diagonal, dS must also be diagonal, i.c. I4 C dS = dS. Therefore,
TU -AT Likewise,
ds I TU dXV)
as was claimed.
Proof of lemma D.3 For W E Mr. let USVT = / be thin SVD of /// =: X.
Lemma D.2 ensures that
L' (' |[] 1 2r
(39)
According to Lemma D.4, the matrix S differentiable and has differential dS = I O (UT dXV). Therefore, the loss L My is differentiable. With the fact that d S tr dS (see e.g. (Magnus & Neudecker 2019 Chap. 8, Eq. 18)), we can compute
d dS = tr (4 o TUT aXv)) = tr (UT dxv)
(UV .xx (UVT, /21/ dw) (2///2V .dW).
Moreover, d 2(W,dW), and so
dL I... (W) a W|| 2dt S = 2(W 2V ,dW), ,WW),
and
LL M. (W) = 2 (W - - -V/
Since matrices (U, V) are continuously differentiable on M₁ VL I,, (W) = 2(W /VV) again continuously differentiable, and [,M,,, is twice continuously differentiable. 17
21

---

Under review as a conference paper at ICLR 2023
We ready to give the proof of Theorem 4.2 We divide the proof into necessary and sufficient conditions for a point to be a critical point of M'
Lemma D.5 (Necessary condition on the critical points of LwlMe) Assume Ey has n distinct eigenvalues. Let W E M2 be critical point of M. Then, with U U"S" Σg = thin SVD of 20 W*, and QAQ Eo an spectral decomposition of Eg (i.e. with Q (())), there exists Jk 5 (1.....m), such that 5* As, and U 52
Proof. Since W E MM, and USS.V 2 3 a thin SVD of E / this means that S E kk Then,
VL' =0 =
W = U-V
by 38)
/² EOUINIT
UUS'V= = EVU
S EqU",
U Ig, v' Ik-
Therefore, U+ EnU* must be diagonal; and since U is semi-orthogonal, this is the case if and only the vectors in U* are eigenvectors for Co. by uniqueness the spectral decomposition of Co Therefore, there exist ji.... .& indices between I and n such that U (W/ Wja) = P in which case
S -1 Eolla
AJe.
Aja
Now we are ready to prove the first part of Theorem 4.2
Proof of Theorem 4.2 first part. Consider the expression for the gradient of [M, given in The necessary condition follows from Lemma D.S, since
38
VL Im (W*) 0 = 2/ W' /2
=lgAA V
W = 1/2 As. VT
-2A -1/2- 2J AJ. VT
S2A 122-1/2A, V
-SAPP3 V
= -SPPA .VP.,AA//2 1/
= 2 VTT Jh
which corresponds to the necessary condition in Theorem
The sufficient condition can be verified follows. With W
lgg A J. 1/2VT one has /
QA1/22T Sz A X1222 Pz Ar, VT, and, as this is correct thin SVD of 1//2 "'" Lemma D.3 gives
VL VL' P'' = 2(W* 210 21/20
Further,
x/ 2//2 = 1211/20 T25
-2A1/223
-SPP 1/2 A//2
-2 1/2
22

---

Under review as a conference paper at ICLR 2023
Hence
VL (m, = 2(W = 2 VT) = A - -P A 1 =
and the sufficient condition is verified.
Now, the loss can be evaluated at the critical points in order to statute on its minimizers.
Corollary D.6 (Value of L the critical points). The value of the loss L at a critical point W = -2y A A1/2TT is L (W*) = tr A - tr Ag, = =Eigh Ai-
Proof. For l > 0, let W be critical point of L IM,- From Theorem 4.2 with Eg QQAR
spectral decomposition of Σg, there exists set Th and a semi-orthogonal matrix V IR such that W 2g J VT One can then compute the value of the loss at W":
L' L'()) t + Eo 2t (2-0 1 (2 w22 T) 1/2
1/2
= =trΩa AJ, 2J + tt - 2tr P2 Jr R3,))
= =ir + - 22tr
-rr -ia--rAs Ar
We now have all need to prove the second part of Theorem 4.2
Proof of Theorem 4.2 second part. The first part of the statement is readily implied by Corol- lary [.6] as the eigenvalues are in decreasing order. The second part is implied by the fact that the minimum I, IM, is indeed achieved for any k & n (by selecting the k largest eigenvalues of Eg) and the optimal value of the loss L is smaller when considering more eigenvalues, i.e. minM, L < minme Lt,
Moreover, it can be shown that only one point per set M₁ is minimizer of the loss L [M all other points are (strict) saddle points. We recall the definition of a strict saddle point: point where there exist descent direction.
Definition D.7 (Strict saddle point). A critical point x of a function f is said to be a strict saddle has strict negative eigenvalue. If all critical points of f are either
point if the Hessian of strict saddle point or the global minimizer, the we say that f satisfies the strict saddle point property.
If the gradient flow can expressed on a manifold, with a Riemannian gradient corresponding to given metric, there is equivalent definition of those saddle points, which will be handy to use. See (Bah et al.) [021]] 16.1) for the details.
Proposition D.8. The loss L satisfies the strict saddle point property.
Proof. Let E. T be the spectral decomposition of Eg with decreasing eigenvalues. For k E N, according to Theorem 4.2] W is critical point of L L' if and only if there exists
Th C (1....,, n). such that W = Ug A Jz 1/2VT with any V e Bm that VTV TV Ig. Jr
(1....,k), W is a global minimum of L MM,, as shown in Corollary D.6 and the proposition holds.
Assume Ts * (l....,k) k}. then there exists jo E Ts such that djo < Ax. and there exists j1 $ Jk but ji E {1, (1.2..... .., k) such that Xj > Ajo- We will show that W* a strict saddle point L M
The critical point W* can equivalently be expressed as
W Eg 1/2
Maaayy
(40)
TEJi
where ue ii are corresponding orthogonal uni-vectors in U and V, and A; are eigenvalues in A.
23

---

Under review as a conference paper at ICLR 2023
The key is, by the following perturbation, for t (-1,1), define UJo u1o(1) = fu / uuuj
and the curve (-1,1) - MA. We look at the perturbative matrix
(() Σ 1/ 2500 "30 (1)
Σ )aw;)
iEJ
Note that 7(0 = W. Recall T. = ir WWT + +50 - 2 (=(²)/ WW 1/2
It is enough
to show that (Bah ct al. 2021 $6.1)):
d²
(7(1))
<0.
dr2
1=0
1=0
We check it term by term,
tr (1(0)((TT) tr E 1/2 (A%
+
Σ
Aeuus (1i0 Uio (t)o
E
iEJ
tr 1 (A 2/ (t)usa
Σ
iEfYio
tr ((
Σ xo uu (A U3o (t)
Σ
Xu;u,))
+
15t5n
iEJMoo
X2
y2, 12 \1o (1 - 122)
Σ x
A51
iEJi(Jo)
and
(12) ((t))(t) ((t)()((t))(() T
1/2)
= tt ((1))) (1)0 Σ
\ratoy (2 (1)3)
Σ
)mne, 1)112)
MEJ
tr (A)), (f)us,
Σ Xurn T )112)
iETViio)
tr (1P2 Ujs uJ +1 - t² ) 23 Us 3 +
Σ
xuruT) )12)
= =[12jo + v1 1 -22 + Σ Pdl.
Thus, since Xj.- > Xjo:
d2
(7(1))
= =2(22 -110) - -Ajool < O.
dt
1==
This completes the proof.
The loss LP satisfies the strict-saddle point property in a similar fashion. Lemma D.9. The loss L.M, satisfies the strict saddle point property.
Proof of Lemma D.9 The proof of Proposition D.8 can be adapted, with the expression of the
critical points as, if En = QANT and with V E Bank any semi-orthogonal matrix, W (E. -l,) 1/2 Σ (X: 114-779 TTw.oT
24

---

Under review as a conference paper at ICLR 2023
D.2 CRITICAL POINTS OF THE PERTURBATIVE LOSS
In this section, we provide the different derivations for Section 4.2 orem 4.2 first the gradient of L, is computed, then the critical points characterized and ranked.
The structure is similar to The-
Lemma D.10 (Differential of L). The differential of L on (++(n) is
AE E (+)( X E S(n), a = (X =/ / Σ>
/ X).
Corollary D.11 (Gradient of L). The gradient of L on S++ (m) is
E vi(E)= = - _ ( 2 1/2 1/2 1//2
Σ
Lemma D.12 (Gradient L2). The loss L has the following gradient
A W E a
LL) = 2(W - [x) (WWT + -T,))22 120 12
11/2 x1 W).
(41)
Proof. This results comes from the chain rule for the loss (W) Lo Pr TW).. With E
T x(W)= = WWT and Σ, = (x) = E + T Tla and since (W) = W + ZWT and dp-(E.) id, one has
(L) d(I o fo o x)
d dip. (2) (T( [dp[(22)(((([)]]] di(E,)(WZ + ZW (V))(W), Z) = (Vi(E,), WT + T
VL.(W) (Vi(E,) vi(E,) >
= =2(W - 1=2
1/3 Σx 1/2 11/2 11/2/////7 W).
Proof of Theorem 4.4 The eigenvectors of WW + are the same as WW and the eigenvalues shifted by Therefore, the expression of the critical points in the original loss be adapted, so that the modified critical points have the same left singular vectors and shifted singular values. This
(Aa. - 1/2
)
leads to having W Sh.. (AJ, - -T(x) 1/2
(v 1)
O₂ On-kxn-k
with V m.n such that VTV Tg. One checks that (W*) =
The value at such critical point W 2Z (Ag. -1+)) V is (W Eier, 2
which is uniquely minimized for [k] when the eigenvalues of Eo are distinct and in descending order.
D.3 PROOF OF PROPOSITION 4.5
We state here the proof on Proposition 4.5 We will transfer the results obtained on the space of linear maps MoR to the space of covaraance matrices S+ (k, n). Borrowing the terminology from|Levin Levin et al. (2022), we introduce the following notations and definitions. Let M be any smooth manifold, linear space, yo: M -> a smooth (overiparametrizztion (or lift) of the search space X (MM) C E. The following problems are considered
min f(x) x
(P)
min f fo pp()), 3EM
(Q)
where we assume that f: E - smooth, and hence so is g == fo p. The following property is relevant for us.
25

---

Under review as a conference paper at ICLR 2023
Definition D.13 (Levin al. 2022 Definition 2.7). The lift pi M - X satisfies the "1 = 1"
property at y if for all differentiable f: X + R, y is a critical point for QQ. then x = ((y) is critical point for PP.
Recall that
(+ : Σ> 0, rank(E) = =
(42)
and let
S_ k, (2 ((n): E 0, rank(2)) & kj.
(43)
We will make use of the following result from Levin et al. (2022).
Proposition D.14 (Levin et al. 2022 Proposition 3.4). Let & n, and let pi Rank + + k. n)
be the parametrization ((R) = RR Then, ip satisfies the "1 = 1" property at R Rx* if and only rank R k.
Said differently, the condition at which a critical point on the space Raxk is such that its image through yp is also critical point on S+ (ki n)) are exactly the points of full rank R The image of the parametrization is g
= Therefore, we need to adapt Proposition D.14 in order to work on
Ba*k
which not smooth - instead of Ba This is performed in the next proposition. Ek
S+ n) be the parametrization (WW)= = WWT. Then, n satisfies the "1 1" property at W E TR ifrank W = k.
Proposition D.15. Let k min(n, m)), and let w Raxm
Proof. Let yo: Baxk S+(( n), h RR and T: RRxm + S+(S k,,, W. F WWT be the covariance parametrizations. Since the manifold R is not smooth, we will focus on the smooth manifold IR IR,x". Therefore, p.: Raxk - (k, n) and u. Raxm + Sk, n) be the parametrization sp, defined on matrices of rank exactly A. We know that y. satisfies the "I + 1" property, and want to show that m, also satisfies it. The idea of the proof is the pass through given quotient spaces which the functions are equivalent.
Let O₂ be the set of x orthogonal matrices, with the dimension omitted when inferred from context. Consider the equivalent relation on I (or RR**) such that X1 X2 = X
X2XT From (Massart & Absil 2020 Proposition 2.1), know that X1 X2 = 3Q
O, X1 X2Q. Denote the equivalent class [X] = {XQ: Q E O} = XO.
Let p: m R R2xm RRxm//On.
W + [W] be the quotient map on Rax= and let
IR R
10m
S (k,n)
(44)
WOm
WWT
be the map on the quotient space, that n. II o Likewise, let q: Raxk - Runk OOk, R + [R] be the quotient map on Raxk and let
Rank O -
(k,
(45)
ROk + RR
be the map on the quotient space, so that y. 4 o
The map T is an diffeomorphism MMssart & Absil 2020
Proposition A.7), and therefore satisfies
the "1 1" property.
For W E R=X", we can find R e Ri*k such that WWT T = RRR.. it unique up to an orthogonal matrix. Therefore, let
TaXm
10 m
R /Ok
(46)
[
R
be the identification map between the quotient spaces.
With the next two lemma, we will be able to finish the proof of Proposition D.15
Lemma D.16. The map Raxm /0 Raxk Deaxk//O. /Ok, [W] + [R] is diffeomorphism.
26

---

Under review as a conference paper at ICLR 2023
Lemma D.17. The map is submersion from x onto R /O,.
So now, to conclude the proof of Proposition D.15 the map n, can be written ग,, = Top = $ Dorop. Since know that Φ being diffeomorphism, satisfies the "1 = 1" property, and since q submersion (Lemma D.17 by (Levin et al. 2022 Proposition 2.42 (b)), the map 4 pocop satisfies the "1 1 property.
remains show Lemmas D.16 and D.17
Proof of Lemma D. 16 From Massart & Absil 2020 Proposition A.7), the mapping $: R
(k, n), + RR is a diffeomorphism. Likewise, the mapping II: Rg Raxm - S (k n) is also a diffeomorphism, since both т, and are submersions. Then, since +((R) = [W], we have [R] - -- [[([W) =: ([W), and is a diffeomorphism.
Proof of Lemma D.17 The map p is a submersion (Massart & Absil 2020 Proposition A.5) and the map is a diffeomorphism (Lemma 6.16,, hence submersion. Therefore, the composition top a submersion.
We are now ready to proof Proposition 4.5
Proposition 5), we know that a critical point in the parameter space W with rank rank(W) W) k will be critical point for L lM Now, from Proposition D.15 critical point W for LL with rank W l is such that I (WW") is a critical point for s___(k,)) (k,m), and the first part of the proposition is proved. For the second part, assume that &o = d = min;(dd). Then, according to (Trager et al. 2020 Proposition 6). W is local minimizer for LN if and only if W p(W) is local minimizer for L. Since W is local minimizer for MM according to Theorem 4.4 there exists V On orthogonal, such that if Eg = ΩAΩ is spectral decomposition of Eo. we have W S4 (A - rIa) 1/V
Proof of Proposition 4.5 From (Trger et al. 2020
so that
[d]
( -)
Σ; W- + +l =
S2 also a minimizer of Lrls (dm)
E PROOFS OF SECTION 5
E. BOUNDS ON THE HESSIAN
In this section, provide bounds the Hessian of the perturbative loss L, We first express the loss as a function of the covariance matrix, in which case the Hessian is known (Kroshnin et al.] 202 Then, simple chain rule for the differential allows express the Hessian in the case the loss is a function of the end-to-end matrix W.
Lemma E. (Second-order differential of L.,, Kroshnin et al. 2021 Lemma A.2). Let W E Raxm and let > O. Define Σ, WWT + TIn, to be the regularized covariance matrix. Given that E, > O, the loss 10 is twice continuously differentiable for any W. Let IQr = = /1/2 E be spectral decomposition of x1/2 ,, with Q diag (q1, -.......9.) For Y E Raxn, define (TT
Let it the linear operator A(Y) E [Raxn to be the matrix with element A(Y) = vort/a) defined as
i(Y) = Then, the second order differential of L, is given by
(47)
v(X.Y) E ( d
(48)
27

---

Under review as a conference paper at ICLR 2023
Proof. We begin by stating the first-order differential for the loss L evaluated the PD matrix Σ,- This is given in lemma D.10
di tr(X- 2 E, 20/
11/2 11/2 210 12 X).
21 2) x)
(1 -1-0 (=:))
1/2 2.1.1////
Let GL(n) = (A E R= det A + #0}, and let f: GL(n) 3 F F-1; then f is differen- tiable with differential [f()) = -F--XE-1 MMagus & Neudecker 2019 Theorem 8.3). Let g: ST+ 3 A + A1/2 be the matrix square root. The function g is differentiable on Sh and its differential can be computed as follows (Kroshnin et al.) 2021 Lemma A. Let A E S and let TQrT be its spectral decomposition, with diag (0:)-1 For X S", define A(X) Roxx
(TXT 1g, Then, the differential of g at A in the direction VØ+vV
to be the matrix with elements A(X); X is dg(A)[X] = IA(())".
Therefore, the chain rule on the differentials gives
d(f (A){)) = off dd((((4))dg = 1/2 dg())))A 1/2 -A -AA/TA A-1/2
and, with A= = - = ///
d' Y] d(E, + ((
= =d(tr(X- - - (= (=//2///////000 -,, / 1/2
11/2 X
--tr(2)/ 1/2 N, ", 1/2 [Y]) X)
-
= (1/2///0001/21110000
(X),
with = (000==//"00---000000-00-000
In order to express the Hessian of the loss as a function of the end-to-end matrix W, we need the chain rule for the second-order differential. We first recall the chain nule for the second-order differential.
Theorem 6.9). Lemma E.2 (Chain rule for second-order differential, Magnus Neudecker 2019
Let R + S and g: S + T be two differentiable functions on open sets, such that h = gof: R + is always well defined. Then, given two directions u, v, the second-order differential of at is
dPM()lu, v] = - (F
+ &g(((e)))O{((()le w]].
(49)
With this computation rule, are able give the second-order differential of I.ţ i., Lemma E.3 (Second-ordes differential of L.}). Let W E RRaxm, For any U, V E RR the second order differential of L. W in the directions U, V is
(L),(W), -
(50)
where
H(V) -(()=2(((VW" + WW + (I E
Σx (0, -1/2 -1/22-1/
21 vV)
(51)
and H is defined in (47
28

---

Under review as a conference paper at ICLR 2023
Proof. Applying the formula 49 to L L, gives, with Σ x(W) and d' VI
UV +
d =
=UUWT -WU MVV +WV tr +VUT)
(20 / " (UV
=2(U, H(VW + WV JW +V - ( -U,
-1 "2/²/)
where we used the symmetry of En
///² (x1// ( Σ, 1/22 11/2
Eo to simplify the expression.
The maximal eigenvalue of will then be computed as Amax (H) supu: ||0||== (U, H(U)) in Lemma E.6
E.2
LIPSCHITZ- OF L,
One can use the bounds of Kroshnin et al. (0021 Lemma A.3) to bound the Hessian of the loss.
Lemma E.4 (Bounds on the second-order differential, Kroshnin et al. 2021 Lemma A.3), Let H HX) be defined as in (4 The second-order differential of L., respects the following bounds (=// 1/2 Amax
1/2
(x,)
1
1/2
(52a)
2
1/2
/ 1/2
(x, H())
min
1 1/2 1/
(52b)
2
Those turn bound the extremal eigenvalues of the Hessian, defined as Amax (X.H(X))
supx300 (X (X)) and Amin(() in [|X||
Lemma E.5 (Bounds on the Hessian HT, Let H be defined as in 47 Then, the extremal eigenvalues of H are bounded as
V CanAmax(EM) 272
Ammin
Amaax H &
Amia (H)
(53)
2C
where Co (0) tr((o)) is initialiaation-deeeneen.
In particular, the loss T.p strongly convex, with parameter K
(Amin(FF
20%
Proof. We first provide the proof for the maximal eigenvalue.
The maximal eigenvalue of the Hessian is defined as
Amax H
sup
(X)).
X
From the upper-bound of (x, H(X)) in 2a
one has
X |D
1/2 E, 20
1/2
sup (X, (x)) < X||==
| 10-1//xx;111
X2; 1/2 x--1/211
sup
X:I|X||
m=1
A
[D Σx
1/2
|Σ 1/2 Xx,, 1/2
2
X r=1
Amax
1/2
(2)
(11/2 11/2
2,20
x 1
sup
X:||X p=1
X (2) 1/2 1/2
20
max (Σ-1)
2
11/2x 1/2 23
1/22
20
272
29

---

Under review as a conference paper at ICLR 2023
The last inequality comes from the definition of E,; if \1 X2 > itive eigenvalues of WWT, then Σ (WWT + =T,)) has eigenvalues
d 0 are the pos-
T2 k times
(Ax +T) > ... (A1 +
For any positive definite matrices A, E S++ (n) with increasing eigenvalues, and for any & E [n]. we know that
(B)A))((B) & (AB) (A (A)A....B).
Therefore, we have the bound Anaxx ( E/ X (E)) Amax (.,). Moreover, ( x Σ,, and from Lemma C.7 we know that tr Σ, E 2(L(2,)) L (oo) Co. Therefore, we obtain
Co CoAmax (Eo)
Amax (H)
272
The proof for the minimal eigenvalue is similar and follows from the bound (2bb. In this case, the term A 1/2 [,,,, can be lower bounded by / Amin (mo)
min
We now turn to the Hessian of L.-
Lemma E.6 (Spectral bound of H). Let H be defined as in (51) The maximal eigenvalue for the Hessian of L2 respects the following bound
Amax (H)
1/2 (
(x1 E 20 )
1/2
202
+ 2(1 Amin (E2 Ex 22
-1/2
x1
(54)
72
Proof. From 52a). one has for any X E SS.
1/2
/ 2/1
(x,
Amax
200
|Σ
1/2 x--
2
Let U E RR With X(U) = UW + TWU the bound becomes
1/2 (5//
1//2E,
1/2
(UWT + WUT, (I(((U)))
|E-1/2X(U)>7)1))
2
1/2
(E))
1/2
1/2
2 ,H(X(U))
2
X
1/25,
D/1
1/2
2 H(X())))
|,,1)3((0));1111
2
Therefore,
(U,H(U)) = ((())) + -/ 2 2// -1/2
-
x///2)0)
1/2 2,, | 2U, / 11/2
-1/2 =1/2 UU). {
2
(55)
We proceed by bounding each of the summands.
First consider the term | 1/- XU)), 121 |Σ - X(U) If is such that || = 1, then [ 3/- = |UWT + 1 < W| We know that ||W]| < C for some constant C,
c.f. 35 Therefore, ||U|| = |X(U]] 2C and
sup U: ==
| <
sup
|2:1X(U)
x |2
40 | X
sup
x: |||==
12 max (I-)) 102
72
30

---

Under review as a conference paper at ICLR 2023
Therefore,
X
E
(E/ 1212 20 72
(()))/)(());; 1/2 &
1222 max
sup Us |
2
The second summation in 55 can be bounded as
sup U: ||0||==1
2(U, (1-20 1/2 (2 / E,
1/2 5/1/2 (0)
201
2dmax (I - 10 (20) Lip 2 1-1/2/ - -1/2 -1/2211 201
//////222
= 2(1 Amin((y) (2/
-1/2 //) /511/2, ", / (E)/
Lemma E.7 (Lipshitz-smoothhee of L2). For > 0, the loss W + (W) is Lipschitz smooth.
Proof. This directly follows from the boundedness of the Hessian showed previously and the con- vexity of L using Taylor approximation.
Once the suaschitz-smoothhess of the loss has been proven, one can turn to showing that the rank is preserved under balanced initial conditions.
Proposition E.8 (aah et al. 2021 Proposition 4.4). Let L R R be a Lipschitz smooth function (i.e., a differentiable function with Lipschitz gradient). Suppose that Wi(t).... WN (t) are solutions of the gradient flow (GF) of T.N with balanced initial values W (0) and define the product (1)) = ((())) (t( ... Wi(t). (0) is contained in Mr for some k. E N, then W(1) is contained in for all t 3 0.
Proof. Let P(t) = T = (w(t) (w()))))))) I/N and Q(t) = Wt) = (W()W(()))) I/N
The proof follows the gradient flow is locally Lipschitz continuous in P. Q, W. so that the curves P. Q, W are uniquely determined by initial datum P(0).Q(0) W(O). From Equations (GF and 21
P - W VZ vC) w,
Q VWW - WV),
N
W = Σ Q~ (W))
Now, with the assumption of Lipschitz continuity of the flow, given solution is uniquely determined by the initial data Po. Qu, Wo. and the proof tools of Bah et al. 2021 Proposition 4.4) can be used here as well.
Remark E.9. The loss I.L satisfies the conditions of Proposition E.8 therefore, the flow Lt remains in the manifold Mr if W (to) E M₁ for some to-
E.3 PROOFS OF GRADIENT FLOW CONVERGENCE
Proof of Theorem 5.6 The idea of the proof to transfer the strong convexity property from î,, to
the evolution of the parameters. Let us start by the inequality which holds due to strong convexity
i(E,) - i(E;) < 2K
Vi(E.))
31

---

Under review as a conference paper at ICLR 2023
where K is the constant from Lemma
E.5
Rearranging the terms in the above equation, have
- -I -2K (L(Z,) L(;;)
(56)
On the covariance space, for the perturbative loss, the gradient flow is written
d [
( E())
dt
-(VxLr W dW aW de
WWWT) W
dt
[I),,,,II), di
The expression of dW is given in Lemma C.1
de
dw
N
WN Ny41 (W We-1-
dt
8=1
Since VL
= ( and from the balancedness assumption we have WN WN++1
(WWT) dx4 and W W W8-111 (WTW)* 4
get
aL(E()) di
Σ
(WWT) *-t
(W f-1
IN) N)
Now, let USVT = W be a (thin) SVD of W, so that WWT = USU and WTW = VS
For one layer & [N], we then have
(VLE)).(WW) N-1 VW f=1 tr (VL())W(WTW)) W) W (WW) N-t
(VL(E)USVVVV - VVSU VICEUSS 2 UT)
-U(VEU( 2IN U VILEUSS
aN
=(TVi(Z)US UN " U VV(EJ)
Let X = VVVÚ R= 5 2(N+-11 and L = S 3N-- We evaluate (XR, LX) for diagonals
R, L as
(XR,LX) Σ Xi XR,
Since I, 8 2(N-F
2(N+/-1)
and R2 = and due the uniform margin deficiency assumption, for
all (i.j) E [ have L, 2N-0) and R2 21N+2-1)
so that
(XR,LX)
20N
Σ x =c
212N
LX
Since X UVVLEE we have that |X [|) that in total
di(2(())
Σ 22N-11
v 4Nc
20N
|
(57)
dt
From the strong convexity of i. (56). we get the bound
aL(x()
-8Nc
2(2-11) N
K(i(E)-L(E')
d(f((5(()) ((D"))
2/N 1)
8Nc
N
K
L(E(t) L(E*))
dt
32

---

Under review as a conference paper at ICLR 2023
Now, by integrating both sides from 0 to
((),()) ((E;) (0) ((E;) S S-NNc
In
Kt.
(58)
Let A = (0) Σ which is the distance to optimality from the initialization. Finally we get the desired exponential rate
((E((0) L(E)) S SNc
which concludes the proof.
E.4 PROOF OF GRADIENT DESCENT CONVERGENCE
We start by proving Lemma 5.3 so that with the uniform margin deficiency assumption on the initial weights. WWT does not degenerate along the gradient descent training algorithms.
Proof of Lemma 5.3 Let (k) == arg ( (K))((R)) - -E/ U1? for each k then as L. (W))) I (W)) for all & 2 0, we have
mia (VW()) = Gmin
(VW(L) (0)0
(2)(2((6)
-2
min
-mmax
(x)/²0(()) 11 W(K)))) -2/20(())11
min
(k)) LL(((())
(59)
Min
[b(2 (W)((())
min
11v WO)) - E ((()]l
Omin
(x1/)) cc=
2 min
m
The cancellation in the last equality works due to the fact that the multiplication with an arbitrary unitary matrix does not change singular values.
Now we are ready to prove the finite step size gradient descent convergence of the BW loss. We consider the perfect balancedness of initial values (O))), 1 < i < N in the remaining proof. The approximation balancedness case can also be carried out but require more complicated auxiliary estimates. We leave the approximate balancedness assumption as future direction.
Proof of Theorem
5.7
Let us start from the gradient descent of the loss with respect to each layer
(k(k + 1)
(k(k) - Vw, L (W, (k), .. W..
(60)
= (k(k) - -,W Vw (13))) I5iSN,
with the boundary conditions (i.(()) Id. and W WN+I:N (l)) Iay for all k 2 0.
With the notations W = (W1, W2 ... Wxv and
Vw, WW,LN(WF
VLN(W)
Vw (
33

---

Under review as a conference paper at ICLR 2023
consider to write the Taylor expansion in the form
Ld 12(WW(k: (K + 1)) = L- (W())))+ - (k)))
(61)
with
(k) = W (k) + E(W (k + 1) - (k)), for some E [0,1].
Recall the relation 21 for 1 S s N,
V Ww, (W....., Wx) W ... WW ... W then the first order term in 61). under (60).
can be written
VL (W (k)),
(A: + (()))
Wk(k)
j=1
W W--1. ( WN....
W ... I W
WWW (W ...
Mmin (WW) miin
(W. ((V
<
(62)
Throughout the computation above, W; = (k for all 1 S N. Moreover, we use the balancedness W,W3. W. WTT W W,+1 j41 for all S i s N so that, in the symmetric structure above,
WN W J+1 WN = NWW
WT WJ_ - ... Wi (W
Therefore, thanks to Lemma 5.3
Jmin (wW (W))
miin (WW)
main (W())W(()7) 2 c²,
from which we get
(VL (VL (k+ -1)++)))))()))))))))))))))))))) S-
- Ne 2NN
V (W
(63)
Let us mention that Arora et al. 2018 Theorem and Claim 1) provide rigorous derivations about the equalities above. The second order term in 61 is more complicated to handle, as we have
LN
N
82 (W)
V
Ow,ow, x.)
(64)
Ow?
1.i
C.4 we have expressions of NL (W) ow?
and LL
Thanks to Corollary
aw.ow, ready.
Note that we have the boundedness (C.)
|| S / 2 22( + ||= (1 /13) S /22 (L)((W()) + 1122 =: M,
(65)
and it is straightforward to see that
2/N Moreover, for all 1 i N. since g E [0,
|| ||
for all 15i5N. N.
(66)
Ak,((k) (kk)) <(Wa(k + 1) W.(k)) (1 ()))()) + <W. + +1),
34

---

Under review as a conference paper at ICLR 2023
then have the uniform upper bound for all A 2 0,
[k)xl 3 ( + 1)1le S M
(67)
Using (k) -V (1,)+ V (WW we can obtain a lower bound in terms of the minimum singular value,
mmin (AL,(H)AA,,(1))
(+)) 2E wwwp:////w/l//////./////////////
2 Omin
2 e² An LTMVTA 2 e² M (T)(())))
(68)
where we utilize 25
(66 and 65
as well as non-increment of L (W) throughout the training.
We denote X, = W W (k: w (15-)))))1(()) We may choose
S
8M (
so that for all l 2 O.
2N 2
Mmin (A).
22
and
(4+))) mmin
2N
(69)
Then combining all estimates above, we have (W (k + (K)) T 22 (EE(K)), (K +11 -)
LL (A(()))
LL (Ac(k)) L
Σ 1
Σ
{x,,
OW,oW, x.)
OW?
X E2/ 1/2 (<():){((k)
T5/
1/2
|| (Ac( [ "1.MARNN-1NN
2
( [XIIll//X,/ll////2244444
1/2 Amax (2
(E (x)(4(())4((()]) "r
2
(
by using 67 E.4 and applying the Cauchy-Schwar inequality for the last term. Notice that ||*|||>
< Mr 1//N Vw (W( (I)))l... Now combining all the bounds we obtained previously,
addition (69). we get that
K + (1) (F))
+1) -k+1)-W(AH)
11/N
S 2n
1/2 (En)
V (W
(70)
(A))
Mmin
+ + M)M (3N 4N V
Moreover, we can use
25 again to get
IV 2
- 2//UUl1r (2)((4((AAA(())))) 2/111/// 11 S 2M//N + +2|1=/ |r--
35

---

Under review as a conference paper at ICLR 2023
Thus, conclude the estimate for the second order term by
2 < N2 2 (20)
+ +8N(N- - 11M (ONTIINN_)NNMMM (M//X + (M -12)//10))
Let us denote the constant
A 2N 22N N M(N-33))
X max (Eo) + 8N(N 1)M(NN (MI/N + +122//111
then, with we can write the iteration as
(W + 1)) = (1 (1-4N) 4Nc x 11 + L (W(()).
20N
If we choose
Ne 21N-12
1 S 2A
then we have
I. (W)(W())) S (1 - Nc
21N 11 & I (0)).
N
For being sufficiently small, we have 1 - 2n.Ne - I exp e-crror for the loss,
(-2m N
Thus, to achieve
L-(W(O))
k
N1,Ne 2N log
36


Reviews:
[{"review_id": "cxq9jKJsJD", "reviewer": "Reviewer_LN7P", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the minimization of the Bures-Wassertein loss on positive symmetric matrices. Seen as a function over matrices of constant ranks, they give a characterisation of its critical points and of its gradient flow when added an overparametrized multiplicative structure.", "summary_of_the_review": "Under this state, the paper is inappropriate for publication.", "strength_and_weaknesses": "### **Strength**\n\nThe paper tries to derive results on what can be rebaptized the Bures-Wasserstein loss that may appear for generative purposes.\n\n-----------------------------------------------\n\n\n### **Weaknesses**\n\nIt is hard to really understand the purpose of the paper: \n- Is it to derive properties of the Bures-Wasserstein distance function restricted to manifolds of constant rank as a theoretical object? In this perspective, not only ICLR may not be a target for such a paper, but also, the results do not seem very conclusive with respect to this precise literature. \n- Is this to theoretically ground generative models? Here again, the overall motivation is not clearly depicted and furthermore, even if this motivation appears in the introduction and the title, there is no comment and/or simulation of this. Which, once again, is clearly out of target when considering this is a ICLR submission.\n- Finally, the mathematical writing as well as the story behind the paper is really hard to follow. Notations are very messy even-though all the reparametrisations used necessitate a careful writing. As a striking example, the beginning of the Section 3  concerning the notations is really hard to follow: $\\pi(W)$ is introduced but is in fact $\\Sigma$, $\\mu$ appears for no reel reason sometimes, and there is no consistency from one line to the other. I advise the authors to ***carefully and explicitly*** write any parameter dependence on the notations used, otherwise it is impossible to follow. Examples: $\\Sigma \\rightarrow \\Sigma(W)$,  $W^* \\rightarrow W^*_{J_k} $, etc...\n- The theorems, the assumptions and the Lemma are not commented. The setup of the Theorems are not recalled. For example, Theorem 5.6 is on the gradient flow in terms of $W$ or $\\Sigma$? Furthermore Theorem 5.6 cannot be really considered as a convergence result as $\\|\\Delta_t^*\\|$ appears in the right term. Finally why do we have a Theorem 5.7 as a linear convergence result whereas Therem 5.6 is polynomial ?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Xtu05JWmbo9", "reviewer": "Reviewer_5WMe", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper shows the convergence of specific first-order gradient-based algorithms on deep linear networks with Bures-Wasserstein loss, under the assumption that the initial weights are balanced. The author studies both the original non-smooth Bures-Wasserstein loss and a smooth perturbative one. They characterize the critical points of those two losses over the rank-constrained matrices and they show the convergence of GF on smooth perturbative BW loss, and GD on BW loss, under the assumption that the initial weights are balanced, and has a uniform deficiency margin.\n\n", "summary_of_the_review": "One assumption is not justified in the paper; \n\nThe results are not very significant compared to prior work.", "strength_and_weaknesses": "Strength: Although this paper considers a deep linear network, the problem setting is relatively new as previous work primarily studies square loss or logistic/exponential loss. The writings are mostly clear.\n\nWeakness: I have several concerns about the quality and significance of the results. Specifically:\n1. The discussion of critical points of BW loss and its smooth version is less relevant to the rest of the paper because the results are shown for the space of $W$, the end-to-end matrix. If one is concerned about the landscape of deep linear networks, he should study the critical points of BW loss in the weight space $W_1, \\cdots, W_N$, as those in [1]. If this paper is mainly about convergence, I don't know why Section 4 is presented.\n2. The convergence requires a uniform deficiency margin at initialization. Although this is very related to the condition proposed in [2], this UDM assumption is much more restrictive than the one in [2]. In fact, it is highly nontrivial to see whether this UDM assumption can be satisfied, but the UDM assumption is essential to the convergence proof in the paper. The author needs to justify the UDM assumption formally by showing there exists an initial $W(0)$ that satisfies UDM.\n3. Regarding the contribution, the new thing is the BW loss. Except for the loss, the convergence analysis relies on a balanced assumption together with a deficiency margin (with minor modification), which is mostly based on the analysis in [2]. Moreover, even the balanced assumption is restrictive and recent work has shown convergence of GF under non-balanced initialization [3,4]. Unless there is a significant difference when analyzing the convergence under the BW loss, the contribution of this paper is minor, in my opinion.\n\nMinor comments:\n1. It's unclear from Theorem 5.6 whether the loss has $\\mathcal{O}(t^{-1})$ convergence or not because $||\\Delta_t^* ||$ depends on time $t$\n\nReference:\n\n[1] Kenji Kawaguchi. Deep learning without poor local minima, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.\n\n[2] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. ICLR, 2018\n\n[3] Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. ICLR, 2020.\n\n[4] Hancheng Min, Salma Tarmoun, René Vidal, and Enrique Mallada. On the explicit role of initialization on the convergence and implicit bias of overparametrized linear networks. ICML, 2021", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "A19pd7HyR-D", "reviewer": "Reviewer_8S82", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the convergence of a deep linear network trained with a Bures-Wasserstein loss. \n\nThe authors first give a description of the critical point of the Bures-Wasserstein loss (and its perturbation) when restricted to the set of rank $k$ matrices $\\mathcal{M}_k$. If I understand correctly, the critical points of the loss as a function of the parameters of the network are not studied.\n\nConvergence of gradient flow and gradient descent on the parameters of the network (with the perturbative Bures-Wasserstein loss) are then showed for initialization that are balanced and have a uniform deficiency margin.", "summary_of_the_review": "The paper adapts the proof of convergence of (Arora et al. 2020) to the Bures-Wasserstein loss. The same strong assumptions are required for their convergence proof.", "strength_and_weaknesses": "The proof of convergence is strong, but I have issues with its novelty (see next section) and with the assumptions:\n - Though the balanced assumption is common due to its power, I find it quite restrictive. The main reason why I am interested in the convergence of deep linear networks is as a stepping stone to later prove convergence of nonlinear networks. And these kind of balanced arguments break down in the non-linear case. I think it is important to focus on non-balanced linear networks.\n - The uniform deficiency margin assumption is very strong and the discussion of this assumption is misleading: (Arora et al., 202) show that a random network at initialization will have a positive deficiency margin with prob. close to 0.5 only when when the input and outputs have dimension 1. In Remark 5.4 they cite the above paper but omit the dimension assumption which is very strong (balanced dimension 1 network are equivalent to scalar networks whose convergence with prob. 0.5 is obvious). In Remark 5.4 the authors write \"in our setting, the uniform deficiency margin can also be satisfied using a similar result\" but I couldn't find the corresponding argument anywhere in the appendix.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "7AB0wzx9zY", "reviewer": "Reviewer_pMGB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the convergence of deep linear networks under the so-called Bures-Wasserstein loss, a loss function on the space of positive semi-definite matrices. Such losses arise as the 2-Wasserstein distance between gaussian distributions, and hence can be considered as approximating a gaussian target in a generative setting. \n\nThe authors characterized the critical points of the loss. Since the Bures-Wasserstein loss is non-smooth, the authors introduced a smooth version of the loss and prove the convergence rates of the gradient flow under an initial condition and a uniform margin condition. Extention to gradient descent with small step-sizes was considered.", "summary_of_the_review": "This paper studies a mathematically interesting problem and provides a detailed analysis. My two major concerns are:\n\n1. The various gaps in the derivation, and \n2. Motivation.\n\nI suggest the authors to improve the paper based in these two directions.", "strength_and_weaknesses": "Strength:\n- Many properties of the Bures-Wasserstein loss in the context of deep linear networks are proved, which might be interesting for future works.\n\n\nWeaknesses:\n- The motivation for using the Bures-Wasserstein loss is not clear. In other words, while I find the math interesting, I do not understand why we should study the problem presented in the paper. Can the authors provide further arguments other than \"the Bures-Wasserstein distance has also attracted much interest\"? Please note that each of the cited literature in the introduction, unlike this work, has its own motivation for the Bures-Wasserstein distance.\n\n- The characterization of the critical points of (non-smooth) Bures-Wasserstein losses is quite disconnected from the rest of the paper.\n\n- Why does the gradient descent (Theorem 5.7) converge faster than the gradient flow, ie exponential versus $O(t^{-1})$? Moreover, the proof of Theorem 5.7 relies on a contraction argument. Why does this not apply to GF, as an \"infinitesimal step-size GD\"?\n\n\n- The bound in Lemma 3.2 does not decrease to 0. In fact, quite unintuitively, it grows with $r$, the rank constraint. These facts suggest that the relation between the smooth and non-smooth Bures-Wasserstein losses is not captured correctly.\n\n- The uniform margin condition in Definition 5.2 is exceedingly strong as it is required to hold for *all* unitary matrix $U$. This is different from Arora et al. 2019a where the \"over all $U$\" part is not present. For instance, let $\\Sigma_\\theta = \\Sigma_0 = I_n$. Then the $\\max_U$ over the left-hand side of (7) would grow with $n$, which is not a property of the deficiency margin condition in (Arora et al. 2019a).\nIn this light, Remark 5.4 seems quite misleading as it suggests that the uniform deficiency margin condition can be easily met with high probability by choosing a specific initialization as in Arora et al. 2019a. \n\n\n- The authors claimed to have proved $O(t^{-1})$ convergence for GF. However, inspecting Theorem 5.6, there is an $\\Delta_t^*$ term in (10) which might increase over $t$ (at least the authors did not provide a rigorous proof that $\\Delta_t^*$ remains bounded). The proof therefore seems incomplete to me.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Reject