B EHIND THE S CENES OF G RADIENT D ESCENT : A T RAJECTORY A NALYSIS VIA B ASIS F UNCTION D ECOMPOSITION
BSTRACT
This work analyzes the solution trajectory of gradient-based algorithms via a novel basis function decomposition .
We show that, although solution trajectories of gradient-based algorithms may vary depending on the learning task, they behave almost monotonically when projected onto an appropriate orthonormal function basis.
Such projection gives rise to a basis function decomposition of the solution trajectory.
Theoretically, we use our proposed basis function decomposition to establish the convergence of gradient descent (GD) on several representative learning tasks.
In particular, we improve the convergence of GD on symmetric matrix factorization and provide a completely new convergence result for the orthogonal symmetric tensor decomposition.
Empirically, we illustrate the promise of our proposed framework on realistic deep neural networks (DNNs) across different architectures, gradient-based solvers, and datasets.
Our key finding is that gradient-based algorithms monotonically learn the coefficients of a particular orthonormal function basis of DNNs defined as the eigenvectors of the conjugate kernel after training. Our code is available at github.com/jianhaoma/function-basis-decomposition .
I NTRODUCTION
Learning highly nonlinear models amounts to solving a nonconvex optimization problem, which is typically done via different variants of gradient descent (GD).
But how does GD learn nonlinear models?
Classical optimization theory asserts that, in the face of nonconvexity, GD and its variants may lack any meaningful optimality guarantee; they produce solutions that—while being first- or second-order optimal (Nesterov, 1998; Jin et al., 2017)—may not be globally optimal.
In the rare event where the GD can recover a globally optimal solution, the recovered solution may correspond to an overfitted model rather than one with desirable generalization.
Inspired by the large empirical success of gradient-based algorithms in learning complex models, recent work has postulated that typical training losses have benign landscapes : they are devoid of spurious local minima and their global solutions coincide with true solutions—i.e., solutions corresponding to the true model.
For instance, different variants of low-rank matrix factorization (Ge et al., 2016; 2017) and deep linear NNs (Kawaguchi, 2016) have benign landscapes.
However, when spurious solutions do exist (Safran & Shamir, 2018) or global and true solutions do not coincide (Ma & Fattahi, 2022b), such a holistic view of the optimization landscape cannot explain the success of gradient-based algorithms.
To address this issue, another line of research has focused on analyzing the solution trajectory of different algorithms.
Analyzing the solution trajectory has been shown extremely powerful in sparse recovery (Vaskevicius et al., 2019), low-rank matrix factorization (Li et al., 2018;
Stöger & Soltanolkotabi, 2021), and linear DNNs (Arora et al., 2018; Ma & Fattahi, 2022a).
However, these analyses are tailored to specific models and thereby cannot be directly generalized.
In this work, we propose a unifying framework for analyzing the optimization trajectory of GD based on a novel basis function decomposition .
We show that, although the dynamics of GD may vary
[Figure]
drastically on different models, they behave almost monotonically when projected onto an appropriate choice of orthonormal function basis.
Motivating example: Our first example illustrates this phenomenon on DNNs.
We study the optimization trajectories of two adaptive gradient-based algorithms, namely AdamW and LARS, on three different DNN architectures, namely AlexNet (Krizhevsky et al., 2017), ResNet-18 (He et al., 2016), and Vision Transformer (ViT)
(Dosovitskiy et al., 2020) with the CIFAR-10 dataset.
The first row of the Figure 1 shows the top- 5 coefficients of the solution trajectory when projected onto a randomly generated orthonormal basis.
We see that the trajectories of the coefficients are highly non-monotonic and almost indistinguishable (they range between -0.04 to 0.06), implying that the energy of the obtained model is spread out on different orthogonal components.
The second row of Figure 1 shows the same trajectory after projecting onto an orthogonal basis defined as the eigenvectors of the conjugate kernel after training (Long, 2021)
(see Section 3.4 and Appendix B for more details).
Unlike the previous case, the top-5 coefficients carry more energy and behave monotonically (modulo the small fluctuations induced by the stochasticity in the algorithm) in all three architectures, until they plateau around their steady state.
In other words, the algorithm behaves more monotonically after projecting onto a correct choice of orthonormal basis.
1. M AIN C ONTRIBUTIONS
The monotonicity of the projected solution trajectory motivates the use of an appropriate basis function decomposition to analyze the behavior of gradient-based algorithms.
In this paper, we show how an appropriate basis function decomposition can be used to provide a much simpler convergence analysis for gradient-based algorithms on several representative learning problems, from simple kernel regression to complex DNNs.
Our main contributions are summarized below:
- Global convergence of GD via basis function decomposition: We prove that GD learns the coefficients of an appropriate function basis that forms the true model.
In particular, we show that GD learns the true model when applied to the expected ℓ 2 -loss under certain gradient independence and gradient dominance conditions.
Moreover, we characterize the convergence rate of GD, identifying conditions under which it enjoys linear or sublinear convergence rates.
Our result does not require a benign landscape for the loss function and can be applied to both convex and nonconvex settings.
- Application in learning problems: We show that our general framework is well-suited for analyzing the solution trajectory of GD on different representative learning problems.
Unlike the existing results, our proposed method leads to a much simpler trajectory analysis of GD for much broader classes of models.
Using our technique, we improve the convergence of
GD on the symmetric matrix factorization and provide an entirely new convergence result for GD on the orthogonal symmetric tensor decomposition.
We also prove that GD enjoys an incremental learning phenomenon in both problems.
- Empirical validation on DNNs
: We empirically show that our proposed framework applies to DNNs beyond GD.
More specifically, we show that different gradient-based algorithms monotonically learn the coefficients of a particular function basis defined as the eigenvectors of the conjugate kernel after training (also known as “after kernel regime”).
We show that this phenomenon happens across different architectures, datasets, solvers, and loss functions, strongly motivating the use of function basis decomposition to study deep learning.
G ENERAL F RAMEWORK : F UNCTION B ASIS D ECOMPOSITION
We study the optimization trajectory of GD on the expected (population) ℓ 2 -loss [Equation] (expected ℓ
[Equation]
Here the input x ∈ R d is drawn from an unknown distribution D , and the output label y is generated as y
= f
⋆ ( x ) + ε , where ε is an additive noise, independent of x , with mean E [ ε ] = 0 and variance E [ ε 2 ]
= σ 2 ε < ∞ .
The model f θ ( x ) is characterized by a parameter vector θ ∈ R m , which naturally induces a set of admissible models (model space for short) F Θ :=
{ f θ : θ ∈ R m } .
We do not require the true model
f ⋆ to lie within the model space; instead, we seek to obtain a model f θ ⋆ ∈ F Θ that is closest to f ⋆ in L 2 ( D ) -distance.
In other words, we consider f ⋆ = f
θ ⋆ ( x ) + f
⋆ ⊥ ( x ) , where θ ⋆ = arg min θ ∥ f θ − f ⋆ ∥
L 2 ( D ) .
1 To minimize the expected ℓ 2 -loss, we use vanilla GD with constant step-size η > 0 :
[Equation]
Definition 1 (Orthonormal function basis) .
A set of functions { ϕ i ( x ) } i ∈I forms an orthonormal function basis for the model space F Θ with respect to the L 2 ( D ) -metric
if
• for any i ∈
I , we have E x ∼D
[ ϕ 2 i ( x )] = 1 ; • for any i, j ∈ I such that i ̸
= j , we have E x ∼D
[ ϕ i ( x ) ϕ j ( x )] = 0 ; • for any f θ ∈ F Θ , there exists a unique sequence of basis coefficients { β i ( θ ) } i ∈I such that f θ ( x ) = (cid:80) i ∈I β
i ( θ ) ϕ
i ( x ) .
Example 1 (Orthonormal basis for polynomials) .
Suppose that F Θ is the class of all univariate real polynomials of degree at most n , that is, F Θ = { (cid:80) n +1 i =
1 θ
i x i − 1 : θ ∈ R n +1 } .
If D is a uniform distribution on [ − 1 , 1] , then the so-called Legendre polynomials form an orthonormal basis for
F Θ with respect to the L 2 ( D ) -metric
(Olver et al., 2010, Chapter 14).
Moreover, if D is a normal distribution, then Hermite polynomials define an orthonormal basis for F Θ with respect to the
L 2 ( D ) -metric
(Olver et al., 2010, Chapter 18).
Example 2 (Orthonormal basis for symmetric matrix factorization) .
Suppose that the true model is defined as f U ⋆ ( X ) = ⟨ U ⋆ U ⋆ ⊤ , X ⟩ with some rank- r matrix
U
⋆ ∈ R d × r , and consider an “overparameterized” function class F Θ = { f U ( X ) : U ∈ R d × r ′ } where r ′ ≥ r is an overestimation of the rank.
Moreover, suppose that the elements of X ∼ D are iid with zero mean and unit variance.
Consider the eigenvalues of U ⋆ U ⋆ ⊤ as σ 1 ≥ · · · ≥ σ d with σ r +1 = · · · = σ d = 0 , and their corresponding eigenvectors z 1 , . .
z d .
It is easy to verify that the functions ϕ ij ( X ) =
⟨ z i z ⊤ j , X ⟩ for 1 ≤ i, j ≤ d define a valid orthogonal basis for F Θ with respect to the L 2 ( D ) -metric.
Moreover, for any f U ( X ) , the basis coefficients can be obtained as β ij ( U ) =
E (cid:2)(cid:10) UU ⊤ , X (cid:11) (cid:10) z i z ⊤ j , X (cid:11)(cid:3)
= (cid:10) z
i z ⊤ j , UU ⊤ (cid:11) .
As will be shown in Section 3.2, this choice of orthonormal basis significantly simplifies the dynamics of GD for symmetric matrix factorization.
Given the input distribution D , we write f θ ⋆ ( x )
= (cid:80) i ∈I β
i ( θ ⋆ ) ϕ
i ( x ) , where { ϕ ( x ) } i ∈I is an orthonormal basis for F Θ with respect to L 2 ( D ) -metric, and { β i ( θ ⋆ ) } i ∈I
are the true basis coefficients .
For short, we denote β ⋆i
=
β i ( θ ⋆ ) .
In light of this, the expected loss can be written as:
[Equation]
Accordingly, GD takes the form
[Equation]
Two important observations are in order based on GD dynamic: first, due to the decomposed nature of the expected loss, the solution trajectory becomes independent of the approximation error and noise.
Second, in order to prove the global convergence of GD, it suffices to show the convergence of β
i ( θ t ) to β ⋆i .
In fact, we will show that the coefficients
β
i ( θ t ) enjoy simpler dynamics for particular choices of orthonormal basis that satisfy appropriate conditions.
Assumption 1 (Boundedness and smoothness) .
There exist constants L f , L g , L H > 0 such that
[Equation]
The above assumptions are common in the optimization literature (Bubeck et al., 2015) and sometimes necessary to ensure the convergence of local-search algorithms (Patel & Berahas, 2022).
Moreover, although these assumptions may not hold globally, all of our subsequent results hold when Assumption 1 is satisfied within any bounded region for θ that includes the solution trajectory.
We will also relax these assumptions for several learning problems.
Proposition 1 (Dynamic of β i ( θ t ) ) .
Under Assumption 1 and based on GD dynamic, we have
[Equation]
The above proposition holds for any valid choice of orthonormal basis { ϕ i ( x ) } i ∈I .
Indeed, there may exist multiple choices for the orthonormal basis, and not all of them would lead to equally simple dynamics for the coefficients.
Examples of “good” and “bad” choices of orthonormal basis were presented for DNNs in our earlier motivating example.
Indeed, an ideal choice of orthogonal basis should satisfy ⟨∇
β i ( θ t ) , ∇ β j ( θ t ) ⟩ ≈ 0 for i ̸
= j , i.e., the gradients of the coefficients remain orthogonal along the solution trajectory.
Under such assumption, the dynamics of β i ( θ t ) almost decompose over different indices:
[Equation]
where the last term accounts for the second-order interactions among the basis coefficients.
If such an ideal orthonormal basis exists, then our next theorem shows that GD efficiently learns the true basis coefficients.
To streamline the presentation, we assume that β ⋆ 1 ≥ · · · ≥ β ⋆k > 0 and β ⋆i = 0 , i > k for some k < ∞ .
We refer to the index set S = { 1 , . . .
, k } as signal and the index set E = I\S
as residual .
When there is no ambiguity, we also refer to β
i ( θ t ) as a signal if i ∈ S .
Theorem 1 (Convergence of GD with finite ideal basis) .
Suppose that the initial point θ 0 satisfies
[Equation]
for all
[Equation]
S , (lower bound on signals at θ 0 )
[Equation]
(upper bound on energy at θ 0 )
[Figure] for C 1 , C 2 > 0 and α ≲ β ⋆k .
Moreover, suppose that the orthogonal function basis is finite, i.e., |I| =
d for some finite d , and the gradients of the coefficients satisfy the following conditions for every
[Equation]
[Equation]
for all i ̸
= j , (gradient independence)
[Equation]
for all i ∈ S , (gradient dominance) for C > 0 and 1 / 2 ≤ γ ≤ 1 .
Then, GD with step-size η ≲
[Equation]
log
[Equation]
satisfies:
• If γ = 12 , then within T
= O (cid:16) 1 C 2 ηβ ⋆k log (cid:16) β
⋆k C 1 α (cid:17)(cid:17) iterations, we have ∥ f θ T − f θ ⋆ ∥
L 2 ( D ) ≲ α. •
If 12 <γ ≤ 1 , then within T = O (cid:16) 1 C 2 ηβ ⋆k α 2 γ − 1 (cid:17) iterations, we have ∥ f θ T − f θ ⋆ ∥
L 2 ( D ) ≲ α.
Theorem 1 shows that, under certain conditions on the basis coefficients and their gradients, GD with constant step-size converges to a model that is at most α -away from the true model.
In particular, to achieve an ϵ -accurate solution for any ϵ > 0 , GD requires O ((1 /ϵ ) log(1 /ϵ ))
iterations for γ = 1 / 2 , and O (1 /ϵ 2 γ ) iterations for 1 / 2 < γ ≤ 1 (ignoring the dependency on other problem-specific parameters).
Due to its generality, our theorem inevitably relies on a small step-size and leads to a conservative convergence rate for GD.
Later, we will show how our proposed approach can be tailored to specific learning problems to achieve better convergence rates in each setting.
How realistic are the assumptions of Theorem 1?
A natural question arises as to whether the conditions for Theorem 1 are realistic.
We start with the conditions on the initial point.
Intuitively, these assumptions entail that a non-negligible fraction of the energy is carried by the signal at the initial point.
We note that these assumptions are mild and expected to hold in practice.
For instance,
We will show in Section 3 that, depending on the learning task, they are guaranteed to hold with fixed, random, or spectral initialization.
3 We have also empirically verified that these conditions are satisfied for DNNs with random or default initialization.
For instance, Figure 2a illustrates the topbasis coefficients at θ 0 for LARS with random initialization on a realistic CNN.
It can be seen that a non-negligible fraction of the energy at the initial point is carried by the first few coefficients.
The conditions on the coefficient gradients are indeed harder to satisfy; as will be shown later, the existence of an ideal orthonormal basis may not be guaranteed even for linear NNs.
Nonetheless, we have empirically verified that, with an appropriate choice of the orthonormal basis, the gradients of the coefficients remain approximately independent throughout the solution trajectory.
Figure 2b shows that, when the orthonormal basis is chosen as the eigenvectors of the conjugate kernel after training, the maximum value of | cos( ∇ β i ( θ t ) , ∇ β j ( θ t )) | remains small throughout the solution trajectory.
Finally, we turn to the gradient dominance condition.
Intuitively, this condition entails that the gradient of each signal scales with its norm.
We prove that this condition is guaranteed to hold for kernel regression, symmetric matrix factorization, and symmetric tensor decomposition.
Moreover, we have empirically verified that the gradient dominance holds across different DNN architectures.
Figure 2c shows that this condition is indeed satisfied for the top-4 basis coefficients of the solution trajectory (other signal coefficients behave similarly).
We also note that our theoretical result on GD may not naturally extend to LARS.
Nonetheless, our extensive simulations suggest that our proposed analysis can be extended to other stochastic and adaptive variants of GD (see
Section 3.4 and Appendix B); a rigorous verification of this conjecture is left as future work.
A PPLICATIONS
In this section, we show how our proposed basis function decomposition can be used to study the performance of GD in different learning tasks, from simple kernel regression to complex DNNs.
We start with the classical kernel regression, for which GD is known to converge linearly Karimi et al.
Our purpose is to revisit GD through the lens of basis function decomposition, where there is a natural and simple choice for the basis functions.
Next, we apply our approach to two important learning problems, namely symmetric matrix factorization and orthogonal symmetric tensor decomposition.
In particular, we show how our proposed approach can be used to improve the convergence of GD for the symmetric matrix factorization and leads to a completely new convergence result for the orthogonal symmetric tensor decomposition.
Finally, through extensive experiments, we showcase the promise of our proposed basis function decomposition on realistic DNNs.
3. K ERNEL R EGRESSION
In kernel regression (KR), the goal is to fit a regression model f θ ( x )
=
(cid:80) di
=
θ
i
ϕ
i ( x ) from the function class F Θ = { f θ ( x ) : θ ∈ R d } to observation y , where { ϕ i ( x ) } di =1 are some known kernel functions.
Examples of KR are linear regression, polynomial regression (including those described in Example 1), and neural tangent kernel (NTK)
(Jacot et al., 2018).
Without loss of generality, we may assume that the kernel functions { ϕ i ( x ) } di =1 are orthonormal.
4 Under this assumption, the basis coefficients can be defined as β i ( θ )
=
θ
i and the expected loss can be written as
[Equation]
Moreover, the coefficients satisfy the gradient independence condition.
Therefore, an adaptation of
Proposition 1 reveals that the dynamics of the basis coefficients are independent of each other.
Proposition 2 (dynamics of β
i ( θ t ) ) .
Consider GD with a step-size that satisfies 0 < η < 1 .
Then,
• for i ∈ S , we have β
i
θ t )
=
β ⋆i − (1 − η ) t ( β ⋆i
−
β
i ( θ 0 )) , • for i ̸∈ S , we have β
i
θ t ) =
(1 − η ) t β i ( θ 0 ) .
Without loss of generality, we assume that 0 < θ ⋆k ≤ · · · ≤ θ ⋆ 1 ≤ 1 and ∥ θ 0 ∥ ∞ ≤ α .
Then, given Proposition 2, we have | β
i ( θ t ) | ≤ 2 + α for every 1 ≤ i ≤ k .
Therefore, the gradient dominance is satisfied with parameters ( C, γ ) =
(1 / √ 2 + α, 1 / 2) .
Since both gradient independence and gradient dominance are satisfied, the convergence of GD can be established with an appropriate initial point.
Theorem 2.
Suppose that θ 0 =
α 1 , where α ≲ k | θ ⋆k | /d .
Then, within T ≲ (1 /η ) log ( k | θ ⋆ 1 | /α ) iterations, GD with step-size 0 < η < 1 satisfies ∥ θ T − θ ⋆ ∥ ≲ α .
Theorem 2 reveals that GD with large step-size and small initial point converges linearly to an ϵ -accurate solution, provided that the initialization scale is chosen as α = ϵ .
This is indeed better than our result on the convergence of GD for general models in Theorem 1.
} di =1 are not orthonormal.
Let { (cid:101) ϕ
i ( x ) } i ∈I be any orthonormal basis for F Θ . Then, there exists a matrix A such that ϕ i ( x ) =
(cid:80)
j A ij (cid:101) ϕ
j ( x ) for every 1 ≤ i ≤ d .
Therefore, upon defining (cid:101) θ
= θ ⊤ A , one can write f (cid:101) θ ⋆ ( x )
= (cid:80) i ∈I (cid:101) θ ⋆i (cid:101) ϕ
i ( x ) which has the same form as the regression model.
3. S YMMETRIC M ATRIX F ACTORIZATION
In symmetric matrix factorization (SMF), the goal is to learn a model f U ⋆ ( X ) = ⟨ U ⋆ U ⋆ ⊤ , X ⟩ with a low-rank matrix U ⋆ ∈ R d × r , where we assume that each element of X ∼ D is iid with E
[ X ij ] =and E [ X 2 ij ] = 1 .
Examples of SMF are matrix sensing Li et al.
and
completion Ge et al.
Given the eigenvectors { z 1 , . . .
, z d } of U ⋆ U ⋆ ⊤ and a function class F Θ = { f U ( X ) : U ∈ R d × r ′ } with r ′ ≥ r , it was shown in Example 2 that the functions ϕ ij ( X ) =
⟨ z i z ⊤ j , X ⟩ define a valid orthogonal basis for F Θ with coefficients β ij ( U ) =
(cid:10) z i z ⊤ j , UU ⊤ (cid:11) .
Therefore, we have
[Equation]
Here, the true basis coefficients are defined as β ⋆ii =
σ i for i ≤ r , and β ⋆ij = 0 otherwise.
Moreover, one can write ∥∇ β ii ( U ) ∥ F = 2 (cid:13)(cid:13)
z i z
⊤ i U (cid:13)(cid:13)
F = 2 (cid:112) β ii ( U ) .
Therefore, gradient dominance holds with parameters ( C, γ ) =
However, gradient independence does not hold for this choice of function basis: given any pair ( i, j ) and ( i, k ) with j ̸
= k , we have ⟨∇ β ij ( U ) , ∇ β ik ( U ) ⟩ = ⟨
z j z ⊤ k , UU ⊤ ⟩ which may not be zero.
Despite the absence of gradient independence, our next proposition characterizes the dynamic of β ij ( U t ) via a finer control over the coefficient gradients.
Proposition 3.
Suppose that γ :=
min 1 ≤ i ≤ r { σ i − σ i +1 } > 0 .
Let U
= α B , where α ≲ min (cid:8) ( ησ 2 r ) Ω( σ 1 /γ ) , ( σ r /d )
Ω( σ 1 /γ ) , ( κ log 2 ( d )) − Ω(1 / ( ησ r )) (cid:9) and the entries of B are independently drawn from a standard normal distribution.
Suppose that the step-size for GD satisfies η ≲ 1 /σ 1 .
Then, with probability of at least 1 − exp( − Ω( r ′ )) :
• For 1 ≤ i ≤ r , we have 0 .
99 σ i ≤ β ii ( U t ) ≤ σ i within O ((1 / ( ησ i )) log ( σ i /α ))
iterations.
• For t ≥
and i ̸
= j or i, j > r , we have | β ij ( U t ) | ≲ poly( α ) .
Proposition 3 shows that GD with small random initialization learns larger eigenvalues before the smaller ones, which is commonly referred to as incremental learning .
Incremental learning for
SMF has been recently studied for gradient flow (Arora et al., 2019a; Li et al., 2020), as well as
GD with identical initialization for the special case r ′ =
d Chou et al.
To the best of our knowledge, Proposition 3 is the first result that provides a full characterization of the incremental learning phenomenon for GD with random initialization on SMF.
Theorem 3.
Suppose that the conditions of Proposition 3 are satisfied.
Then, with probability of at least 1 − exp( − Ω( r ′ )) and within T ≲ (1 / ( ησ r )) log ( σ r /α ) iterations, GD satisfies
[Equation]
It has been shown in (Stöger & Soltanolkotabi, 2021, Thereom 3.3) that GD with small random initialization satisfies (cid:13) (cid:13)
U T U ⊤ T − M ⋆ (cid:13) (cid:13)
F ≲ (cid:16) d 2 /r ′ 15 / 16 (cid:17) α 21 / 16 within the same number of iterations.
Theorem 3 improves the dependency of the final error on the initialization scale α .
3. O RTHOGONAL S YMMETRIC T ENSOR D ECOMPOSITION
We use our approach to provide a new convergence guarantee for GD on the orthogonal symmetric tensor decomposition (OSTD).
In OSTD, the goal is to learn f U ⋆ ( X ) = ⟨ T U ⋆ , X ⟩ , where U ⋆ =
[ u ⋆ 1 , . . .
, u ⋆r ] ∈ R d × r and T U ⋆ = (cid:80) ri =1 u ⋆i ⊗ l = (cid:80) di
=
σ i
z i ⊗ l is a symmetric tensor with order l and rank r .
Here, σ 1 ≥ · · · ≥ σ d are tensor eigenvalues with σ r +1 = · · · = σ d = 0 , and z 1 , . . .
, z d are the corresponding tensor eigenvectors.
The notation u ⊗ l refers to the l -time outer product of u .
We assume that X ∼ D is an l -order tensor whose elements are iid with zero mean and unit variance.
Examples of OSTD are tensor regression (Tong et al., 2022) and completion (Liu et al., 2012).
When the rank of T U ⋆ is unknown, it must be overestimated.
Even when the rank is known, its overestimation can improve the convergence of gradient-based algorithms (Wang et al., 2020).
This leads to an overparameterized model f U ( X ) = ⟨ T U , X ⟩ , where T U =
(cid:80) r ′ i =
1 u i ⊗ l with an overestimated rank r ′ ≥ r .
Accordingly, the function class is defined as F Θ = (cid:110)
f U ( X ) : U =
[ u 1 , · · · , u r ′ ] ∈ R d × r ′ (cid:111) .
Upon defining a multi-index Λ = ( j 1 , · · · , j l ) , the functions ϕ Λ ( X ) =
(cid:10) ⊗ lk =
1 z j k , X (cid:11) for 1 ≤ j 1 , . . .
, j l ≤ d form an orthonormal basis for F Θ with basis coefficients defined as
[Equation]
and the expected loss can be written as
[Equation]
where the true basis coefficients are β ⋆ Λ i
= σ i for Λ i
=
( i, . . .
, i ) , 1 ≤ i ≤ r , and β ⋆ Λ =otherwise.
Unlike KR and SMF, neither gradient independence nor gradient dominance are satisfied for OSTD with a random or equal initialization.
However, we show that these conditions are approximately satisfied throughout the solution trajectory, provided that the initial point is nearly aligned with the eigenvectors z 1 , . . .
, z r ′ ; in other words, cos( u i (0) , z i )
≈
for every 1 ≤ i ≤ r ′ .
Assuming that the initial point satisfies this alignment condition, we show that the entire solution trajectory remains aligned with these eigenvectors, i.e., cos( u i ( t ) , z i )
≈
for every 1 ≤ i ≤ r ′ and 1 ≤ t ≤ T .
Using this key result, we show that both gradient independence and gradient dominance are approximately satisfied throughout the solution trajectory.
We briefly explain the intuition behind our approach for gradient dominance and defer our rigorous analysis for gradient independence to the appendix.
Note that if cos( u i ( t ) , z i ) ≈ 1 , then β Λ i ( U t )
≈ ⟨ u i ( t ) , z i ⟩ l and ∥∇ β Λ i ( U t ) ∥
F ≈ ∥∇ u i
β Λ i ( U t )
∥
≈ l ⟨ u i ( t ) ,
z ⟩ l − 1 .
Therefore, gradient dominance holds with parameters ( C, γ ) =
( l, ( l − 1) /l ) .
We will make this intuition rigorous in Appendix G.
Proposition 4.
Suppose that the initial point U 0 is chosen such that ∥ u i (0) ∥
=
α 1 /l and cos( u i (0) , z i ) ≥ √ 1 − γ , for all 1 ≤ i ≤ r ′ , where α ≲ d − l 3 and γ ≲ ( lκ ) − l/ ( l − 2) .
Then,
GD with step-size η ≲ 1 / ( lσ 1 ) satisfies:
• For 1 ≤ i ≤ r , we have 0 .
99 σ i ≤ β Λ i ( U t ) ≤ 1 .
σ i within O (cid:16) (1 / ( ηlσ r ))
α − l − 2 l (cid:17) iterations.
• For t ≥ 0 and Λ ̸ = Λ i , we have | β Λ ( U t )
| = poly( α ) .
Proposition 4 shows that, similar to SMF, GD learns the tensor eigenvalues incrementally.
However, unlike SMF, we require a specific alignment for the initial point.
We note that such initial point can be obtained in a pre-processing step via tensor power method within a number of iterations that is almost independent of d (Anandkumar et al., 2017, Theorem 1).
We believe that Proposition 4 can be extended to random initialization; we leave the rigorous verification of this conjecture to future work.
Equipped with this proposition, we next establish the convergence of GD on OSTD.
Theorem 4.
Suppose that the conditions of Proposition 4 are satisfied.
Then, within T ≲ (1 / ( ηlσ r ))
α − ( l − 2) /l iterations, GD satisfies
[Equation]
Theorem 4 shows that, with appropriate choices of η and α , GD converges to a solution that satisfies ∥ [Equation] ⋆ ∥ 2 F ≤ ϵ within O ( d l ( l − 2) /ϵ
l − 2 ) iterations.
To the best of our knowledge, this is the first result establishing the convergence of GD with a large step-size on OSTD.
3. E MPIRICAL V ERIFICATION ON N EURAL N ETWORKS
In this section, we numerically show that the conjugate kernel after training (A-CK) can be used as a valid orthogonal basis for DNNs to capture the monotonicity of the solution trajectory of different optimizers on image classification tasks.
To ensure consistency with our general framework, we use ℓ 2 -loss, which is shown to have a comparable performance with the commonly-used cross-entropy loss Hui & Belkin (2020).
In Appendix B, we extend our simulations to cross-entropy loss.
The conjugate kernel (CK) is a method for analyzing the generalization performance of DNNs that uses [Figure] the second to last layer (the layer before the last linear layer) at the initial point as the feature map (Daniely et al., 2016; Fan & Wang, 2020; Hu & Huang, 2021).
Recently, Long (2021) shows that A-CK, a variant of CK that is evaluated at the last epoch, better explains the generalization properties of realistic DNNs.
Surprisingly, we find that A-CK can be used not only to characterize the generalization performance but also to capture the underlining solution trajectory of different gradient-based algorithms.
To formalize the idea, note that any neural network whose last layer is linear can be characterized as f θ ( x )
=
W ψ ( x ) , where x ∈ R d is the input drawn from the distribution D , ψ ( x ) ∈ R m is the feature map with number of features m , and
W ∈ R k
×
m is the last linear layer with k referring to the number of classes.
We denote the trained model, i.e., the model in the last epoch, by f θ ∞ ( x ) =
W ∞ ψ ∞ ( x ) .
To form an orthogonal basis, we use SVD to obtain a series of basis functions ϕ
i ( x ) =
W ∞ ,i ψ ∞ ( x ) that satisfy E x ∼D [ ∥ ϕ i ( x ) ∥ 2 ] = 1 and E x ∼D [ ⟨ ϕ i ( x ) , ϕ j ( x ) ⟩ ] =
δ ij where δ ij is the delta function.
Hence, the coefficient β i ( θ t ) at each epoch t can be derived as β i ( θ t ) =
E x ∼D [ ⟨ f θ t ( x ) , ϕ j ( x ) ⟩ ] , where the expectation is estimated by its sample mean on the test set.
More details on our implementation can be found in Appendix B.
Performance on convolutional neural networks: We use LARS to train CNNs with varying depths on MNIST dataset.
These networks are trained such that their test accuracies are above 96% .
Figures 3a-3c illustrate the evolution of the top-5 basis coefficients after projecting LARS onto the orthonormal basis obtained from A-CK.
It can be observed that the basis coefficients are consistently monotonic across different depths, elucidating the generality of our proposed basis function decomposition.
In the appendix, we discuss the connection between the convergence of the basis functions and the test accuracy for different architectures and loss functions.
Performance with different optimizers: The monotonic behavior of the projected solution trajectory is also observed across different optimizers.
Figures 3d-3f show the solution trajectories of three optimizers, namely LARS, SGD, and AdamW, on AlexNet with the CIFAR-10 dataset.
It can be seen that all three optimizers have a monotonic trend after projecting onto the orthonormal basis obtained from A-CK.
Although our theoretical results only hold for GD, our simulations highlight the strength of the proposed basis function decomposition in capturing the behavior of other gradient-based algorithms on DNN.
We thank Richard Y. Zhang and Tiffany Wu for helpful feedback. We would also like to thank Ruiqi Gao and Chenwei Wu for their insightful discussions.
This research is supported, in part, by NSF Award DMS-2152776, ONR Award N00014-22-1-2127, MICDE Catalyst Grant, MIDAS PODS grant and Startup Funding from the University of Michigan. C ONTENTS A Discussion and Future Direction
B Additional Experiments
B.1 Numerical Verification of our Theoretical Results . . . . . . . . . . . . . . . . . .
B.2 Derivation of Basis Functions for DNNs . . . . . . . . . . . . . . . . . . . . . . .
B.3 Further Details on the Experiments . . . . . . . . . . . . . . . . . . . . . . . . . .
B.4 Experiments for CIFAR-100 . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.5 Experiments for Different Losses . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.6 Experiments for Different Batch Size . . . . . . . . . . . . . . . . . . . . . . . . .
B.7 Experiments for ResNet-18 with SGD on CIFAR-10 . . . . . . . . . . . . . . . .
C Related Work
D Proofs for General Framework
D. Proof of Proposition 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D. Proof of Theorem
E Proofs for Kernel Regression
E.1 Proof of Proposition 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.2 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F Proofs for Symmetric Matrix Factorization
F. Initialization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F. One-step Dynamics . .
F. Proofs of Proposition 3 and Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . F.
Analysis of the Logistic Map . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G Proof for Tensor Decomposition
G.1 Proof of Lemma 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.2 Proof of Proposition 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.3 Proof of Proposition 8 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H Auxiliary Lemmas
A D ISCUSSION AND F UTURE D IRECTION
Extension to empirical loss and implication on generalization.
The most significant future direction is to extend our analysis of expected loss to the finite sample regime.
It has been shown that incremental learning can drive generalization (Gissin et al., 2019).
Hence, our proposed framework is likely to explain the puzzling generalization ability of overparameterized machine learning models.
For such an extension, the main technical difficulty is the relaxation of gradient independence.
In the finite-sample regime, gradient independence does not hold exactly due to the randomness of samples.
Therefore, an alternative approach would be to establish gradient independence approximately and with high probability.
We have successfully applied our framework to matrix factorization and tensor decomposition, where gradient independence only holds approximately.
Therefore, we believe that more general guarantees on the finite sample regime are not out of reach.
Extension to other optimization algorithms.
In this paper, we mainly focused on GD as a representative of various optimization algorithms.
Nonetheless, we believe that our analysis can be adapted to investigate other local-search optimization algorithms, such as GD with momentum (Nesterov, 1983), SGD (Robbins & Monro, 1951), and Adam (Kingma & Ba, 2014).
Overall, our approach can offer a unified framework for examining the implicit bias and incremental learning phenomena for different local-search algorithms.
B A DDITIONAL E XPERIMENTS
In this section, we provide more details on our simulation and further explore the empirical strength of the proposed basis function decomposition on different datasets, optimizers, loss functions, and batch sizes; see Table 1 for a summary of our simulations in this section.
[Table]
B. N UMERICAL V ERIFICATION OF OUR T HEORETICAL R ESULTS
In this section, we provide experimental evidence to support our theoretical results on kernel regression (KR), symmetric matrix factorization (SMF), and orthogonal symmetric tensor decomposition (OSTD).
The results are presented in Figure 4.
Kernel regression.
We randomly generate 20 orthonormal kernel functions.
The true model is comprised of 4 signal terms with basis coefficients 10 , 5 , 3 , 1 .
Figure 4a shows the trajectories of the top-4 basis coefficients of GD with initial point θ 0 = 5 × 10 − 7 × 1 and step-size η = 0 .
It can be seen that GD learns different coefficients at the same rate, which is in line with Proposition 2.
Symmetric matrix factorization.
In this simulation, we aim to recover a rank- 4 matrix M ⋆ =
V Σ V ⊤ ∈ R 20 × 20 .
In particular, we assume that V ∈ R 20 × 4 is a randomly generated orthonormal matrix and Σ
= Diag { 10 , 5 , 3 , 1 } .
We consider a fully over-parameterized model where U ∈ R 20 ×(i.e., r ′ = 20 ).
Figure 4b illustrates the incremental learning phenomenon that was proved in
Proposition 3 for GD with small Gaussian initialization U ij i.i.d.
∼ N (0 , α 2 ) , α = 5 × 10 − 7 and step-size η = 0 .
Orthogonal symmetric tensor decomposition.
Finally, we present our simulations for OSTD.
We aim to recover a rank-4 symmetric tensor of the form T ⋆ = (cid:80) 4 i =
σ i z ⊗ 4 i where σ i are the nonzero eigenvalues with values { 10 , 5 , 3 , 1 } and z i ∈ R 10 are the corresponding eigenvectors.
We again consider a fully over-parameterized model with r ′ = 10 .
Figure 4c shows the incremental learning phenomenon for GD with an aligned initial point that satisfies cos( u i (0) , z i ) ≥ 0 . 9983 , 1 ≤ i ≤ r ′ and step-size η = 0 . [Figure]
B. D ERIVATION OF B ASIS F UNCTIONS FOR DNN S
In this section [Figure] we provide more details on how we evaluate our proposed orthogonal basis induced by
A-CK and calculate the corresponding coefficients
β
i ( θ t ) for an arbitrary neural network.
First, recall that any neural network whose last layer is linear can be characterized as f θ ( x ) =
W ψ ( x ) , where x ∈ R d is the input drawn from the distribution D , ψ ( x ) ∈ R m is the feature map with number of features m , and
W ∈ R k
×
m is the last linear layer with k referring to the number of classes.
We denote the trained model, i.e., the model in the last epoch, by f θ ∞ ( x ) =
W ∞ ψ ∞ ( x ) .
To form an orthogonal basis, we use SVD to obtain a series of basis functions ϕ
i ( x ) =
W ∞ ,i ψ ∞ ( x ) that satisfy
E x ∼D [ ∥ ϕ i ( x ) ∥ 2 ] = 1 and E x ∼D [ ⟨ ϕ i ( x ) , ϕ j ( x ) ⟩ ] =
δ ij where δ ij is the delta function.
Hence, the coefficient β i ( θ t ) at each epoch t can be derived as β i ( θ t ) =
E x ∼D [ ⟨ f θ t ( x ) , ϕ j ( x ) ⟩ ] .
In all of our implementation, we use the test dataset to approximate the population distribution.
Step 1: Obtaining the orthogonal basis ϕ
i ( x ) .
We denote Ψ =
[ ψ ( x 1 ) , · · · , ψ ( x N )] ∈ R m × N as the feature matrix where N is the number of the test data points.
We write the SVD of Ψ as Ψ = U Σ V ⊤ .
The right singular vectors collected in V can be used to define the desired orthogonal basis of { ψ i ( x ) } mi =1 .
To this goal, we write the prediction matrix as F = W Ψ =
(cid:102) W (cid:101) Ψ
where (cid:102) W = W U Σ and (cid:101) Ψ = V ⊤ .
Our goal is to define a set of matrices A i such that ϕ
i ( x ) =
A i (cid:101) ψ
i ( x ) form a valid orthonormal basis for F = W Ψ =
(cid:102) W (cid:101) Ψ .
Before designing such A i , first note that, due to the orthogonality of { (cid:101) ψ i ( x ) } , we have
[Equation]
Therefore, it suffices to ensure that { A i } are orthonormal.
Consider the SVD of (cid:102) W as (cid:102) W =
(cid:80) i σ i
u i v ⊤ i .
We define A i
= u
i v ⊤ i .
Clearly, defined { A i } are orthonormal.
Moreover, it is easy to see that the basis coefficients (treated as the true basis coefficients) are exactly the singular values of (cid:102) W .
Step 2: Obtaining the basis coefficients
β
i ( θ t ) .
After obtaining the desired orthonormal basis ϕ
i ( x ) , we can calculate the coefficient β
i ( θ t ) for each epoch.
Given the linear layer W t and the feature matrix Φ t at epoch t , we can obtain the coefficients for the signal terms by projecting the prediction matrix F t = W t Ψ t onto ˜ Ψ
= V ⊤ .
In particular, we write the prediction matrix as
P (cid:101) ψ F t = P (cid:101)
ψ W t Ψ t =
W t Ψ t V V ⊤ = (cid:102) W t (cid:101) Ψ where (cid:102) W t
=
W t Ψ t V .
Hence, the basis coefficients can be easily calculated as β i ( θ t ) =
(cid:68)(cid:102) W t , u i v ⊤ i (cid:69) .
B. F URTHER D ETAILS ON THE E XPERIMENTS
In this section, we provide more details on our experiments presented in the main body of the paper and compare them with other DNN architectures.
All of our experiments are implemented in Python 3.9, Pytorch 1.12.1 environment and run through a local server SLURM using NVIDIA Tesla with V100-PCIE-16GB GPUs.
We use an additional NNGeometry package for calculating batch gradient, and our implemention of ViT is adapted from https://juliusruseckas.github.io/ml/cifar10-vit.html .
To ensure consistency with our theoretical results, we drop the last softmax operator and use the ℓ 2 -loss throughout this section.
All of our training data are augmented by RandomCrop and
RandomHorizontalFlip , and normalized by mean and standard deviation.
Experimental details for Figure 1.
Here, we describe our implementation details for Figure 1, and present additional experiments on VGG-11, ResNet-34, and ResNet-50 with the CIFAR-10 dataset.
The results can be seen in Figure 5.
We use standard data augmentation for all architectures except for ViT. For ViT, we only use data normalization.
To obtain a stable A-CK, we trained the above models for 300 epochs.
For ResNet-18, we used
LARS with a learning rate of 0 . 5 and applied small initialization with α = 0 .
3 , i.e., we scale the default initial point by α = 0 .
For ResNet-34 and ResNet-50, we choose the default learning rate and apply small initialization with α = 0 .
For ViT, we use AdamW with a learning rate of 0 . 01 .
The remaining parameters are set to their default values.
Experiments details for the first row of Figure 3.
Here we conduct experiments on MNIST dataset with different CNN architectures.
The CNNs are composed of l blocks of layers, followed by a single fully connected layer.
A block of a CNN consists of a convolutional layer, an activation layer, and a pooling layer.
In our experiments, we use ReLU activation and vary the depth of the network.
For the first block, we used identity pooling.
For the remaining blocks, we used max-pooling.
For 2 -block CNN, we set the convolutional layer width to 256 and 64 , respectively.
For 3 -block CNN, we set the convolutional layer width to 256 , 128 , and 64 , respectively.
And for 4 -block CNN, we set the convolutional layer width to 256 , 128 , 128 , and 64 , respectively.
To train these networks, we used
LARS with the learning rate of 0 .
The remaining parameters are set to their default values.
We run 20 epochs to calculate A-CK.
Experiments details for second row of Figure 3.
We conduct experiments to compare the performance of different optimizers on the CIFAR-10 dataset.
In particular, we use AlexNet to compare the performance of three optimizers, i.e., SGD, AdamW, LARS.
For SGD, we set the base learning rate to be 2 with Nesterov momentum of 0 .
9 and weight decay of 0 . 0001 , together with the “linear warm-up” technique.
For AdamW, we set the learning rate to 0 .
01 and keep the remaining parameters unchanged.
For LARS, we set the learning rate to 2 with Nesterov momentum of 0 .
9 and weight decay of 0 . 0001 .
The remaining parameters are set to the default setting.
B. E XPERIMENTS FOR CIFAR-
In this section, we conduct experiments using the CIFAR-100 dataset which is larger than both
CIFAR-10 and MNIST.
Our simulations are run on AlexNet, VGG-11, ViT, ResNet-18, ResNet-34, and ResNet-50.
In particular, we use the “loss scaling trick” (Hui & Belkin, 2020) defined as follows: consider the datapoint ( x , y )
where x ∈ R
d is the input and
y ∈ R k is a one-hot vector with 1 at In linear warm-
, we linearly increase the learning rate the epochs.
More precisely we set the initial learning rate to 1 × and linearly increase it the selected learning rate
5 epochs
After the firstepochs, the learning rate follows a regular decay scheme.
[Figure] position i . Then, the scaled ℓ 2 -loss is defined as ℓ 2 ,
[Equation]
for some constants k, M > 0 .
We set these parameters to k = 1 , M = 4 .
For AlexNet and VGG-11, we use LARS with a base learning rate of η = 1 .
For ViT, we use AdamW with a base learning rate of 0 .
01 and batch size of 256 .
For ResNet architectures, we use SGD with a base learning rate of 0 .
3 and batch size of 64 .
We also add 5 warm-up epochs for ResNets.
All the remaining parameters for the above architectures are set to their default values.
The results can be seen in Figure 6.
Our experiments highlight a trade-off between the monotonicity of the projected solution trajectories and the test accuracy: in order to obtain a higher test accuracy, one typically needs to pick a larger learning rate, which in turn results in more sporadic behavior of the solution trajectories.
Nonetheless, even with large learning, the basis coefficients remain relatively monotonic after the first few epochs and converge to meaningful values.
B. E XPERIMENTS FOR D IFFERENT L OSSES
Next, we compare the projected solution trajectories on two loss functions, namely ℓ 2 -loss and cross-entropy (CE) loss.
We use LARS to train AlexNet on the CIFAR-10 dataset with both ℓ 2 -loss and CE loss.
In particular, we add the softmax operator before training the CE loss.
For CE loss, we set the base learning rate of LARS to 1 .
For ℓ 2 -loss, we use the base learning rate of 2 .
The remaining parameters are set to their default values.
The results can be seen in Figure 7.
We observe that, similar to the ℓ 2 -loss, the solution trajectory of the CE loss behaves monotonically after projecting onto the orthogonal basis induced by A-CK.
Inspired by these observations, another venue for future research would be to extend our framework to general loss functions.
Interestingly, the convergence of the basis coefficients is much slower than those of the ℓ 2 -loss.
This is despite the fact that CE loss can learn slightly faster than ℓ 2 -loss in terms of test accuracy as shown by Hui & Belkin (2020).
B. E XPERIMENTS FOR D IFFERENT B ATCH S IZE
Next, we study the effect of different batch sizes on the solution trajectory.
We train AlexNet on the CIFAR-10 dataset.
When testing for different batch sizes, we follow the “linear scaling” rule, i.e., [Figure] [Figure] the learning rate scales linearly with the batch size.
For batch size of 32 , we used SGD with a base learning rate η = 0 .
1 and 5 warm-up epochs.
For batch size of 64 , we used SGD with a base learning rate of η = 0 .
2 and 5 warm-up epochs.
For batch size of 256 , we used LARS with a base learning rate of 1 .
The remaining hyperparameters are set to their default values.
The results are reported in Figure 8.
We see that the projected solution trajectories share a similar monotonic behavior for different batch sizes.
B. E XPERIMENTS FOR R ES N ET -18 WITH SGD ON CIFAR-
Lastly, we train ResNet-18 with SGD on CIFAR-10. The results can be seen in Figure 9. To achieve good generalization, we use a large learning rate η = 0 .
3 (with 5 warm-up epochs starting at 1 × 10 − 5 ), and decrease the learning rate by 0 .
33 for every 50 epochs.
Moreover, We use a large batch size 512 to better imitate the trajectory of GD.
Based on the simulation result, we can find an approximately monotonic (modulo the fluctuations caused by the randomness of the gradients) behavior of the dynamics of the top- 5 components, which again validifies our theoretical results.
[Figure]
[Figure]
βββββ
C R ELATED W ORK
GD for general nonconvex optimization.
Gradient descent and its stochastic or adaptive variants are considered as the “go-to” algorithms in large-scale (unconstrained) nonconvex optimization.
Because of their first-order nature, they are known to converge to first-order stationary points (Nesterov, 1998).
Only recently it has been shown that GD (Lee et al., 2019; Panageas et al., 2019) and its variants, such as perturbed GD (Jin et al., 2017) and SGD
(Fang et al., 2019; Daneshmand et al., 2018), can avoid saddle points and converge to a second-order stationary point.
However, these guarantees do not quantify the distance between the obtained solution and the globally optimal and/or true solutions.
To the best of our knowledge, the largest subclass of nonconvex optimization problems for which GD or its variants converge to meaningful solutions are those with benign landscapes.
These problems include different variants of low-rank matrix optimization with exactly parameterized rank, namely matrix completion (Ge et al., 2016), matrix sensing (Ge et al., 2017; Zhang et al., 2021), dictionary learning (Sun et al., 2016), and robust PCA (Fattahi & Sojoudi, 2020), as well as deep linear neural networks (Kawaguchi, 2016).
However, benign landscape is too restrictive to hold in practice; for instance,Zhang (2021) shows that spurious local minima are ubiquitous in the low-rank matrix optimization, even under fairly mild conditions.
Therefore, the notion of benign landscape cannot be used to explain the success of local search algorithms in more complex learning tasks.
GD for specific learning problems.
Although there does not exist a unifying framework to study the global convergence of GD for general learning tasks, its convergence has been established in specific learning problems, such as kernel regression (which includes neural tangent kernel (Jacot et al., 2018)), sparse recovery (Vaskevicius et al., 2019), matrix factorization (Li et al., 2018), tensor decomposition (Wang et al., 2020;
Ge et al., 2021), and linear neural network (Arora et al., 2018).
In what follows, we review specific learning tasks that are most related to our work.
The convergence of GD on kernel regression was studied far before the emergence of deep learning.
Bauer et al.
; Raskutti et al.
(2014) establish the convergence of gradient descent on a special class of nonparametric kernel regression called reproducing kernel Hilbert space (RKHS).
Recently,
Jacot et al.
(2018) discovered that under some conditions, neural networks can be approximated by a specific type of kernel models called neural tangent kernel (NTK).
Later on, a series of papers studied the optimization (Allen-Zhu et al., 2019b) and generalization (Allen-Zhu et al., 2019a;
Arora et al., 2019b) properties of NTK.
As for the matrix factorization, Li et al.
; Stöger &
Soltanolkotabi (2021) studied the global convergence of GD on the symmetric matrix sensing with noiseless measurements and overestimated rank.
Later, these results were extended to noisy (Zhuo et al., 2021), asymmetric (Ye & Du, 2021), and robust (Ma & Fattahi, 2022b) variants.
Wang et al.
Ge et al.
(2021) studied the dynamic of a modified GD for overcomplete nonconvex tensor decomposition.
Moreover, Razin et al.
(2021; 2022) analyzed the implicit regularization and the incremental learning of gradient flow in hierarchical tensor decomposition and showed its connection to neural networks.
Conjugate kernel.
Conjugate kernel (CK) at the initial point has been considered as one of the promising methods for studying the generalization properties of DNNs (Daniely et al., 2016;
Hu &
Huang, 2021; Fan & Wang, 2020).
However, similar to NTK, a major shortcoming of CK is that it cannot fully characterize the behavior of the practical neural networks (Vyas et al., 2022).
Recent results have suggested that the conjugate kernel evaluated after training (for both NTK and CK) can better describe the generalization properties of DNNs (Fort et al., 2020; Long, 2021).
In our work, we show that such “after kernel regime” can also be adapted to study the optimization trajectory of practical DNNs.
D P ROOFS FOR G ENERAL F RAMEWORK D. P ROOF OF P ROPOSITION
To prove this proposition, we first combine (1) and (GD):
[Equation]
For notational simplicity, we denote E ( θ t )
=
12 (cid:80) i ∈E β
i ( θ t ) .
Then, one can write
[Equation]
[Equation]
Due to the Mean-Value Theorem, there exists a ξ ∈ R m such that
[Equation]
[Equation]
On the other hand, one can write
[Equation]
For sup θ (cid:13) (cid:13) ∇ 2 β
i ( θ ) (cid:13) (cid:13) , we further have sup
[Equation]
Here, we used Cauchy-Schwartz inequality for (a).
Moreover, for (b), we used Assumption 1 and the definition of the orthonormal basis.
Hence, we have
[Equation]
Now [Equation] it suffices to bound ∥∇L ( θ t ) ∥ 2 .
Using Cauchy-Schwarz inequality [Equation] we have
[Equation]
Therefore
[Equation]
we conclude that
[Equation]
which completes the proof
[Equation]
□
D. P ROOF OF T HEOREM
Proof.
Invoking the gradient independence condition, Proposition 1 can be simplified as
[Equation]
[Equation]
We next provide upper and lower bounds for the residual and signal terms.
Recall that S
= { i ∈ I :
β ⋆i ̸ = 0 } , and E = I\S .
We first consider the dynamic of the signal term
β
i ( θ t ) , i ∈ S .
Without loss of generality, we assume β ⋆i > 0 .
Then, due to the gradient dominance condition, we have the following lower bound
[Equation]
Next, for the dynamic of the residual term β i ( θ t ) , i ∈ E , we have
[Equation]
Next, we show that ∥∇ β
i ( θ t ) ∥
≤ L g .
One can write,
[Equation] [Equation]
Due to our choice of the step-size, we have η ≲ 1 L 2 g , which in turn implies (cid:12)(cid:12)(cid:12) [Equation] β i [Equation]
θ
[Equation]
∥
[Equation]
≤ [Equation]
Therefore, we have
[Equation]
Now, we are ready to prove the theorem.
We divide it into two cases.
Case 1: γ = 12 .
In this case, since we set the step-size η ≲
[Equation]
, we can simplify the dynamics of both signal and residual terms in Equation 19 and Equation 20 as follows
[Equation]
We first analyze the dynamic of signal
β i ( θ t ) for i ∈ S .
To this goal, we further divide this case into two phases.
In the first phase, we assume C 1 α ≤ β
i ( θ t ) ≤ 12 β ⋆i .
Under this assumption, we can simplify the dynamic of β
i ( θ t ) as
[Equation]
Therefore, within T 1 = O (cid:16)
[Equation]
iterations, β i ( θ t ) becomes larger than 12 β ⋆i .
In the second phase, we assume that β
i ( θ t ) ≥ β
⋆i / 2 and define y t
=
β ⋆i − β
i ( θ t ) .
One can write
[Equation]
Hence, with additional T 2 = O (cid:16)
[Equation]
dβ ⋆i α (cid:17)(cid:17) , we have y t ≤ α √ d which implies β i ( θ t ) ≥ β
⋆i − α √ d .
Next, we show that there exists a time t ⋆ such that β ⋆i − 1 √ d α ≤ β
i ( θ t ⋆ ) ≤ β ⋆i + 1 √ d α .
Without loss of generality, we assume that t ⋆ is the first time that β
i ( θ t ) ≥ β ⋆i − 1 √ d α .
Due to the dynamic of β i ( θ t ) , the distance between two adjacent iterations can be upper bounded as
[Equation]
In particular, for t = t ⋆ − 1 , we have β
i
θ t ⋆ − 1 ) ≤ β ⋆ i − 1 √ d α , which in turn implies
[Equation]
Therefore, for each i ∈ S , we have | β
i
θ t ) − β ⋆i | ≤ 1 √ d α within T i
= O (cid:16)
[Equation]
iterations.
Meanwhile, we can show that the residual term | β i θ t ) | , ∀ i ∈ E remains small for
[Equation]
Therefore, we have that within T
= O (cid:16) 1 ηβ [Equation]
[Equation]
[Equation]
Case 2: 12 < γ ≤ 1 .
In this case, we have the following bounds for the signal and residual terms
[Equation]
[Equation]
We first analyze the dynamic of the signal term
β
i ( θ t ) for i ∈ S .
We will show that | β i ( θ t ) − β ⋆i | ≤ α √ k within T
= O (cid:16)
[Equation]
iterations.
Due to η ≲
[Equation]
k , we can further simplify the dynamic of β
i ( θ t ) as
[Equation]
Next, we divide our analysis into two phases.
In the first phase, we have β
i
θ t ) ≤ 12 β ⋆i .
We denote the number of iterations for this phase as T i, 1 .
We further divide this period into ⌈ log ( β ⋆i / 2 α ) ⌉ substages.
In each Substage k , we have C 1 2 k − 1 α ≤ β
i ( θ t ) ≤ C 1 2 k α .
Let t k be the number of iterations in Substage k .
We first provide an upper bound for t k .
To this goal, note that at this substage
[Equation] [Equation]
Hence, we have
[Equation]
Summing over t k , we obtain an upper bound for T i,
T i, 1 =
[Equation]
Via a similar argument, we can show that in the second phase, we have | β i ( θ t ) − β ⋆i | ≲ 1 √ d α within additional T i, 2 = O (cid:16)
[Equation]
(cid:17) iterations.
Therefore, for each i ∈ S , we conclude that | β i ( θ t ) − β ⋆i | ≤ α √ d
within T i
= T i, 1 + T i, 2 = O (cid:16)
[Equation]
(cid:17) iterations.
Meanwhile, for the residual term β i ( θ t ) , i ∈ E , we have
[Equation]
Therefore,
[Equation]
within T
= O (cid:16)
[Equation]
(cid:17) iterations.
This completes the proof.
E P ROOFS FOR K ERNEL R EGRESSION E. P ROOF OF P ROPOSITION
Note that θ t +
=
θ t − η ( θ t − θ ⋆ ) , and β i ( θ )
= θ i .
Hence, for every 1 ≤ i ≤ k , we have
[Equation]
[Equation]
This in turn implies
[Equation]
For i > k , we have
[Equation]
which completes the proof.
□
E. P ROOF OF T HEOREM
Due to our choice of initial point θ 0 =
α 1 , α ≲ | β ⋆k | , we have | β ⋆i − β
i ( θ 0 ) | ≤ 2 | β ⋆i | .
Hence, by Proposition 2, we have
[Equation]
[Equation]
Therefore, to prove | β ⋆i − β
i
( θ t ) | ≤ α √ k , i ∈ S , it suffices to have
[Equation]
On the other hand, to ensure | β i ( θ t ) | ≤ α √ d , i ∈ E , it suffices to have
[Equation]
Recall that α ≲ k | θ k | d and | θ 1 | ≥ | θ 2 | ≥ · · · ≥ | θ k | > 0 .
Therefore, within T =
O [Equation] iterations, we have
[Equation]
which completes the proof.
□
F P ROOFS FOR S YMMETRIC M ATRIX F ACTORIZATION F. I NITIALIZATION
We start by proving that both lower bound on signals at θ 0 and upper bound on energy at θ 0 are satisfied with high probability.
Recall that each element of U 0 is drawn from N (0 , α 2 ) .
The following proposition characterizes the upper and lower bounds for different coefficients β ij ( U 0 ) .
Proposition 5 (Initialization) .
With probability at least 1 − e − Ω( r ′ ) , we have
[Equation]
(45) and [Equation] for all i
[Equation]
j
[Equation]
i
[Equation]
r
[Equation]
Proof.
First note that U ⊤
z i ∼ N (0 , α 2 I r ′ × r ′ ) .
Hence, a standard concentration bound on Gaussian random vectors implies
[Equation]
Via a union bound, we have that with probability of at least 1 − e − Cr ′ :
[Equation]
Given these bounds, one can write
[Equation]
and
[Equation]
which completes the proof.
F. O NE - STEP D YNAMICS
In this section, we characterize the one-step dynamics of the basis coefficients.
To this goal, we first provide a more precise statement of Proposition 3 along with its proof.
Proposition 6.
For the diagonal element β ii ( U t ) , we have
[Equation]
where σ i = 0 for r < i ≤ d .
Moreover, for every i ̸ =
j
[Equation]
we have
[Equation]
Proof.
The iterations of GD on SMF take the form
[Equation]
This leads to
[Equation]
Recall that β ij ( U t ) =
⟨ U t U ⊤ t , z i z ⊤ j ⟩ , U t
U ⊤ t = (cid:80)
i,j β ij ( U t )
z i z ⊤ j , and M ⋆ = (cid:80) ri =
σ i
z
i z ⊤ i .
Based on these definitions, one can write
[Equation]
In light of the above equality and the orthogonality of { z i } i ∈ [ d ] , the RHS of Equation 55 can be written in terms of β kl ( U t ) , 1 ≤ k, l ≤ d .
In particular
[Equation]
Other terms in Equation 55 can be written in terms of β ij ( U t ) in an identical fashion.
Substituting these derivations back in Equation 55, we obtain
[Equation]
Note that the above equality holds for any 1 ≤ i, j ≤ r ′ .
In particular, for i =
j , we further have β ii ( U t +1 ) = (1 + [Equation] ( σ i − β ii ( U t )))
β ii ( U t ) − 2 η (cid:88)
j ̸ = i β 2 ij ( U t )
+ [Equation]
 (cid:88)
