Title: Efficient Dictionary Learning with Gradient Descent

Abstract: Randomly initialized first-order optimization algorithms are the method of choice for solving many high-dimensional nonconvex problems in machine learning, yet general theoretical guarantees cannot rule out convergence to critical points of poor objective value. For some highly structured nonconvex problems however, the success of gradient descent can be understood by studying the geometry of the objective. We study one such problemcomplete orthogonal dictionary learning, and provide converge guarantees for randomly initialized gradient descent to the neighborhood of a global optimum. The resulting rates scale as low order polynomials in the dimension even though the objective possesses an exponential number of saddle points. This efficient convergence can be viewed as a consequence of negative curvature normal to the stable manifolds associated with saddle points, and we provide evidence that this feature is shared by other nonconvex problems of importance as well.

1. Introduction:

Many central problems in machine learning and signal processing are most naturally formulated as optimization problems. These problems are often both nonconvex and highdimensional. High dimensionality makes the evaluation of second-order information prohibitively expensive, and thus randomly initialized first-order methods are usually employed instead. This has prompted great interest in recent years in understanding the behavior of gradient descent on nonconvex objectives (18; 14; 17; 11) . General analysis of first-and second-order methods on such problems can provide guarantees for convergence to critical points but these may be highly suboptimal, since nonconvex optimization is in general an NP-hard probem (4) . Outside of a convex setting (28) one must assume additional structure in order to make statements about convergence to optimal or high quality solutions. It is a curious fact that for certain classes of problems such as ones that involve sparsification (25; 6) or matrix/tensor recovery (21; 19; 1) first-order methods can be used effectively. Even for some highly nonconvex problems where there is no ground truth available such as the training of neural networks first-order methods converge to high-quality solutions (40) .

Dictionary learning is a problem of inferring a sparse representation of data that was originally developed in the neuroscience literature (30) , and has since seen a number of important applications including image denoising, compressive signal acquisition and signal classification (13; 26) . In this work we study a formulation of the dictionary learning problem that can be solved efficiently using randomly initialized gradient descent despite possessing a number of saddle points exponential in the dimension. A feature that appears to enable efficient optimization is the existence of sufficient negative curvature in the directions normal to the stable manifolds of all critical points that are not global minima 1 . This property ensures that the regions of the space that feed into small gradient regions under gradient flow do not dominate the parameter space. Figure 1 illustrates the value of this property: negative curvature prevents measure from concentrating about the stable manifold. As a consequence randomly initialized gradient methods avoid the "slow region" of around the saddle point.

Figure 1 : Negative curvature helps gradient descent. Red: "slow region" of small gradient around a saddle point. Green: stable manifold associated with the saddle point. Black: points that flow to the slow region. Left: global negative curvature normal to the stable manifold. Right: positive curvature normal to the stable manifold -randomly initialized gradient descent is more likely to encounter the slow region.

The main results of this work is a convergence rate for randomly initialized gradient descent for complete orthogonal dictionary learning to the neighborhood of a global minimum of the objective. Our results are probabilistic since they rely on initialization in certain regions of the parameter space, yet they allow one to flexibly trade off between the maximal number of iterations in the bound and the probability of the bound holding.

While our focus is on dictionary learning, it has been recently shown that for other important nonconvex problems such as phase retrieval (8) performance guarantees for randomly initialized gradient descent can be obtained as well. In fact, in Appendix C we show that negative curvature normal to the stable manifolds of saddle points (illustrated in Figure 1 ) is also a feature of the population objective of generalized phase retrieval, and can be used to obtain an efficient convergence rate.

2. Related Work:

Easy nonconvex problems. There are two basic impediments to solving nonconvex problems globally: (i) spurious local minimizers, and (ii) flat saddle points, which can cause methods to stagnate in the vicinity of critical points that are not minimizers. The latter difficulty has motivated the study of strict saddle functions (36; 14) , which have the property that at every point in the domain of optimization, there is a large gradient, a direction of strict negative curvature, or the function is strongly convex. By leveraging this curvature information, it is possible to escape saddle points and obtain a local minimizer in polynomial time. 2 Perhaps more surprisingly, many known strict saddle functions also have the property that every local minimizer is global; for these problems, this implies that efficient methods find global solutions. Examples of problems with this property include variants of sparse dictionary learning (38) , phase retrieval (37) , tensor decomposition (14) , community detection (3) and phase synchronization (5) .

Minimizing strict saddle functions. Strict saddle functions have the property that at every saddle point there is a direction of strict negative curvature. A natural approach to escape such saddle points is to use second order methods (e.g., trust region (9) or curvilinear search (15) ) that explicitly leverage curvature information. Alternatively, one can attempt to escape saddle points using first order information only. However, some care is needed: canonical first order methods such as gradient descent will not obtain minimizers if initialized at a saddle point (or at a point that flows to one) -at any critical point, gradient descent simply stops. A natural remedy is to randomly perturb the iterate whenever needed. A line of recent works shows that noisy gradient methods of this form efficiently optimize strict saddle functions (24; 12; 20) . For example, (20) obtains rates on strict saddle functions that match the optimal rates for smooth convex programs up to a polylogarithmic dependence on dimension. 3Randomly initialized gradient descent? The aforementioned results are broad, and nearly optimal. Nevertheless, important questions about the behavior of first order methods for nonconvex optimization remain unanswered. For example: in every one of the aforemented benign nonconvex optimization problems, randomly initialized gradient descent rapidly obtains a minimizer. This may seem unsurprising: general considerations indicate that the stable manifolds associated with non-minimizing critical points have measure zero (29) , this implies that a variety of small-stepping first order methods converge to minimizers in the large-time limit (23) . However, it is not difficult to construct strict saddle problems that are not amenable to efficient optimization by randomly initialized gradient descent -see (12) for an example. This contrast between the excellent empirical performance of randomly initialized first order methods and worst case examples suggests that there are important geometric and/or topological properties of "easy nonconvex problems" that are not captured by the strict saddle hypothesis. Hence, the motivation of this paper is twofold: (i) to provide theoretical corroboration (in certain specific situations) for what is arguably the simplest, most natural, and most widely used first order method, and (ii) to contribute to the ongoing effort to identify conditions which make nonconvex problems amenable to efficient optimization.

3 Dictionary Learning over the Sphere Suppose we are given data matrix Y = y 1 , . . . y p ∈ R n×p . The dictionary learning problem asks us to find a concise representation of the data (13) , of the form Y ≈ AX, where X is a sparse matrix. In the complete, orthogonal dictionary learning problem, we restrict the matrix A to have orthonormal columns (A ∈ O(n)). This variation of dictionary learning is useful for finding concise representations of small datasets (e.g., patches from a single image, in MRI ( 32)).

To analyze the behavior of dictionary learning algorithms theoretically, it useful to posit that Y = A 0 X 0 for some true dictionary A 0 ∈ O(n) and sparse coefficient matrix X 0 ∈ R n×p , and ask whether a given algorithm recovers the pair (A 0 , X 0 ). 4 In this work, we further assume that the sparse matrix X 0 is random, with entries i.i.d. Bernoulli-Gaussian5 . For simplicity, we will let A 0 = I; our arguments extend directly to general A 0 via the simple change of variables q → A * 0 q. (34) showed that under mild conditions, the complete dictionary recovery problem can be reduced to the geometric problem of finding a sparse vector in a linear subspace (31) . Notice that because A 0 is orthogonal, row(Y ) = row(X 0 ). Because X 0 is a sparse random matrix, the rows of X 0 are sparse vectors. Under mild conditions (34) , they are the sparsest vectors in the row space of Y , and hence can be recovered by solving the conceptual optimization problem min q * Y 0 s.t. q * Y = 0. This is not a well-structured optimization problem: the objective is discontinuous, and the constraint set is open. A natural remedy is to replace the 0 norm with a smooth sparsity surrogate, and to break the scale ambiguity by constraining q to the sphere, giving

EQUATION

Here, we choose h µ (t) = µ log(cosh(t/µ)) as a smooth sparsity surrogate. This objective was analyzed in (35) , which showed that (i) although this optimization problem is nonconvex, when the data are sufficiently large, with high probability every local optimizer is near a signed column of the true dictionary A 0 , (ii) every other critical point has a direction of strict negative curvature, and (iii) as a consequence, a second-order Riemannian trust region method efficiently recovers a column of A 0 . 6 The Riemannian trust region method is of mostly theoretical interest: it solves complicated (albeit polynomial time) subproblems that involve the Hessian of f DL . Note the similarity to the dictionary learning objective. Right: The objective for complete orthogonal dictionary learning (discussed in section 6) for n = 3.

In practice, simple iterative methods, including randomly initialized gradient descent are also observed to rapidly obtain high-quality solutions. In the sequel, we will give a geometric explanation for this phenomenon, and bound the rate of convergence of randomly initialized gradient descent to the neighborhood of a column of A 0 . Our analysis of f DL is probabilistic in nature: it argues that with high probability in the sparse matrix X 0 , randomly initialized gradient descent rapidly produces a minimizer.

To isolate more clearly the key intuitions behind this analysis, we first analyze the simpler separable objective

min f Sep (q) ≡ n i=1 h µ (q i ) s.t. q ∈ S n-1 .

