Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learning

Introduction

There are usually far more learnable parameters in deep neural nets than the number of training data, but still deep learning works well on real-world tasks. Even with explicit regularization, the model complexity of state-of-the-art neural nets is so large that they can fit randomly labeled data easily [Zhang et al., 2017]. Towards explaining the mystery of generalization, we must understand what kind of implicit regularization does Gradient Descent (GD) impose during training. Ideally, we are hoping for a nice mathematical characterization of how GD constrains the set of functions that can be expressed by a trained neural net.

With sufficiently small initialization, GF converges to the minimum nuclear norm solution of matrix sensing.

Subsequently, Arora et al. [2019a] challenged this view by arguing that a simple mathematical norm may not be a sufficient language for characterizing implicit regularization. One example illustrated in Arora et al. [2019a] is regarding matrix sensing with a single observation. They showed that GD with small initialization enhances the growth of large singular values of the solution and attenuates that of smaller ones. This enhancement/attenuation effect encourages low-rank, and it is further intensified with depth in deep matrix factorization (i.e., GD optimizes $f({\bm{U}}_{1}\cdots{\bm{U}}_{L})$ for $L\geq 2$). However, these are not captured by the nuclear norm alone. Gidel et al. , Gissin et al. further exploited this idea and showed in the special case of full-observation matrix sensing that GF learns solutions with gradually increasing rank. Razin and Cohen showed in a simple class of matrix completion problems that GF decreases the rank along the trajectory while any norm grows towards infinity. More aggressively, they conjectured that the implicit regularization can be explained by rank minimization rather than norm minimization.

In this paper, we move one further step towards resolving the implicit regularization in the matrix factorization problem. Our theoretical results show that GD performs rank minimization via a greedy process in a broader setting. Specifically, we provide theoretical evidence that GF with infinitesimal initialization is in general mathematically equivalent to another algorithm called Greedy Low-Rank Learning (GLRL). At a high level, GLRL is a greedy algorithm that performs rank-constrained optimization and relaxes the rank constraint by $1$ whenever it fails to reach a global minimizer of $f(\,\cdot\,)$ with the current rank constraint. As a by-product, we refute 1.1 by demonstrating an counterexample (Example 5.9).

We also extend our results to deep matrix factorization Section 6, where we prove that the trajectory of GF with infinitesimal identity initialization converges to a deep version of GLRL, at least in the early stage of the optimization. We also use this result to confirm the intuition achieved on toy models [Gissin et al., 2020], that benefits of depth in matrix factorization is to encourage rank minimization even for initialization with a relatively larger scale, and thus it is more likely to happen in practice. This shows that describing the implicit regularization using GLRL is more expressive than using the language of norm minimization. We validate all our results with experiments in Appendix C.

Related Works

The view of norm minimization, or the closely related view of margin maximization, has been explored in different settings. Besides the nuclear norm minimization for matrix factorization [Gunasekar et al., 2017] discussed in the introduction, previous works have also studied the norm minimization/margin maximization for linear regression [Wilson et al., 2017, Soudry et al., 2018a, b, Nacson et al., 2019b, c, Ji and Telgarsky, 2019b], deep linear neural nets [Ji and Telgarsky, 2019a, Gunasekar et al., 2018], homogeneous neural nets [Nacson et al., 2019a, Lyu and Li, 2020], ultra-wide neural nets [Jacot et al., 2018, Arora et al., 2019b, Chizat and Bach, 2020].

The initialization scale can greatly influence the implicit regularization. A sufficiently large initialization can make the training dynamics fall into the lazy training regime defined by Chizat et al. and diminish test accuracy. Using small initialization is particularly important to bias gradient descent to low-rank solutions for matrix factorization, as empirically observed by Gunasekar et al. . Arora et al. [2019a], Gidel et al. , Gissin et al. , Razin and Cohen studied how gradient flow with small initialization encourages low-rank in simple settings, as discussed in the introduction. Li et al. proved recovery guarantees for gradient flow solving matrix sensing under Restricted Isometry Property (RIP), but the proof cannot be generalized easily to the case without RIP. Belabbas made attempts to prove that gradient flow is approximately rank-1 in the very early phase of training, but it does not exclude the possibility that the approximation error explodes later and gradient flow is not converging to low-rank solutions. Compared to these works, the current paper studies how GF encourages low-rank in a much broader setting.

Background

In this paper, we particularly focus on the overparameterized case, where $r=d$, to understand the implicit regularization of GF when there is no rank constraint for the matrix ${\bm{W}}$.

Warmup Examples

Before introducing our main results, we illustrate how GD performs greedy learning using two warmup examples.

In general, for a loss function $\mathcal{L}({\bm{U}})=\frac{1}{2}f({\bm{U}}{\bm{U}}^{\top})$, we can always apply Taylor expansion $f({\bm{W}})\approx f({\bm{0}})+\left<{\bm{W}},\nabla f({\bm{0}})\right>$ around the origin to approximate it with a linear function. This motivates us to study the linear case: $f({\bm{W}}):=f_{0}-\left<{\bm{W}},{\bm{Q}}\right>$ for some symmetric matrix ${\bm{Q}}$. In this case, the matrix ${\bm{U}}$ follows the ODE, $\frac{\textup{{d}}{\bm{U}}}{\textup{{d}}t}={\bm{Q}}{\bm{U}}$, which can be understood as a continuous version of the classical power iteration method for solving the top eigenvector. Let ${\bm{Q}}:=\sum_{i=1}^{d}\mu_{i}{\bm{v}}_{i}{\bm{v}}_{i}^{\top}$ be the eigendecomposition of ${\bm{Q}}$, where $\mu_{1}\geq\mu_{2}\geq\cdots\geq\mu_{d}$ and ${\bm{v}}_{1},\dots,{\bm{v}}_{d}$ are orthogonal to each other. Then we can write the solution as:

When $\mu_{1}>\mu_{2}$, the ratio between $e^{\mu_{1}t}$ and $e^{\mu_{i}t}$ for $i\neq 1$ increases exponentially fast. As $t\to+\infty$, ${\bm{U}}(t)$ and ${\bm{W}}(t)$ become approximately rank-1 as long as ${\bm{v}}_{i}^{\top}{\bm{U}}(0)\neq{\bm{0}}$, i.e.,

The analysis for the simple linear case reveals that GD encourages low-rank through a process similar to power iteration. However, $f({\bm{W}})$ is non-linear in general, and the linear approximation is close to $f({\bm{W}})$ only if ${\bm{W}}$ is very small. With sufficiently small initialization, we can imagine that GD still resembles the above power iteration in the early phase of the optimization. But what if ${\bm{W}}(t)$ grows to be so large that the linear approximation is far from the actual $f({\bm{W}})$?

Let ${\bm{W}}^{*}:=\sum_{i=1}^{d}\mu_{i}{\bm{v}}_{i}{\bm{v}}_{i}^{\top}$ be the eigendecomposition of ${\bm{W}}^{*}$. Our previous analysis shows that the dynamics is approximately $\frac{\textup{{d}}{\bm{U}}}{\textup{{d}}t}={\bm{W}}^{*}{\bm{U}}$ in the early phase and thus encourages low-rank.