[Equation]
U
[Equation]
U
[Equation]
(cid:88)
j β
[Equation]
U
[Equation]
U t
[Equation]
  , (58) which completes the proof.
F. P ROOFS OF P ROPOSITION 3 AND T HEOREM
To streamline the presentation, we prove Proposition 3 and Theorem 3 simultaneously.
The main idea behind our proof technique is to divide the solution trajectory into r substages: in Substage i , the basis coefficient β ii ( U t ) converges linearly to σ i while all the remaining coefficients remain almost unchanged.
More precisely, suppose that Substage i lasts from iteration t i,s to t i,e .
We will show that β ii ( U t i,e )
≈ σ i and β ij (
U t
i,e ) ≈
β ij ( U t i,s ) .
Recall that γ = min 1 ≤ i ≤ r σ i − σ i +1 is the eigengap of the true model, which we assume is strictly positive.
Substage 1 .
In the first stage, we show that β
( U t ) approaches σ 1 and | β ij ( U t )
| , i, j ≥ 2 remains in the order of poly( α ) within T
= O (cid:16) 1 ησ 1 log (cid:0) σ 1 α (cid:1)(cid:17) iterations.
To formalize this idea, we further divide this substage into two phases.
In the first phase (which we refer to as the warm-up phase), we show that β
( U t ) will quickly dominate the remaining terms β ij ( U t ) , ∀ ( i, j ) ̸
= (1 , 1) within
O (cid:16) 1 ηγ log log( d ) (cid:17) iterations.
This is shown in the following lemma.
Lemma 1 (Warm-up phase) .
Suppose that the initial point satisfies Equation 45 and Equation 46.
Then, within O (cid:16) 1 ηγ log log( d ) (cid:17) iterations, we have
[Equation]
Proof.
To show this, we use an inductive argument.
Due to our choice of the initial point, we have | β ij ( U 0 ) | ≲ r ′
α 2 log( d ) 1+ σ 1 /γ , 1 ≤ i, j ≤ d .
Now, suppose that at time t , we have | β ij ( U t ) | ≲ r ′
α 2 log( d ) 1+ [Equation] 1 /γ , 1 ≤ i, j ≤ d .
Then, by Proposition 6, we have
[Equation]
Similarly, for the remaining coefficients β ij ( U t ) , ∀ ( i, j ) ̸
=
[Equation] we have [Equation]
[Equation]
Note that the stepsize satisfies η ≲ 1 σ 1 , and σ 1 − σ 2 ≥ γ ≳ dr ′ α 2 log( d ) 1+ σ 1 /γ .
Hence, we have β
[Equation]
t +
[Equation]
This further implies
[Equation]
On the other hand [Equation] Equation 49 and Equation 50 imply that
[Equation]
Hence, within O (cid:16) 1 ηγ log log( d ) (cid:17) iterations [Equation] we have β ij ( U t ) ≥ | β 11 ( U t ) | , ∀ ( i, j ) ̸
= (1 , 1) .
Moreover, we have that during this phase,
[Equation]
(65) which completes the proof.
After the warm-up phase [Equation] we show that β 11 ( U t ) quickly approaches σ 1 while the remaining coefficients remain small.
Lemma 2 (Fast growth) .
After the warm-up phase followed by O (cid:16) 1 ησ 1 log (cid:0) σ 1 α (cid:1)(cid:17) iterations, we have
[Equation]
Moreover, for | β ij ( U t ) | , ∀ ( i, j ) ̸ =
(1 , 1) , we have
[Equation]
Before providing the proof of Lemma 2 we analyze an intermediate logistic map which, as will be shown later, closely resembles the dynamic of β ij ( U t ) :
[Equation]
The following two lemmas characterize the dynamic of a single logistic map, as well as the dynamic of the ratio between two different logistic maps.
Lemma 3 (Iteration complexity of logistic map) .
Suppose that α ≤ ε ≤ 0 .
1 σ .
Then, for the logistic map, we have x T ≥ σ − ε within T =
[Equation]
Lemma 4 (Separation between two logistic maps) .
Let σ 1 , σ 2 be such that σ 1 − σ 2 > 0 , and
[Equation]
Then, within T =
[Equation]
iterations, we have
[Equation]
The proofs of Lemmas 3 and 4 are deferred to Appendix F.4.
We are now ready to provide the proof of Lemma 2.
Proof of Lemma 2.
Similar to the proof of Lemma 1, we use an inductive argument.
Suppose that t 0 is when the second phase starts.
According to Lemma 1, we have | β ij ( U t 0 ) | ≤ β 11 ( U t 0 ) ≲
[Equation]
Therefore, the base case of our induction holds.
Next, suppose that at some time t within the second phase, we have | β ij ( U t ) | ≲ σ 1 α 2 σ 1 − σi − σj 2 σ, ∀
( i, j ) ̸
= (1 , 1) .
Our goal is to show that | β ij ( U t +1 ) | ≲ σ 1 α 2 σ 1 − σi − σj 2 σ.
To this goal, we consider two cases.
Case I: i ≤ r or j ≤ r and ( i, j ) ̸
=
We have
[Equation]
Here in (a) we used the assumption | β ij ( U
[Equation]
, ∀ ( i, j ) ̸
= (1 , 1) .
Hence, by Lemma 4, we have
[Equation]
, where 1 ≤ i ≤ r or 1 ≤ j ≤ r.
Case II: i, j ≥ r + 1 .
For β ij ( U t ) such that i, j ≥ r + 1 , its dynamic is characterized by
[Equation]
+
[Equation]
Hence, for t ≲ 1 ησ 1 log (cid:0) σ 1 α (cid:1) , we have | β ij ( U t ) | ≤ (cid:16) 1 + η O (cid:16) dα γσ 1 (cid:17)(cid:17) O (cid:16) 1 ησ 1 log ( σ 1 α ) (cid:17) | β ij ( U 0 ) | ≲ | β ij ( U 0 ) | since we assume α ≲ (cid:0) σ 1 d (cid:1) σ 1 /γ .
This completes our inductive proof for | β ij ( U t ) | ≲
[Equation]
, ∀ ( i, j ) ̸
= (1 , 1) in the second phase.
Finally, we turn to β 11 ( U t ) .
One can write
[Equation]
Here in (a) we used the fact that | β ij ( U t ) | ≲ σ 1 r ′ log
[Equation]
In (b) [Equation] we used the assumption that α ≲ σ r d 2 σ 1 /γ .
The above inequality together with Lemma 3 entails that within
[Equation]
iterations, we have β 11 ( U t ) [Equation]
[Equation]
This completes the proof of Lemma 2 and marks the end of Substage 1.
□
Next, we move on to Substage 2 .
Substage 2 .
In Substage 2 , we show that the second component β 22 ( U t ) converges to σ 2 within
O (cid:16) 1 ησ 2 log (cid:0) σ 2 α (cid:1)(cid:17) iterations while the other coefficients remain small.
To this goal, we first study the one-step dynamic of β 22 ( U t ) :
[Equation]
Different from the dynamic of β 11 ( U t ) , not all the coefficients β ij ( U t ) with i = 2 or j = 2 are smaller than β 22 ( U t ) at the beginning of Substage 2 .
In particular, the basis coefficient | β 12 ( U t ) | may be much larger than β 22 ( U t ) at the beginning of Substage 2 .
To see this, note that, according to Equation 68, we have β 12 ( U t ) ≲ α σ 1 − σ 2 2 σ 1 and β 22 ( U t ) ≲
α 2( σ 1 − σ 2) 2 σ.
Hence, it may be possible to have | β 12 ( U t ) | ≍ (cid:112) β 22 ( U t )
≫ β 22 ( U t ) .
Therefore, the term 2 ηβ 212 ( U t ) in Equation 76 must be handled with extra care.
Note that if we can show σ 2 β 22 ( U t ) ≫ β 212 ( U t ) , then 2 ηβ 212 ( U t ) can be combined with the first term in the RHS of Equation 76 and the argument made in Substagecan be repeated to complete the proof of Substage 2.
However, our provided bound in Equationcan only imply β 12 ( U t ) 2 ≍ β 22 ( U t ) .
Therefore, we need to provide a tighter analysis to show that β 212 ( U t ) ≪ σ 2 β 22 ( U t ) along the trajectory.
Upon controlling β 212 ( U t ) , we can then show the convergence of β 22 ( U t ) similar to our analysis for β 11 ( U t ) in Substage 1 .
To control the behavior of β 212 ( U t ) , we study the ratio ω ( t ) :
=
β 212 ( U t ) β 22 ( U t ) .
We will show that ω ( t ) ≪ σalong the trajectory.
To this goal, we will show that ω ( t ) can only increase for O ( 1 ησ 1 log( σ 1 α )) iterations.
Therefore, its maximum along the solution trajectory happens at T
= O ( 1 ησ 1 log( σ 1 α )) .
Therefore, by bounding the maximum, we can show that ω ( t ) remains small throughout the solution trajectory.
First, at the initial point, we have ω (0) ≤ 64 log 2 ( d ) r ′
α 2 , which satisfies our claim.
We next provide an upper bound for ω ( t + 1) based on ω ( t ) .
Note that
[Equation]
Due to the first inequality, ω ( t ) can be increasing only until β 11 ( U t ) ≥ σ 1 2 ± poly( α ) .
On the other hand, due to the dynamic of β 11 in Substage 1 , we can show that β 11 ( U t ) ≥ σ 1 2 ± poly( α ) in at most
T =
O (cid:16) 1 ησ 1 log (cid:0) σ 1 α (cid:1)(cid:17) iterations.
Therefore, ω ( t ) takes its maximum at T
= O (cid:16) 1 ησ 1 log (cid:0) σ 1 α (cid:1)(cid:17) .
On the other hand, we know that β
11 ( U t ) satisfies
[Equation]
Hence, we can bound ω ( t ) as
[Equation]
Here in (a) we used the fact that α ≲ (cid:0) ησ 2 r (cid:1) σ 1 /γ .
Hence [Equation] we have
[Equation]
Due to our assumption α ≲ (cid:0) κ log 2 ( d ) (cid:1) − Ω ( 1 ησr ) [Equation] we conclude that ω ( t ) ≤ 0 .
01 σ 2 .
Therefore, equation 76 can be lower bounded as
[Equation]
The rest of the proof is a line by line reconstruction of Substage 1 and hence omitted for brevity.
Substage 3 ≤ k ≤ r .
Via an identical argument to Substage 2 , we can show that for each Substage 3 ≤ k ≤ r , we have
[Equation]
within T k
= O (cid:16) 1 ησ k log (cid:0) σ 1 α (cid:1)(cid:17) iterations.
This completes the proof of the first statement of Proposition 3.
To prove the second statement of Proposition 3 as well as Theorem 3, we next control the residual terms.
First, we consider the residual term β ij ( U t ) where either i ≤ r or j ≤ r .
Note that β ij ( U t )
=
β ji ( U t ) and hence we can assume i ≤ r without loss of generality.
We will show that β ij ( U t ) decreases linearly once the corresponding signal β ii ( U t ) converges to the vicinity of σ i .
To this goal, it suffices to control the largest component β max ( U t ) :
= max i ̸
= j,i ≤ r | β ij ( U t ) | .
Without loss of generality, we assume that the index ( i, j ) attains the maximum at time t , i.e., β max ( U t ) =
| β ij ( U t ) | .
One can write
[Equation]
[Equation]
Here in (a) we used the fact that 0 .
99 σ i ≤ β ii ( U t ) ≤ σ i , 1 ≤ i ≤ r and the fact that β max ( U t ) =
| β ij ( U t ) | .
Hence, we conclude that β max ( U t +1 ) ≤ (1 − 0 .
9 ησ r ) β max ( U t ) .
Therefore, within additional O (cid:16) 1 ησ r log (cid:0) 1 α (cid:1)(cid:17) iterations, we have | β ij ( U t ) | ≤ r ′ log( d ) α 2 for all ( i, j ) such that i ≤ r, i ̸
= j .
The remaining residual terms, i.e., those coefficients β ij ( U t ) for which i, j > r , can be bounded via the same approach in Case II of substage 1 .
In particular, we can show that | β ij ( U t ) | ≲ | β ij ( U 0 ) | ≲ r ′ log( d ) α 2 , ∀ i, j > r .
For brevity, we omit this step.
This completes the proof of the second statement of Proposition 3.
Finally, to prove Theorem 3, we show that once | β ij ( U t ) | ≤ r ′ log( d ) α 2 , ∀ ( i, j ) ̸
=
( k, k ) , 1 ≤ k ≤ r , the signals β kk ( U t ) will further converge to σ k ± O (cid:0) α 2 (cid:1) within O (cid:16) 1 ησ k log (cid:0) σ k α (cid:1)(cid:17) iterations.
To see this, we simplify the dynamic of β kk ( U t ) as
[Equation]
which leads to
[Equation]
Hence, within additional O (cid:16) 1 ησ k log (cid:0) σ k α (cid:1)(cid:17) iterations, we have | σ k − β kk
( U t +1 )
| =
O (cid:16) 1 σ k dr ′ 2 log 2 ( d ) α 4 (cid:17) for every 1 ≤ k ≤ r .
In conclusion, we have
[Equation]
within O (cid:16) 1 ησ r log (cid:0) σ r α (cid:1)(cid:17) iterations.
This completes the proof of Theorem 3.
□
F. A NALYSIS OF THE L OGISTIC M AP
In this section, we provide the proofs of Lemmas 3 and 4.
U PPER B OUND OF I TERATION C OMPLEXITY
Recall the logistic map
[Equation]
Here the initial value satisfies 0 <
α ≪ σ .
Vaskevicius et al.
(2019) provide both upper and lower bounds for x t that follows the above logistic map.
However, their bounds are not directly applicable to our setting.
Hence, we need to develop a new proof for Lemma 3.
Proof of Lemma 3.
We divide the dynamic into two stages: (a) x t ≤ 12 σ , and (b) x t ≥ 12 σ .
Stage 1: x t ≤ 12 σ .
We consider K = ⌈ log (cid:0) 12 σ/α (cid:1) ⌉ substages, where in each Substage k , we have αe [Equation] ≤ [Equation]
t ≤ αe k +1 .
Suppose that t k is the number of iterations in Substage k .
One can write
[Equation]
[Equation]
Hence, it suffices to find the smallest t =
t min such that
[Equation]
Solving this inequality leads to
[Equation]
Based on the above equality, we provide an upper bound for t min :
[Equation]
Here in (a) we used the fact that log(1 + x ) ≤ x 1+ x , ∀ x > − 1 and in (b) we used the fact that x t ≤ 12 σ .
Hence, we have t k ≤ 1+ ηαe k +1 log(1+ ησ ) .
Therefore, the total iteration complexity of Stage 1 is upper bounded by
[Equation]
Stage 2: x t ≥ 12 σ .
In this stage, we rewrite the equation 82 as
[Equation]
Via a similar trick, we can show that within additional T 2 = log(4 σ/ε ) log(1+ ησ ) iterations, we have x t ≥ σ − ε , which can be achieved in the total number of iterations T
= T 1 + T 2 =log(1+ ησ ) (log(4 σ/α ) + log(4 σ/ε )) iterations.
This completes the proof of Lemma□
S EPARATION BETWEEN T WO I NDEPENDENT S IGNALS
In this section, we show that there is a sharp separation between two logistic maps with signals σ 1 , σprovided that σ 1 ̸ = σ 2 .
In particular, suppose that σ 1 − σ 2 ≥ γ > 0 and
[Equation]
Proof of Lemma 4.
By Lemma 3, we have x T ≥ σ 1 − ε within T =
[Equation]
Therefore, it suffices to show that y t remains small for t ≤ T .
To this goal, note that
[Equation]
Hence, we need to bound Γ =
(1 + ησ 2 ) T .
Taking logarithm of both sides, we have
[Equation]
Now, we provide a lower bound for the ratio log(1 + ησ 1 ) / log(1 + ησ 2 ) :
[Equation] [Equation]
[Equation]
where (a) follows from the assumption η ≤ 14 σ 1 .
Therefore, we have
[Equation]
which implies that
[Equation]
This completes the proof of Lemma 4.
□
G P ROOF FOR T ENSOR D ECOMPOSITION
In this section, we prove our results for the orthonormal symmetric tensor decomposition (OSTD).
Different from matrix factorization, we use a special initialization that aligns with the ground truth.
In particular, for all 1 ≤ i ≤ r ′ , we assume that sin( u i (0) , z i ) ≤ γ for some small γ .
We will show that u i ( t ) aligns with z i along the whole optimization trajectory.
To this goal, we define v ij ( t ) = ⟨ u i ( t ) , z j ⟩ for every 1 ≤ i ≤ r ′ and 1 ≤ j ≤ d .
Recall that Λ is a multi-index with length l .
We define | Λ | k as the number of times index k appears as one of the elements of Λ .
Evidently, we have 0 ≤ | Λ | k ≤ l .
Based on these definitions, one can write
[Equation]
Now, it suffices to study the dynamic of v ij ( t ) .
In particular, we will show that v ij ( t ) remains small except for the top- r diagonal elements v jj ( t ) , 1 ≤ j ≤ r , which will approach σ 1 /lj .
To make this intuition more concrete, we divide the terms { v ij ( t ) } into three parts:
• signal terms defined as v jj ( t ) , 1 ≤ j ≤ r , • diagonal residual terms defined as v jj ( t ) , r + 1 ≤ j ≤ d , and • off-diagonal residual terms defined as v ij ( t ) , ∀ i ̸
= j .
Moreover, we define V ( t ) = max i ̸
=
j | v ij ( t ) | as the maximum element of the off-diagonal residual terms at every iteration t .
When there is no ambiguity, we will omit the dependence on iteration t .
For example, we write v ij = v ij ( t ) and V = V ( t ) .
Similarly, when there is no ambiguity, we write β Λ ( t ) or β Λ in lieu of β Λ ( U ( t )) .
Our next lemma characterizes the relationship between β Λ and v ij .
Lemma 5.
Suppose that max j ≥ r +1 (cid:12)(cid:12)
[Equation]
V , and V ≤ σ
[Equation]
Then,
• For β Λ j with Λ j = ( j, . . .
, j ) , we have
[Equation]
• For β Λ with at least two different indices in Λ , we have
[Equation]
The proof of this lemma is deferred to Appendix G.1.
Lemma 5 reveals that the magnitude of β Λ can be upper bounded by max i ̸
=
j | v ij | .
Next, we control v ij by providing both lower and upper bounds on its dynamics.
Proposition 7 (One-step dynamics for v ij ( t ) ) .
Suppose that we have V ( t ) ≲ σ 1 /l 1 d − 1 l − 1 and v ii ( t ) ≤ σ 1 /l 1 , ∀ 1 ≤ i ≤ d .
Moreover, suppose that the step-size satisfies η ≲ 1 lσ 1 .
Then,
• For the signal term v ii ( t ) , 1 ≤ i ≤ r , we have
[Equation]
• For the diagonal residual term v ii ( t ) , r + 1 ≤ i ≤ d , we have
[Equation]
• For the off-diagonal term V ( t ) , we have
[Equation]
The proof of this proposition is deferred to Appendix G.2.
Equipped with the above one-step dynamics, we next provide a bound on the growth rate of v ij .
Proposition 8.
Suppose that the initial point satisfies sin( u i (0) , z i ) ≤ γ and ∥ u i (0) ∥
=
α 1 /l with α ≲ 1 d l 3 , γ ≲ 1 lκ l l − 2 .
Moreover, suppose that the step-size satisfies η ≲ 1 lσ 1 .
Then, within t ⋆ = 8 ηlσ r α − l − 2 l iterations,
• For the signal term v ii ( t ) , 1 ≤ i ≤ r , we have
[Equation]
• For the diagonal residual term v ii ( t ) , r + 1 ≤ i ≤ d , we have
[Equation]
• For the off-diagonal term V ( t ) , we have
[Equation]
The proof of this proposition is deferred to Appendix G.3.
With the above proposition, we are ready to prove Theorem 4.
Proof of Thereom 4.
We have the following decomposition
[Equation]
Hence [Equation] it suffices to bound each | β Λ − β ⋆ Λ | .
Combining Lemma 5 and Proposition 8, we have for every
[Equation]
where in the last inequality
[Equation]
we used σ j ≥
[Equation]
α ≤
[Equation]
and γ ≤ 1 .
For the remaining diagonal elements β Λ j , r + 1 ≤ j ≤ d , we have
[Equation]
For the general β Λ with at least two different indices in the multi-index Λ
[Equation]
we have
[Equation]
Hence [Equation] we conclude
[Equation]
[Equation]
which completes the proof of the theorem.
G. P ROOF OF L EMMA
Proof.
We first analyze β Λ j .
Note that
[Equation]
Hence,
[Equation] [Equation]
where we used the definition of V .
For general β Λ where there are at least two different elements in the multi-index Λ , we have
[Equation]
where in the last inequality, we used the assumption that V ≲ σ 1 /l 1 d − 1 l − 1 and max j ≥ r +1 (cid:12)(cid:12) v ljj (cid:12)(cid:12) ≲ σ l − 1 l
V .
This completes the proof.
G. P ROOF OF P ROPOSITION
In this section, we provide the proof for Proposition 7.
For simplicity and whenever there is no ambiguity, we omit the iteration t and show iteration t + 1 with superscript ‘ + ’.
For instance, we write v ij = v ij ( t ) and v + ij = v ij ( t + 1) .
Recall that v ij = ⟨ u i , z j ⟩ .
For simplicity, we denote µ = σ 1 /l 1 .
Hence, by our assumption, we have v ii ≤ µ, ∀ 1 ≤ i ≤ r .
We first provide the exact dynamic of v ij in the following lemma.
Lemma 6.
The one-step dynamic of v ij takes the following form
[Equation]
Proof.
Recall that L = 12 (cid:80) Λ ( β Λ − β ⋆ Λ ) 2 and β Λ = (cid:80) r ′ i =1 (cid:81) k ∈ Λ ⟨ u i , z k ⟩ | Λ | k .
Moreover, we have β ⋆ Λ j = σ j for 1 ≤ j ≤ r , and β ⋆ Λ = 0 otherwise.
We first calculate
[Equation]
Hence, the partial derivative of L ( t ) with respect to u i is
[Equation]
Note that { z j } j ∈ [ d ] are unit orthogonal vectors.
Hence, we have
[Equation]
By the definition of v ij
[Equation]
its update rule can be written in the following way
[Equation]
Here in ( a ) , we used the update rule for u i .
In ( b ) , we applied equation 114 to exclude those Λ without j .
In ( c ) , we simply rearranged the above equation according to the cardinality | Λ | j .
We further isolate the term that only has v ij :
[Equation]
Here in ( a ) , we rearranged terms and isolated the term with | Λ | j = l .
Note that the remaining terms must satisfy 1 ≤ | Λ | j ≤ l − 1 , which indicates that there must be at least 2 different indexes in Λ which in turn implies β ⋆ Λ = 0 .
In ( b ) , we used the definition of β Λ j .
This completes the proof of
Lemma 6.
Equipped with Lemma 6, we are ready to prove Proposition 7.
Proof of Proposition 7.
The proof is divided into three parts:
[Equation]
We first consider the signal terms v ii ( t ) , 1 ≤ j ≤ r .
First, upon setting i =
j in Lemma 6, we have
[Equation]
v lki 
 v l − 1 ii − η l − 1 (cid:88) s = s (cid:88) Λ: | Λ | i = s β Λ   (cid:89) k ∈ Λ ,k ̸ = i v | Λ | k ik 
 v s − 1 ii ≥ v ii + ηl (cid:0) σ i − v lii − mV l (cid:1) v l − 1 ii − η l − 1 (cid:88) s
= s (cid:88) Λ: | Λ | i =
s β Λ   (cid:89) k ∈ Λ , k ̸
=
i v | Λ | k ik 
 v s − 1 ii . Now we aim to control ( A ) = β Λ (cid:16)(cid:81) k ∈ Λ ,k ̸ =
i v | Λ | k ik (cid:17) v s − 1 ii for | Λ | i = s ∈ { 1 , . . .
, l − 1 } .
We have ( A ) = r ′ (cid:88) j = (cid:32)(cid:89) h ∈ Λ v | Λ |
h jh (cid:33)   (cid:89) k ∈ Λ ,k ̸
=
i v | Λ | k ik 
 v | Λ | i − 1 ii ( a ) ≤ r ′ (cid:88)
j = v | Λ | j jj V l −| Λ | j V l − s v s − 1 ii ( b ) ≤ v 2 s − 1 ii
V 2 l − 2 s + r ′ max
j ̸
= i (cid:12)(cid:12)(cid:12) v | Λ | j jj V 2 l − s −| Λ | j v s − 1 ii (cid:12)(cid:12)(cid:12) ( c ) ≤ v 2 l − 3 ii
V 2 + r ′ max j ̸ =
i µ | Λ | j + s − 1 V 2 l − s −| Λ | j ≤ v 2 l − 3 ii
V 2 +
[Equation]
In ( a ) , we used the fact that | v ij | ≤ V, ∀ i ̸
=
j .
In ( b ) , we isolated the term with j = i and bounded the remaining terms with their maximum value.
In ( c ) , we used the fact that V ≤ v ii ≤ µ .
After substituting equation 118 into equation [Equation] we have
[Equation]
where C sl = (cid:0) ls (cid:1) .
Note that (cid:80) l − 1 s =1 sC sl ( d − 1) l − s = ld l − 1 − l ≤ ld l − 1 .
Therefore, we have
[Equation]
where in the last inequality, we used the fact that r ′ ≤ d .
Diagonal Residual Term: v ii , r + 1 ≤ i ≤ d
In this case we consider the terms v ii with r + 1 ≤ i ≤ d , which is similar to the case 1 ≤ i ≤ r .
Without loss of generality, we assume that v ii ≥ 0 .
The case v ii ≤ 0 can be argued in an identical fashion.
By equation 117, we have
[Equation]
i v | Λ | k ik 
 v s − 1 ii ≤ v ii − ηl (cid:0) v lii − r ′
V l (cid:1) v l − 1 ii − η l − 1 (cid:88) s
= s (cid:88) Λ: | Λ | i =
s β Λ   (cid:89) k ∈ Λ ,
k ̸
=
i v | Λ | k ik 
 v s − 1 ii .
For ( A ) =
β Λ (cid:16)(cid:81) k ∈ Λ ,k ̸
=
i v | Λ | k ik (cid:17) v s − 1 ii , we further have ( A ) =
r ′ (cid:88)
j = (cid:32)(cid:89) h ∈
Λ v | Λ |
h jh (cid:33)   (cid:89) k ∈ Λ ,k ̸
= i v | Λ | k ik 
 v s − 1 ii ≥ (cid:32)(cid:89) h ∈ Λ v | Λ | h ih (cid:33)   (cid:89) k ∈ Λ ,k ̸ =
i v | Λ | k ik 
 v s − 1 ii − (cid:88)
j ̸ = i (cid:32)(cid:89) h ∈ Λ v | Λ |
h jh (cid:33)   (cid:89) k ∈ Λ ,k ̸
=
i v | Λ | k ik 
 v s − 1 ii ≥   (cid:89) k ∈ Λ ,k ̸
= i v 2 | Λ | k ik 
 v 2 s − 1 ii (cid:124) (cid:123)(cid:122) (cid:125) ≥− r ′ max
j ̸
= i (cid:12)(cid:12)(cid:12) v | Λ | j jj V 2 l − s −| Λ | j v s − 1 ii (cid:12)(cid:12)(cid:12) ≥
[Equation]
Therefore, we obtain
[Equation]
Off-diagonal Residual Term: V ( t )
Finally, we characterize the dynamic of V ( t ) = max i ̸
=
j | v ij ( t ) | .
To this goal, we first consider the dynamic of each v ij such that i ̸
=
j .
Without loss of generality, we assume that v ij ≥ 0 .
One can write
[Equation]
where in the last inequality we use the assumption V ≲ σ 1 /l 1 d − 1 l − 1 .
Similar to the previous case, we next bound ( A ) =
β Λ (cid:16)(cid:81) k ∈ Λ ,k ̸
= j v | Λ | k ik (cid:17) v s − 1 ij .
First note that | Λ | j = s ∈
[ l − 1] .
Hence, we have that
[Equation]
We further note that | Λ | h + | Λ | i ≤ l − s .
Hence, it can be lower bounded by
[Equation]
Pluggin our estimation for ( A ) into equation 124, we finally have
[Equation]
The last inequality comes from the assumption that V ≤ 1 d l .
This completes the proof of Proposition 7.
G. P ROOF OF P ROPOSITION
In this section, we provide the proof of Proposition 8.
We will show that the signal terms v ii ( t ) , 1 ≤ i ≤ r quickly converge to σ 1 /l i , and the residual terms remain small.
To this goal, we first study the dynamic of V ( t ) .
Iteration complexity of the off-diagonal residual term V ( t ) .
To start with the proof, we first study the time required for the off-diagonal term V ( t ) to go from
V (0) to 2 1 /l V (0) , i.e., T 1 = min t ≥ 0 { V ( t ) ≥ 2 1 /l V (0) } .
By Proposition 7, we know
[Equation]
Hence, for 0 ≤ t ≤ T 1 , we have
[Equation]
[Equation]
Note that V ( T 1 ) ≥ 2 1 /l V (0) .
Therefore,
[Equation]
[Equation]
Solving the above inequality for T 1 , we obtain
[Equation]
On the other hand, our initial point satisfies sin( u i (0) , z i ) ≤ γ and ∥ u
i (0) ∥
=
α 1 /l .
Hence, we have
V (0) ≤ α 1 /l γ .
Substituting this into the above equation, we conclude that
[Equation]
[Equation] Note that T 1 ≥ [Equation] t ⋆ .
Hence, we have V ( t ⋆ ) ≤ 2 1 /l V (0) ≤ [Equation] 2 α [Equation]
1 /l γ . Iteration complexity of the diagonal residual term v ii ( t ) , r + 1 ≤ i ≤ d .
Next, we show that v ii ( t ) , r + 1 ≤ i ≤ d will remain small during 0 ≤ t ≤ t ⋆ = 8 ηlσ r α − l − 2 l .
First by equation 123 in the proof of Lemma 6, we have for every 0 ≤ t ≤ t ⋆
[Equation]
Note that V ( t ) ≤ 2 1 /l V (0) , ∀ t ≤ t ⋆ , which leads to
[Equation]
for 0 ≤ t ≤ t ⋆ .
Here we used the assumption that α ≲ 1 d l 3 and γ ≲ 1 κl 1 l − 2 .
Iteration complexity of the signal term v ii ( t ) , 1 ≤ i ≤ r .
As the last piece of the proof, we show that by iteration t ⋆ = 8 ηlσ r α − l − 2 l , the signal v ii , 1 ≤ i ≤ r will converge to the eigenvalue σ 1 /l i .
First, recall that
[Equation]
where in the last inequality we used the fact that V ( t ) ≤ 2 1 /l V (0) for 0 ≤ t ≤ t ⋆ .
In light of the above inequality, we characterize the convergence of v ii using a similar method as in (Ma & Fattahi, 2022a).
In particular, we divide our analysis into two phases.
Phase 1.
In the first phase, we have v ii ≤ (0 .
5 σ i ) 1 /l .
First, since v ii (0) ≥ α 1 /l (cid:112) 1 − γ 2 , we can easily conclude that v ii ( t + 1) ≥ v ii ( t ) by induction.
Hence, we can simplify the dynamic as
[Equation] [Equation]
Next, we further split the interval I = (cid:104) 0 , 0 .
5 σ 1 /l i (cid:105) into N = O (cid:16) log (cid:16) 0 .
5 σ 1 /l i /α (cid:17)(cid:17) sub-intervals {I 0 , · · · , I N − 1 } , where I k
=
[2 k v ii (0) , 2 k +1 v ii (0)) .
Let T k collect the iterations that v ii spends in
I k .
Accordingly, let |T k | =
t k be the number of iterations that v ii spends within I k .
First note that v ii ( t ) ≥ 2 k v ii (0) for every t ∈ T k .
Hence, we have
[Equation]
[Equation]
which implies
[Equation]
[Equation]
By summing over k = 0 , · · · , N − 1 , we can upper bound the required number of iterations T T
[Equation]
where the last inequality is due to our assumption γ ≲ 1 κl 1 l − 2 .
Phase 2.
In the second phase, we have v ii ≥ 0 .
5 σ 1 /l i .
We further simplify equation 134 as
[Equation]
where we denote ˜ σ i
=
σ i −
[Equation]
d
[Equation] σ [Equation] i α 2 /l γ 2 .
Then, via a similar trick, within additional T 4 ≤ [Equation] iterations, we have v lii [Equation]
t ) ≥ σ i − 8 d l
[Equation]
1 σ l − 2 l i α 2 /l
γ 2 .
A similar argument on the upper bound shows v l ii ( t ) ≤ [Equation] α 2 /l γ 2 , which completes the proof.
□
H A UXILIARY L EMMAS
Lemma 7 (Bernoulli inequality) .
For 0 ≤ x < 1 r − 1 , and r > 1 , we have
[Equation]