(

EQUATION

Figure 2 plots both f Sep and f DL as functions over the sphere. Notice that many of the key geometric features in f DL are present in f Sep ; indeed, f Sep can be seen as an "ultrasparse" version of f DL in which the columns of the true sparse matrix X 0 are taken to have only one nonzero entry. A virtue of this model function is that its critical points and their stable manifolds have simple closed form expressions (see Lemma 1).

4. Outline of Important Geometric Features:

Our problems of interest have the form min f (q) s.t. q ∈ S n-1 , where f : R n → R is a smooth function. We let ∇f (q) and ∇ 2 f (q) denote the Euclidean gradient and hessian (over R n ), and let grad [f ] (q) and Hess [f ] (q) denote their Riemannian counterparts (over S n-1 ). We will obtain results for Riemannian gradient descent defined by the update q → exp q (-η grad[f ](q)) for some step size η > 0, where exp q : T q S n-1 → S n-1 is the exponential map. The Riemannian gradient on the sphere is given by grad[f ](q) = (I -qq * )∇f (q).

We let A denote the set of critical points of f over S n-1 -these are the points q s.t. grad [f ] (q) = 0. We let Ȃ denote the set of local minimizers, and " A its complement. Both f Sep and f DL are Morse functions on S n-1 ,7 we can assign an index α to every q ∈ A, which is the number of negative eigenvalues of Hess [f ] (q). Our goal is to understand when gradient descent efficiently converges to a local minimizer. In the small-step limit, gradient descent follows gradient flow lines γ : R → M, which are solution curves of the ordinary differential equation

γ(t) = -grad [f ] (γ(t))

To each critical point α ∈ A of index λ, there is an associated stable manifold of dimension dim(M) -λ, which is roughly speaking, the set of points that flow to α under gradient flow: In the light blue region, there is negative curvature normal to ∂C, while in the orange region the gradient norm is large, as illustrated by the arrows. There is a single global minimizer in the yellow region. For the separable objective, the stable manifolds of the saddles and maximizers all lie on ∂C (the black circles denote the critical points, which are either maximizers " ", saddles " ", or minimizers " "). The red dots denote ∂C ζ with ζ = 0.2.

W s (α) ≡ q ∈ M lim t→∞ γ(t) = α

Our analysis uses the following convenient coordinate chart

EQUATION

where w ∈ B 1 (0). We also define two useful sets:

EQUATION

Since the problems considered here are symmetric with respect to a signed permutation of the coordinates we can consider a certain C and the results will hold for the other symmetric sections as well. We will show that at every point in C aside from a neighborhood of a global minimizer for the separable objective (or a solution to the dictionary problem that may only be a local minimizer), there is either a large gradient component in the direction of the minimizer or negative curvature in a direction normal to ∂C. For the case of the separable objective, one can show that the stable manifolds of the saddles lie on this boundary, and hence this curvature is normal to the stable manifolds of the saddles and allows rapid progress away from small gradient regions and towards a global minimizer8 . These regions are depicted in Figure 3 .

In the sequel, we will make the above ideas precise for the two specific nonconvex optimization problems discussed in Section 3 and use this to obtain a convergence rate to a neighborhood of a global minimizer. Our analysis are specific to these problems. However, as we will describe in more detail later, they hinge on important geometric characteristics of these problems which make them amenable to efficient optimization, which may obtain in much broader classes of problems.

5. Separable Function Convergence Rate:

In this section, we study the behavior of randomly initialized gradient descent on the separable function f Sep . We begin by characterizing the critical points:

Lemma 1 (Critical points of f Sep ). The critical points of the separable problem (2) are

EQUATION

For every α ∈ A and corresponding a(α), for µ < c √ n log n the stable manifold of α takes the form

W s (α) = P S n-1 [ a(α) + b ] supp(a(α)) ∩ supp(b) = ∅, b ∞ < 1 (6)

where c > 0 is a numerical constant.

Proof. Please see Appendix A By inspecting the dimension of the stable manifolds, it is easy to verify that that there are 2n global minimizers at the 1-sparse vectors on the sphere ± e i , 2 n maximizers at the least sparse vectors and an exponential number of saddle points of intermediate sparsity. This is because the dimension of W s (α) is simply the dimension of b in 6, and it follows directly from the stable manifold theorem that only minimizers will have a stable manifold of dimension n -1. The objective thus possesses no spurious local minimizers.

When referring to critical points and stable manifolds from now on we refer only to those that are contained in C or on its boundary. It is evident from Lemma 1 that the critical points in " A all lie on ∂C and that α∈ " A W s (α) = ∂C , and there is a minimizer at its center given by q(0) = e n .

5.1. The effect of negative curvature on the gradient:

We now turn to making precise the notion that negative curvature normal to stable manifolds of saddle points enables gradient descent to rapidly exit small gradient regions. We do this by defining vector fields u (i) (q), i ∈ [n -1] such that each field is normal to a continuous piece of ∂C ζ and points outwards relative to C ζ defined in 4. By showing that the Riemannian gradient projected in this direction is positive and proportional to ζ, we are then able to show that gradient descent acts to increase ζ(q(w)) = qn w ∞ -1 geometrically. This corresponds to the behavior illustrated in the light blue region in Figure 3 .

Lemma 2 (Separable objective gradient projection). For any w:

∈ C ζ , i ∈ [n -1], we define a vector u (i) ∈ T q(w) S n-1 by u (i) j =      0 j / ∈ {i, n}, sign(w i ) j = i, -|wi| qn j = n. (7) If µ log 1 µ ≤ w i and µ < 1 16 , then u (i) * grad[f Sep ](q(w)) ≥ c w ∞ ζ, where c > 0 is a numerical constant. Proof. Please see Appendix A.

Since we will use this property of the gradient in C ζ to derive a convergence rate, we will be interested in bounding the probability that gradient descent initialized randomly with respect to a uniform measure on the sphere is initialized in C ζ . This will require bounding the volume of this set, which is done in the following lemma:

Lemma 3 (Volume of C ζ ). For C ζ defined as in (4) we have Vol(C ζ ) Vol(S n-1 ) ≥ 1 2n - log(n) n ζ

Proof. Please see Appendix D.3.

5.2. Convergence rate:

Using the results above, one can obtain the following convergence rate:

Theorem 1 (Gradient descent convergence rate for separable function). For any

0 < ζ 0 < 1, r > µ log 1 µ , Riemannian gradient descent with step size η < min c1 n , µ 2 on the separable objective (2) with µ < c2 √ n log n , enters an L ∞ ball of radius r around a global minimizer in T < C η √ n r 2 + log 1 ζ 0 iterations with probability P ≥ 1 -2 log(n)ζ 0 , where c i , C > 0 are numerical constants. Proof. Please see Appendix A.

We have thus obtained a convergence rate for gradient descent that relies on the negative curvature around the stable manifolds of the saddles to rapidly move from these regions of the space towards the vicinity of a global minimizer. This is evinced by the logarithmic dependence of the rate on ζ. As was shown for orthogonal dictionary learning in (38) , we also expect a linear convergence rate due to strong convexity in the neighborhood of a minimizer, but do not take this into account in the current analysis.

6. Dictionary Learning Convergence Rate:

The proofs in this section will be along the same lines as those of Section 5. While we will not describe the positions of the critical points explicitly, the similarity between this objective and the separable function motivates a similar argument. It will be shown that initialization in some C ζ will guarantee that Riemannian gradient descent makes uniform progress in function value until reaching the neighborhood of a global minimizer. We will first consider the population objective which corresponds to the infinite data limit

EQUATION

)

and then bounding the finite sample size fluctuations of the relevant quantities. We begin with a lemma analogous to Lemma 2:

Lemma 4 (Dictionary learning population gradient).

For w ∈ C ζ , r < |w i |, µ < c 1 r 5/2 √ ζ the dictionary learning population objective 8 obeys u (i) * grad[f pop DL ](q(w)) ≥ c θ r 3 ζ

where c θ depends only on θ, c 1 is a positive numerical constant and u (i) is defined in 7.

Proof. Please see Appendix B:

Using this result, we obtain the desired convergence rate for the population objective, presented in Lemma 11 in Appendix B. After accounting for finite sample size fluctuations in the gradient, one obtains a rate of convergence to the neighborhood of a solution (which is some signed basis vector due to our choice A 0 = I) Under review as a conference paper at ICLR 2019 Theorem 2 (Gradient descent convergence rate for dictionary learning). For any

1 > ζ 0 > 0, s > µ 4 √

2 , Riemannian gradient descent with step size η < c5θs n log np on the dictionary learning objective 1 with µ < c6 √ ζ0 n 5/4 , θ ∈ (0, 1 2 ), enters a ball of radius c 3 s from a target solution in

T < C 2 ηθ 1 s + n log 1 ζ 0

iterations with probability

P ≥ 1 -2 log(n)ζ 0 -P y -c 8 p -6

where y = c7θ(1-θ)ζ0

n 3/2

, P y is given in Lemma 10 and c i , C i are positive constants.

Proof. Please see Appendix B

The two terms in the rate correspond to an initial geometric increase in the distance from the set containing the small gradient regions around saddle points, followed by convergence to the vicinity of a minimizer in a region where the gradient norm is large. The latter is based on results on the geometry of this objective provided in (38) .

7. Discussion:

The above analysis suggests that second-order properties -namely negative curvature normal to the stable manifolds of saddle points -play an important role in the success of randomly initialized gradient descent in the solution of complete orthogonal dictionary learning. This was done by furnishing a convergence rate guarantee that holds when the random initialization is not in regions that feed into small gradient regions around saddle points, and bounding the probability of such an initialization. In Appendix C we provide an additional example of a nonconvex problem that for which an efficient rate can be obtained based on an analysis that relies on negative curvature normal to stable manifolds of saddles -generalized phase retrieval. An interesting direction of further work is to more precisely characterize the class of functions that share this feature.

