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