Understanding Deflation Process in Over-parametrized Tensor Decomposition
Year: 2021

Abstract:
In this paper we study the training dynamics for gradient flow on over-parametrized tensor decomposition problems. Empirically, such training process often first fits larger components and then discovers smaller components, which is similar to a tensor deflation process that is commonly used in tensor decomposition algorithms. We prove that for orthogonally decomposable tensor, a slightly modified version of gradient flow would follow a tensor deflation process and recover all the tensor components. Our proof suggests that for orthogonal tensors, gradient flow dynamics works similarly as greedy low-rank learning in the matrix setting, which is a first step towards understanding the implicit regularization effect of over-parametrized models for low-rank tensors.

heading: 1 Introduction
text: Recently, over-parametrization has been recognized as a key feature of neural network optimization. A line of works known as the Neural Tangent Kernel (NTK) showed that it is possible to achieve zero training loss when the network is sufficiently over-parametrized (Jacot et al., 2018; Du et al., 2018; Allen-Zhu et al., 2018b). However, the theory of NTK implies a particular dynamics called lazy training where the neurons do not move much (Chizat et al., 2019), which is not natural in many settings and can lead to worse generalization performance (Arora et al., 2019b). Many works explored other regimes of over-parametrization (Chizat and Bach, 2018; Mei et al., 2018) and analyzed dynamics beyond lazy training (Allen-Zhu et al., 2018a; Li et al., 2020a; Wang et al., 2020).
Over-parametrization does not only help neural network models. In this work, we focus on a closely related problem of tensor (CP) decomposition. In this problem, we are given a tensor of the form
T ∗ = r∑ i=1 ai(U [:, i]) ⊗4,
where ai ≥ 0 and U [:, i] is the i-th column of U ∈ Rd×r. The goal is to fit T ∗ using a tensor T of a similar form:
T = m∑ i=1 (W [:, i])⊗4 ‖W [:, i]‖2 .
∗Alphabetical order.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Here W is a d × m matrix whose columns are components for tensor T . The model is overparametrized when the number of components m is larger than r. The choice of normalization factor of 1/‖W [:, i]‖2 is made to accelerate gradient flow (similar to Li et al. (2020a); Wang et al. (2020)). Suppose we run gradient flow on the standard objective 12‖T −T
∗‖2F , that is, we evolve W according to the differential equation:
dW
dt = −∇
( 1
2 ‖T − T ∗‖2F
) ,
can we expect T to fit T ∗ with good accuracy? Empirical results (see Figure 1) show that this is true for orthogonal tensor T ∗2 as long as m is large enough. Further, the training dynamics exhibits a behavior that is similar to a tensor deflation process: it finds the ground truth components one-by-one from larger component to smaller component (if multiple ground truth components have similar norm they might be found simultaneously).
In this paper we show that with a slight modification, gradient flow on over-parametrized tensor decomposition is guaranteed to follow this tensor deflation process, and can fit any orthogonal tensor to desired accuracy3(see Section 4 for the algorithm and Theorem 1 for the main theorem). This shows that for orthogonal tensors, the trajectory of modified gradient-flow is similar to a greedy low-rank process that was used to analyze the implicit bias of low-rank matrix factorization (Li et al., 2020b). We emphasize that our goal is not to propose another tensor decomposition algorithm. Instead, we hope our results can serve as a first step in understanding the implicit bias of over-parameterized gradient descent for low-rank tensor problems.