The effect of curvature can be seen in the dependence of the maximal number of iterations T on the parameter ζ 0 . This parameter controlled the volume of regions where initialization would lead to slow progress and the failure probability of the bound 1 -P was linear in ζ 0 , while T depended logarithmically on ζ 0 . This logarithmic dependence is due to a geometric increase in the distance from the stable manifolds of the saddles during gradient descent, which is a consequence of negative curvature. Note that the choice of ζ 0 allows one to flexibly trade off between T and 1 -P. By decreasing ζ 0 , the bound holds with higher probability, at the price of an increase in T . This is because the volume of acceptable initializations now contains regions of smaller minimal gradient norm. In a sense, the result is an extrapolation of works such as (23) that analyze the ζ 0 = 0 case to finite ζ 0 .

Our analysis uses precise knowledge of the location of the stable manifolds of saddle points.

For less symmetric problems, including variants of sparse blind deconvolution (41) and overcomplete tensor decomposition, there is no closed form expression for the stable manifolds. However, it is still possible to coarsely localize them in regions containing negative curvature.

Understanding the implications of this geometric structure for randomly initialized first-order methods is an important direction for future work.

One may hope that studying simple model problems and identifying structures (here, negative curvature orthogonal to the stable manifold) that enable efficient optimization will inspire approaches to broader classes of problems. One problem of obvious interest is the training of deep neural networks for classification, which shares certain high-level features with the problems discussed in this paper. The objective is also highly nonconvex and is conjectured to contain a proliferation of saddle points (11) , yet these appear to be avoided by first-order methods (16) for reasons that are still quite poorly understood beyond the two-layer case (39) .

A Proofs -Separable Objective:

Proof of Lemma 1: (Critical point structure of separable objective) . Denoting by tanh( q µ ) a vector in R n elements tanh( q µ ) i = tanh( qi µ ) we have

grad[f Sep ](q) i = (I -qq * ) tanh( q µ )

. Thus critical points are ones where either tanh( q µ ) = 0 (which cannot happen on S n-1 ) or tanh( q µ ) is in the nullspace of (I -qq * ), which implies tanh( q µ ) = cq for some constant b. The equation tanh( x µ ) = bx has either a single solution at the origin or 3 solutions at {0, ±r(b)} for some r(b). Since this equation must be solves simultaneously for every element of q, we obtain ∀i ∈ [n] : q i ∈ {0, ±r(b)}. To obtain solutions on the sphere, one then uses the freedom we have in choosing b (and thus r(b)) such that q = 1. The resulting set of critical points is thus

A = P S n-1 {-1, 0, 1} n \ {0} .

To prove the form of the stable manifolds, we first show that for q i such that |q i | = q ∞ and any q j such that |q j | + ∆ = |q i | and sufficiently small ∆ > 0, we have

EQUATION

For ease of notation we now assume q i , q j > 0 and hence ∆ = q i -q j , otherwise the argument can be repeated exactly with absolute values instead. The above inequality can then be written as

(q i -q j ) n k=1 tanh( q k µ )q k -tanh( q i µ ) + tanh( q j µ ) ≡h > 0.

If we now define

s 2 = n-1 k = 1 k = i, n q 2 k and q n = 1 -s 2 -(q j + ∆) 2 we have h = ∆ tanh( qj +∆ µ ) (q j + ∆) + tanh( √ 1-s 2 -(qj +∆) 2 µ ) 1 -s 2 -(q j + ∆) 2 +∆ k =i,n tanh( q k µ )q k -tanh( qj +∆ µ ) + tanh( qj µ ) = ∆            k =i,n tanh( q k µ )q k + tanh( qj µ )q j + tanh( √ 1-s 2 -q 2 j µ ) 1 -s 2 -q 2 j ≡h1 -sech 2 ( q j µ ) 1 µ ≡h2            + O(∆ 2 )

where the O(∆ 2 ) term is bounded. Defining a vector r ∈ R n by

k = i, n : r k = q k , r i = tanh( q j µ )q j , r n = 1 -s 2 -q 2 j

we have r 2 = 1. Since tanh(x) is concave for x > 0, and |r i | ≤ 1, we find

h 1 = n k=1 tanh( r k µ )r k ≥ tanh( 1 µ ) n k=1 r 2 k = tanh( 1 µ ).

From

|q i | = q ∞ it follows that q i ≥ 1 √ n and thus q j ≥ 1 √ n -∆.

Using this inequality and properties of the hyperbolic secant we obtain

h 2 ≤ 4 exp(-2 q j µ -log µ) ≤ exp( 2∆ µ - 2 µ √ n -log µ + log 4)

and plugging in µ

= c √ n log n for some c < 1 ≤ exp( 2∆ µ - 2 log n c -log c + 1 2 log n + log log n + log 4).

We can bound this quantity by a constant, say h 2 ≤ 1 2 , by requiring

A ≡ 2∆ µ -log c + ( 1 2 - 2 c ) log n + log log n ≤ -log 8

and for and c < 1, using -log n + log log n < 0 we have

A < 2∆ µ -log c -( 2 c -1) log n.

Since ∆ can be taken arbitrarily small, it is clear that c can be chosen in an n-independent manner such that A ≤ -log 8. We then find

h 1 -h 2 ≥ tanh( 1 µ ) - 1 2 ≥ tanh( √ n log n) - 1 2 > 0 since this inequality is strict, ∆ can be chosen small enough such that O(∆ 2 ) < ∆(h 1 -h 2 )

and hence

EQUATION

It follows that under negative gradient flow, a point with

|q j | < ||q|| ∞ cannot flow to a point q such that |q j | = ||q || ∞ .

From the form of the critical points, for every such j, q must thus flow to a point such that q j = 0 (the value of the j coordinate cannot pass through 0 to a point where |q j | = ||q || ∞ since from smoothness of the objective this would require passing some q with q j = 0, at which point grad [f Sep ] (q ) j = 0).

As for the maximal magnitude coordinates, if there is more than one coordinate satisfying

|q i1 | = |q i2 | = q ∞ ,

it is clear from symmetry that at any subsequent point q along the gradient flow line q i1 = q i2 . These coordinates cannot change sign since from the smoothness of the objective this would require that they pass through a point where they have magnitude smaller than 1/ √ n, at which point some other coordinate must have a larger magnitude (in order not to violate the spherical constraint), contradicting the above result for non-maximal elements. It follows that the sign pattern of these elements is preserved during the flow. Thus there is a single critical point to which any q can flow, and this is given by setting all the coordinates with |q j | < q ∞ to 0 and multiplying the remaining coordinates by a positive constant to ensure the resulting vector is on S n . Denoting this critical point by α, there is a vector b such that q = P S n-1 [a(α) + b] and supp(a(α)) ∩ supp(b) = ∅, b ∞ < 1 with the form of a(α) given by 5 . The collection of all such points defines the stable manifold of α.

Proof of Lemma 2: (Separable objective gradient projection). i) We consider the sign(w i ) = 1 case; the sign(w i ) = -1 case follows directly.

Recalling that

EQUATION

for some c > 0 whose form will be determined later. The inequality clearly holds for w i = q n . To verify that it holds for smaller values of w i as well, we now show that

∂ ∂w i tanh w i µ -tanh q n µ w i q n -c(q n -w i ) < 0

which will ensure that it holds for all w i . We define s 2 = 1 -||w|| 2 + w 2 i and denote q n = s 2 -w 2 i to extract the w i dependence, giving

∂ ∂w i tanh w i µ -tanh q n µ w i q n -c(q n -w i ) = 1 µ sech 2 wi µ + 1 µ sech 2 √ s 2 -w 2 i µ w 2 i s 2 -w 2 i -tanh √ s 2 -w 2 i µ s 2 (s 2 -w 2 i ) 3/2 + c( wi √ s 2 -w 2 i + 1) ≤ 4 µ e -2 w i µ + e -2 √ s 2 -w 2 i µ -tanh √ s 2 -w 2 i µ s 2 (s 2 -w 2 i ) 3/2 + 2c

Where in the last inequality we used properties of the sech function and q n ≥ w i . We thus want to show

EQUATION

we have

EQUATION

n and it follows that 10 holds. For µ < 1 16 we are guaranteed that c > 0.

From examining the RHS of 10 (and plugging in q n = s 2 -w 2 i ) we see that any lower bound on the gradient of an element w j applies also to any element |w i | ≤ |w j |. Since for |w j | = ||w|| ∞ we have q n -w j = w j ζ, for every log( 1 µ )µ ≤ w i we obtain the bound

u (i) * grad[f Sep ](q(w)) ≥ c w ∞ ζ

Proof of Theorem 1: (Gradient descent convergence rate for separable function).

We obtain a convergence rate by first bounding the number of iterations of Riemannian gradient descent in C ζ0 \C 1 , and then considering

C 1 \B ∞ r . From Lemma 16 we obtain C ζ0 \C 1 ⊆ C ζ0 \B ∞ 1/ √ n+3 .

Choosing c 2 so that µ < 1 2 , we can apply Lemma 2, and for u defined in 7, we thus have

|w i | > µ log( 1 µ ) ⇒ u (i) * grad[f Sep ](q(w)) > c||w|| ∞ ζ 0 .

Since from Lemma 7 the Riemannian gradient norm is bounded by