However, it is unclear how and why this sequential learning/incremental learning can occur in general. Through the first warmup example, we may understand why GD learns a rank-1 matrix in the early phase, but does GD always learn solutions with rank $2,3,4,\dots$ sequentially? If true, what is the mechanism behind this? The current paper answers the questions by providing both theoretical and empirical evidence that the greedy learning behavior does occur in general with a similar reason as for the first warmup example.

Greedy Low-Rank Learning (GLRL)

In this section, we present a trajectory-based analysis for the implicit bias of GF on matrix factorization. Our main result is that GF with infinitesimal initialization is generically the same as that of a simple greedy algorithm, Greedy Low-Rank Learning (GLRL, Algorithm 1).

We define the (limiting) trajectory of GLRL by taking the learning rate $\eta\to 0$. The goal is to show that the trajectory of GLRL is close to that of GF with infinitesimal initialization. Recall that $\phi({\bm{W}}_{0},t)$ stands for the solution ${\bm{W}}(t)$ in (2) when ${\bm{W}}(0)={\bm{W}}_{0}$.

The most related one to GLRL (Algorithm 1) is probably Rank-1 Matrix Pursuit (R1MP) proposed by Wang et al. for matrix completion, which was later generalized to general convex loss in [Yao and Kwok, 2016]. R1MP maintains a set of rank-1 matrices as the basis, and in phase $r$, R1MP adds the same ${\bm{u}}_{r}{\bm{u}}_{r}^{\top}$ as defined in Algorithm 1 into its basis and solve $\min_{{\bm{\alpha}}}f(\sum_{i=1}^{r}\alpha_{i}{\bm{u}}_{i}{\bm{u}}_{i}^{\top})$ for rank-$r$ estimation. The main difference between R1MP and GLRL is that the optimization in each phase of R1MP is performed on the coefficients ${\bm{\alpha}}$, while the entire ${\bm{U}}_{r}$ evolves with GD in each phase of GLRL. In Figure 3, we provide empirical evidence that GLRL generalizes better than R1MP when ground truth is low-rank, although GLRL may have a higher computational cost depending on $\eta,\epsilon$.

Similar to R1MP, Greedy Efficient Component Optimization (GECO, Shalev-Shwartz and Singer 2010) also chooses the $r$-th component of its basis as the top eigenvector of $-\nabla f({\bm{W}}_{r})$, while it solves $\min_{{\bm{\beta}}}f(\sum_{1\leq i,j\leq r}\beta_{ij}{\bm{u}}_{i}{\bm{u}}_{j}^{\top})$ for the rank-$r$ estimation. Khanna et al. provided convergence guarantee for GECO assuming strong convexity. Haeffele and Vidal proposed a local-descent meta algorithm, of which GLRL can be viewed as a specific realization.

1 The Limiting Trajectory: A General Theorem for Dynamical System

To prove the equivalence between GF and GLRL, we first introduce our high-level idea by analyzing the behavior of a more general dynamical system around its critical point, say ${\bm{0}}$. A specific example is (2) if we set ${\bm{\theta}}$ to be the vectorization of ${\bm{W}}$.

The key observation is that if the initialization is infinitesimal, the trajectory is almost uniquely determined. To be more precise, we need the following definition:

A special case is that the direction of ${\bm{\theta}}_{\alpha}-{\bm{\theta}}_{0}$ converges, i.e., $\bar{{\bm{\theta}}}:=\lim_{\alpha\to 0}\frac{{\bm{\theta}}_{\alpha}-{\bm{\theta}}_{0}}{\|{\bm{\theta}}_{\alpha}-{\bm{\theta}}_{0}\|_{2}}$ exists. In this case, $\{{\bm{\theta}}_{\alpha}\}$ has positive alignment with either ${\bm{u}}$ or $-{\bm{u}}$ except for a zero-measure subset of $\bar{{\bm{\theta}}}$. This means any convergent sequence generically falls into either of these two categories.

2 Equivalence Between GD and GLRL: Rank-One Case

Now we establish the equivalence between GF and GLRL in the first phase. The main idea is to apply Theorem 5.3 on (2). For this, we need the following lemma on the eigenvalues and eigenvectors.

Let ${\bm{g}}({\bm{W}}):=-{\bm{W}}\nabla f({\bm{W}})-\nabla f({\bm{W}}){\bm{W}}$ and ${\bm{J}}({\bm{W}})$ be its Jacobian. Then ${\bm{J}}({\bm{0}})$ is symmetric and thus diagonalizable. Let $-\nabla f({\bm{0}})=\sum_{i=1}^{d}\mu_{i}{\bm{u}}_{1[i]}{\bm{u}}_{1[i]}^{\top}$ be the eigendecomposition of the symmetric matrix $-\nabla f({\bm{0}})$, where $\mu_{1}\geq\mu_{2}\geq\cdots\geq\mu_{d}$. Then ${\bm{J}}({\bm{0}})$ has the form:

where ${\bm{J}}({\bm{0}})[{\bm{\Delta}}]$ stands for the resulting matrix produced by left-multiplying ${\bm{J}}({\bm{0}})$ to the vectorization of ${\bm{\Delta}}$. For every pair of $1\leq i\leq j\leq d$, $\mu_{i}+\mu_{j}$ is an eigenvalue of ${\bm{J}}({\bm{0}})$ and ${\bm{u}}_{1[i]}{\bm{u}}_{1[j]}^{\top}+{\bm{u}}_{1[j]}{\bm{u}}_{1[i]}^{\top}$ is a corresponding eigenvector. All the other eigenvalues are .

$\mu_{1}>\max\{\mu_{2},0\}$, where $\mu_{1}:=\lambda_{1}(-\nabla f({\bm{0}})),\mu_{2}:=\lambda_{2}(-\nabla f({\bm{0}}))$.

$f({\bm{W}})$ is locally analytic at each point.

5.7 is a natural assumption, since $f(\,\cdot\,)$ in most cases of matrix factorization is a quadratic or polynomial function (e.g., matrix sensing, matrix completion). In general, it is unlikely for a gradient-based optimization process to get stuck at saddle points [Lee et al., 2017, Panageas et al., 2019]. Thus, we should expect to see in general that GLRL finds the rank-1 solution if the problem is feasible with rank-1 matrices. This means at least for this subclass of problems, the implicit regularization of GD is rather unrelated to norm minimization. Below is a concrete example:

3 Equivalence between GD and GLRL: General Case

Benefits of Depth: A View from GLRL

In this section, we consider matrix factorization problems with depth $L\geq 3$. Our goal is to understand the effect of the depth-$L$ parametrization ${\bm{W}}={\bm{U}}_{1}{\bm{U}}_{2}\cdots{\bm{U}}_{L}$ on the implicit bias — how does depth encourage GF to find low rank solutions? We take the standard assumption in existing analysis for the end-to-end dynamics that the weight matrices have a balanced initialization, i.e. ${\bm{U}}_{i}^{\top}(0){\bm{U}}_{i}(0)={\bm{U}}_{i+1}(0){\bm{U}}_{i+1}^{\top}(0),\ \forall 1\leq i\leq L-1$. Arora et al. showed that if $\{{\bm{U}}_{i}\}_{i=1}^{L}$ is balanced at initialization, then we have the following end-to-end dynamics. Similar to the depth-2 case, we use $\phi({\bm{W}}(0),t)$ to denote ${\bm{W}}(t)$, where