heading: 1.1 Our approach and technique
text: To understand the tensor deflation process shown in Figure 1, intuitively we can think about the discovery and fitting of a ground truth component in two phases. Consider the beginning of the gradient flow as an example. Initially all the components in T are small, which makes T negligible compared to T ∗. In this case each component w in W will evolve according to a simpler dynamics that is similar to tensor power method, where one updates w to T ∗(w⊗3, I)/
∥∥T ∗(w⊗3, I)∥∥ (see Section 3 for details).
For orthogonal tensors, it’s known that tensor power method with random initializations would be able to discover the largest ground truth components (see Anandkumar et al. (2014)). Once the largest ground truth component has been discovered, the corresponding component (or multiple components) w will quickly grow in norm, which eventually fits the ground truth component. The flat regions in the trajectory in Figure 1 correspond to the period of time where the components w’s are small and
2We say T ∗ is an orthogonal tensor if the ground truth components U [:, i]’s are orthonormal. 3Due to some technical challenges, we actually require the target accuracy to be at least exp(−o(d/ log d)). This is only a very mild restriction since the dependence is exponential in d, and in practice, d is usually large and this lower bound can easily drop below the numerical precision.
T − T ∗ remains stable, while the decreasing regions correspond to the period of time where a ground truth component is being fitted.
However, there are many challenges in analyzing this process. The main problem is that the gradient flow would introduce a lot of dependencies throughout the trajectory, making it harder to analyze the fitting of later ground truth components, especially ones that are much smaller. We modify the algorithm to include a reinitialization step per epoch, which alleviates the dependency issue. Even after the modification we still need a few more techniques:
Local stability One major problem in analyzing the dynamics in a later stage is that the components used to fit the previous ground truth components are still moving according to their gradients, therefore it might be possible for these components to move away. To address this problem, we add a small regularizer to the objective, and give a new local stability analysis that bounds the distance to the fitted ground truth component both individually and on average. The idea of bounding the distance on average is important as just assuming each component w is close enough to the fitted ground truth component is not sufficient to prove that w cannot move far. While similar ideas were considered in Chizat (2021), the setting of tensor decomposition is different.
Norm/Correlation relation A key step in our analysis establishes a relationship between norm and correlation: we show if a component w crosses a certain norm threshold, then it must have a very large correlation with one of the ground truth components. This offers an initial condition for local stability and makes sure the residual T ∗− T is almost close to an orthogonal tensor. Establishing this relation is difficult as unlike the high level intuition, we cannot guarantee T ∗ − T remains unchanged even within a single epoch: it is possible that one ground truth component is already fitted while no large component is near another ground truth component of same size. In previous work, Li et al. (2020a) deals with a similar problem for neural networks using gradient truncation that prevents components from growing in the first phase (and as a result has super-exponential dependency on the ratio between largest and smallest ai). We give a new technique to control the influence of ground truth components that are fitted within this epoch, so we do not need the gradient truncation and can characterize the deflation process.

heading: 1.2 Related works
text: Neural Tangent Kernel There is a recent line of work showing the connection between Neural Tangent Kernel (NTK) and sufficiently wide neural networks trained by gradient descent (Jacot et al., 2018; Allen-Zhu et al., 2018b; Du et al., 2018, 2019; Li and Liang, 2018; Arora et al., 2019b,c; Zou et al., 2020; Oymak and Soltanolkotabi, 2020; Ghorbani et al., 2021). These papers show when the width of a neural network is large enough, it will stay around the initialization and its training dynamic is close to the dynamic of the kernel regression with NTK. In this paper we go beyond the NTK setting and analyze the trajectory from a very small initialization.
Mean-field analysis There is another line of works that use mean-field approach to study the optimization for infinite-wide neural networks (Mei et al., 2018; Chizat and Bach, 2018; Nguyen and Pham, 2020; Nitanda and Suzuki, 2017; Wei et al., 2019; Rotskoff and Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020). Chizat et al. (2019) showed that, unlike NTK regime, the parameters can move away from its initialization in mean-field regime. However, most of the existing works need width to be exponential in dimension and do not provide a polynomial convergence rate.
Beyond NTK There are many works showing the gap between neural networks and NTK (AllenZhu and Li, 2019; Allen-Zhu et al., 2018a; Yehudai and Shamir, 2019; Ghorbani et al., 2019, 2020; Dyer and Gur-Ari, 2019; Woodworth et al., 2020; Bai and Lee, 2019; Bai et al., 2020; Huang and Yau, 2020; Chen et al., 2020). In particular, Li et al. (2020a) and Wang et al. (2020) are closely related with our setting. While Li et al. (2020a) focused on learning two-layer ReLU neural networks with orthogonal weights, they relied on the connection between tensor decomposition and neural networks (Ge et al., 2017) and essentially worked with tensor decomposition problems. In their result, all the ai’s are within a constant factor and all components are learned simultaneously. We allow ground truth components with very different scale and show a deflation phenomenon. Wang et al. (2020) studied learning a low-rank non-orthogonal tensor, but they only showed the learned tensor T will eventually be close to the ground truth tensor T ∗ and does not guarantee the components of T will
align with the components of T ∗. On the other hand, we fully characterize the training trajectory and the components of the learned tensor.
Implicit regularization Many works recently showed that different optimization methods tend to converge to different optima and have different optimization trajectories in several settings (Saxe et al., 2014; Soudry et al., 2018; Nacson et al., 2019; Ji and Telgarsky, 2018a,b, 2019, 2020; Gunasekar et al., 2018a,b; Moroshko et al., 2020; Arora et al., 2019a; Lyu and Li, 2019; Chizat and Bach, 2020). In particular, Saxe et al. (2014) related the dynamics of gradient descent to the magnitude of the singular values of the target weight matrices for linear networks with orthogonal inputs. The phenomenon there is qualitatively similar to our results, but the settings and the proof techniques are very different. The more related and recent works are Li et al. (2020b) and Razin et al. (2021). Li et al. (2020b) studied matrix factorization problem and showed gradient descent with infinitesimal initialization is similar to greedy low-rank learning, which is a multi-epoch algorithm that finds the best approximation within the rank constraint and relax the constraint after every epoch. Razin et al. (2021) studied the tensor factorization problem and showed that it biases towards low rank tensor. Both of these works considered partially observable matrix or tensor and are only able to fully analyze the first epoch (i.e., recover the largest direction). We focus on a simpler setting with fully-observable ground truth tensor and give a complete analysis of learning all the ground truth components.