√ n, we can choose c 1 , c 2 such that µ log( 1 µ ) < 1 2 √ n+3 , η < 1 6 √ n 2 +3n

. This choice of η then satisfies the conditions of Lemma 17 with r = µ log

( 1 µ ), b = 1 √ n+3 , M =

√ n, which gives that after a gradient step

EQUATION

for some suitably chosen c > 0. If we now define by w (t) the t-th iterate of Riemannian gradient descent and

ζ (t) ≡ q (t) n w (t) ∞ -1, ζ (0) ≡ ζ 0 , for iterations such that w (t) ∈ C ζ \C 1 we find ζ (t) ≥ ζ (t-1) (1 + cη) ≥ ζ 0 (1 + cη)

t and the number of iterations required to exit C ζ0 \C 1 is

EQUATION

)

To bound the remaining iterations, we use Lemma 2 to obtain that for every

w ∈ C ζ0 \B ∞ r , grad[f Sep ](q(w)) 2 ≥ u (i) * grad[f Sep ](q(w)) 2 ||u (i) || 2 ≥ ζ 2 0 c 2 r 2

where we have used ||u (i) || 2 = 1 +

EQUATION

)

Choosing η < 1 2L where L is the gradient Lipschitz constant of f s , from Lemma 5 we obtain

2 f Sep (q (0) ) -f * Sep η > T -1 i=0 grad[f Sep ](q (i) ) 2 .

According to Lemma B, L = 1/µ and thus the above holds if we demand η < µ 2 . Combining 12 and 13 gives

T < 2 f Sep (q (0) ) -f * Sep ηc 2 r 2 + 1 - ζ 2 0 (n+3)r 2 log( 1 ζ0 ) log(1 + cη)

.

To obtain the final rate, we use in g(w 0 ) -g * ≤ √ n and cη < 1

EQUATION

From Lemma 1 the ball B ∞ r contains a global minimizer of the objective, located at the origin.

The probability of initializing in:

Ȃ C ζ0 is simply given from Lemma 3 and by summing over the 2n possible choices of C ζ0 , one for each global minimizer (corresponding to a single signed basis vector).

Lemma 5 (Riemannian gradient descent iterate bound). For a Riemannian gradient descent algorithm on the sphere with step size t k < 1 2L , where L is a lipschitz constant for ∇f (q), one has

f (q 1 ) -f (q ) ≥ f (q 1 ) -f (q T ) ≥ t k 2 grad [f ] (q k ) 2 .

Proof. Just as in the euclidean setting, we can obtain a lower bound on progress in function values of iterates of the Riemannian gradient descent algorithm from a lower bound on the Riemannian gradient. Consider f : S n-1 → R, which has L-lipschitz gradient. Let q k denote the current iterate of Riemannian gradient descent, and let t k > 0 denote the step size. Then we can form the Taylor approximation to f • Exp q k (v) at 0 q k :

f : B 1 (0 q k ) ∩ T q k S n-1 → R : v → f (q k ) + v, ∇f (q k ) .

From Taylor's theorem, we have for any

v ∈ B 1 (0 q k ) ∩ T q k S n-1 | f (v) -f • Exp q k (v)| ≤ 1 2 Hess[f ](q k ) v -0 q k 2 ,

where the matrix norm is the operator norm on R n×n . Using the gradient-lipschitz property of f , we readily compute

Hess[f ](q k ) ≤ ∇ 2 f (q k ) + | ∇f (q k ), q k | ≤ 2L,

since ∇f (0) = 0 and q k ∈ S n-1 . We thus have

f • Exp q k (v) ≤ f (q k ) + v, ∇f (q k ) + L v 2 .

If we put v = -t k grad[f ](q k ) and write q k+1 = Exp q k (-t k grad [f ] (q k )), the previous expression becomes

f (q k+1 ) ≤ f (q k ) -t k grad [f ] (q k ) 2 + t 2 k L grad [f ] (q k ) 2 ≤ f (q k ) - t k 2 grad [f ] (q k ) 2 if t k < 1 2L

. Thus progress in objective value is guaranteed by lower-bounding the Riemannian gradient.

B.1 Bounds for E sech 2 (Y ):

We already have a lower bound in Lemma 20 of (38) that we can use for the second term, so we need an upper bound for the first term. Following from p. 865, we define Y ∼ N (0, σ 2 Y ) , Z = exp -2Y µ , and defining β = 1 -1 √ T for some T > 1 we have

sech 2 (Y /µ) = 4Z (1 + Z) 2 ≤ 4Z (1 + βZ) 2 = ∞ k=0 b k Z k+1

Where b k = (-β) k (k + 1). Using B.3 from Lemma 40 in (38) we have

E ∞ k=0 b k Z k+1 1 Y >0 = ∞ k=0 b k E e -2(k+1)Y /µ 1 Y >0 = ∞ k=0 b k exp 1 2 2(k + 1) µ 2 σ 2 Y Φ c 2(k + 1) µ σ Y

Where Φ c (x) is the complementary Gaussian CDF (The exchange of summation and expectation is justified since Y > 0 implies Z ∈ [0, 1], see proof of Lemma 18 in (38) for details).

Using the following bounds 1

√ 2π 1 x -1 x 3 e -x 2 /2 ≤ Φ c (x) ≤ 1 √ 2π 1 x -1 x 3 + 3

x 5 e -x 2 /2 by applying the upper (lower) bound to the even (odd) terms in the sum, and then adding a non-negative quantity, we obtain

≤ 1 √ 2π ∞ k=0 (-β) k (k + 1)    1 2(k+1) µ σ Y - 1 2(k+1) µ σ Y 3    + 1 √ 2π ∞ k=0 β k (k + 1)    3 2(k+1) µ σ Y 5    and using ∞ k=0 (-β) k = 1 1+β , ∞ k=0 b k (k+1) 3 ≥ 0, ∞ k=0 |b k | (k+1) 5 ≤ 2 (from Lemma 17 in (38)) and taking T → ∞ so that β → 1 we have ∞ k=0 b k E Z k+1 1 Y >0 ≤ 1 2 √ 2π 1 2 µ σ Y + 1 √ 2π 6 2 µ σ Y 5

giving the upper bound

EQUATION

Y while the lower bound (Lemma 20 in (38)) is

2 π µ σ Y - 2µ 3 √ 2πσ 3 Y - 3µ 5 2 √ 2πσ 5 Y ≤ E sech 2 (Y )

B.2 Gradient bounds:

After conditioning on J \{n, i} the variables X + q n v n , X + q i v i are Gaussian. We can thus plug the bounds into 16 to obtain

u (i) * grad[f pop DL ](q(w)) ≥ 2 π w i θ(1 -θ) * E J \{n,i}    1 √ σ 2 X +w 2 i - µ 2 (σ 2 X +w 2 i ) 3/2 - 3µ 4 4(σ 2 X +w 2 i ) 5/2 - 1 √ σ 2 X +q 2 n - 3µ 4 4(σ 2 X +q 2 n ) 5/2    ≥ 2 π w i θ(1 -θ)    E J \{n,i} √ σ 2 X +q 2 n - √ σ 2 X +w 2 i √ σ 2 X +q 2 n √ σ 2 X +w 2 i -µ 2 w 3 i -3µ 4 2w 5 i   

the term in the expectation is positive since

q n > ||w|| ∞ (1 + ζ) > w i giving ≥ 2 π w i θ(1 -θ)    E J \{n,i} σ 2 X + q 2 n -σ 2 X + w 2 i -µ 2 w 3 i -3µ 4 2w 5 i   

. To extract the ζ dependence we plug in q n > w i (1 + ζ) and develop to first order in ζ (since the resulting function of ζ is convex) giving

≥ 2 π w i θ(1 -θ)    E J \{n,i} w 2 i ζ √ σ 2 X +w 2 i -µ 2 w 3 i -3µ 4 2w 5 i    ≥ 2 π θ(1 -θ) w 3 i ζ - µ 2 w 2 i - 3µ 4 2w 4 i

Given some ζ and r such that w i > r, if we now choose µ such that µ <

√ 1+ 3 4 r 3 ζ-1 3

r we have the desired result. This can be achieved by requiring µ < c 1 r 5/2 √ ζ for a suitably chosen c 1 > 0.

Lemma 8 (Point-wise concentration of projected gradient). For u (i) defined in 7, the gradient of the objective 1 obeys

P u (i) * grad[f DL ](q) -E u (i) * grad[f DL ](q) ≥ t ≤ 2 exp - pt 2 4 + 2 √ 2t

Proof of Lemma 8: (Point-wise concentration of projected gradient). If we denote by x i a column of the data matrix with entries x i j ∼ BG(θ), we have

u (i) * grad[f DL ](q(w)) = 1 p p k=1 tanh q * (w)x k µ x k i -x k n w i q n ≡ 1 p p k=1 Z k . Since tanh(x) is bounded by 1, |Z k | ≤ x k i -x k n w i q n ≡ u T x k

. Invoking Lemma 21 from (38) and u 2 = 1 +

EQUATION

and using Lemma 36 in (38) with R = √ 2, σ = √ 2 we have

P [|∇g DL (w) i -E [∇g DL (w) i ]| ≥ t] ≤ 2 exp - pt 2 4 + 2 √ 2t

Lemma 9 (Projection Lipschitz Constant). The Lipschitz constant for

u (i) * grad[f DL ](q(w)) is L = 2 √ n X ∞ X ∞ µ + 1