The lemma below is the foundation of our analysis for the deep case, which greatly simplifies (11). Due to the space limit, we defer its derivations and applications into Appendix I.

If ${\bm{W}}(t)$ is a symmetric solution of (11), then for ${\bm{M}}(t):={\bm{W}}(t)^{2/L}$, we have

Let $P=\frac{L}{2}$, $L\geq 3$. Suppose $\left\|\nabla f({\bm{0}})\right\|_{2}=\lambda_{1}(-\nabla f({\bm{0}}))>\max\{\lambda_{2}(-\nabla f({\bm{0}})),0\}$,$\left\|\nabla f({\bm{0}})\right\|_{2}=\lambda_{1}(-\nabla f({\bm{0}}))$ is a technical assumption which we believe could be removed with a more refined analysis.

When the ground truth is low-rank, say rank-$k$, our experiments (Figure 2) suggest that GF with small initialization deep matrix sensing finds solutions with smaller $k$-low-rankness compared to the depth-2 case, thus achieving better generalization. At first glance, this is contradictory to what Theorem 6.2 suggests, i.e., the convergence rate of deep GLRL at a constant time gets slower as the depth increases. However, it turns out the uniform upper bound for the distance between GF and GLRL is not the ideal metric for the eventual $k$-low-rankness of learned solution. Below we will illustrate why the $r$-low-rankness of GF within each phase $r$ is a better metric and how they are different.

For depth-2 GLRL, the low-rankness is raised to some power less then $1$ per phase (depending on the eigengap). For deep GLRL, we show the low-rankness is only multiplied by some constant for the first phase and speculate it to be true for later phases. This conjecture is supported by our experiments; see Figure 2. Interestingly, our theory and experiments (Figure 5) suggest that while being deep is good for generalization, being much deeper may not be much better: once $L\geq 3$, increasing the depth does not improve the order of low-rankness significantly. While this theoretical result is only for identity initialization, Theorem D.1 and Corollary D.2 further show that the dynamics of GF (11) with any initialization pointwise converges as $L\to\infty$, under a suitable time rescaling. See Figure 6 for experimental verification.

Conclusion and Future Directions

In this work, we connect gradient descent to Greedy Low-Rank Learning (GLRL) to explain the success of using gradient descent to find low-rank solutions in the matrix factorization problem. This enables us to construct counterexamples to the implicit nuclear norm conjecture in [Gunasekar et al., 2017]. Taking the view of GLRL can also help us understand the benefits of depth.

Our result on the equivalence between gradient flow with infinitesimal initialization and GLRL is based on some regularity conditions that we expect to hold generically. We leave it a future work to justify these condition, possibly through a smoothed analysis on the objective $f(\,\cdot\,)$. Another interesting future direction is to find the counterpart of GLRL in training deep neural nets. This could be one way to go beyond the view of norm minimization in the study of the implicit regularization of gradient descent.

The authors thank Sanjeev Arora and Jason D. Lee for helpful discussions. The authors also thank Runzhe Wang for useful suggestions on writing. ZL and YL acknowledge support from NSF, ONR, Simons Foundation, Schmidt Foundation, Mozilla Research, Amazon Research, DARPA and SRC. ZL is also supported by Microsoft PhD Fellowship.

References

Appendix A Preliminary Lemmas

${\bm{U}}_{0}$ is a stationary point of $\mathcal{L}({\bm{U}})=\frac{1}{2}f({\bm{U}}{\bm{U}}^{\top})$;

$\nabla f({\bm{W}}_{0}){\bm{W}}_{0}={\bm{0}}$;

${\bm{W}}_{0}:={\bm{U}}_{0}{\bm{U}}_{0}^{\top}$ is a critical point of (2).

(2) $\Rightarrow$ (3) is trivial. We only prove (1) $\Rightarrow$ (2), (3) $\Rightarrow$ (1).

If ${\bm{U}}_{0}$ is a stationary point, then ${\bm{0}}=\nabla\mathcal{L}({\bm{U}}_{0})=\nabla f({\bm{W}}_{0}){\bm{U}}_{0}$. So

If ${\bm{W}}_{0}$ is a critical point, then

which implies $\nabla\mathcal{L}({\bm{U}}_{0})={\bm{0}}$. ∎

Note that $\left<\nabla f({\bm{W}}_{0}),{\bm{W}}_{0}\right>=\operatorname{Tr}(\nabla f({\bm{W}}_{0}){\bm{W}}_{0})$. By Lemma A.1, $\left<\nabla f({\bm{W}}_{0}),{\bm{W}}_{0}\right>=0$. Combining this with (16), we know that ${\bm{W}}_{0}$ is a global minimizer iff

It is easy to check that this condition is equivalent to $\nabla f({\bm{W}}_{0})\succeq{\bm{0}}$. ∎

Appendix B Proofs for Counter-example

where $\Omega=\{(1,3),(1,4),(2,3),(3,1),(3,2),(4,1)\}$.

Let ${\bm{M}}_{0}:=\nabla f({\bm{0}})$, then

Let ${\bm{W}}_{\epsilon}(t)$ be the following matrix:

Since $g(x(t),y(t))$ is non-increasing overtime, and $\lim\limits_{t\to-\infty}g(x(-t),y(-t))=g(x(-\infty),y(-\infty))=g(0,0)=R^{2}+0.5$, we know $\lvert x(t)y(t)\rvert\leq 3R$ for all $t$. So whenever $y^{2}(t)-x^{2}(t)\geq 9R^{2}$, we have $x^{2}(t)\leq\frac{9R^{2}}{y^{2}(t)}\leq\frac{9R^{2}}{y^{2}(t)-x^{2}(t)}\leq 1$. In this case, $\frac{\textup{{d}}(y^{2}(t)-x^{2}(t))}{\textup{{d}}t}=2x^{2}(t)(x^{2}(t)-1)\leq 0$. Combining this with $y(-\infty)^{2}-x(-\infty)^{2}=0\leq 9R^{2}$, we have $y^{2}(t)-x^{2}(t)\leq 9R^{2}$ for all $t$, which also implies that $y(t)$ is bounded. Noticing that $9R^{2}\geq g(x(t),y(t))\geq(x^{2}(t)-1)^{2}$, we know $x^{2}(t)$ is also bounded. Therefore, ${\bm{W}}_{1}^{G}(t)$ is bounded.

Let $m_{ij}$ be $(i,j)$th element of ${\bm{M}}$. Suppose ${\bm{M}}\succeq{\bm{0}}$, we have

Thus $4R=\min_{{\bm{W}}\succeq\bm{0},f({\bm{W}})=0}\left\|{\bm{W}}\right\|_{\ast}$, where the equality is only attained at $m_{ii}=R,i=1,2,3,4$. Otherwise, either $\begin{bmatrix}m_{11}&m_{14}\\ m_{41}&m_{44}\end{bmatrix}$ or $\begin{bmatrix}m_{22}&m_{23}\\ m_{32}&m_{33}\end{bmatrix}$ will have negative eigenvalues. Contradiction to that ${\bm{M}}\succeq{\bm{0}}$.