heading: 1.3 Outline
text: In Section 2 we introduce the basic notations and problem setup. In Section 3 we review tensor deflation process and tensor power method. We then give our algorithm in Section 4. Section 5 gives the formal main theorem and discusses high-level proof ideas. We conclude in Section 6 and discuss some limitations of the work. The detailed proofs and additional experiments are left in the appendix.

heading: 2 Preliminaries
text: Notations We use upper-case letters to denote matrices and tensors, and lower-case letters to denote vectors. For any positive integer n, we use [n] to denote the set {1, 2, · · · , n}. We use Id to denote d× d identity matrix, and omit the subscript d when the dimension is clear. We use δ0Unif(Sd−1) to denote the uniform distribution over (d− 1)-dimensional sphere with radius δ0. For vector v, we use ‖v‖ to denote its `2 norm. We use vk to denote the k-th entry of vector v, and use v−k to denote vector v with its k-th entry removed. We use v̄ to denote the normalized vector v̄ = v/ ‖v‖, and use v̄k to denote the k-th entry of v̄. For a matrix A, we use A[:, i] to denote its i-th column and col(A) to denote the set of all column vectors of A. For matrix M or tensor T , we use ‖M‖F and ‖T‖F to denote their Frobenius norm, which is equal to the `2 norm of their vectorization.
For simplicity we restrict our attention to symmetric 4-th order tensors. For a vector v ∈ Rd, we use v⊗4 to denote a d × d × d × d tensor whose (i, j, k, l)-th entry is equal to vivjvkvl. Suppose T = ∑ w w ⊗4, we define T (v⊗4) as ∑ w 〈w, v〉 4, T (v⊗3, I) as ∑ w 〈w, v〉 3 w, and T (v⊗2, u, I) =∑
w 〈w, v〉 2 〈w, u〉w.
For clarity, we always call a component in T ∗ as ground truth component and call a component in our model T simply as component.
Problem setup We consider the problem of fitting a 4-th order tensor. The components of the ground truth tensor is arranged as columns of a matrix U ∈ Rd×r, and the tensor T ∗ is defined as
T ∗ = r∑ i=1 ai(U [:, i] ⊗4),
where a1 ≥ a2 ≥ · · · ≥ ar ≥ 0 and ∑r i=1 ai = 1. For convenience in the analysis, we assume
ai ≥ / √ d for all i ∈ [r]. This is without loss of generality because the target accuracy is and we
can safely ignore very small ground truth components with ai < / √ d. In this paper, we focus on the case where the components are orthogonal—that is, the columns U [:, i]’s are orthonormal. For
simplicity we assume without loss of generality that U [:, i] = ei where ei is the i-th standard basis vector4. To reduce the number of parameters we also assume r = d, again this is without loss of generality because we can simply set ai = 0 for i > r.
There can be many different ways to parametrize the tensor that we use to fit T ∗. Following previous works (Wang et al., 2020; Li et al., 2020a), we use an over-parameterized and two-homogeneous tensor
T = m∑ i=1 W [:, i]⊗4 ‖W [:, i]‖2 .
Here W ∈ Rd×m is a matrix with m columns that corresponds to the components in T . It is overparametrized when m > r.
Since the tensor T only depends on the set of columns W [:, i] instead of the orderings of the columns, for the most part of the paper we will instead write the tensor T as
T = ∑
w∈col(W )
w⊗4 ‖w‖2 ,
where col(W ) is the set of all the column vectors in W . This allows us to discuss the dynamics of coordinates for a component w without using the index for the component. In particular, wi always represents the i-th coordinate of the vector w. This representation is similar to the mean-field setup (Chizat and Bach, 2018; Mei et al., 2018) where one considers a distribution on w, however since we do not rely on analysis related to infinite-width limit we use the sum formulation instead. For the ease of presentation, we choose to restrict our setting to fourth-order tensor decomposition, but our results can be easily generalized to tensor with order at least three.