Proof of Lemma 9: (Projection Lipschitz Constant). We have

|u (j) * grad[f DL ](q(w)) -u (j) * grad[f DL ](q(w ))| = 1 p p i=1   tanh( q * (w)x i µ ) x i j - x i n qn(w) w j -tanh( q * (w )x i µ ) x i j - x i n qn(w ) w j   ≡ 1 p p i=1 tanh( q * (w)x i µ )s(w) -tanh( q * (w )x i µ )s(w )

where we have defined s(w) = x i j -xn qn(w) w j . Using q(w), q(w ) ∈ C ⇒ q n (w), q n (w )

≥ 1 2 √ n we have |s(w) -s(w )| = x i n w j q n (w) - w j q n (w ) ≤ |x n | 2 √ n w -w

Lemma 25 in (38) gives

tanh( q * (w)x µ ) -tanh( q * (w )x µ ) ≤ 2 √ n µ x w -w

We also use the fact that tanh is bounded by 1 and s(w) is bounded by X ∞ . We can then use Lemma 23 in (38) to obtain

|u (j) * grad[f DL ](q(w)) -u (j) * grad[f DL ](q(w ))| ≤ 2 √ n p p i=1 ( 1 µ x i 2 ∞ + x i ∞ ) w -w ≤ 2 √ n X ∞ X ∞ µ + 1 w -w we thus have L = 2 √ n X ∞ X ∞ µ + 1 .

Lemma 10 (Uniformized gradient fluctuations). For all w ∈ C ζ , i ∈ [n], with probability P > P y we have

u (i) * grad[f DL ](q(w)) -E u (i) * grad[f DL ](q(w)) ≤ y(θ, ζ)

where

P y ≡ 2 exp    -1 4 py(θ,ζ) 2 4+ √ 2y(θ,ζ) + log(n) +n log 48 √ n 4 log(np) µ + √ log(np) y(θ,ζ)    Proof: B

Proof of Lemma 10:(Uniformized gradient fluctuations). For X ∈ R n×p with i.i.d. BG(θ) entries, we define the event E ∞ ≡ {1 ≤ X ∞ ≤ 4 log(np)}. We have

P[E c ∞ ] ≤ θ(np) -7 + e -0.3θnp For any ε ∈ (0, 1) we can construct an ε-net N for C ζ \B 2 1/20 √ 5(n-1) (0) with at most (3/ε) n points. Using Lemma 9, on E ∞ , grad[f DL ](q) i is L-Lipschitz with L = 8 √ n 4 log(np) µ + log(np) . If we choose ε = y(θ,ζ) 2L we have |N | ≤ ( 6L y(θ, ζ) ) n

. We then denote by E g the event

max w∈N,i∈[n] u (i) * grad[f DL ](q(w)) -E u (i) * grad[f DL ](q(w)) ≤ y(θ, ζ) 2 and obtain that on E g ∩ E ∞ sup w∈C ζ ,i∈[n] |∇g DL (w) i -E [∇g DL (w) i ]| ≤ y(θ, ζ) . Setting t = b(θ)

2 in the result of Lemma 8 gives that for all

w ∈ C ζ , i ∈ [n], P u (i) * grad[f DL ](q(w)) -E u (i) * grad[f DL ](q(w)) ≥ y(θ, ζ) 2 ≤ 2 exp - 1 4 py(θ, ζ) 2 4 + 2 √ 2y(θ, ζ)

and thus

P E c g ≤ 2 exp   -1 4 py(θ,ζ) 2 4+ √ 2y(θ,ζ) 2 +n log 6L b(θ) + log(n)  

Lemma 11 (Gradient descent convergence rate for dictionary learning -population). For any 1 > ζ 0 > 0 and s > µ 4 √ 2 , Riemannian gradient descent with step size η < c2s n on the dictionary learning population objective 8 with µ < c4 √ ζ0 n 5/4 , θ ∈ (0, 1 2 ), enters a ball of radius c 3 s from a target solution in

T < C 1 ηθ 1 s + n log 1 ζ 0 iterations with probability P ≥ 1 -2 log(n)ζ 0 where the c i , C i are positive constants.

Proof of Lemma 11: (Gradient descent convergence rate for dictionary learning -population). The rate will be obtained by splitting C ζ0 into three regions. We consider convergence to B 2 s (0) since this set contains a global minimizer. Note that the balls in the proof are defined with respect to w.

B.3 C ζ0 \B 2 1/20 √ 5 (0)

The analysis in this region is completely analogous to that in the first part of the proof of Lemma 1. For every point in this set we have

w ∞ > 1 20 5(n -1)

. From Lemma 16 we know that n-1

(2+ζ (t) )ζ (t) +n < 1 20 √ 5 ⇒ w (t) ∈ B 2 1/20 √ 5 (0) hence in this set ζ < 8. If we choose r = 1 40 √ 5(n-1)

, since for every point in this region r 3 ζ < 1, we have . We now require η < , M = √ θn (since the maximal norm of the Riemannian gradient is √ θn from Lemma 12) , obtaining that at every iteration in this region