Below we will show the rest unknown off-diagonal entries must be $1$. Let ${\bm{V}}=\begin{bmatrix}1&-1&0&0\\ 0&0&1&0\\ 0&0&0&1\end{bmatrix}$, then

which implies $m_{13}=m_{23}$, $m_{14}=m_{24}$.

Appendix C Experiments

The code is written in Julia [Bezanson et al., 2012] and PyTorch [Paszke et al., 2019].

In Figures 1, 2, 3 and 4, the GLRL’s trajectory is obtained by running Algorithm 1 with $\epsilon=10^{-7}$ and $\eta=10^{-3}$. The stopping criterion is that if the loop has been iterated for $10^{7}$ times.

C.2 Experimental Equivalence between GLRL and Gradient Descent

Here we provide experimental evidence supporting our theoretical claims about the equivalence between GLRL and GF for both cases, $L=2$ and $L\geq 3$.

C.3 How well does GLRL work?

We compare GLRL with gradient descent (with not-so-small initialization), nuclear norm minimization and R1MP [Wang et al., 2014]. We use CVXPY [Diamond and Boyd, 2016, Agrawal et al., 2018] for finding the nuclear norm solution. The results are shown in Figure 3. GLRL can fully recover the ground truth, while others have difficulty doing so.

C.4 How does initialization affect the convergence rate to the rank-1 GLRL trajectory?

The result is shown at Figure 4. We observe that GLRL trajectories are closer to the reference matrix ${\bm{W}}_{\text{ref}}$ by magnitudes. Thus the take home message here is that GLRL is in general a more computational efficient method to simulate the trajectory of GF (GD) with infinitesimal initialization, as one can start GLRL with a much larger initialization, while still maintaining high precision.

C.5 Benefit of Depth: polynomial vs exponential dependence on initialization

To verify the our theory in Section 6, we run gradient descent with different depth and initialization. The results are shown in Figure 5. We can see that as the initialization becomes smaller, the final solution gets closer to the ground truth. However, a depth-2 model requires exponentially small initialization, while deeper models require polynomial small initialization, though it takes much longer to converge.

Appendix D The marginal value of being deeper

Theorem D.1 shows that the end-to-end dynamics (19) converges point-wise while $L\to\infty$ if the product of learning rate and depth, $\eta L$, is fixed as constant. Interestingly, (19) also allows us to simulate the dynamics of ${\bm{W}}(t)$ for all depths $L$ while the computation time is independent of $L$. In Figure 6, we compare the effect of depth while fixing the initialization and $\eta L$. We can see that deeper models converge faster. The difference between $L=1,2$, and $4$ is large, while difference among $L\geq 16$ is marginal.

where $K^{(L)}_{i,i}=\sigma_{i}^{2-2/L}$, $K^{(L)}_{i,j}=\frac{\sigma_{i}^{2}-\sigma_{j}^{2}}{L\sigma_{i}^{2/L}-L\sigma_{j}^{2/L}}$ for $i\neq j$.

where ${\bm{H}}^{(l)}_{i,j}=\sigma_{i}^{\frac{2l}{L}}\sigma_{j}^{\frac{2(L-1-l)}{L}}$. Therefore,

where ${\bm{K}}^{(L)}=L^{-1}\sum_{l=0}^{L-1}{\bm{H}}^{(l)}$. Hence,

The entries of ${\bm{K}}^{(L)}$ can be directly calculated by

As $L\to\infty$, ${\bm{K}}^{(L)}$ converges to ${\bm{K}}^{*}$, where $K^{*}_{i,i}=\sigma_{i}^{2}$, $K^{*}_{i,j}=\frac{\sigma_{i}^{2}-\sigma_{j}^{2}}{\ln\sigma_{i}^{2}-\ln\sigma_{j}^{2}}$ for $i\neq j$.

Appendix E Proofs for Dynamical System

In this section, we prove Theorem 5.3 in Section 5.1. In Section E.1, we show how to reduce Theorem 5.3 to the case where ${\bm{J}}({\bm{0}})$ is exactly a diagonal matrix, then we prove this diagonal case in Section E.2. Finally, in Section E.3, we discuss how to extend it to the case where ${\bm{J}}({\bm{0}})$ is non-diagonalizable.

E.2 Proof for the Diagonal Case

Let $R>0$. Since ${\bm{g}}({\bm{\theta}})$ is $\mathcal{C}^{2}$-smooth, there exists $\beta>0$ such that

for all $\|{\bm{\theta}}\|_{2},\|{\bm{\theta}}+{\bm{h}}\|_{2}\leq R$. Then the following can be proved by integration:

For ${\bm{\theta}}(t)=\phi({\bm{\theta}}_{0},t)$ with $\|{\bm{\theta}}_{0}\|_{2}\leq\alpha$ and $t\leq T_{\alpha}(r)$,

Expending $F(\alpha)$ proves the lemma. ∎

For ${\bm{\theta}}(t)=\phi({\bm{\theta}}_{0},t)$ with $\|{\bm{\theta}}_{0}\|_{2}\leq\alpha$ and $t\leq T_{\alpha}(r)$, we have

Let $\hat{{\bm{\theta}}}(t)=e^{t{\bm{J}}({\bm{0}})}{\bm{\theta}}_{0}$. Then we have

where the last inequality is due to Lemma E.2. By (24) and Lemma E.3, we have

Let ${\bm{\theta}}(t)=\phi({\bm{\theta}}_{0},t),\hat{{\bm{\theta}}}(t)=\phi(\hat{{\bm{\theta}}}_{0},t)$. If $\max\{\|{\bm{\theta}}_{0}\|_{2},\|\hat{{\bm{\theta}}}_{0}\|_{2}\}\leq\alpha$, then for $t\leq T_{\alpha}(r)$,

For every $t\in(-\infty,+\infty)$, ${\bm{z}}(t)$ exists and ${\bm{z}}_{\alpha}(t)$ converges to ${\bm{z}}(t)$ in the following rate:

where $O$ hides constants depending on $g({\bm{\theta}})$ and $t$.

Now it is only left to prove (7). WLOG we can assume that $\|{\bm{\delta}}_{\alpha}\|_{2}$ is decreasing and $\frac{\alpha}{2}\leq\|{\bm{\delta}}_{\alpha}\|_{2}\leq\alpha$ (otherwise we can do reparameterization). Then our goal becomes proving

Let $q_{\alpha}:=\left<\frac{{\bm{\delta}}_{\alpha}}{\alpha},{\bm{e}}_{1}\right>$. By Definition 5.2, there exists $q>0$ such that $q_{\alpha}\geq q$ for all sufficiently small $\alpha$. Then we have

Combining this with the convergence rate for ${\bm{z}}_{\alpha_{1}}(t)$, we have

E.3 Extension to Non-Diagonalizable Case

Assume that ${\bm{\theta}}={\bm{0}}$ is a critical point and the following regularity conditions hold:

${\bm{g}}({\bm{\theta}})$ is $\mathcal{C}^{2}$-smooth;

$\phi({\bm{\theta}}_{0},t)$ exists for all ${\bm{\theta}}_{0}$ and $t$;