heading: 3 Tensor deflation process and tensor power method
text: In this section we will first discuss the basic tensor deflation process for orthogonal tensor decomposition. Then we show the connection between the tensor power method and gradient flow.
Tensor deflation For orthogonal tensor decomposition, a popular approach is to first fit the largest ground truth component in the tensor, then subtract it out and recurse on the residual. The general process is given in Algorithm 1. In this process, there are multiple ways to find the best rank-1 approximation. For example, Anandkumar et al. (2014) uses tensor power method, which picks many random vectors w, and update them as w = T ∗(w⊗3, I)/
∥∥T ∗(w⊗3, I)∥∥. Algorithm 1 Tensor Deflation Process
Input: Tensor T ∗ Output: Components W such that T ∗ ≈ ∑ w∈col(W ) w
⊗4/‖w‖2 Initially let the residual R be T ∗. while ‖R‖F is large do
Find the best rank 1 approximation w⊗4/‖w‖2 for R. Add w as a new column in W , and let R = R− w⊗4/‖w‖2.
end while
Tensor power method and gradient flow If we run tensor power method using a tensor T ∗ that is equal to ∑d i=1 aie ⊗4 i , then a component w will converge to the direction of ei where i is equal to arg maxi aiw̄ 2 i . If there is a tie (which happens with probability 0 for random w), then the point will be stuck at a saddle point. Let’s consider running gradient flow on W with objective function 12 ‖T − T ∗‖2F as T :=∑
w∈col(W ) w ⊗4/‖w‖2. If T does not change much, the residual R := T ∗ − T is close to a
4This is without loss of generality because gradient flow (and our modifications) is invariant under rotation of the ground truth parameters.
constant. In this case the trajectory of one component w is determined by the following differential equation:
dw dt = 4R(w̄⊗2, w, I)− 2R(w̄⊗4)w. (1)
To understand how this process works, we can take a look at dw 2 i /dt
w2i (intuitively this corresponds to
the growth rate for w2i ). If R ≈ T ∗ then we have:
dw2i /dt
w2i ≈ 8aiw̄2i − 4 ∑ j∈[d] ajw̄ 4 j .
From this formula it is clear that the coordinate with larger aiw̄2i has a faster growth rate, so eventually the process will converge to ei where i is equal to arg maxi aiw̄2i , same as the tensor power method. Because of their similarity later we refer to dynamics in Eqn. (1) as tensor power dynamics.