r 5/2 √ ζ 2 √ 3 < √ 1+ 3 4 r 3 ζ-1 3 r = z(

ζ ≥ ζ 1 + √ nc DL 2(8000(n -1)) 3/2 η

and the maximal number of iterations required to obtain ζ > 8 and exit this region is given by

t 1 = log(8/ζ 0 ) log 1 + √ nc DL 2(8000(n-1)) 3/2 η (17) B.4 B 2 1/20 √ 5 (0)\B 2 s (0)

According to Proposition 7 in (38) , which we can apply since s ≥ µ 4 √ 2 , µ < 9 50 , in this region we have

w * ∇ w g pop DL (w) w ≥ cθ

A simple calculation shows that ∇ w g pop DL (w) = ∂ϕ ∂w * grad[f pop DL ](q(w)) where ϕ is the map defined in 3, and thus

EQUATION

. Defining h(q) = w 2 2 , and denoting by q an update of Riemannian gradient descent with step size η, we have (using a Lagrange remainder term)

h(q ) = h(q) + ∂h(q ) ∂η η + η 0 dt ∂ 2 h(q ) ∂η 2 η=t (η -t) ≡R = w 2 2 -grad[f pop DL ](q), ∂h(q) ∂q + R

where in the last line we used q = cos(gη)q -sin(gη)

grad[f pop DL ](q) g

where g ≡ grad[f pop DL ](q) . Since grad[f pop DL ](q), ∂h(q) ∂q = grad[f pop DL ](q), (I -qq * ) ∂h(q) ∂q and

(I -qq * ) ∂h(q) ∂q = (I -qq * ) w -q n = w -q n -( w 2 -q 2 n )q = 2(1 -w 2 ) w -w 2 qn

we obtain (using 18)

w 2 2 = w 2 2 + 2(1 -w 2 )η grad[f pop DL ](q), w -w 2 qn + R < w 2 2 -2(1 -w 2 ) w θcη + R It remains to bound R. Denoting r = w -q n * grad[f ](q) we have ∂ 2 h(q ) ∂η 2 η=t = ∂q ∂η * ∂ 2 h(q) ∂q∂q ∂q ∂η η=t + ∂h(q) ∂q * ∂ 2 q ∂η 2 η=t = cos 2 (gt) grad[f pop DL ](q) 2 -grad[f pop DL ](q) 2 n +g 2 sin 2 (gt) -cos(gt) w 2 -q 2 n +g sin(gt)r(1 + 2 cos(gt)) hence for some C > 0, if grad[f pop DL ](q) < M we have R < CM 2 η 2

and thus choosing η < (1-w 2 ) w θc CM 2 we find

w 2 < w 2 -2(1 -w 2 ) w cθη

and in our region of interest w 2 < w 2 -csθη for some c > 0 and thus summing over iterations, we obtain for some C2 > 0

EQUATION

From Lemma 12, M = √ θn and thus with a suitably chosen c 2 > 0, η < c2s n satisfies the above requirement on η as well as the previous requirements, since θ < 1.

B.5 Final rate and distance to minimizer:

Combining these results gives, we find that when initializing in C ζ0 , the maximal number of iterations required for Riemannian gradient descent to enter B 2 s (0) is

T ≤ t 1 + t 2 < C 1 ηθ n log 1 ζ 0 + 1 s

for some suitably chosen C 1 , where t 1 , t 2 are given in 17,19. The probability of such an initialization is given by the probability of initializing in one of the 2n possible choices of C ζ , which is bounded in Lemma 3.

Once w ∈ B 2 s (0), the distance in R n-1 between w and a solution to the problem (which is a signed basis vector, given by the point w = 0 or an analog on a different symmetric section of the sphere) is no larger than s, which in turn implies that the Riemannian distance between ϕ(w) and a solution is no larger than c 3 s for some c 3 > 0. We note that the conditions on µ can be satisfied by requiring µ < c4

√ ζ0 n 5/4 .

Lemma 12 (Dictionary learning gradient upper bound). The dictionary learning population gradient obeys

∇ w g pop DL (w) ≤ √ 2θn grad[f pop DL ](q) ≤ √ θn

while in the finite sample case

∇ w g DL (w) 2 ≤ √ 2n X ∞ grad[f DL ](q) ≤ √ n X ∞

where X is the data matrix with i.i.d. BG(θ) entries.

Proof. Denoting x ≡ (x, x n ) we have

∇ w g pop DL (w) 2 = E tanh q * x µ x -x n w q n 2

and using Jensen's inequality, convexity of the L 2 norm and the triangle inequality to obtain

≤ E tanh q * x µ x 2 + tanh q * x µ x n w q n 2 ≤ E x 2 + x n w q n 2 ≤ 2θn while grad[f pop DL ](q) ≤ ∇f pop DL (q) = E tanh q * x µ x ≤ √ θn

Similarly, in the finite sample size case one obtains

∇ w g DL (w) 2 ≤ 1 p p i=1 x i 2 + x i n w q n 2 ≤ 2n X 2 ∞ grad[f DL ](q) ≤ 1 p p i=1 tanh q * x i µ x i ≤ √ n X ∞

Proof of Theorem 2: (Gradient descent convergence rate for dictionary learning). The proof will follow exactly that of Lemma 11, with the finite sample size fluctuations decreasing the guaranteed change in ζ or ||w|| at every iteration (for the initial and final stages respectively) which will adversely affect the bounds.

B.6 C ζ0 \B 2 1/20 √ 5 (0)

To control the fluctuations in the gradient projection, we choose

y(θ, ζ 0 ) = ζ 0 c DL 2(8000(n -1)) 3/2 which can be satisfied by choosing y(θ, ζ 0 ) = c7θ(1-θ)ζ0 n 3/2

for an appropriate c 7 > 0 . According to Lemma 10, with probability greater than P y we then have

u (i) * grad[f DL ](q(w)) -E u (i) * grad[f DL ](q(w)) ≤ y(θ, ζ)

With the same condition on µ as in Lemma 11, combined with the uniformized bound on finite sample fluctuations, we have that at every point in this set

u (i) * grad[f pop DL ](q(w)) ≥ c DL 2(8000(n -1)) 3/2

. According to Lemma 12 the Riemannian gradient norm is bounded by

M = √ n X ∞ . Choosing r, b as in Lemma 11, we require η < 1 360 X ∞ √ 5n(n-1) = b-r

3M and obtain from Lemma 17

EQUATION

B.7 B 2 1/20 √ 5 (0)\B 2 s (0)

From Theorem 2 in (38) there are numerical constants c b , c such that in this region

w * ∇ w g DL (w) w = w * ∂ϕ ∂w * grad[f ](q(w)) w ≥ c θ with probability P > 1 -c b p -6

. Following the same analysis as in Lemma 11, since from Lemma 12 the norm of the gradient gradient is bounded by

√ n||X|| ∞ we require η < (1-w 2 ) w θc Cn||X|| 2 ∞ which is satisfied by requiring η < cθs n||X|| 2 ∞

for some chosen c > 0. We then obtain

EQUATION

for a suitably chosen C 2 > 0.

B.8 Final rate and distance to minimizer:

The final bound on the rate is obtained by summing over the terms for the three regions as in the population case, and convergence is again to a distance of less than c 3 s from a local minimizer. The probability of achieving this rate is obtained by taking a union bound over the probability of initialization in C ζ0 (given in Lemma 3) and the probabilities of the bounds on the gradient fluctuations holding (from Lemma 10 and ( 38)). Note that the fluctuation bound events imply by construction the event E ∞ = {1 ≤ X ∞ ≤ 4 log(np)} hence we can replace X ∞ in the conditions on η above by 4 log(np). The conditions on η, µ can be satisfied by requiring η < c5θs n log np , µ < c6 √ ζ0 n 5/4 for suitably chosen c 5 , c 6 > 0. The bound on the number of iterations can be simplified to the form in the theorem statement as in the population case.

C Generalized Phase Retrieval:

We show below that negative curvature normal to stable manifolds of saddle points in strict saddle functions is a feature that is found not only in dictionary learning, and can be used to obtain efficient convergence rates for other nonconvex problems as well, by presenting an analysis of generalized phase retrieval that is along similar lines to the dictionary learning analysis. We stress that this contribution is not novel since a more thorough analysis was carried out by (8) . The resulting rates are also suboptimal, and pertain only to the population objective.

Generalized phase retrieval is the problem of recovering a vector x ∈ C n given a set of magnitudes of projections y k = |x * a k | onto a known set of vectors a k ∈ C n . It arises in numerous domains including microscopy (27) , acoustics (2), and quantum mechanics (10) (see (33) for a review). Clearly x can only be recovered up to a global phase. We consider the setting where the elements of every a k are i.i.d. complex Gaussian, (meaning (a k ) j = u + iv for u, v ∼ N (0, 1/ √ 2)). We analyze the least squares formulation of the problem (7) given by

min z∈C n f (z) = 1 2p p k=1 y 2 k -|z * a k | 2 2

.

Taking the expectation (large p limit) of the above objective and organizing its derivatives using Wirtinger calculus (22) , we obtain

E[f ] = x 4 + z 4 -x 2 z 2 -|x * z| 2 (22) ∇E[f ] = ∇ z E[f ] ∇ z E[f ] =   (2 z 2 -x 2 )I -xx * z (2 z 2 -x 2 )I -xx T z   .

For the remainder of this section, we analyze this objective, leaving the consideration of finite sample size effects to future work.

C.1 The geometry of the objective:

In (37) it was shown that aside from the manifold of minima Ȃ ≡ xe iθ , the only critical points of E[f ] are a maximum at z = 0 and a manifold of saddle points given by

EQUATION

where W ≡ {z|z * x = 0}. We decompose z as

EQUATION

where ζ > 0, w ∈ W . This gives z 2 = w 2 + ζ 2 . The choice of w, ζ, φ is unique up to factors of 2π in φ, as can be seen by taking an inner product with x. Since the gradient decomposes as follows:

EQUATION

the directions e iφ x x , w w are unaffected by gradient descent and thus the problem reduces to a two-dimensional one in the space (ζ, w ). Note also that the objective for this twodimensional problem is a Morse function, despite the fact that in the original space there was a manifold of saddle points. It is also clear from this decomposition of the gradient that the stable manifolds of the saddles are precisely the set W .

It is evident from 24 that the dispersive property does not hold globally in this case. For z / ∈ B ||x|| we see that gradient descent will cause ζ to decrease, implying positive curvature normal to the stable manifolds of the saddles. This is a consequence of the global geometry of the objective. Despite this, in the region of the space that is more "interesting", namely B ||x|| , we do observe the dispersive property, and can use it to obtain a convergence rate for gradient descent.

We define a set that contains the regions that feeds into small gradient regions around saddle points within B ||x|| by

Q ζ0 ≡ {z(ζ, w )|ζ ≤ ζ 0 }.

We will show that, as in the case of orthogonal dictionary learning, we can both bound the probability of initializing in (a subset of) the complement of Q ζ0 and obtain a rate for convergence of gradient descent in the case of such an initialization. 9We now define four regions of the space which will be used in the analysis of gradient descent:

S 1 ≡ z z 2 ≤ 1 2 x 2 S 2 ≡ z 1 2 x 2 < z 2 ≤ (1 -c) x 2 S 3 ≡ z (1 -c) x 2 < z 2 ≤ x 2 S 4 ≡ z x 2 < z 2 ≤ (1 + c) x 2

defined for some c < 1 4 . These are shown in Figure 4 . We now define

EQUATION

and using 24 obtain

EQUATION

)

Figure 4 : The projection of the objective of generalized phase retrieval on the ( ζ x , w x ) plane. The full red curves are the boundaries between the sets S 1 , S 2 , S 3 , S 4 used in the analysis. The dashed red line is the boundary of the set Q ζ0 that contains small gradient regions around critical points that are not minima. The maximizer and saddle point are shown in dark green, while the minimizer is in pink.

These are used to find the change in ζ, w at every iteration in each region:

EQUATION

w ≤ w (27d)

On S 3 :

EQUATION

On S 4 : 

1 -(1 + 2c)η x 2 w ≤ w ≤ 1 -η x 2 w (27g) (1 -2cη x 2 )ζ ≤ ζ ≤ ζ (27h) C.2

+ log(2c) log( 4 √ 7 ) log(1-(1-2c)η x 2 ) log(1+2cη x 2 )

iterations with probability

P ≥ 1 - 8 π erf √ 2n x ζ , Proof. Please see Appendix C.3.

We find that in order to prevent the failure probability from approaching 1 in a high dimensional setting, if we assume that x does not depend on n we require that ζ scale like 1

√ n . This is simply the consequence of the well-known concentration of volume of a hypersphere around the equator. Even with this dependence the convergence rate itself depends only logarithmically on dimension, and this again is a consequence of the logarithmic dependence of ζ due to the curvature properties of the objective. Lemma 13. For any iterate z of gradient descent on 22, assuming η < √ c 4 x 2 , c < 1 4 and defining ζ as in 25, we have i)

EQUATION

Proof of Lemma 13. i) From 27 we see that in 

EQUATION

and since ζ 2 ≥ 0 this implies

1 + εη x 2 (2 -ε) (1 -ε) ≥ 1

by considering the product of these two factors, this in turn implies

1 2b (2 -ε) ≥ η x 2 (2 -ε) ≥ 1

where we have used η < 

ζ 2 ≥ x 2 2 (1 -2η x 2 c) 2 ≥ x 2 2 (1 -4η x 2 c) ≥ (1 - 1 2b ) x 2 2

where in the last inequality we used c < 

) z ∈ 4 i=2 S i ⇒ z ∈ 4 i=2 S i ii) z ∈ S 1 ⇒ z ∈ S 1 ∪ S 2

Proof of Lemma 14. We use the fact that for the next iterate we have