The top eigenvalue of ${\bm{J}}({\bm{0}})$ is unique and is a positive real number, i.e.,

We only need to select a parameter $\delta>0$ and prove the theorem in the case of ${\bm{J}}({\bm{0}})={\bm{J}}_{\delta}$ since we can change the basis in a similar way as we have done in Section E.1. By scrutinizing the proof for Theorem E.1, we can find that we only need to reprove Lemma E.2. However, Lemma E.2 may not be correct since ${\bm{J}}({\bm{0}})$ is not diagonal anymore. Instead, we prove the following:

Let ${\mathcal{K}}$ be the set of pairs $(k_{1},k_{2})$ such that $k_{1}\neq k_{2}$ and the entry of ${\bm{J}}_{\delta}$ at the $k_{1}$-th row and the $k_{2}$-th column is non-zero. Then we have

where the second equality uses the fact that ${\bm{I}}$ and ${\bm{N}}$ are commutable. So we have

where the second equality uses the fact that ${\bm{D}}$ and ${\bm{N}}^{2}$ are commutable. Note that $e^{t{\bm{C}}}=e^{at}\left[\begin{smallmatrix}\cos(bt)&-\sin(bt)\\ \sin(bt)&\cos(bt)\end{smallmatrix}\right]$, which implies $\|e^{t{\bm{D}}}\|_{2}=\|e^{t{\bm{C}}}\|_{2}=e^{at}$. So we have

Appendix F Eigenvalues of Jacobians and Hessians

In this section we analyze the eigenvalues of the Jacobian ${\bm{J}}({\bm{W}})$ at critical points of (2).

for ${\bm{W}}_{0}:={\bm{U}}_{0}{\bm{U}}_{0}^{\top}$, and thus ${\bm{W}}_{0}$ is a critical point of (2).

For a real-valued or vector-valued function $F({\bm{\theta}})$, we use $DF({\bm{\theta}})[{\bm{\delta}}],D^{2}F({\bm{\theta}})[{\bm{\delta}}_{1},{\bm{\delta}}_{2}]$ to denote the first- and second-order directional derivatives of $F(\,\cdot\,)$ at ${\bm{\theta}}$.

We also write $DF({\bm{\theta}})[{\bm{\Delta}}_{1},{\bm{\Delta}}_{2}]:=\left<DF({\bm{\theta}})[{\bm{\Delta}}_{1}],{\bm{\Delta}}_{2}\right>$.

Define ${\bm{J}}({\bm{W}}):=D{\bm{g}}({\bm{W}})$. By simple calculus, we can compute the formula for ${\bm{J}}({\bm{W}}_{0})$:

We can also compute the formula for $D^{2}\mathcal{L}({\bm{U}}_{0})$:

The eigenvalues of ${\bm{J}}({\bm{0}})$ is given in Lemma 5.4. Now we provide the proof.

It is easy to see from the second equation that ${\bm{J}}({\bm{0}})$ is symmetric.

Let $-\nabla f({\bm{0}})=\sum_{i=1}^{d}\mu_{i}{\bm{u}}_{1[i]}{\bm{u}}_{1[i]}^{\top}$ be the eigendecomposition of the symmetric matrix $-\nabla f({\bm{0}})$. Then we have

For ${\bm{\Delta}}={\bm{u}}_{1[i]}{\bm{u}}_{1[j]}^{\top}+{\bm{u}}_{1[j]}{\bm{u}}_{1[i]}^{\top}$, we have

So ${\bm{u}}_{1[i]}{\bm{u}}_{1[j]}^{\top}+{\bm{u}}_{1[j]}{\bm{u}}_{1[i]}^{\top}$ is an eigenvector of ${\bm{J}}({\bm{0}})$ associated with eigenvalue $\mu_{i}+\mu_{j}$. Note that $\{{\bm{u}}_{1[i]}{\bm{u}}_{1[j]}^{\top}+{\bm{u}}_{1[j]}{\bm{u}}_{1[i]}^{\top}:i,j\in[d]\}$ spans all the symmetric matrices, so these are all the eigenvectors in the space of symmetric matrices.

For every antisymmetric matrix ${\bm{\Delta}}$ (i.e., ${\bm{\Delta}}=-{\bm{\Delta}}^{\top}$), we have

So ${\bm{J}}({\bm{0}})[{\bm{\Delta}}]={\bm{0}}$ and every antisymmetric matrix is an eigenvector associated with eigenvalue .

Since every matrix can be expressed as the sum of a symmetric matrix and an antisymmetric matrix, we have found all the eigenvalues. ∎

F.2 Eigenvalues at Second-Order Stationary Points

Let ${\bm{\Delta}}={\bm{v}}{\bm{q}}^{\top}$. Then we have

So $D^{2}\mathcal{L}({\bm{U}}_{0})[{\bm{\Delta}},{\bm{\Delta}}]<0$, which leads to a contradiction. ∎

By (28), the symmetric matrices $-\nabla f({\bm{W}}_{0})$ and ${\bm{W}}_{0}$ commute, so they can be simultaneously diagonalizable. Since (28) also implies that they have different column spans, we can have the following diagonalization:

First we prove the following lemma on the eigenvalues and eigenvectors of the linear operator $-D^{2}\mathcal{L}({\bm{U}}_{0})$:

Replacing ${\bm{\Delta}}$ with ${\bm{U}}_{0}{\bm{R}}$ in (30) gives ${\bm{U}}_{0}({\bm{R}}+{\bm{R}}^{\top}){\bm{U}}_{0}^{\top}={\bm{0}}$, which is equivalent to ${\bm{R}}=-{\bm{R}}^{\top}$ since ${\bm{U}}_{0}$ is full-rank. Since the dimension of $r\times r$ antisymmetric matrices is $\frac{r(r-1)}{2}$, the span spanned by the solutions of (30) also has dimension $\frac{r(r-1)}{2}$. ∎

The eigenvalues of ${\bm{J}}({\bm{W}}_{0})$ can be fully classified into the following $3$ types:

$\mu_{i}+\mu_{j}$ is an eigenvalue for every $1\leq i\leq j\leq d-r$, and $\hat{{\bm{U}}}_{ij}:={\bm{v}}_{i}{\bm{v}}_{j}^{\top}+{\bm{v}}_{j}{\bm{v}}_{i}^{\top}$ is an associated left eigenvector.

is an eigenvalue, and any antisymmetric matrix is an associated right eigenvector, which spans a linear space of dimension $\frac{d(d-1)}{2}$.

We first prove each item respectively, and then prove that these are all the eigenvalues of ${\bm{J}}({\bm{W}}_{0})$.

For $\hat{{\bm{U}}}_{ij}={\bm{v}}_{i}{\bm{v}}_{j}^{\top}+{\bm{v}}_{j}{\bm{v}}_{i}^{\top}$, it is easy to check:

which shows that $\hat{{\bm{U}}}_{ij}$ is a left eigenvector associated with eigenvalue $\lambda_{i}+\lambda_{j}$.

By definition of eigenvector, we have $-D^{2}\mathcal{L}({\bm{U}}_{0})[{\bm{E}}_{p}]=\xi_{p}{\bm{E}}_{p}$, so

Right-multiplying both sides by ${\bm{U}}_{0}^{\top}$, we get