heading: 4 Our algorithm
text: Our algorithm is a modified version of gradient flow as described in Algorithm 2. First, we change the loss function to
L(W ) = 1
2 ‖T − T ∗‖2F +
λ 2 ‖W‖2F .
The additional small regularization λ2 ‖W‖ 2 F allows us to prove a local stability result that shows if there are components w that are close to the ground truth components in direction, then they will not move too much (see Section 5.1).
Our algorithm runs in multiple epochs with increasing length. We use W (s,t) to denote the weight matrix in epoch s at time t. We use similar notation for tensor T (s,t). In each epoch we try to fit ground truth components with ai ≥ β(s). In general, the time it takes to fit one ground truth direction is inversely proportional to its magnitude ai. The earlier epochs have shorter length so only large directions can be fitted, and later epochs are longer to fit small directions.
At the middle of each epoch, we reinitialize all components that do not have a large norm. This serves several purposes: first we will show that all components that exceed the norm threshold will have good correlation with one of the ground truth components, therefore giving an initial condition to the local stability result; second, the reinitialization will reduce the dependencies between different epochs and allow us to analyze each epoch almost independently. These modifications do not change the dynamics significantly, however they allow us to do a rigorous analysis.
Algorithm 2 Modified Gradient Flow Input: Number of components m, initialization scale δ0, re-initialization threshold δ1, increasing rate of epoch length γ, target accuracy , regularization coefficient λ Output: Tensor T satisfying ‖T − T ∗‖F ≤ . Initialize W (0,0) as a d×m matrix with each column w(0,0) i.i.d. sampled from δ0Unif(Sd−1). β(0) ← ∥∥T (0,0) − T ∗∥∥ F
; s← 0 while ∥∥T (s,0) − T ∗∥∥ F > do
Phase 1: Starting from W (s,0), run gradient flow for time t(s)1 = O( d β(s) log(d) ). Reinitialize all components that have `2 norm less than δ1 by sampling i.i.d. from δ0Unif(Sd−1).
Phase 2: Starting from W (s,t (s) 1 ), run gradient flow for t(s)2 − t (s) 1 = O( log(1/δ1)+log(1/λ) β(s) ) time W (s+1,0) ←W (s,t (s) 2 ); β(s+1) ← β(s)(1− γ); s← s+ 1
end while

heading: 5 Main theorem and proof sketch
text: In this section we discuss the ideas to prove the following main theorem5
Theorem 1. For any ≥ exp(−o(d/ log d)), there exists γ = Θ(1), m = poly(d), λ = min{O(log d/d), O( /d1/2)}), α = min{O(λ/d3/2), O(λ2), O( 2/d4)}, δ1 = O(α3/2/m1/2), δ0 = Θ(δ1α/ log
1/2(d)) such that with probability 1 − 1/poly(d) in the (re)-initializations, Algorithm 2 terminates in O(log(d/ )) epochs and returns a tensor T such that
‖T − T ∗‖F ≤ .
Intuitively, epoch s of Algorithm 2 will try to discover all ground truth components with ai that is at least as large as β(s). The algorithm does this in two phases. In Phase 1, the small components w will evolve according to tensor power dynamics. For each ground truth component with large enough ai that has not been fitted yet, we hope there will be at least one component in W that becomes large and correlated with ei. We call such ground truth components “discovered”. Phase 1 ends with a check that reinitilizes all components with small norm. Phase 2 is relatively short, and in Phase 2 we guarantee that every ground truth component that has been discovered become “fitted”, which means the residual T − T ∗ becomes small in this direction. However, there are still many difficulties in analyzing each of the steps. In particular, why would ground truth components that are fitted in previous epochs remain fitted? How to guarantee only components that are correlated with a ground truth component grow to a large norm? Why wouldn’t the gradient flow in Phase 2 mess up with the initialization we require in Phase 1? We discuss the high level ideas to solve these issues. In particular, in Section 5.1 we first give an induction hypothesis that is preserved throughout the algorithm, which guarantees that every ground truth component that is fitted remains fitted. In Section 5.2 we discuss the properties in Phase 1, and in Section 5.3 we discuss the properties in Phase 2.

heading: 5.1 Induction hypothesis and local stability
text: In order to formally define what it means for a ground truth component to be “discovered” or “fitted”, we need some more definitions and notations.
Definition 1. Define S(s,t)i ⊆ [m] as the subset of components that satisfy the following conditions: the k-th component is in S(s,t)i if and only if there exists some time (s
′, t′) that is no later than (s, t) and no earlier than the latest re-initialization of W [:, k] such that∥∥∥W (s′,t′)[:, k]∥∥∥ = δ1 and [W (s′,t′)[:, k]i]2 ≥ 1− α2. We say that ground truth component i is discovered in epoch s at time t, if S(s,t)i is not empty.
Intuitively, S(s,t)i is a subset of components in W such that they have large enough norm and good correlation with the i-th ground truth component. Although such components may not have a large enough norm to fit ai yet, their norm will eventually grow. Therefore we say ground truth component i is discovered when such components exist.
For convenience, we shorthand w(s,t) ∈ {W (s,t)[:, j]|j ∈ S(s,t)i } by w(s,t) ∈ S (s,t) i . Now we will discuss when a ground truth component is fitted, for that, let
â (s,t) i = ∑ w(s,t)∈S(s,t)i ∥∥∥w(s,t)∥∥∥2 . Here â(s,t)i is the total squared norm for all the components in S (s,t) i . We say a ground truth component is fitted if ai − â(s,t)i ≤ 2λ.
5In the theorem statement, we have a parameter α that is not used in our algorithm but is very useful in the analysis (see for example Definition 1). Basically, α measures the closeness between a component and its corresponding ground truth direction (see more in Section 5.1).
Note that one can partition the columns in W using sets S(s,t)i , giving d groups and one extra group that contains everything else. We define the extra group as S(s,t)∅ := [m] \ ⋃ k∈[d] S (s,t) k .
For each of the non-empty S(s,t)i , we can take the average of its component (weighted by ∥∥w(s,t)∥∥2):
E(s,t)i,w f(w (s,t)) :=
1
â (s,t) i ∑ w(s,t)∈S(s,t)i ∥∥∥w(s,t)∥∥∥2 f(w(s,t)). If S(s,t)i = ∅, we define E (s,t) i,w f(w (s,t)) as zero. Now we are ready to state the induction hypothesis:
Proposition 1 (Induction hypothesis). In the setting of Theorem 1, for any epoch s and time t and every k ∈ [d], the following hold. (a) For any w(s,t) ∈ S(s,t)k , we have [ w̄ (s,t) k ]2 ≥ 1− α.
(b) If S(s,t)k is nonempty, E (s,t) k,w
[ w̄
(s,t) k ]2 ≥ 1− α2 − 4smδ21 .
(c) We always have ak−â(s,t)k ≥ λ/6−smδ21; if ak ≥ β(s) 1−γ , we further know ak−â (s,t) k ≤ λ+smδ21 .
(d) If w(s,t) ∈ S(s,t)∅ , then ‖w(s,t)‖ ≤ δ1.
We choose δ21 small enough so that smδ 2 1 is negligible compared with α 2 and λ. Note that if Proposition 1 is maintained throughout the algorithm, all the large components will be fitted, which directly implies Theorem 1. Detailed proof is deferred to Appendix D.
Condition (c) shows that for a ground truth component k with large enough ak, it will always be fitted after the corresponding epoch (recall from Theorem 1 that λ = O(ε/ √ d)). Condition (d) shows that components that did not discover any ground truth components will always have small norm (hence negligible in most parts of the analysis). Conditions (a)(b) show that as long as a ground truth component k has been discovered, all components that are in S(s,t)k will have good correlation, while the average of all such components will have even better correlation. The separation between individual correlation and average correlation is important in the proof. With only individual bound, we cannot maintain the correlation no matter how small α is. Here is an example below:
Claim 2. Suppose T ∗ = e⊗4k and T = v ⊗4/ ‖v‖2 + w⊗4/ ‖w‖2 with ‖w‖2 + ‖v‖2 ∈ [2/3, 1]. Suppose v̄2k = 1−α and v̄k = w̄k, v̄−k = −w̄−k. Assuming ‖v‖ 2 ≤ c1 and α ≤ c2 for small enough constants c1, c2, we have ddt v̄ 2 k < 0.
In the above example, both v̄ and w̄ are close to ek but they are opposite in other directions (v̄−k = w̄−k). The norm of v is very small compared with that of w. Intuitively, we can increase v−k so that the average of v and w is more aligned with ek. See the rigorous analysis in Appendix A.6.
The induction hypothesis will be carefully maintained throughout the analysis. The following lemma guarantees that in the gradient flow steps the individual and average correlation will be maintained.
Lemma 3. In the setting of Theorem 1, suppose Proposition 1 holds in epoch s at time t, we have
d dt [w̄(s,t)]2 ≥ 8
( ak − â(s,t)k )( 1− [w̄(s,t)k ] 2 ) −O ( α1.5 ) ,
d dt E(s,t)k,w [w̄ (s,t) k ]
2 ≥ 8 ( ak − â(s,t)k )( 1− E(s,t)k,w [w̄ (s,t) k ] 2 ) −O(α3).
In particular, when ak − â(s,t)k ≥ Ω(λ) = Ω( √ α), we have ddt [w̄ (s,t) k ] 2 > 0 when [w̄(s,t)k ] 2 = 1− α and ddtE (s,t) k,w [w̄ (s,t) k ] 2 > 0 when E(s,t)k,w [w̄ (s,t) k ] 2 = 1− α2.
The detailed proof for the local stability result can be found in Appendix A. Of course, to fully prove the induction hypothesis one needs to talk about what happens when a component enters S(s,t)i , and what happens at the reinitialization steps. We discuss these details in later subsections.

heading: 5.2 Analysis of Phase 1
text: In Phase 1 our main goal is to discover all the components that are large enough. We also need to maintain Proposition 1. Formally we prove the following: Lemma 4 (Main Lemma for Phase 1). In the setting of Theorem 1, suppose Proposition 1 holds at (s, 0). For t(s)1 := t (s)′ 1 + t (s)′′ 1 + t (s)′′′ 1 with t (s)′ 1 = Θ(d/(β (s) log d)), t(s)′′1 = Θ(d/(β (s) log3 d)), t (s)′′′ 1 = Θ(log(d/α)/β (s)), with probability 1− 1/poly(d) we have
1. Proposition 1 holds at (s, t) for any 0 ≤ t < t(s)1 , and also for t = t (s) 1 after reinitialization.
2. If ak ≥ β(s) and S(s,0)k = ∅, we have S (s,t
(s) 1 )
k 6= ∅ and â (s,t
(s) 1 )
k ≥ δ21 .
3. If S(s,0)k = ∅ and S (s,t
(s) 1 )
k 6= ∅, we have ak ≥ Cβ(s) for universal constant 0 < C < 1.
Property 2 shows that large enough ground truth components are always discovered, while Property 3 guarantees that no small ground truth components can be discovered. Our proof relies on initial components being “lucky” and having higher than usual correlation with one of the large ground truth components. To make this clear we separate components into different sets (here we use v to denote a component in W ): Definition 2 (Partition of (re-)initialized components). For each direction i ∈ [d], define the set of good components S(s)i,good and the set of potential components S (s) i,pot as follow, where Γ (s) i := 1/(8ait (s)′ 1 ) if S (s,0) i = ∅, and Γ (s) i := 1/(8λt (s)′ 1 ) otherwise. Here ρ (s) i := cρΓ (s) i and cρ is a small enough absolute constant.
S (s) i,good := {k | [v̄ (s,0) i ] 2 ≥ Γ(s)i + ρ (s) i , [v̄ (s,0) j ] 2 ≤ Γ(s)j − ρ (s) j ,∀j 6= i and v (s,0) = W (s,0)[:, k]},
S (s) i,pot := {k | [v̄ (s,0) i ] 2 ≥ Γ(s)i − ρ (s) i and v (s,0) = W (s,0)[:, k]}.
Let S(s)good := ∪iS (s) i,good and S (s) pot := ∪iS (s) i,pot. We also define the set of bad components S (s) bad.
S (s) bad := {k | ∃i 6= j s.t. [v̄ (s,0) i ] 2 ≥ Γ(s)i − ρ (s) i , [v̄ (s,0) j ] 2 ≥ Γ(s)j − ρ (s) j and v (s,0) = W (s,0)[:, k]}.
For convenience, we shorthand v(s,t) ∈ {W (s,t)[:, j]|j ∈ Si,good} by v(s,t) ∈ Si,good (same for Si,pot and Sbad). Intuitively, the good components will grow very quickly and eventually pass the norm threshold. Since both good and potential components only have one large coordinate, they will become correlated with that ground truth component when their norm is large. The bad components are correlated with two ground truth components so they can potentially have a large norm while not having a very good correlation with either one of them. In the proof we will guarantee with probability at least 1 − 1/poly(d) that good components exists for all large enough ground truth components and there are no bad components. The following lemma characterizes the trajectories of different type of components: Lemma 5. In the setting of Lemma 4, for every i ∈ [d]
1. (Only good/potential components can become large) If v(s,t) 6∈ S(s)pot, ∥∥v(s,t)∥∥ = O(δ0) and
[v̄ (s,t) i ] 2 = O(log(d)/d) for all i ∈ [d] and t ≤ t(s)1 .
2. (Good components discover ground truth components) If S(s)i,good 6= ∅, there exists v(s,t (s) 1 ) such that ∥∥∥v(s,t(s)1 )∥∥∥ ≥ δ1 and S(s,t(s)1 )i 6= ∅.
3. (Large components are correlated with ground truth components) If ∥∥v(s,t)∥∥ ≥ δ1 for some
t ≤ t(s)1 , there exists i ∈ [d] such that v(s,t) ∈ S (s,t) i .
The proof of Lemma 5 is difficult as one cannot guarantee that all the ground truth components that we are hoping to fit in the epoch will be fitted simultaneously. However we are able to show that T − T ∗ remains near-orthogonal and control the effect of changing T − T ∗ within this epoch. The details are in Appendix B.

heading: 5.3 Analysis of Phase 2
text: In Phase 2 we will show that every ground truth component that’s discovered in Phase 1 will become fitted, and the reinitialized components will preserve the desired initialization conditions.
Lemma 6 (Main Lemma for Phase 2). In the setting of Theorem 1, suppose Proposition 1 holds at (s, t
(s) 1 ), we have for t (s) 2 − t (s) 1 := O( log(1/δ1)+log(1/λ) β(s) )
1. Proposition 1 holds at (s, t) for any t(s)1 ≤ t ≤ t (s) 2 .
2. If S(s,t (s) 1 )
k 6= ∅, we have ak − â (s,t
(s) 2 )
k ≤ 2λ. 3. For any component v that was reinitialized at t(s)1 , we have ∥∥∥v(s,t(s)2 ∥∥∥2 = Θ(δ20) and[
v̄ (s,t
(s) 2 )
i
]2 = [ v̄ (s,t (s) 1 )
i
]2 ± o ( log d d ) for every i ∈ [d].
The main idea is that as long as a direction has been discovered, the norm of the corresponding components will increase very fast. The rate of that is characterized by the following lemma.
Lemma 7 (informal). In the setting of Theorem 6, for any t(s)1 ≤ t ≤ t (s) 2 ,
d dt â (s,t) k ≥
( 2(ak − â(s,t)k )− λ−O ( α2 )) â (s,t) k .
In particular, after O( log(1/δ1)+log(1/λ)ak ) time, we have ak − â (s,t) k ≤ λ.
By the choice of δ1 and λ, the length of Phase 2 is much smaller than the amount of time needed for the reinitialized components to move far, allowing us to prove the third property in Lemma 6. Detailed analysis is deferred to Appendix C.

heading: 6 Conclusion
text: In this paper we analyzed the dynamics of gradient flow for over-parametrized orthogonal tensor decomposition. With very mild modification to the algorithm (a small regularizer and some reinitializations), we showed that the trajectory is similar to a tensor deflation process and the greedy low-rank procedure in Li et al. (2020b). These modifications allowed us to prove strong guarantees for orthogonal tensors of any rank, while not changing the empirical behavior of the algorithm. We believe such techniques would be useful in later analysis for the implicit bias of tensor problems.
A major limitation of our work is that it only applies to orthogonal tensors. Going beyond this would require significantly new ideas—we observed that for general tensors, overparametrized gradient flow may have a very different behavior compared to the greedy low-rank procedure, as it is possible for two large component in the same direction to split into two different directions (see more details in Appendix E). We leave that as an interesting open problem.

heading: Acknowledgements
text: Rong Ge, Xiang Wang and Mo Zhou are supported in part by NSF Award CCF-1704656, CCF1845171 (CAREER), CCF-1934964 (Tripods), a Sloan Research Fellowship, and a Google Faculty Research Award.

Review summaries:
The paper analyzes gradient descent on the orthogonal tensor decomposition problem. More specifically, they consider a ground 4th order tensor with r orthonormal factors and an overparametrized estimate with m>r factors, and analyze the performance of (a variant of) gradient descent on the problem.  The paper establishes a close relationship between the dynamics of gradient descent and deflation based tensor decomposition. In deflation based tensor decomposition, the factors are recovered one at a time, and the factors recovered so far are subtracted from the original tensor to recover the next factor. The paper shows that the dynamics of a modified gradient descent algorithm are similar: the algorithm also recovers factors one at a time.   The analysis uses the following key ideas, some of which are new and some are variants of previous techniques. First, the paper states a straightforward relation between gradient descent and the well-known tensor power method. This relation ensures that a non-negligible correlation with one of the original factors will get amplified on performing the gradient descent updates (similar to the tensor or matrix power methods which are well-analyzed). Next, the paper shows a local stability condition to ensure that once an iterate converges to one of the original factors, then it does not move away. This is done by adding a regularizer. Finally, the paper needs to do a periodic re-randomization of the iterates which have not converged so far to remove dependencies. To identify which factors have not converged so far, the paper establishes a relation between the norm of a factor and its correlation with any of the true factors. It appears that this is the most novel and important part of the proof. While periodic re-randomization has been introduced in the context of tensor power method and alternating least squares (ALS) and their variants, establishing a link between norm and correlation appears novel in this context.  The paper also does an empirical validation of the claims, but this is more to support the theory than a contribution in itself.

This work studied a special training behavior of over-parametrized tensor decomposition by gradient-based methods. The paper pointed out this phenomena as the tensor deflation process and rigorously proved its convergence (Thm. 1) and the behaviors of factors (Prop. 1 and etc.) during the training. The implicit regularization by gradient method is an interesting and important issue in deep learning, and this work shows it would also happen in tensor decomposition. It may help the tensor researchers to develop more effective methods to seek for low-tensor-rank approximations, especially when the tensor model owns complex structures like tensor networks. However, unlike the closed work (li et al., 2020b) on matrices, several additional conditions and assumptions are imposed in the paper (norm regularization and re-initialization shown in Algo. 2), and the theoretical results seem to be only held for the decomposition task (the Fro. objective) and to be difficult for people to extend it into more general forms like tensor learning. It therefore limits the impact on the both tensor and machine learning community. 

The paper analyzes the dynamics of (a slightly modified) gradient flow for overparameterized symmetric CP tensor decomposition of order 4, where the minimized objective is the L2 distance from an orthogonally decomposable target tensor. The main technical result (Theorem 1) establishes that the modified gradient flow follows a trajectory similar to that of a tensor deflation process, in which components are learned one after the other, recovering the target tensor. Empirical evaluations corroborate the theoretical analysis.