EQUATION

We will also repeatedly use η <

√ c b x 2 , c < 1 4 and z ∈ 4 i=1 S i ⇒ w 2 ≤ x 2 2

which is a shown in Lemma 13.

EQUATION

1) We have z ∈ S 3 ⇒ z 2 = (1 -ε) x 2 for some ε ≤ c and using 28 we must show

x 2 2 ≤ 1 -η x 2 (1 -2ε) 2 w 2 + 1 + 2η x 2 ε 2 ζ 2 or equivalently A ≡ ε - x ≤ η x 2   -2(1 -2ε) + (1 -2ε) 2 η x 2 w 2 +4 ε + ε 2 η x 2 ζ 2   ≡ B and using η < √ c b x 2 , c < 1 4 -x 2 b < -2 x √ c b < -2η x 4 ≤ B

while on the other hand

EQUATION

thus picking b = 4 guarantees the desired result.

2) By a similar argument, z

2 ≤ (1 + c) x 2 is equivalent to A ≡ η x 2   -2(1 -2ε) + η x 2 (1 -2ε) 2 w 2 +4 ε + η x 2 ε 2 ζ 2   ≤ x 2 (c + ε) ≡ B . Since w 2 ≤ x 2 2 and z 2 ≤ x 2 ⇒ ζ 2 ≤ x 2 2 we obtain A ≤ η η x 4 + 4 x 2 ε + η x 4 ε 2 x 2 2 < 1 2b 1 b + 2 1 + 1 8b c x 2 .

If we choose b = 4 we thus have A < B which implies

z 2 < (1 + c) x 2 C.5 z ∈ S 4 ⇒ z ∈ 4 i=2 S i We have z ∈ S 4 ⇒ z 2 = w 2 + ζ 2 = (1 + ε) x 2 for some ε ≤ c . 1) x 2 2 < z 2 is equivalent to A ≡ -(ε + 1 2 ) x 2 ≤ η x 2   -4(1 + 2ε) + η x 2 (1 + 2ε) 2 w 2 +4 -ε + η x 2 ε 2 ζ 2   ≡ B . We have B ≥ -4η x 2 (1 + 2ε) w 2 + εζ 2 ≥ - 15 8b x 2

where the last inequality used w

2 ≤ x 2 2 and z 2 ≤ x 2 (1 + c) ⇒ ζ 2 ≤ x 2 ( 1 2 + c).

The choice b = 4 gaurantees A ≤ B which ensures the desired result.

2) This is trivial since z 2 ≤ (1 + c) x 2 and in S 4 both ζ and w decay at every iteration (ref eq).

C.6 z ∈ S 2 ⇒ z ∈ 4 i=2 S i 1) We use z ∈ S 2 ⇒ z 2 = w 2 + ζ 2 = ( 1 2 + ε) x 2 for some ε ≤ 1 2 -c

. Using a similar argument as in the previous section, we are required to show

-ε x 2 < η x 2   4 -ε + ε 2 η x 2 w 2 + 2(1 -2ε) + (1 -2ε) 2 η x 2 ζ 2   ≡ B where B ≥ -ε x 2 b

implies that b = 4 gives the desired result.

2) The condition is equivalent to

A ≡ η x 2   4 -ε + ε 2 η x 2 w 2 + 2(1 -2ε) + (1 -2ε) 2 η x 2 ζ 2   + ε x 2 ≤ ( 1 2 + c) x 2 ≡ B

One can show by looking for critical points of A(ε) in the range 0 ≤ ε ≤ 1 2 that A is maximized at ε = 0, since there is only one critical point at ε

EQUATION

and in both cases b = 4 ensures A ≤ B.

C.7 z ∈ S 1 ⇒ z ∈ S 1 ∪ S 2 We must show z ≤ (1 -c) x 2 using z 2 = (1 -ε) x 2 2 for 0 ≤ ε ≤ 1. z 2 = 1 + εη x 2 2 w 2 + 1 + 2(ε + 1)η x 2 2 2 ζ 2 A ≡ η x 2   2ε + ε 2 η x 2 w 2 + 2(ε + 1) + (ε + 1) 2 η x 2 4 ζ 2   -ε x 2 ≤ ( 1 2 -c) x 2 ≡ B and since A ≤ 1 2b 2 + 1 b x 2 2 and B ≥ x 2 4

once again b = 4 suffices to obtain the desired result.

Lemma 15. For z parametrized as in 23,

w 2 < c x 2 ∨ ζ 2 > (1 -c) x 2 ⇒ dist(z, Ȃ) < √ 5c x

which gives the desired result. In particular, if we choose c = 1 35 we converge to dist 2 (z, Ȃ) <

x 2

7 , a region which is strongly convex according to (38) .

Proof of Theorem 3: (Gradient descent convergence rate for generalized phase retrieval)

. We now bound the number of iterations that gradient descent, after random initialization in S 1 , requires to reach a point where one of the convergence criteria detailed in Lemma 15 is fulfilled. From Lemma 14, we know that after initialization in S 1 we need to consider only the set 

C.7.1 Iterations in S 1:

Assuming we initialize with some ζ = ζ 0 . Then the maximal number of iterations in this region is

ζ 0 (1 + η x 2 ) t1 = x √ 2 t 1 = log x ζ0 √ 2 log(1 + η x 2 )

since after this many iterations z 

S i:

The convergence criteria are w 2 < c x 2 or ζ 2 > (1 -c) x 2 .

After exiting S 1 and assuming the next iteration is in S 2 , the maximal number of iterations required to reach S 3 ∪ S 4 is obtained using

ζ ≥ (1 + 2η x 2 c)ζ

and is given by The only concern is that after an iteration in S 3 ∪ S 4 the next iteration might be in S 2 .

x √ 2 (1 + 2η x 2 c) t2 = (1 -c) x

To account for this situation, we find the maximal number of iterations required to reach S 3 ∪ S 4 again. This is obtained from the bound on ζ in Lemma 13.

Using this result, and the fact that for every iteration in S 

V ζ = (1 + ζ) √ π ∞ 0 e -(1+ζ) 2 x 2 erf n-1 (x)dx

This integral admits no closed form solution but one can construct a linear approximation around small ζ and show that it is convex. Thus the approximation provides a lower bound for V ζ and an upper bound on the failure probability.

From symmetry considerations the zero-order term is V 0 = 1 2n . The first-order term is given by

∂V ζ ∂ζ ζ=0 = 1 n - 2 √ π ∞ 0

x 2 e -x 2 erf n-1 (x)dx

We now require an upper bound for the second integral since we are interested in a lower bound for V ζ . We can express it in terms of the second moment of the L ∞ norm of a Gaussian vector as follows:

1 √ π ∞ 0 x 2 e -x 2 erf n-1 (x) = 1 √ π ∞ 0 x 2 e -x 2 Π i 1 √ π x -x e -t 2 i dt i dx = 1 √ 2π ∞ 0 x 2 2 e -x 2 /2 Π i 1 √ 2π x -x e -t 2 i /2 dt i dx = 1 4n X 2 ∞ dµ(X) = 1 4n Var [ X ∞ ] + (E [ X ∞ ]) 2

where µ(X) is the Gaussian measure on the vector X ∈ R n . We can bound the first term using

Var [ X ∞ ] ≤ max i Var [|X i |] = Var [|X i |] < Var [X i ] = 1

To bound the second term, we use the fact that for a standard Gaussian vector X (X i ∼ N (0, 1)) and any λ > 0 we have Combining these bounds, the leading order behavior of the gradient is

exp (λE [ X ∞ ]) ≤ E exp λmax i |X i | ≤ E i exp (λ |X i |) = nE [exp (λ |X i |)]

∂V ζ ∂ζ ζ=0 ≥ 3 -4 log(2n) 4n ≥ - log(n) n .

This linear approximation is indeed a lower bound, since using integration by parts twice we have where the last inequality holds for any n > 2 since the integrand is non-negative everywhere. This gives Proof. Given the surface of some L ∞ ball for w , we can ask what is the minimal ζ such that ∂C ζm intersects this surface. This amounts to finding the minimal q n given some w ∞ . Yet this is clearly obtained by setting all the coordinates of w to be equal to w ∞ (this is possible since we are guaranteed

∂ 2 V ζ ∂ζ 2 = 1 √ π ∞ 0 e -(1+ζ) 2 x 2 -6(1 + ζ)x 2 +4(1 + ζ) 3 x 4 erf n-1 (x)dx = - 2(n -1) π ∞ 0 e -(1+ζ) 2 x 2 1 -2(1 + ζ) 2 x 2 e -x 2

V ζ ≥ 1 2n - log(n) n ζ Lemma 16. B ∞ s(ζ) (0) ⊆ C ζ ⊆ B 2

q n ≥ w ∞ ⇒ w ∞ ≤ 1 √ n ), giving 1 -(n -1) w 2 ∞ w ∞ = 1 + ζ m w ∞ = 1 (1 + ζ m ) 2 + n -1

:

As well as a lack of spurious local minimizers, and the existence of large gradients or strong convexity in the remaining parts of the space

This statement is nontrivial: finding a local minimum of a smooth function is NP-hard.

This work also proves convergence to a second-order stationary point under more general smoothness assumptions.

This problem exhibits a sign permutation symmetry: A0X0 = (A0Γ)(Γ * X0) for any signed permutation matrix Γ. Hence, we only ask for recovery up to a signed permutation.