Since $\nabla f({\bm{W}})$ is symmetric, ${\bm{g}}({\bm{W}})$ is also symmetric. For any ${\bm{\Delta}}=-{\bm{\Delta}}^{\top}$,

So ${\bm{J}}({\bm{W}}_{0})[{\bm{\Delta}}]={\bm{0}}$ and ${\bm{\Delta}}$ is an eigenvector associated with eigenvalue .

Appendix G Proofs for the Depth-2 Case

G.2 Proof for Theorem 5.8

The proof for Theorem 5.8 relies on the following Lemma on the gradient flow around a local minimizer:

If $\bar{{\bm{\theta}}}$ is a local minimizer of $\mathcal{L}({\bm{\theta}})$ and for all $\|{\bm{\theta}}-\bar{{\bm{\theta}}}\|_{2}\leq r$, ${\bm{\theta}}$ satisfies Łojasiewicz inequality:

for some $\mu\in[1/2,1)$, then the gradient flow ${\bm{\theta}}(t)=\phi({\bm{\theta}}_{0},t)$ converges to a point ${\bm{\theta}}_{\infty}$ near $\bar{{\bm{\theta}}}$ if ${\bm{\theta}}_{0}$ is close enough to $\bar{{\bm{\theta}}}$, and the distance can be bounded by $\|{\bm{\theta}}_{\infty}-\bar{{\bm{\theta}}}\|_{2}=O(\|{\bm{\theta}}_{0}-\bar{{\bm{\theta}}}\|_{2}^{2(1-\mu)})$.

For every $t\geq 0$, if $\|{\bm{\theta}}(t)-\bar{{\bm{\theta}}}\|_{2}\leq r$,

Therefore, $\|{\bm{\theta}}(t)-{\bm{\theta}}_{0}\|_{2}\leq\int_{0}^{t}\left\|\frac{\textup{{d}}{\bm{\theta}}}{\textup{{d}}t}\right\|_{2}\textup{{d}}t\leq\frac{1}{(1-\mu)c}\mathcal{L}({\bm{\theta}}_{0})^{1-\mu}=O(\|{\bm{\theta}}_{0}-\bar{{\bm{\theta}}}\|_{2}^{2(1-\mu)})$. If we choose $\|{\bm{\theta}}(t)-\bar{{\bm{\theta}}}\|_{2}$ small enough, then $\|{\bm{\theta}}(t)-\bar{{\bm{\theta}}}\|_{2}\leq\|{\bm{\theta}}(t)-{\bm{\theta}}_{0}\|_{2}+\|{\bm{\theta}}_{0}-\bar{{\bm{\theta}}}\|_{2}=O(\|{\bm{\theta}}_{0}-\bar{{\bm{\theta}}}\|_{2}^{2(1-\mu)})<r$, and thus $\int_{0}^{+\infty}\left\|\frac{d{\bm{\theta}}}{dt}\right\|_{2}\textup{{d}}t$ is convergent and finite. This implies that ${\bm{\theta}}_{\infty}:=\lim_{t\to+\infty}{\bm{\theta}}(t)$ exists and $\|{\bm{\theta}}_{\infty}-\bar{{\bm{\theta}}}\|_{2}=O(\|{\bm{\theta}}_{0}-\bar{{\bm{\theta}}}\|_{2}^{2(1-\mu)})$. ∎

Combining these together we have $\left\|\phi({\bm{W}}_{\alpha},T({\bm{W}}_{\alpha})+t_{\epsilon})-\overline{{}{\bm{W}}}_{1}\right\|_{2}\leq\epsilon$.

It is easy to construct a factorization $\phi({\bm{W}}_{\alpha},T({\bm{W}}_{\alpha})+t_{\epsilon}):={\bm{U}}_{\alpha,\epsilon}{\bm{U}}_{\alpha,\epsilon}^{\top}$ such that $\left\|{\bm{U}}_{\alpha,\epsilon}-\overline{{\bm{U}}}\right\|_{2}=O(\epsilon)$, e.g., we can find an arbitrary factorization and then right-multiply an orthogonal matrix so that the row vector with the largest norm aligns with the direction of $\bar{{\bm{u}}}$. Applying Lemma G.1, we know that gradient flow starting with ${\bm{U}}_{\alpha,\epsilon}$ converges to a point that is only $O(\epsilon^{2(1-\mu)})$ far from $\bar{{\bm{u}}}$. So we have

Taking $\epsilon\to 0$ complete the proof. ∎

G.3 Proof for Theorem 5.11

Moreover, there exists a constant $C>0$ such that

G.4 Gradient Flow only finds minimizers (Proof for Theorem 5.10)

The proof for Theorem 5.10 is based on the following two theorems from the literature.

Let ${\bm{g}}$ be a ${\mathcal{C}}^{1}$ mapping from ${\mathcal{X}}\to{\mathcal{X}}$ and $\det(D{\bm{g}}(x))\neq 0$ for all ${\bm{x}}\in{\mathcal{X}}$. Then the set of initial points that converge to an unstable fixed point has measure zero, $\mu\left(\{{\bm{x}}_{0}:\lim_{k\to\infty}{\bm{g}}^{k}({\bm{x}}_{0})\in{\mathcal{A}}_{\bm{g}}^{*}\}\right)=0$, where ${\mathcal{A}}^{*}_{\bm{g}}=\{{\bm{x}}:{\bm{g}}({\bm{x}})={\bm{x}},\max_{i}|\lambda_{i}(D{\bm{g}}({\bm{x}}))|>1\}$.

Then the set of initial points that converge to a unstable critical point has measure zero, $\mu\left(\left\{{\bm{x}}_{0}:\lim_{t\to\infty}\phi({\bm{x}}_{0},t)\in{\mathcal{U}}^{*}_{\bm{f}}\right\}\right)=0$, where ${\mathcal{U}}^{*}_{\bm{f}}=\{{\bm{x}}:{\bm{f}}({\bm{x}})={\bm{0}},\lambda_{1}(D{\bm{f}}({\bm{x}}))>0\}$ and $D{\bm{f}}$ is the Jacobian matrix of ${\bm{f}}$.

By Theorem 1 in Section 2.3, Perko , we know $\phi(\cdot,\cdot)$ is $\mathcal{C}^{1}$-smooth for both $x,t$. We let ${\bm{g}}(x)=\phi(x,1)$, then we know ${\bm{g}}^{-1}(x)=\phi(x,-1)$ and both ${\bm{g}},{\bm{g}}^{-1}$ are $\mathcal{C}^{1}$-smooth. Note that $D{\bm{g}}^{-1}(x)$ is the inverse matrix of $D{\bm{g}}(x)$. So both of the two matrices are invertible. Thus we can apply Theorem G.4 and we know $\mu\left(\{x_{0}:\lim_{k\to\infty}{\bm{g}}^{k}(x_{0})\in{\mathcal{A}}_{\bm{g}}^{*}\}\right)=0$.

