Understanding Deflation Process in Over-parametrized Tensor Decomposition

1 Introduction . 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 . 1.1 Our approach and technique . 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 can not 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 can not 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 . 1.2 Related works . 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 .

Summary: 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.