[X0]ij = V ij Ωij, with V ij ∼ N (0, 1), Ωij ∼ Bern(θ) independent.

Combining with a deflation strategy, one can then efficiently recover the entire dictionary A0.

Strictly speaking, fDL is Morse with high probability, due to results of(38).

The direction of this negative curvature is important here, and it is this feature that distinguishes these problems from other problems in the strict-saddle class where this direction may be arbitrary

Q ζ 0 is equivalent to the complement of the set C ζ used in the analysis of the separable objective and dictionary learning.

(in R 2n , using the natural identification of C n with R 2n ). This volume can be calculated by integrating over 2n -1 dimensional balls of varying radius.

Figures and Tables:

Figure 2: Left: The separable objective for n = 3.Note the similarity to the dictionary learning objective. Right: The objective for complete orthogonal dictionary learning (discussed in section 6) for n = 3.

a gradient flow line s.t. γ(0) = q .

Figure3: Negative curvature and efficient gradient descent. The union of the light blue, orange and yellow sets is the set C. In the light blue region, there is negative curvature normal to ∂C, while in the orange region the gradient norm is large, as illustrated by the arrows. There is a single global minimizer in the yellow region. For the separable objective, the stable manifolds of the saddles and maximizers all lie on ∂C (the black circles denote the critical points, which are either maximizers " ", saddles " ", or minimizers " "). The red dots denote ∂C ζ with ζ = 0.2.

r, ζ) and we thus demand µ < 4 that for |w i | > r u (i) * grad[f pop DL ](q(w)) ≥ c DL (8000(n -1))3/2

cannot increase, hence this can only happen in S 1 . We show that for some z ∈ S 1 , a point with w = (1 -ε) x √ 2 , ε < 1 cannot reach a point with w = x √ 2 by a gradient descent step. This would mean 1

i . The number of iterations in each set will be determined by the bounds on the change in ζ, ||w|| detailed in 27.

since after this many iterationsz 2 ≥ ζ 2 ≥ (1 -c) x 2 .For every iteration in S 3 ∪ S 4 we are guaranteedw ≤ 1 -(1 -2c)ηx thus using Lemmas 13.i and 15 the number of iterations in S 3 ∪ S 4 required for convergence is given by x (1 -2c)η x 2

log(1-(1-2c)η x 2 ) log(1+2cη x 2 )C.9 Probability of the bound holdingThe bound applies to an initialization with ζ ≥ ζ 0 , hence in S 1 \Q ζ0 . Assuming uniform initialization in S 1 , the set Q ζ0 is simply a band of width 2ζ 0 around the equator of the changing variables to x = π ε x (1+ζ)

using convexity and non-negativity of the exponent respectively)nE [exp (λ |X i |)] = 2n ∞ 0 exp (λX i ) dµ(X i ) ≤ 2nE [exp (λX i )] = 2n exp λ 2 is minimized for λ = 2 log(2n) giving E max i |X i | ≤ 2 log(2n) ∼ 2 log(n)

1+ζ) 2 +2)x 2 erf n-3 (x)dx > 0

n-1s(ζ) (0) where s(ζ) = (2+ζ)ζ+n . B ∞ s(ζ) (0) is the largest L ∞ ball contained in C ζ , and B 2 √ n-1s(ζ) (0) is the smallest L 2 ball containing C ζ (where these balls are defined in terms of the w vector). All three intersect only at the points where all the coordinates of w have equal magnitude. Additionally,C ζ ⊆ B ∞ 1/ √ 2+ζ(0) and this is the smallest L ∞ ball containing C ζ .

Behavior of gradient descent in ∪ 4 i=1 S i We now show that gradient descent initialized in S 1 \Q ζ0 cannot exit ∪ 4 i=1 S i or enter Q ζ0 . Lemma 14 guarantees that gradient descent initialized in ∪ 4 i=1 S i remains in this set. From equation 27 we see that a gradient descent step can only decrease ζ if z ∈ S 4 . Under the mild assumption ζ 2 0 < 7 16 x 2 we are guaranteed from Lemma 13 that at every iteration ζ ≥ ζ 0 . Thus the region with ζ < ζ 0 can only be entered if gradient descent is initialized in it. It follows that initialization in S 1 \Q ζ0 rules out entering Q ζ0 at any future iteration of gradient descent. Since this guarantees that regions that feed into small gradient regions are avoided, an efficient convergence rate can again be obtained.Theorem 3 (Gradient descent convergence rate for generalized phase retrieval). Gradient descent on 22 with step size η <

Table content: <table><tr><td>z such that dist(z, Ȃ) &lt;</td><td>√</td><td colspan="2">5c x in</td></tr><tr><td/><td/><td>T &lt;</td><td>log log(1+η x 2 ) + x ζ √ 2</td><td>log(2) 2 log(1+2cη x 2 )</td></tr><tr><td colspan="3">C.3 Convergence rate</td></tr><tr><td/><td/><td colspan="2">√ 4 x 2 , c &lt; 1 c 4 , initialized uniformly in S 1 converges to a point</td></tr></table>

Additionally, if we initialize in S 1 ∩ Q ζ0 then we cannot initialize at a point where w = x Since only a step from S 4 can decrease ζ, we have that for the initial point z 2 > x 2 .

Table content: <table><tr><td>ii) Combined with w</td><td>2 ≤ x 2 2</td><td colspan="2">this gives</td></tr><tr><td/><td/><td/><td>ζ 2 ≥</td><td>x 2</td><td>2</td></tr><tr><td colspan="3">and using the lower bound (1 -2η x</td><td>2 c)ζ ≤ ζ we obtain</td></tr><tr><td/><td/><td/><td>√</td><td>2</td></tr><tr><td colspan="3">and hence the inequality is strict.</td></tr></table>

this also ensures that the next iterate cannot lie in the small gradient regions around the stable manifolds of the saddles. Lemma 14. Defining z as in 25, under the conditions of Lemma 13 and we have i

Table content: <table><tr><td/><td colspan="2">1 4 , η &lt;</td><td colspan="2">√ b x 2 . Choosing b = 4 gives c</td></tr><tr><td/><td>ζ 2 ≥</td><td>7 16</td><td>x</td><td>2</td></tr><tr><td>If we require ζ 2 0 &lt; 7 16 x</td><td>2</td><td/><td/></tr></table>

2 ≥ ζ 2 ≥ x 2 2 .

Table content: <table><tr><td>4</td></tr><tr><td>C.7.2 Iterations in</td></tr><tr><td>i=2</td></tr></table>

2 we are guaranteed ζ ≥ (1 + 2η x 2 c)ζ the number of iterations required to reach S 3 ∪ S 4 again is given by

Table content: <table><tr><td colspan="2">√ 4 7</td><td colspan="3">x (1 + 2η x</td><td colspan="2">2 c) tr =</td><td>√</td><td>1 -c x</td></tr><tr><td>t r =</td><td colspan="3">log 4 √ √ 1-c 7 log(1 + 2η x</td><td colspan="2">2 c)</td><td>≤</td><td>log( 4 √ 7 ) log(1 + 2η x</td><td>2 c)</td></tr><tr><td>C.8 Final rate</td><td/><td/><td/><td/><td/></tr><tr><td colspan="3">The final rate to convergence is</td><td/><td/><td/></tr><tr><td/><td/><td/><td colspan="4">T &lt; t 1 + t 2 + t 3+4 t r</td></tr><tr><td/><td colspan="2">=</td><td colspan="3">log log(1+η x 2 ) + x ζ √ 2</td><td>log(2) 2 log(1+2cη x 2 )</td></tr><tr><td/><td>+</td><td/><td colspan="4">log(2c) log( 4</td></tr></table>

Peer review (iclr2019):
 The authors analyze the convergence performance of Riemannian gradient descent based algorithm for the dictionary learning problem with orthogonal dictionary and sparse factors. They demonstrate a polynomial time convergence from a random initialization for a smooth surrogate objective for the original non-smooth one. The problem and the analysis are of interest, but I have several questions regarding the paper as follows. My first concern is that the analysis is on a smooth surrogate of the non-smooth sparse minimization for solving the dictionary learning problem, so it is not clear what is relationship between the global minimizer of the smooth problem to the underlying true dictionary. More specifically, how far is the global minimizer of problem (1) or (2) to the true dictionary parameter, and whether they share (approximately) the space or components regarding the sparse factors. Without clarifying this, it is not well motivated why we are interested in studying the problem considered in this paper at the beginning. Intuitively, since the recovered factors are not sparse anymore, it will impact the dictionary accordingly due to the linear mapping, which may lead to a very different set of dictionary components. Thus, explicit explanation is necessary here to avoid such degenerate case. My second concern is the eligibility of assuming the dictionary A to be an identity matrix and extending it to the general orthogonal matrix case. The analysis uses the fact that rows of A are canonical basis, i.e., each row only has one non-zero entry. I do not see a straightforward extension by replacing A to be an orthogonal matrix as the authors claimed on page 3, since then the inner product of one row of A and one column of X is not just the corresponding entry of the column of X. It will be helpful if the authors can explain this explicitly or adjust the analysis accordingly to make this valid. Another issue is the clarity of the paper. Some statements in the paper are not very clear. Form example, on page 3, third paragraph of Section 3, what does row(Y) = row(X_0) mean? Also, in eqn (1), y_k means k-th column of Y, and in eqn (2), q_i means i-th entry of q? Since both are bold lower case letters, clear distinction will help. Moreover, the reference use (. ) instead [ .], which can be confusing sometimes.