Note that if $\lim_{t\to\infty}\phi(x,t)$ exists, then $\lim_{k\to\infty}{\bm{g}}^{k}({\bm{x}})=\lim_{t\to\infty}\phi({\bm{x}},t)$. It remains to show that ${\mathcal{U}}^{*}_{\bm{f}}\subseteq{\mathcal{A}}^{*}_{\bm{g}}$. For ${\bm{f}}({\bm{x}}_{0})={\bm{0}}$, we have $\phi({\bm{x}}_{0},t)={\bm{x}}_{0}$ and thus ${\bm{g}}({\bm{x}}_{0})={\bm{x}}_{0}$. Now it suffices to prove that $\lambda_{1}(D{\bm{g}}({\bm{x}}_{0}))>1$. For every $t\in$, by Corollary of Theorem 1 in Section 2.3, Perko , we have $\frac{\partial}{\partial t}D\phi({\bm{x}},t)=D{\bm{f}}(\phi({\bm{x}},t))D\phi({\bm{x}},t)$, $\forall{\bm{x}},t$. Thus,

Solving this ODE gives $D{\bm{g}}({\bm{x}}_{0})=D\phi({\bm{x}},1)=e^{D{\bm{f}}(x_{0})}D\phi({\bm{x}},0)=e^{D{\bm{f}}(x_{0})}$, where the last equality is due to $D\phi({\bm{x}},0)\equiv{\bm{I}},\ \forall{\bm{x}}$. Combining this with $\lambda_{1}(D{\bm{f}}(x_{0}))>0$, we have $\lambda_{1}(D{\bm{g}}(x_{0}))>1$.

Thus we have ${\mathcal{U}}^{*}_{\bm{f}}:=\{{\bm{x}}_{0}:{\bm{f}}({\bm{x}}_{0})={\bm{0}},\lambda_{1}(D{\bm{f}}({\bm{x}}_{0}))>0\}\subseteq{\mathcal{A}}^{*}_{\bm{g}}$, which implies that $\{{\bm{x}}_{0}:\lim_{t\to\infty}\phi({\bm{x}}_{0},t)\in{\mathcal{U}}^{*}\}\subseteq\{{\bm{x}}_{0}:\lim_{k\to\infty}{\bm{g}}^{k}({\bm{x}}_{0})\in{\mathcal{A}}_{\bm{g}}^{*}\}$ ∎

For (1), by Theorem G.3, we immediately know all the stationary points of $\mathcal{L}(\,\cdot\,)$ are either global minimizers or strict saddles. (2) is just a direct consequence of Theorem G.5 by setting ${\bm{f}}$ in the above proof to $-\nabla\mathcal{L}$. ∎

Appendix H Equivalence Between GF and GLRL

In this section we elaborate on the theoretical evidence that GF and GLRL are equivalent generically, including the case where GLRL does not end in the first phase. The word “generically” used when we want to assume one of the following regularity conditions:

We want to assume that GF converges to a local minimizer (i.e., GF does not get stuck on saddle points);

We say that GF aligns well with GLRL in the beginning of phase $r$ if there exists $T^{(r)}_{\alpha}$ for every $\alpha>0$ such that $\phi({\bm{W}}_{\alpha},T^{(r)}_{\alpha})$ converges to $\overline{{}{\bm{W}}}_{r-1}$ with positive alignment with ${\bm{u}}_{r}{\bm{u}}_{r}^{\top}$ as $\alpha\to 0$.

If the initialization satisfies that ${\bm{W}}_{\alpha}$ converges to ${\bm{0}}$ with positive alignment with ${\bm{u}}_{1}{\bm{u}}_{1}^{\top}$ as $\alpha\to 0$, then GF aligns well with GLRL in the beginning of phase $1$, which can be seen by taking $T_{\alpha}^{(1)}=0$. Now assume that GF aligns well with GLRL in the beginning of phase $r-1$, then the above argument shows that GF should generically align well with GLRL in the beginning of phase $r$, if GLRL does not exit in phase $r-1$. In the other case, we can use a similar argument as in Theorem 5.8 to show that GF converges to a solution near the minimizer $\overline{{}{\bm{W}}}_{r}$ of $f(\,\cdot\,)$ as $t\to\infty$, and the distance between the solution and $\overline{{}{\bm{W}}}_{r}$ converges to as $\alpha\to 0$. By this induction we prove that GF with infinitesimal initialization is equivalent to GLRL generically.

Appendix I Proofs for Deep Matrix Factorization

Let $f(x)=P(1-x)-(1-x^{P})$. Since $f^{\prime}(x)=-P+Px^{P-1}<0$ for all $x\in[0,1)$, $f(x)\geq f(0)=0$. Then substituting $x$ by $\frac{b}{a}$ completes the proof. ∎

Recall we use $D{\bm{F}}({\bm{N}})[{\bm{M}}]$ to denote the directional derivative along ${\bm{M}}$ of ${\bm{F}}$ at ${\bm{N}}$.

where $D{\bm{F}}({\bm{N}})[{\bm{M}}]:=\lim_{t\to 0}\frac{{\bm{F}}({\bm{N}}+t{\bm{M}})-{\bm{F}}({\bm{N}})}{t}$ is the directional derivative of ${\bm{F}}$ along ${\bm{M}}$.

Therefore, it suffices to prove the lemma for the case where ${\bm{N}}$ is diagonal, i.e., ${\bm{N}}={\bm{\Sigma}}$.

Let ${\bm{H}}({\bm{G}})={\bm{G}}^{q}$. With the same argument, we know

Note that ${\bm{H}}({\bm{G}}({\bm{\Sigma}}))={\bm{F}}({\bm{\Sigma}})$. By chain rule, we have

When $\sigma_{i}=\sigma_{j}$, clearly $[D{\bm{F}}({\bm{\Sigma}})[{\bm{M}}]]_{ij}=m_{ij}\cdot\frac{q}{p}\cdot\sigma_{i}^{\frac{q-p}{p}}=Pm_{ij}\sigma_{i}^{P-1}$. Otherwise, we assume WLOG that $\sigma_{i}>\sigma_{j}$, we multiply $\sigma_{i}-\sigma_{j}$ to both numerator and denominator and we have

where the first inequality is by Lemma I.2. Thus we conclude the proof. ∎

$\left\|{\bm{N}}(t)\right\|_{2}\leq\rho$, since $\left\|\cdot\right\|_{2}$ is convex.

For a locally Lipschitz function $f(\,\cdot\,)$, the Clarke subdifferential [Clarke, 1975, 1990, Clarke et al., 2008] of $f$ at any point ${\bm{x}}$ is the following convex set

Clarke subdifferential generalize the standard notion of gradients in the sense that, when $f$ is smooth, $\frac{\partial^{\circ}f({\bm{x}})}{\partial{\bm{x}}}=\{\nabla f({\bm{x}})\}$. Clarke subdifferential satisfies the chain rule:

I.2 Proof of Lemma 6.1

Since ${\bm{W}}(t)\succeq{\bm{0}}$ by Lemma I.1, (11) can be rewritten as the following:

Suppose ${\bm{W}}(t)$ is a symmetric solution of (11). By Lemma I.1, we know ${\bm{W}}(t)$ also satisfies (31). Now we let ${\bm{R}}(t)$ be the solution of the following ODE with ${\bm{R}}(0):=({\bm{W}}(0))^{\frac{1}{L}}$. Note we don’t define ${\bm{R}}(t)$ by $({\bm{W}}(t))^{\frac{1}{L}}$.

The calculation below shows that ${\bm{R}}^{L}(t)$ also satisfies (31).

I.3 Proof for Theorem 6.2

Now we turn to prove Theorem 6.2. Let $P=L/2$. Then (12) can be rewritten as

The following lemma about the growth rate of $\lambda_{k}({\bm{M}})$ is used later in the proof.

Suppose ${\bm{M}}(t)$ satisfies (33), we have for any $T^{\prime}>T$, and $k\in[d]$,

Since $\lambda_{k}({\bm{M}}(t))$ is locally Lipschitz in $t$, by Rademacher’s theorem, we know $\lambda_{k}({\bm{M}}(t))$ is differentiable almost everywhere, and the following holds

When $\frac{\textup{{d}}\lambda_{k}({\bm{M}}(t))}{\textup{{d}}t}$ exists, we have

Let $R>0$. Since $f(\,\cdot\,)$ is $\mathcal{C}^{3}$-smooth, there exists $\beta>0$ such that

for all ${\bm{W}}_{1},{\bm{W}}_{2}$ with $\left\|{\bm{W}}_{1}\right\|_{2},\left\|{\bm{W}}_{2}\right\|_{2}\leq R$.

Let $\kappa=\beta/\mu_{1}$. We assume WLOG that $R\leq\frac{1}{\kappa(P-1)}$. Let $F_{\hat{\alpha}}(x):=\int_{x^{-(P-1)}}^{\hat{\alpha}^{-(P-1)}}\frac{dz}{1+\kappa z^{-P/(P-1)}}$. Then $F^{\prime}_{\hat{\alpha}}(x)=\frac{(P-1)x^{-P}}{1+\kappa x^{P}}=\frac{P-1}{(1+\kappa x^{P})x^{P}}$. We will use this function to bound norm growth. Let $g_{\hat{\alpha},c}(t)=\frac{1}{\hat{\alpha}^{-(P-1)}-\kappa(P-1)c-2\mu_{1}(P-1)t}$. Define $T_{\hat{\alpha}}(r)=\frac{\hat{\alpha}^{-(P-1)}-\kappa(P-1)r-r^{-(P-1)}}{2\mu_{1}(P-1)}$. It is easy to verify that $g_{\hat{\alpha},r}(T_{\hat{\alpha}}(r))=r^{P-1}$.

Let ${\bm{M}}_{0}$ be a PSD matrix with $\|{\bm{M}}_{0}\|_{2}\leq 1$. For ${\bm{M}}(t):={\phi_{m}}(\hat{\alpha}{\bm{M}}_{0},t)$ and $t\leq T_{\hat{\alpha}}(c)$,

Since $\|\nabla f({\bm{M}}^{P})\|_{2}\leq\|\nabla f({\bm{0}})\|_{2}+\beta\|{\bm{M}}\|_{2}^{P}\leq\mu_{1}+\beta(\lambda_{1}({\bm{M}}))^{P}$, by Lemma I.7, we have

If $\|{\bm{M}}(t)\|_{2}<\hat{\alpha}$, then $\|{\bm{M}}(t)\|_{2}\leq g_{\hat{\alpha},c}(t)^{\frac{1}{P-1}}$. If $\|{\bm{M}}(t)\|_{2}\geq\hat{\alpha}$, then by Lemma I.8,

so $\|{\bm{M}}(t)\|_{2}\leq c$ for all $t\leq T_{\hat{\alpha}}(c)$. Applying Lemma I.8 again, we have

which implies $\|{\bm{M}}(t)\|_{2}\leq g_{\hat{\alpha},c}(t)^{\frac{1}{P-1}}$ by definition. ∎

We use $\hat{\phi}_{m}(\widehat{{\bm{M}}}_{0},t)$ to denote the solution of $\widehat{{\bm{M}}}(t)$ when $\widehat{{\bm{M}}}(0)=\widehat{{\bm{M}}}_{0}$. For diagonal matrix $\widehat{{\bm{M}}}_{0}$, $\widehat{{\bm{M}}}(t)$ is also diagonal for any $t$, and it is easy to show that

Unlike depth-2 case, the closed form solution, $\widehat{{\bm{M}}}(t)$ is only tractable for diagonal initialization, i.e., (36) (note that the identity matrix is diagonal). And this is the main barrier for extending our two-phase analysis to the case of general initialization when $L\geq 3$. In Appendix J, we give a more detailed discussion on this barrier.

The following lemma shows that the trajectory of ${\bm{M}}(t)$ is close to $\widehat{{\bm{M}}}(t)$.

Let ${\bm{M}}_{0}$ be a diagonal PSD matrix with $\|{\bm{M}}_{0}\|_{2}\leq 1$. For ${\bm{M}}(t):={\phi_{m}}(\hat{\alpha}{\bm{M}}_{0},t)$ and $\widehat{{\bm{M}}}(t):=\hat{\phi}_{m}(\hat{\alpha}{\bm{M}}_{0},t)$, we have

We bound the difference ${\bm{D}}:={\bm{M}}-\widehat{{\bm{M}}}$ between ${\bm{M}}$ and $\widehat{{\bm{M}}}$.

where the last step is by Lemma I.4. This implies that

Let ${\bm{M}}(t)={\phi_{m}}(\hat{\alpha}{\bm{M}}_{0},t),\widetilde{{\bm{M}}}(t)={\phi_{m}}(\hat{\alpha}\widetilde{{\bm{M}}}_{0},t)$. If $\max\{\|{\bm{M}}_{0}\|_{2},\|\widetilde{{\bm{M}}}_{0}\|_{2}\}\leq 1$. For $t\leq T_{\hat{\alpha}}(r)$, we have

Define ${\bm{D}}(t)={\bm{M}}(t)-\widetilde{{\bm{M}}}(t)$. Then we have

Let $c$ be a sufficiently small constant. Let $\bar{T}:=\frac{-\kappa(P-1)c-c^{-(P-1)}}{2\mu_{1}(P-1)}$. We prove this lemma in the cases of $t\in(-\infty,\bar{T}]$ and $t>\bar{T}$ respectively.

For $t=\bar{T}+\tau$ with $\tau>0$, ${\phi_{m}}({\bm{M}},\tau)$ is locally Lipschitz with respect to ${\bm{M}}$. So

which proves the lemma for $t>\bar{T}$. ∎

For every $t\in(-\infty,+\infty)$, as $\alpha\to 0$, we have:

Let ${\bm{M}}_{\hat{\alpha}}(t):={\phi_{m}}\left(\hat{\alpha}{\bm{I}},\frac{\hat{\alpha}^{-(P-1)}}{2\mu_{1}(P-1)}+t\right)$. Again we let $c$ be a sufficiently small constant and $\bar{T}:=\frac{-\kappa(P-1)c-c^{-(P-1)}}{2\mu_{1}(P-1)}$. We prove in the cases of $t\in(-\infty,\bar{T}]$ and $t>\bar{T}$ respectively.

By Lemma I.11, $\lambda_{1}({\bm{M}}_{\hat{\alpha}}(\bar{T}))=\left\|{\bm{M}}_{\hat{\alpha}}(\bar{T})\right\|_{2}=c+O(c^{P+1})$. For $k\geq 2$,