Behind the Scenes of Gradient Descent: A Trajectory Analysis via Basis Function Decomposition
Authors: Jianhao Ma
Lingjun Guo
Salar Fattahi
Venue: ICLR 2023 poster (2023)
Keywords: nonconvex optimization
trajectory analysis
neural network optimization

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

Published as a conference paper at ICLR 2023
BEHIND THE SCENES OF GRADIENT DESCENT: A TRAJECTORY ANALYSIS VIA BASIS FUNCTION DE- COMPOSITION
Jianhao Ma, Lingjun Guo, Salar Fattahi
Department of Industrial and Operational Engineering, University of Michigan, Ann Arbor (jianhac, glingiun, fattahi/ @uuichh edu
ABSTRACT
This work analyzes the solution trajectory of gradient-based algorithms via a novel basis function decomposition. We show that, although solution trajecto- algorithms may vary depending on the learning task, they
ries of gradient-base behave almost monotonically when projected onto appropriate orthonormal function basis. Such projection gives rise to a basis function decomposition of the solution trajectory. Theoretically, we decomposition to establish the convergence gradient descent (GD) on sev- eral representative learning tasks. In particular, we improve the convergence of GD on symmetric matrix factorization and provide a completely new con- vergence result for the orthogonal symmetric tensor decomposition. Empiri- cally, 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 funct i ceon-baasii baw.s-decoooossseeeeeeeee decomposition
proposed basis function
INTRODUCTION
Learning highly nonlinear models amounts to solving nonconvex optimization problem, which 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.) (017]- may not be globally optimal. In the rare event where the GD can recover a globally optimal solution, the recovered solution may correspond overfitted model rather than with desirable generalization.
Inspired by the large empirical success of gradient-based algorithms learning complex models, recent work has postulated that typical training losses have benign landscapes: they devoid of spurious local minima and their global solutions coincide with true solutions-i.e. solutions corresponding 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) global and true solutions do not coincide (Ma)
& Fattahi 2022b). such 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 al. 2019), low-rank matrix factorization (Li et al. 2018 Stögerr & Soltanolkotabi 2021), and linear DNNs (Arora 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 novel basis function decomposition. We show that, although the dynamics of GD may vary

---

Published as a conference paper at ICLR 2023
3.004
0.044
01032
3.00
8.00
D1000
-2004
0.04
10002
opech
opcch
(a) AlexNet
(b) ResNet-18
(c) ViT
10.0
0.:30
0.31
11:3
000
045
Reration
(d) AlexNet
(c) ResNet-18
(f) VIT
The solution trajectories of on AlexNet and ResNet-18 Adam on ViT with E2-loss after projecting onto two different orthonormal bases. The first row shows the trajectories of the top-5 coefficients after projecting onto randomly generated orthonormal basis. The second row shows the trajectories of the top-5 coefficients after projecting onto the eigenvectors of the conjugate kernel evaluated at the last epoch. More detail on implementation he found Appendix B
Figure 1:
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 optimization trajectories of two adaptive gradient-based algorithms, namely Adam W and LARS, on three different DNN architectures, namely AlexNet (Krizhevsky et al. 2017). ResNet-18 18 (He) et al. 2016), and Vision Transformer (ViT) (Dosovitskiy et al.) (2020) with the CIFAR-10 dataset. The first row of the T the top-5 coefficients of the solution trajectory when projected onto a randomly generated orthonormal basis. We see that the trajectories of the coefficients 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] shows the same trajectory after projecting onto an orthogonal basis defined as the eigenvectors of the conjugate kernel after training (Long) 202T) (see Section 3.4 and Appendix 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.
MAIN CONTRIBUTIONS
The monotonicity of the projected solution trajectory motivates the of an appropriate basis function decomposition to analyze the behavior of gradient-basee algorithms. In this paper, we show how appropriate basis function decomposition analysis for gradient-basee algorithms on several representative learning problems, from simple kernel regression complex DNNs. Our main contributions are summarized below:
be used to provide much simpler convergence
Global convergence of GD via basis function decomposition: We prove that GD learns the coefficients an appropriate function basis that forms the true model. In particular, show that GD learns the true model when applied the expected loss under certain gradient independence and gradient dominance conditions. Moreover, we characterize the convergence rate of GD, identifying conditions under which it enjoys linear sublinear
2

---

Published as a conference paper at ICLR 2023
convergence rates. Our result does not require 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 much simpler trajectory analysis of GD for much broader classes of models. Using technique, we improve the convergence of GD on the symmetric matrix factorization and provide an entirely new convergence result for GD the orthogonal symmetric tensor decomposition. We also prove that GD enjoys an incremental learning phenomenon in both problems.
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.
Empirical validation on DNNs:
GENERAL FRAMEWORK' FUNCTION BASIS DECOMPOSITION
We study the optimization trajectory of GD the expected (population) f2-loss
min L (0) : 2 En, [fo(xx) v)² OEG
(expected E2-loss)
Here the input x E IR is drawn from an unknown distribution D, and the output label is generated as y (x) + e, where is additive noise, independent of æ. with mean E[c] and variance E[="]= o² < oc. The model fe(x) is characterized by a parameter vector E R which naturally induces set of admissible models (model space for short) Fg == {Jo 0
E Rm R." We do not
require the true model f to lie within the model space; instead, we seek to obtain model fe- E Fe that is closest to ft in diss)ddiitt
other words, we consider f = fe- (x) fi (æ). where To minimize the expected fo-loss we use vanilla GD with
0r arg ming ll.e f
constant step-size 17 > 0:
0.+1 = 0, - -)
(GD)
Definition 1 (Orthonormal function basis). A set of functions (d:()))eep forms an orthonormal function basis for the model space Fe with respect to the (Dm)mmmmi if
for any E I, we have (x= =1;
for any i,j E I such that i + j. have EE (=)=)))(( = 0:
for any fe Fe, there exists a unique sequence of basis coefficients (8)(0))vE such that
fo(æ)= = =iET (i (,)))x((x)
Example (Orthonormal basis for polynomials). Suppose that Fe is the class of all univariate real polynomials of degree at most n, that is, Fe = {Ei+) 0,r²=1 0 E BR+t]). If is uniform distribution on [-1,1], then the so-called Legendre polynomials form an orthonormal basis for Fg with respect to the distribution, then Hermite polynomials define an orthonormal basis for Fo with respect to the
(Olver et al.) [210 Chapter 14). Moreover, if D is normal
1,2 (D) (D)-mtrric (Olver el al. 2010 Chapter 18)
Example (Orthonormal basis for symmetric matrix factorization). Suppose that the true model is defined as fu-(X) (U+U+T, X) with some rank-r matrix U* E Rdxr, and consider an "vverparameterizea" function class Fe = (fv((X U E Rdxx' where r' 2 r is an overestimation of the rank. Moreover, suppose that the elements of ~ are iid with zero mean and unit variance. Consider the eigenvalues of U* as a1 2 ... 2 aa with ar+1
ad = O. and their
'Given a probability distrihution D, we define the L2 (D)) as [If EX-D [x))
2 'otth Legendre and Hermite polynomials can be derived sequentially using Gram-Schmidt procedure. For instance, the first three Legendre polynomials are defined (x) 1//2, Pa(x) v3/2x, and Pa(x) /5/8(322- 11)
3

---

Published as a conference paper at ICLR 2023
corresponding eigenvectors 21....,2u' Zd- It is easy to verify that the functions Øis (X) (z12] X)
for 1 I i.j S d define a valid orthogonal basis for Fe with respect the L (D) -mttrc.... Moreover
for any fü (X), the basis coefficients can be obtained as 3.y = E [ X) { x]]
UU :U ) 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 fo- (a:) = -iez Bi((0 )0,(()) where (((z)) )ee is an or- thonormal basis for Fe with respect to (D and {3:(0*) Jiez are the true basis coefficients. For short, denote _ (O*)). In light of this, the expected loss can be written
(0) = 2
Σ REI
- + 2 l + +o2/2.
(1)
optimization error
approximation entor
noise
Accordingly, GD takes the form
0z+1 = 0: -7 Σ
(8((0.)- (.,)
(GD dynamic)
KEI
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 (O.) to 8r. In fact, we will show that the coefficients (0,) enjoy simpler dynamics for particular choices of orthonormal basis that satisfy appropriate conditions.
Assumption (Boundedness and smoothness). There exist constants L1, Lyg Lu > 0 such that
|| [ool S L,,, [Volll S Lgg I ( Ln
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 (0022). Moreover, although these assumptions may not hold globally, all of our subsequent results hold when Assump- tion is satisfied within any bounded region for 0 that includes the solution trajectory. We will also relax these assumptions for several learning problems.
Proposition 1 (Dynamic of (o,)))) Under Assumption 7 aand based on GD dynamic
we have
3.(0++1) = (0,) = Σ (8,(0,) - (V., + (1) L 12 JEI
(3)
The above proposition holds for any valid choice of orthonormal basis (d)( 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 (0,). 3 (0.)) ss 0 for / j. i.e., the gradients of the coefficients remain orthogonal along the solution trajectory. Under such assumption, the dynamics of (e,) almost decompose over different indices:
(4) 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 81 2 ... 2 B > O and 3 = 0, i for some k < oc. We refer the index set S = (1,..., k) as signal and the index set غ = I as
(6++1) R (0.) 17 (3.(0.) (3) (V) = 0(r)),
,.,
residual. When there is no ambiguity, we also refer to 3,(0,) as a signal if i E S.
Theorem (Convergence of GD with finite ideal basis). Suppose that the initial point Op satisfies
(:0))) 2 Cia,
for all E S.
(lower bound on signals at fp)
1/2
f1 11 ((pD)) = Σ (Ba)
S Cza.
(upper bound on energy at 0p)
Vie
4

---

Published as a conference paper at ICLR 2023
0.10
0.8
00
13:00)
0/00
-10
0114
-20
-15
000
index
(
(a) Small initialization
Gradient independence
(c) Gradient dominance
The conditions of Theorem I are approximately satisfied for LARS 2-layer CNN with MNIST dataset. Here, the coefficients obtaindd by projecting the solution trajectory onto the eigenvectors of the conjugate kernel after training. (a) The top-20 basis coefficients a small random initial point along with residual energy on the remaining coefficients. (b) The maximum value of cos( 3, 3 (0)) for < < j & 10 along the solution trajectory. (c) The scaling of | (V.)) with respect to 18, (0.) for the top-4 coefficients.
Figure 2:
for C1, F1,C2 0 and S. By- Moreover, suppose that the orthogonal function basis is finite, i.e. [ d for some finite d, and the gradients of the coefficients satisfy the following conditions for every 0 S s T:
(V.: (0))) =0
for all i
(gradient independence)
[V,) II 2 C(3((0.))*
for all i E S,
(gradient dominance)
for C >> 0 and 1/2 < 1. Then, GD with step-size 1) Vac= 2222 log -1
43L Cia
satisfies:
If y then within T=O
log The
2 Cja
iterations, we have \l -o (D()) Sa.
< then within T=O (?"mora"--)
iterations, we have [[ - fo ( Sa.
Theorem [soww that, under certain conditions the basis coefficients and their gradients, GD with constant step-size converges to model that is at most o-away from the true model. In particular, to achieve e-accurate solution for any >> 0, GD requires (((1/e) (og(]/))) iterations for 7 = 1/2, and (1))21) iterations for 1/2 < < 1 (ignoring the dependency on other problem-specifee parameters). Due generality, our theorem inevitably relies on small step-size and leads to conservative convergence rate for GD. Later, we will show how our proposed approach can be tailored to specific learning problems achieve better convergence rates in each setting.
How realistic are the assumptions of Theorem TP A natural question arises as to whether the conditions for Theorem [ are realistic. We start with the conditions on the initial point. Intuitively, these assumptions entail that 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 hold with fixed, random, or spectral initialization 3 We have also empirically verified that these conditions are satisfied for DNNs with random default initialization. For instance, Figure 2a illustrates the top-20 basis coefficients at 0g for LARS with random initialization on a realistic CNN. It can be seen that a non-negligible fraction of the energy 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( 3 (0,)), ), V B (0.)) remains small throughout the solution trajectory. Finally, we turn to the gradient dominance condition. Intuitively, this condition entails
IIf Po is selected from isotropic Gaussian distribution, then C1 and C2 may scale with & and d. However, streamline the presentation, keep this dependency implicit.
5

---

Published as a conference paper at ICLR 2023
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 Showws 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 other stochastic and adaptive variants of GD (sce Section and Appendix [): rigorous verification of this conjecture is left as future work.
APPLICATIONS
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 complex DNNs. We start with the classical kernel regression, for which GD is known to converge linearly Karimi jtt al. (2016). 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 Ieads 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.1 KERNEL REGRESSION
In kernel regression (KR), the goal is fit regression model fo(x) El= (1))(x) from the function class Fe (Jo(x) 0 Rdd to observation y, where {0)())) some known kernel functions. Examples of KR are linear regression, polynomial regression (including those described in Example []. and neural tangent kernel (NTK) (Jacot al. 2018). Without loss of generality, we may assume that the kernel functions {ør ) K are orthonormal Under this assumption, the basis coefficients can be defined as (0) 0g and the expected loss can be written as
1 F. (f) (oo)--fo-(a (x))) [(0,)= =
1 100 1-
Σ (B.(O.) 6,)
(5)
Moreover, the coefficients satisfy the gradient independence condition. Therefore, an adaptation of Proposition [ reveals that the dynamics of the basis coefficients are independent of each other.
Proposition 2 (dynamics of (,.)))). Consider GD with step-size that satisfies n 1. Then,
for E S, we have (,,))) 8 (1 (i- - - (,())))
for & S, we have (0,)) = (1 - -,)
Without loss of generality, we assume that 0 < 0% < ... < 0 S and 100l < a. Then, given Proposition 2 we have (8.)) < 2 for every 1 < i k. Therefore, the the gradient dominance is
satisfied with parameters (C, 7) (1/v2 + a, 1/2). Since both gradient independence and and gradient dominance are satisfied, the convergence of GD can be established with an appropriate initial point.
Theorem 2. Suppose that 0 = a1, where a S k /dd. Then, within T S (1/n) log (A/011/o) iterations, GD with step-size 0 < r < satisfies 0r 0 S a.
Theorem 2 reveals that GD with large step-size and small initial point converges linearly to an e-accurate solution, provided that the initialization scale is chosen as = e. This is indeed better
than our result on the convergence of GD for general models Theorem I
*Suppose that (#)) are not orthonormal. Let (x) be any orthonormal basis for Fe. Then, there exists matrix A such that Φi(x) = E, A for every 1 < i < d. Therefore, upon defining 0 A, one write fg. (æz) Exer 0: (.))) which has the same form as the regression model.

---

Published as a conference paper at ICLR 2023
3.2 SYMMETRIC MATRIX FACTORIZATION
In symmetric matrix factorization (SMF), the goal is to learn a model fix (X)
UUT x) with
low-rank matrix U E Rdxx, where we assume that each element of X D is iid with E E and E[X3]] = 1. Examples of SMF are matrix sensing et al. (0118) and completion Ge et al. 2016))
Given the eigenvectors {.),..., _z, of UU T and a function class Fe = (uu(X) U E RRd
with r' 2 r, it was shown in Example 2 that the functions $15 (X) = (2,2, x) define a valid
orthogonal basis for Fe with coefficients (U) = (ziz,.UK _UT). Therefore, we have
((U.) -(U.)= ((x 1/ (// (UU.U,T-UU.... x)) xx)
(U)(
Here, the true basis coefficients are defined as 9 = oi for i < r, and 1 0 otherwise. Moreover, one can write [U) = - Ull 2/ 3is (U). Therefore, gradient dominance holds with parameters (C.7) = 1/2.. However, gradient independence does not hold for this choice Igradiett
of function basis: given any pair (i.j) and k) with j * k. we have (V,,,(U), VB,((U)) (zj=ī UU") which may not be zero. Despite the absence of [graddient independence our next proposition characterizes the dynamic of (U,) via a finer control over the coefficient gradients.
Let Uo aB, where a S min { (noo))lo./ (a, (or/d)((e//)) (Ir log" ld)) (d))-R(/((n-))) and the entries of are independently drawn from standard normal distribution. Suppose that the step-size for GD satisfies v S 1/01. Then, with probability of at least 1 - -xppp (s)r))
Proposition 3. Suppose that 7 == minisier o+11
For 1 Si we have 0.99I; S (u,(U)) S o, within O ((1/(10:)) log (ai/a)) iterations.
For t 11220 0 and i + j or i,j > r, we have |8t (U,) poly(o).
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 al.] [201aa] Li et al. 2020). as well as GD with identical initialization for the special case r' = d d'Chou et al. (2020) 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(-S((r)) and within T % (1/(no,) log (a./ai) iterations, GD satisfies
U - 1 S log((ddda?
(6)
It has been shown in (Stöger & Soltanolkotabi 2021 Thereom 3.3) that GD with small random ini-
tialization satisfies |UJUT - 11 & (d² /,
21/16 within the same number of iterations.
Theorem improves the dependency of the final error on the initialization scale a.
3.3 ORTHOGGNNAL SYMMETRIC TENSOR DECOMPOSITION
We our approach provide a new convergence guarantee for GD on the orthogonal symmetric tensor decomposition (OSTD). In OSTD, the goal is to learn fu- (X) = (T X). where U* [ui,... ux]] E dd and Tu. Σ u St Σd= i= o,2, @l is symmetric tensor with order ( a = 0, and 2i...., Zd
and rank r. Here, 01 2 ... 2 aa are tensor eigenvalues with 0,+1 = the corresponding tensor eigenvectors. The notation 3 refers to the I-time outer product of We assume that X ~ 'D is an 1-order tensor whose elements are iid with zero mean and unit variance. Examples of OSTD are tensor regression (Tong al. 2022) and completion (Liu et al. 2012).
When the rank of Tu is unknown, it must be overestimated. Even when the rank is known, its overes- timation can improve the convergence of gradient-based algorithms (Wang al. 2020). This leads to an overparameterized model fo (x) = (Tv, X), where To = Σt= u. @l with an overestimated rank
r' 2 r. Accordingly, the function class is defined as Fe = [x) : UU= U[(ur,"" ur]] E RRdxr''
7

---

Published as a conference paper at ICLR 2023
Upon defining multi-index A = (f1. j.. the functions xx ( 1 X) for 1 fi...., Jt I d form an orthonormal basis for Fe with basis coefficients defined as
(A = E (Tu. X)
x)
21 (u
II k=
(us 2j²}
and the expected loss can be written as
((U) =B Su. (0000)
1 ml
1 Σ 2
-
2
A
where the true basis coefficients are 37, oj for A; (i..... < < and 3 otherwise. Unlike KR and SMF, neither gradient independence nor show that these conditions approximately satisfied throughout the solution trajectory, provided that the initial point is nearly aligned with the eigenvectors 21,-... 2r'; in other words, cos(u;(0), 2,) for every s sish Assuming that the initial point satisfies this alignment condition, we show that the entire solution trajectory remains aligned with these eigenvectors, i.e., (t) z:) for every 1 and 1 < & < T. Using this key result, we show that both gradieet independence and gradient dominance are approximately satisfied throughout the solution trajectory. We briefly explain the intuition behind our approach for forlgradient dominance]and defer our rigorous analysis for gradient independence to the appendix. Note that if cos( u. (O), z;) 1, then (U,) (t, 2,) and
for OSTD with random equal initialization. However, nor gradient dominance are satisfied
r'
BV (U+) R 3A, (U,) s I (t,(t), z)" Therefore, |gradient dominance holds with parameters (C,7) (1, 1)/1). We will make this intuition rigorous Appendix
Proposition 4. Suppose that the initial point Uo chosen such that [|,](0)]| a/ and cos(u,(0), (0). 2:) 2 71,,, for all 1 S i S r', where a S d= and y S (lk) -1/0-2) Then, GD with step-siza nj I 1/((o)) satisfies:
For 1sisr, we have 0.990, S B,,, (U.) within ((I()lo,))) 4) iterations.
Fort 2 and A * A,, we have |,xA (U.) poly(a).
Proposition 4 shows that, similar SMF, GD learns the tensor eigenvalues incrementally. However, unlike SMF, we require 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 number of iterations that is almost independent of d (Anandkumar et al. 2017 Theorem 1). We believe that Proposition can be extended to random initialization; we leave the rigorous verification this conjecture to future work. Equipped with this proposition,
next establish the convergence of GD OSTD.
Theorem 4. Suppose that the conditions of Proposition f are satisfied. Then, within T x (1/(nla,))e (1-2)/1 iterations, GD satisfies
[ Tu- S rd rd - at.
Theorem shows that, with appropriate choices of and a, GD converges solution that satisfies |Tu, T 2 within ((d(((2)) /er 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.4 EMPIRICAL VERIFICATION ON NEURAL NETWORKS
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 losso which is shown to have 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 method for analyzing the generalization performance of DNNs that uses
We use the notations u, ut) interchangeably denote the solution iteration
8

---

Published as a conference paper at ICLR 2023
00:3
10:0
31
91
10.00
9
3/2
1
(a) 2-block CNN
(b) 3-block CNN
(c) 4-hlock CNN
0.33
1020
0.31
0:35
36
86
000
0.8
(d) LARS
(c) SGD
(f) Adam W
accuracies for 2- 2-Block. 3- B-Block, and 4 block CNN are 96. 96,10%, 96 48%, 96. 58%. (Second row) the projected trajectories of different optimizers 90.54%. 91. .3%.. and 90.26%. respectively. We use the following settings for each optimizer: (d) LARS: learning of Nesterov momentum of 0.9, and weight decay of 10- (e) SGD: learning rate of 2 with "linear warm-up", Nesterov momentum of 0.9, weight decay of x 10 (f) AdamW: learning rate 0.01.
Figure 3: (First row) the projected trajectory of LARS on CNNs MNIST dataset.
AlexNet with the CIFAR- 10 dataset. The test accuracies for LARS, SGD, and Adam W
the second to last layer (the layer before the last linear layer) at the initial point as the feature map (Daniely 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,
find that A-CK can be used not only 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 fe(x) = W (xx),, where x IRd is the input drawn from the distribution D, (xx) E R" is the feature map with number of features m. and W E p is the last linear layer with & referring to the number classes. We denote the trained model, i.e., the model the last epoch,
by fo (æ) = Woc (oo) To form an orthogonal basis, we use SVD to obtain a series of basis functions (xxx) 5.j where dij is the delta function. Hence, the coefficient (0,) at each epoch can be derived as (.(0.) ET [Ko, (1=), (z))), where the expectation estimated by its sample mean the test set. More details on our implementation can be found in Appendix B
W We, (z:) that satisfy Eæ<D [loo 11
and Ez~D (x) Φ (x)))
Performance convolutional neural networks: We use LARS train CNNs with varying depths on MNIST dataset. These networks are trained such that their test accuracies are above 96%. Figures Ba 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 also observed across different optimizers. Figures 3f show the solution trajectories 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 theoretical results only hold for GD, simulations highlight the strength of the proposed basis function decomposition in capturing the behavior of other gradient- based algorithms DNN.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
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 supported. in part, by NSF Award DMS- 2MM-2152276, ONR Award 000014-22-1-2127, MICDE Catalyst Grant, MIDAS PODS grant and Startup Funding from the University Michigan.
REFERENCES
Zeyuan Allen-Zhiee ZZu, Yuanzhi Li, and Yingyu Liang, Learning and generalization in overparameterized neural networks, going beyond two layers. Advances neural information processing systems, 32, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp- 242-252. PMLR, 2019b.
Animashree Anandkumar, Rong Ge, and Majid Janzamin. Analyzing tensor power method dynamics in overcomplete regime. Journal of Machine Learning Research, 18(22):1-40) 2017.
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. arXiv preprint arXiv:1810.02281, 2018. 2018.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine- suue-graied analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332. PMLR, 2019b.
Frank Bauer, Sergei Pereverzev, and Lorenzo Rosasco. On regularization algorithms in learning theory. Journal of complexity, 23(1):52-72 2007.
Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 8(3-4):231-357, 2015.
Hung- Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep trix factorization: Dynamics and implicit bias towards low rank. arXiv preprint arXiv:2011.13772 2020.
Hadi Daneshmand, Jonas Kohler, Aurelien Lucchi, and Thomas Hofmann. Escaping saddles with stochastic gradients. International Conference on Machine Learning, pp. 1155-1164. PMLR, 2018.
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding neural networks: The power of initialization and dual view on expressivity. Advances in neural information processing systems, 29, 2016.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint 1rXiv:2010.11929, 2020.
Zhou Fan and Zhichao Wang. Spectra of the conjugate kernel and neural tangent kernel for linear- width neural networks. Advances neural information processing systems, 33:7710-7721, 2020.
Cong Fang, Zhouchen Lin, and Tong Zhang. Sharp analysis for nonconvex sgd escaping from saddle points. In Conference on Learning Theory, pp. 192-1234, PMLR, 2019.
Salar Fattahi and Somaych Sojoudi. Exact guarantees on the absence of spurious local minima for non-negative rank-1 robust principal component analysis. Journal of machine learning research, 2020.
10

---

Published as a conference paper at ICLR 2023
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems, 33:5850-5861, 2020.
Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has no spurious local minimum. Advances in neural information processing systems, 29, 2016.
Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. In International Conference on Machine Learning, pp. 1233-1242. PMLR, 2017.
Rong Ge, Yunwei Ren, Xiang Wang, and Mo Zhou. Understanding deflation process in over- parametrized tensor decomposition. Advances in Neural Information Processing Systems, 34, 2021.
Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv:I909, 12051, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778 2016.
Zhengmian Hu and Heng Huang, On the random conjugate kernel and neural tangent kernel. In International Conference on Machine Learning, pp. 4359-4368. PMLR, 2021.
Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss vs cross- entropy in classification tasks. arXiv preprint DrXiv:2000600322. 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How to escape saddle points efficiently. In International Conference on Machine Learning, pp. 1724-1733 PMLR, 2017.
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewice condition. In Joint European conference machine learning and knowledge discovery in databases, pp. 115-811.55 Springer, 2016.
Kenji Kawaguchi. Deep learning without poor local minima. Advances in neural information processing systems, 29, 2016.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXhv:1412.6980, 2014.
Alex Krizhevsky, llya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Communications of the ACM, 6016):84-90, 2017.
Jason D Lee, Ioannis Panageas, Georgios Piliouras, Max Simchowitz, Michael I Jordan, and Benjamin Recht. First-order methods almost always avoid strict saddle points. Mathematical programming. 176(1):311-337) 2019.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang, Algorithmic regularization in over- parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning, arXiv preprint arXiv:2012.09839 2020.
Ji Liu, Przemyslaw Musialski, Peter Wonka, and Jieping Ye. Tensor completion for estimating missing values in visual data. IEEE transactions on pattern analysis and machine intelligence, 35 (1):208-2200 2012.
Philip M Long. Properties of the after kernel. arXiv preprint arXiv:2105.10588. 2021.
11

---

Published as a conference paper at ICLR 2023
Jianhao Ma and Salar Fattahi. Blessing of nonconvexity in deep linear models: Depth flattens the optimization landscape around the true solution. arXiv preprint arXiv:2207.07612, 2022a.
Jianhao Ma and Salar Fattahi. Global convergence of sub-gradient method for robust matrix re- covery: Small initialization, noisy measurements, and over-parameterization. arXiv preprint arXiv:2202.08888, 2022b.
Yu Nesterov. Introductory lectures on convex programming, 1998.
Yurii Evgen' evich Nesterov. A method of solving a convex programming problem with convergence (bbig)). In Doklady Akademii Nauk, volume 269, pp. 543-547. Russian Academy of Sciences, 1983.
rale o olbigl((k2)) bigr).
Frank WJ Olver, Daniel W Lozier, Ronald F Boisvert, and Charles W Clark. NIST handbook of mathematical functions functioos hardback and CD-ROM. Cambridge university press, 2010.
Ioannis Panageas, Georgios Piliouras, and Xiao Wang. First-order methods almost always avoid saddle points: The case of vanishing step-sizes. Advances in Neural Information Processing Systems, 32, 2019.
Vivak Patel and Albert Berahas. Gradient descent the absence of global lipschitz continuity of the gradients: Convergence, divergence and limitations of its continuous approximation. arXiv preprint arXiv:2210.02418 2022.
Garvesh Raskutti, Martin Wainwright, and Bin Yu. Early stopping and non-parametric regression: an optimal data- fata-dependent stopping rule. The Journal of Machine Learning Research, 15(1): 335-366, 2014.
Noam Razin, Asaf Maman, and Nadav Cohen. Implicit regularization in tensor factorization. In International Conference on Machine Learning, pp. 8913-8924. PMLR, 2021.
Noam Razin, Asaf Maman, and Nadav Cohen. Implicit regularization in hierarchical tensor factoriza- tion and deep convolutional neural networks. arXiv preprint arXiv:200..1172) 2022.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
Itay Safran and Ohad Shamir. Spurious local minima are common in two-layer relu neural networks. In International conference on machine learning, pp. 4433 4441. PMLR, 2018
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterizeed low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 34:23831-23844 2021.
Ju Sun, Qing Qu, and John Wright. Complete dictionary recovery the sphere i: Overview and the geometric picture. IEEE Transactions on Information Theory, 63(2):853-88) 884, 2016.
Tian Tong, Cong Ma, Ashley Prater- Biiter--BBnette, Erin Tripp, and Yuejie Chi. Scaling and scalability: Provable nonconvex low-rank tensor estimation from incomplete measurements. Journal of Machine Learning Research, 23(163):1-77, 2022.
2022.
Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. Advances in Neural Information Processing Systems, 32, 2019.
Nikhil Vyas, Yamini Bansal, and Preetum Nakkiran. Limitations of the ntk for understanding generalization in deep learning. arXiv preprint arXiv:2206. 100/2, 2022.
Xiang Wang, Chenwei Wu, Jason D Lee, Tengyu Ma, and Rong Ge. Beyond lazy training for over paraaeterzee tensor decomposition. Advances in Neural Information Processing Systems, 33:21934-21944 2020.
Tian Ye and Simon S Du. Global convergence of gradient descent for asymmetric low-rank matrix factorization. Advances in Neural Information Processing Systems, 34:1429-1444 1439, 2021.
12

---

Published as a conference paper at ICLR 2023
Jialun Zhang, Salar Fattahi, and Richard Y Zhang. Preconditioned gradient descent for over- parameterized nonconvex matrix factorization. Advances in Neural Information Processing Systems, 34:5985-5996, 2021.
Richard Y Zhang. Sharp global guarantees for nonconvex low-rank matrix recovery in the overpa- rameterized regime. arXiv preprini arXiv:2104. 10790, 2021.
Jiacheng Zhuo, Jeongyeol Kwon, Nhat Ho, and Constantine Caramanis. On the computational and statistical complexity of over-parameterized matrix sensing. arXiv preprint arXiv:2102.02756. 2021.
13

---

Published as a conference paper at ICLR 2023
CONTENTS
Discussion and Future Direction
15
Additional Experiments
15
B.
Numerical Verification of our Theoretical Results
15
2
Derivation Basis Functions for DNNs
16
B.3
Further Details on the Experiments
17
B.4
Experiments for CIFAR-I00
17
B.5
Experiments for Different Losses
18
B.6
Experiments for Different Batch Size
18
B.
Experiments for ResNet-18 with SGD on CIFAR-10 10
19
Related Work
D
Proofs for General Framework
21
D.I Proof of Proposition
D.2 Proof of Theorem
22
Proof' for Kernel Regression
25
Proof of Proposition
25
Proof of Theorem
25
Proofs for Symmetric Matrix Factorization
25
Initialization
25
F:2
One-step Dynamics
26
Proofs of Proposition and
Theorem
27
Analysis of the Logistic Map
32
Proof for Tensor Decompasition
34
Proof of Lemmal3
36
G.2 Proof of Proposition
37
Proof of Proposition
40
Auxiliary Lemmas
42
14

---

Published as a conference paper at ICLR 2023
A DISCUSSION AND FUTURE DIRECTION
Extension to empirical loss and implication on generalization. The most significant future direction extend our analysis of expected loss to the finite sample regime. It has been shown that incremental learning can drive generalization (Gissin et al.] (011)) 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 offgradient 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 establishg gradient independence approximately and with high probability. We have successfully applied our framework matrix factorization and tensor decomposition,
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 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 [95I], and Adam (Kingma & Ba 2014). Overall, our approach offer unified framework for examining the implicit bias and incremental learning phenomena for different local-search algorithms.
B ADDITIONAL EXPERIMENTS
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 [ for a summary our simulations this section.
Architectures
CNN
AlexNet
VGG11
ResNet-18 ResNet-34
ResNet-50
VIT
Datasets
MNIST
CIFAR-10
CIFAR-100
Optimizers
SGD
AdamW
LARS
Losses
Eg-loss
CE loss
Table I: The summary of our experiments.
B. NUMERICAL VERIFICATION OF OUR THEORETICAL RESULTS
In this section, we provide experimental evidence to support our theoretical results on kernel regres- sion (KR), symmetric matrix factorization (SMF), and orthogonal symmetric tensor decomposition (OSTD). The results are presented Figure
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 4alsshos the trajectories the top-4 basis coefficients of GD with initial point O, = 5 x 10- 1 and step-size 7 0.4. It can be
seen that GD learns different coefficients at the same rate, which is line with Proposition 2
Symmetric matrix factorization. In this simulation, aim recover a rank-4 matrix M VEVT E 20 In particular, we assume that V E 2020x4 is randomly generated orthonormal matrix and Σ- Diag 5. 1}. We consider fully over-parameterized model where U e 20x20
(i.e., r' 20). Figure 4b illustrates the incremental learning phenomenon that was proved in
Proposition for GD with small Gaussian initialization Uis step-size 7) 0.04.
i.i.d. NO (20,a²), = 5 x 10-7 and
Orthogonal symmetric tensor decomposition. Finally, we present our simulations for OSTD. We aim to recover a rank-4 symmetric tensor the form T Σ'= o1 where o; are the nonzero eigenvalues with values {10,5, 3. 1} and z; E I. are the corresponding eigenvectors. We again consider fully pee--parameterieed model with r' 10. Figure 4c shows the incremental learning phenomenon for GD with an aligned initial point that satisfies cos( (0), 2,) 2 0.9983, i and step-size y 0.001.
15

---

Published as a conference paper at ICLR 2023
8d0
13
8u10
8,00
8u(0)
9,
6-00
Polt
303
420
10005
30000
Iteration
Iteration
Iteration
(a) Solution trajectory for KR
(b) Solution trajectory for SMF
(c) Solution trajectory for OSTD
121
122
20"
20-4
1010
10
20-2
10
10
20
20-
0
35
120
1.5
200
250
300
120
200
200
402
300
5000
12020
.3000
20000
25030 20186
Iteration
Iteration
(d) Estimation error for KR
(e) Estimation error for SMF
(f) Estimation error for OSTD
Figure 4: Experimental verification to support Theorems and in Section The first row shows the projected trajectories of GD onto the specific basis functions we defined for each problem. The second row shows the estimation error.
B.2 DERIVATION OF BASIS FUNCTIONS FOR DNNS
In this section, we provide more details on how we evaluate our proposed orthogonal basis induced by A-CK and calculate the corresponding coefficients B; (0,) for arbitrary neural network. First, recall that any neural network whose last layer is linear can be characterized as fo(x) = W*) where Id is the input drawn from the distribution D, vx(x) R" is the feature map with number of features and W E R 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 fe.. (x) = W (xx). To form an orthogonal basis,
W Wo. (x) that satisfy EgzDl pb (x) ][²] = and Eu>D [( (x)), de (x))) diy where dig is the delta function. Hence, the coefficient 3, (O.) at each epoch / can be derived as (O,)) ER [[o, (ær), di (x))). In all of our implementation, we use the test dataset to approximate the population distribution.
SVD to obtain a series of basis functions (x
Step 1: Obtaining the orthogonal basis (x)). We denote v [=(x)),... (EN)) R as the feature matrix where is the number the test data points. We write the SVD of v as I UEVT. The right singular vectors collected in V can be used define the desired orthogonal basis { To this goal, we write the prediction matrix as F W W where W WUΣ and V VT. Our goal is to define a set of matrices A; such that (x) = ( form valid orthonormal basis for F W W Before designing such A,, first note that, due the orthogonality (x)), we have
E [x)]
|| [((d((z), (x) (A.,
(8)
Therefore, it suffices ensure that (A,) are orthonormal. Consider the SVD of W W E; u.u.U," We define A, = u.u;" Clearly, defined {A,} are orthonormal. Moreover, it is easy to that the basis coefficients (treated as the basis coefficients) are exactly the singular values of W.
Step 2: Obtaining the basis coefficients (,.). After obtaining the desired orthonormal basis (x),, we can calculate the coefficient (0,) for each epoch. Given the linear layer W. and the feature matrix & epoch can obtain the coefficients for the signal terms by projecting the prediction matrix F. W,Y, onto v VT. In particular, write the prediction matrix as
16

---

Published as a conference paper at ICLR 2023
P.F.PP -W. W where W. W V. Hence, the basis coefficients can be easily calculated as (3:(O,) = (W,, u.
B.3 FURTHER DETAILS ON THE EXPERIMENTS
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-I6GBB GPUs. We use an additional NNGeometry package for calculating batch gradient, and our implemention of ViT is adapted from http::///jjiiuuuueeeeek ensure consistency with our theoretical results, drop the last softmax operator and use the f2-loss throughout this section. All of our training data are augmented by RandomCrop and RandomHorizontal 1i, and normalized by mean and standard deviation.
github. io mwl/cifaall
html To
Experimental details for Figure 1 Here, we describe our implementation details for Figure I and present additional experiments on VGG-11, ResNet-34, and ResNet-50 with the CIFAR-10 10 dataset. We use standard data augmentation for all architectures except
The results can be seen in Figure 5 for VIT. For VIT, we only use data normalization.
To obtain stable A-CK, we trained the above models for 300 epochs. For ResNet-18, used LARS with a learning rate of 0.5 and applied small initialization with a = 0.3, i.e., we scale the default initial point by a 0.3. For ResNet-34 and ResNet-50, we choose the default learning rate and apply small initialization with a = 0.3. For ViT, we use Adam with a learning rate of 0.01. The remaining parameters are their default values.
Experiments details for the first row of Figure Here we conduct experiments on MNIST dataset with different CNN architectures. The CNNs are composed of blocks of layers, followed by single fully connected layer. A block of a CNN consists of a convolutional layer, 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, used max-pooling.
For 2-block CNN, we set the convolutional layer width to 256 and 64, respectively. For 3-block CNN, 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.05. The remaining parameters set to their default values. We run 20 epochs to calculate A-CK.
We conduct experiments to compare the perfor- mance of different optimizers 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 @ For AdamW, we set the learning rate to 0.01 and keep the remaining parame- ters unchanged. For LARS, we the learning rate with Nesterov momentum of 0.9 and weight decay of 0.0001. The remaining parameters are set to the default setting.
Experiments details for second row of Figure 3
warm-up" technique
B.4 EXPERIMENTS FOR -IFAR-100
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 the "loss scaling trick" (Hui & Belkin 2020) defined follows: consider the datapoint (x, y) where æ nd is the input and E R2 is a one-hot vector with at
'In "linear warm-up" we linearly increase the learning rate in the first 5 epochs. More precisely, we set the initial learning rate to x 10-5 and linearly increase it to the selected learning rate in 5 epochs. After the first 5 epochs, the learning rate follows regular decay scheme.
17

---

Published as a conference paper at ICLR 2023
10.30
10.0
3,
5.
0.35
Reration
(a) AlexNet
(b) VGG-11
(c) VIT
0.:0
10:0
11:3
9.
043
008
0.5
Reration
(d) ResNet-18
(c) ResNet-34
(l) ResNet-50
Figure 5: Solution trajectories different architectures trained on CIFAR-1O. test accuracies Alexnet, VGG- 11 and VIT are 90.54%, 91.11%, 77 7.99% respectively. The test accuracies for ResNet-18, ResNet- and Re are 88 38,,1%, 93.93% and 94.06%, respectively.
position Then, the scaled Ez-loss is defined
la scalingg (x) = ks (fo(z)])] M)² + (Vo0)))))
for some constants k, Al > O. We set these parameters to 1, = For AlexNet and VGG-I we use LARS with base learning rate of 1. For VIT, of 0.01 and batch size 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 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 meaningful values.
AdamW with base learning rate
B.5 EXPERIMENTS FOR DIFFERENT LOSSES
Next, compare the projected solution trajectories two loss functions, namely Enlloss and cross-entropy (CE) loss. We use LARS to train AlexNet on the CIFAR- dataset with both loss and CE loss. In partieular, set the base learning rate LARS to 1. For fosso we use the base learning rate of The remaining parameters are to their default values. The results to the folsoss, 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 much slower than those the fosss.. This is despite the fact that CE loss can learn slightly faster than foss in terms of test accuracy shown by Hui & Belkin (2020).
add the softmax operator before training the CE loss. For CE loss,
be seen in Figure 7 We observe that, similar
B.6 EXPERIMENTS FOR DIFFERENT BATCH SIZE
Next, we study the effect of different batch sizes on the solution trajectory. We train AlexNet on the CIFAR-10 10 dataset. When testing for different batch sizes, we follow the "linear scaling" rule, i.e.,
18

---

Published as a conference paper at ICLR 2023
86-
00
15
60
(a) AlexNet
(b) VGG-11
(c) VIT
057
0G-
10
84
8,
61
05
20
00
110
121
141
01
20
100
141
20
31
30
10:0
120
142
(d) ResNet-18
(c) ResNet-34
(f) ResNet-50
Figure 6: Solution trajectories for different architectures trained on CIFAR-1O0 100. The test accuracies for AlexNet, VG,-11, and VIT are 61. 31..0% 50. 64%, and 50.41 50.41%, respectively. The test accuracies for ResNet- ResNet-34 and ResNet-50 78%1 74.26% and 80 30.05%, respectively.
25
0.30
B1
8
0.20
03
0.10
84
14
85
epoch
iteration
(a) Egss
(b) CE loss
Figure 7: Solution trajectories of LARS on the CIFAR-10 dataset with loss and CE loss. The test accuracies 90.54% for fo-soss and 91.09% for CE loss.
the learning rate scales linearly with the batch size. For batch size of 32, we used SGD with a base learning rate 1) 0.1 and 5 warm-up epochs. For batch size of 64, used SGD with base learning rate of = 0.2 and 5 warm-up epochs. For batch size of 256, we used LARS with base learning rate of 1. The remaining hyperparameters are their default values. The results reported in Figure g We see that the projected solution trajectories share a similar monotonic behavior for different batch sizes.
B.7 EXPERIMENTS FOR RESNET- 18 WITH SGD ON CIFAR-10
Lastly, we train ResNet-18 with SGD on CIFAR-10. The results can be seen in Figure 9 To achieve good generalization, use a large learning rate 1 0.3 (with warm-up epochs starting at x 10 -), and decrease the learning rate by 0.33 for every 50 epochs. Moreover, We use large batch size 512 better imitate the trajectory of GD. Based the simulation result, an approximately monotonic (modulo the fluctuations caused by the randomness the gradients) behavior of the dynamics of the top-5 components, which again validifies theoretical results.
find
19

---

Published as a conference paper at ICLR 2023
1030
11.0
0.10
0g
0115
6:15
85
000
600
Seeratio
opoch
Itoration
(a) Batch size 32
(b) Batch size 64
(c) Batch size 256
Figure 8: Solution trajectories for AlexNet the CIFAR-10 10 dataset with different batch sizes. The test accuracies for size 32, 256 are 91.50% 91. 99.79% 90.12%, respectively.
0.30
0.25
B
0.20
scale
82
0.15
B3
0.10
B
0.05
B5
0.00
40
100
120
140
iteration
Figure 9: Solution trajectories ResNet-18 on the CIFAR- 10 dataset with SGD solver. The best test accuracy 94.45%.
20

---

Published as a conference paper at ICLR 2023
C RELATED WORK
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. Be- cause of their first-order nature, they 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 (Tin al. 2017) and SGD (Fang et al. 2019; Daneshmand 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., (016)) and robust PCA (Fattahi & Sojoudi (020), as well as deep linear neural networks (Kawaguchi 2016). However, benign landscape is too restrictive to hold in practice; for instance 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 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 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.] [02 ), and linear neural network (Arora et al. 2018). In what follows, review specific learning tasks that most related to work.
The convergence of GD on kernel regression was studied far before the emergence of deep learning. Bauer et al. (2007); Raskutti et al. (2014) establish the convergence of gradient descent a special class of nonparametric kernel regression called reproducing kernel Hilbert space (RKHS). Recently,
Jacot et al. (218)) 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., 20195) properties of NTK. As for the matrix factorization, [i et al. (2018); Stöger & Soltanolkotabi (202T) 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 al. 2021> , asymmetric (Ye Du 202T). and robust (Ma Fattahi 20226) variants. Wang et al. (2020): Ge et al. (2021) studied the dynamic of 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 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., 20200] Long 202T). In our work, we show that such "after kernel regime" can also be adapted study the optimization trajectory of practical DNNs.
D PROOFS FOR GENERAL FRAMEWORK
D. PROOF OF PROPOSITION I
To prove this proposition, we first combine ! and GD
0. 0 V. 0,
(B,(0.) -8) (8.)
(10)
iez
21

---

Published as a conference paper at ICLR 2023
For notational simplicity, denote E(0;) Eies (0,). Then, one write
01+1 0: - 0, 77
Σ (B. - (VE()))
(11)
i
Due to the Mean-Vaie Value Theorem, there exists & R such that
(,(0++1) =B
0: 1 Σ
(3,(0.) - 83) (.,(
JEI
(3, 53) (.), ()) + T (VC())), Σ
= =3,(0,)
JET
(12)
On the other hand, one can write
(V V(
sup (0) /
(13)
For sup, V (0) we further have
sup sup V B.))|
0
sup
|
fo())
sup
[
fo [x))
(14)
sup E [1v2
1/2
(E (1)6)
0
LH
Here, we used Cauchy- Schwartz inequality for (a). Moreover, for (b), we used Assumption and the definition of the orthonormal basis. Hence, have
8.(0. = = (0,) - Σ
(3, 9) (V.( (3,( + (1/2) [VC(0)))
jEI
(15)
Now, it suffices to bound VV((e)) 2 Using
inequality, we have
VE [
|V((0))
(So, - fo-
||E| [(fe. fo.) fe.
(16)
llSo. Sor l(E=(D) Vo ((DD
- 4L213.
Therefore, we conclude that
B,(0.-1) (3,
(3,(0.) B3) (B)( =2n L..
(17)
which completes the proof.
D.2 PROOF OF THFOREM
Proof. Invoking the gradient independence condition, Prooddiiiooooooootttiioooon can be simplified as
3,(0++1) = (,((.)) Σ (8,(0.) - 83) (O.) v3, + 2n LH L jEz
(18)
= B.(.)) 77 (3, Bi) (3)) 27 LHL
22

---

Published as a conference paper at ICLR 2023
We next provide upper and lower bounds for the residual and signal terms. Recall that S = {i E I 3 * 0). and & = S. We first consider the dynamic of the signal term (,(0,), i E S. Without loss of generality, we assume 3 > O. Then, due the gradient dominance condition, we have the following lower bound
((+++))) 2 (,(O.) 17 ( 21 -n
(19)
2 (1+ C1+ - (.(0.)) B.(O.) 2nj² LH L8g iES
Next, for the dynamic of the residual term (0))), E &, we have
B,(0++1) (1 (1- ( B:(O,) 21j² LH LH Next, we show that [.:)( < Lg. One can write.
(20)
[V)())))1 = VE fee (x))
< III So, (z)
[[
1/2 (E(o
(21)
Lg
Due to our choice of the step-size, we have 1) Therefore, we have
S IZ which in turn implies 1- [V
(8+(0++)| < |3, (0.) + 27 L.
(22)
Now, we are ready to prove the theorem. We divide it into two cases.
Case 1: 7 7 1.- In this case, since we the step-size 1/ VdceLyLEL B log simplify the dynamics of both signal and residual terms in Equation 19 and Equation
d3g
Gja
can
20 as follows
(.(0+)) 2 (1 (1+0 (5, - (.,))) (0,)) 00++1l) s 2n
Vie S,
(23)
ViE
We first analyze the dynamic of signal (0,) f S. To this goal, we further divide this case into two phases. In the first phase, we assume Cjax S (0,) S 13 Under this assumption, we can simplify the dynamic of (,,))) as
(.(++)) 2 + (1+0,2CC2 8.().).
(24)
Therefore, within Ti = O ट log Cia
iterations, (,(O,) becomes larger than 1 In the second phase, we assume that (0)) 2 33 /2 and define y: = S (,,)).. One can write
Vt+1 (1 0, (0.) yt
(25)
< (1 0.250 n18 )v-.
Hence, with additional T3- O
2 log
d0 we have V1 S Va which implies (,(0.) >
3 -a Next, show that there exists time such that 83 Ja S (0. S 0 + aaa. Without loss of generality, we assume that t* is the first time that (0,) 2 8 Jaa. Due to the dynamic of (,)). the distance between two adjacent iterations can be upper bounded as
V
(3.40)+1) (o.)) I 1 |8] 3, (, (0)) + 2nj Lu L2
(26)
< n |3 - (3.())) + 2n)² LH L. In particular, for t t* - we have (,(0:-1)) s , J/a which in turn implies
(..O..) S (0, 1) + 12 18 - (0 - + 2n Ln 1,
< 3 + 2m LH 13
(27)
$ 3 + Vi
23

---

Published as a conference paper at ICLR 2023
Therefore, for each i E S, we have [3,(0.) - - & va within T, = O ب
1
log Cia
di
iterations. Meanwhile, we can show that the residual term 13,0,), Vi & remains small for maxies Ti O クタ
log
dBg
iterations:
Cia
[B,(O.)) - (3,(6u)
max T₁ T:221) 22
= [3,(0u) +
11
log
Cia
Lu 11
(28)
= 13:(0o)) +
Vdd a
Therefore, we have that within = 1/
log
[9g
iterations:
C1a
[for fo- (()=(
(8)
(.(()))²
(0) S
(29)
SES
SEE
K/51. 1. In this case, we have the following bounds for the signal and residual terms Case 2: < < $
3,(0++1) 2 (1 C2n (8) (,)) (0.) (,(0.) - 2n Ln L Vi E S (0(41)l (3)(0.)) + 2n
(30)
Vi E E. We first analyze the dynamic of the signal term (0,))) for i S. We will show that (3,)
87 S
iterations. Due to n VacaLa 2 we can further simplify VE within T 0 (mm8f
a²
the dynamic of (,(0,) as
0,(+1) 2 (1+ (1+0 (3) - (,)))) 8,- (0.) (,())).
(31)
Next, we divide our analysis into two phases. In the first phase, we have (0,) S 1 We denote the number of iterations for this phase as Ti,1- We further divide this period into [log (8//2a) [2)] substages. In each Substage k. we have 2 a S 3, (0,) I 2 Let t. be the number of iterations in Substage We first provide an upper bound for 4. To this goal, note that this substage
(.(0+))) 2
0.50 (B: (0.) (0.)) (..).
(32)
> + 0.255 (3- (0.) (0))
Hence, we have
log 2)
to 5 log (1 2 3 (C 26 20)--11
(33)
Summing over the we obtain an upper bound for T. [log(9//2a) 12a)
T.
(34)
C2np 1a-2-
(Cia)
k=1
Via similar argument, we can show that in the second phase, we have (B.():) 81 S Ja" a within additional Ti O (mm))(C1a))5
iterations. Therefore, for each i E S, we conclude that
(8,) - - < within T, Tú, Ti, = =0 +
iterations. Meanwhile, for the (11))111) (c)
residual term (0,),, E E, we have
[3,(0.) < [Br(00)) max Ti 2nj L3 iES
1
(35)
= [B,(0)) + O
Therefore,
[.fer for 11AA()) Σ (8) B,(OT))²
Σ (7)r)) Sa²
(36)
within 0 (Cmpat(h.app11)) iterations. This completes the proof.
24

---

Published as a conference paper at ICLR 2023
E PROOFS FOR KERNEL REGRESSION
E. PROOF OF PROPOSITION
2
Note that 0.+1 0: (0, 0*), and (,(0) 8,. Hence, for every 1 S i S k, we have
3.(0++1) (O, + n (O,)).
(37)
This in turn implies
B -
=(1 (3) -:((0.) = - = (1- - 11)² (3, - Bi(0o))
(38)
For i > k we have
(.(0+1)) (1 (),)(())) = = 3,(0,) = (1 - 1,)33,((0)).
(39)
which completes the proof.
E.2
PROOF OF THEOREM 2
Due to our choice of initial point 0g = al, a S |82]. we have 187 - 3;(0,) - 2| Hence, by
Proposition 2
we have
|8t - Bc((.)) (1 - 1)) - -:0())) 2 - -)) ies
1
(40)
[3.(0,)) 5
E
(41)
Therefore, to prove - -3,(00)) S .
iE E ESS, suffices to have
2(1 - 1) <
a
t ビ log
kt
(42)
V
On the other hand, ensure 13:(0.)) < Sg E E, suffices to have
1
(1 (1-n)aa v Vd
t log (d)
(43)
7
Recall that S Klaal and 101| 2 102| 2 ... > 10ll > O. Therefore, within T 0 log
kl0:
d
iterations, we have
a
a SV2a,
0 - 0-11 <
(44)
d
i=1
++1
which completes the proof.
F PROOFS FOR SYMMETRIC MATRIX FACTORIZATION
F.1 INITIALIZATION
We start by proving that both lower bound on signals at 0g and jupper bound on energy at Og are satisfied with high probability. Recall that each element of Uo is drawn from N NVO,, a"). The following proposition characterizes the upper and lower bounds for different coefficients Biy (Uo).
Proposition 5 (Initialization). With probability at least 1 - 1e eQ('', we have
1 tr'a, for all ssisT.
Bu) = =<xI, -xx1T,UU,,) 2
(45)
and
[B,)(Un) (212]. (oU < log(d)r' for all # j or i.j
(46)
25

---

Published as a conference paper at ICLR 2023
Proof. First note that Ug 2, a2 Hence, a standard concentration bound Gaussian random vectors implies
P (1110 zill c// ad) 2exp(_c8²).
(47)
Via union bound, have that with probability of at least 1 - -e-or'
U zi s 20//1" r' log(d); isisd, Given these bounds, one can write
U 2; 0.50N sisisr
(48)
(n((U)) (ziz, U,U)) U 2: =| 1 ra2 visisr,
(49)
4
and
18,( - U00% $ 100 2 2,11 < 4log(d)r'a LA Misijsd d,
(50)
which completes the proof.
F.2 ONE-STEP DYNAMICS
In this section, we characterize the one-step dynamics of the basis coefficients. To this goal, we first provide more precise statement of Proposition 3 along with its proof.
Proposition 6. For the diagonal element B& (Ut) we have
(u(U++1) (1+ (1+22(66 (o. Bu((U))) (U.) 2n Σ (,))
jti
(51)
+ 1
Σ
(U)(())) 20, (3)(( +
j.k
where == 0ffor < cis d. for every # we have
(eu(U)+)) ( -2,a(U.) 28,, (U))) Big (U.) + 2n Σ (U,)8a) (U.)
+ n²
Ba -U (o. Bix (0 Σ k,
(52)
Proof. The iterations of GD on SMF take the form
U4i = U, - nU.U - My
(53)
This leads
U4 U.U, - NU.UT M.UUU - U,U (U.U, M²)
(54)
+ (U.UUT - -MJUUU (U.U. - MM).
Recall that = (U,U,,z,z])) U,U, Based on these definitions, one can write
and M El= O.
=
(U++)) = (U++1 U1112111))22)))
-U.UT [(z]) - -U(UUUTT - (T)JUUT,Z22]] -,(,)_"(U.UU] (U,U, -M*),
(UUUUUUU -M.)UUU, - -M"),,2]]1
(55)
26

---

Published as a conference paper at ICLR 2023
In light of the above equality and the orthogonality of (zilve(d) the RHS of Equation 55 written in terms of BRa (U,),, S k, 1 < d. In particular
can be
(
(UU MU 2z,7)
Σ (B)(U)) - Big) xiz]
Σ b
Biy (Ui)zizz 21z]
(Bia(Us) - -Bie) Bkg (U,) 2iz]
1, b
Σ (3x)(U)) -Pic) Bui(U.)
(3,(Ui) 01) Big (U,) + Σ (Ba(U,) - 8%x) 3ky (Ut).
kpi
(56)
Other terms in Equation 53 can be written in terms of (U, in an identical fashion. Substituting these derivations back in Equation 55] we obtain
(U,+1)))) (1
(U,(U) - 2n Σ Bix (v,(),))
+ r2
Σ Rd
Bu (il
(o. + to)) Σ Bux (U.) + t
(57)
Note that the above equality holds for any 1 5 i r'. In particular, for i j. we further have
B: = =11 (1+2m (a, -B((U)))) (U, - 2n Σ (U,))
(58)
+ x2
Σ (B(( 20; Σ +
j.k
which completes the proof.
F.3 PROOFS OF PROPOSITION AND THEOREM 3
To streamline the presentation, we prove Proposition 3 and Theorem simuutannss........sssy. The main idea behind our proof technique is to divide the solution trajectory into r substages: in Substage i, the basis coefficient; 3 (U,) converges linearly a, while all the remaining coefficients remain almost unchanged. More precisely, suppose that Substage lasts from iteration ti to ti,e- We will show that 3a (Ut..) o; and Biy (U...) (U..). Recall that 7 minigigr G, 0++1 is the eigengap of the true model, which we assume is strictly positive.
Substage 1. In the first stage, we show that 311 (U,) approaches on and (3,y(U)) (U,),, 2 2 remains in the order poly(a) within Ti O yor log (a) iterations. To formalize this idea, we further divide this substage into phases. In the first phase (which we refer to the warm-up phase), we show that (U,) will quickly dominate the remaining terms B.;(U.) V(i,j) * (1,1) within
log log(d) iterations. This is shown in the following lemma.
Lemma (Warm-up phase). Suppose that the initial point satisfies Equation 45 and Equation 46
Then, within log log(d)) iterations, we have
my
(U.) (l.(U) r'a² 200177
(59)
Proof. To show this, we use an inductive argument. Due to our choice of the initial point, we have [iy(Uo)) $ r'a² ( s i, < d. Now, suppose that at time I, have (B.)(..) $
27

---

Published as a conference paper at ICLR 2023
r'a loo'logl) +017 ,15 (, V d. Then, by Proposition 6 have
B11(U+11) = (1+ (0 - 3.. Bn)))) - -2n Σ (U,)
j&l
+ n² Σ (1,((U))51((U)33))(1)1 201 E +
(60)
i.
2 1 (101 2no
- log(d log(d) (0))))) (ui(U.).
Similarly, for the remaining coefficients Big (U,) V(i.j) * (1, 1), have
(B,((++1)] S (1+ 01 + 02 0 dr' a² log(d) (o) 18, Note that the stepsize satisfies
(61)
S o. and 01 02 ス こ dr a² log () 101/11 Hence, we have
(1+())
1 + 2m01 - (ndr' a² log d) +01
(U,)
(U++ll
+ 1) (01 + 0 + (dr'
(U)(()))
n 02 O a² log dd (d)
()((U)
1 + 1 ( + 02 O(dr'a log(d)
[.,(
(62)
(U,)
+ 1+ 0.5m(a 01 + +2)
BBg (U,)
(U,(
(1 + 0.5m%)
[B:}(U;)
This further implies
(U,)
(U,)
2 (1 + 0.5m7)
(63)
(U,)
(U)(
On the other hand, Equation 49 and Equation 50 imply that
B1 (U))
(64)
(U,) 16 log()
Hence, within O 1 log (d) iterations, we have B,y 2 (312(U,),, V(i,j) * (1, More- over, have that during this phase,
(U,((U.) < (U,) < a' log(d) (1 + 21111) ( log log(d)) S I''a"log(d)
(65)
which completes the proof.
After the warm-up phase, we show that 311 (U,) quickly approaches oi while the remaining coeffi- cients remain small.
Lemma (Fast growth). After the warm-up phase followed by O y00 log (a)) iterations, we have
0.99a S (U,(U,) 50.
(66)
Moreover, for (B,((U)),, v(i.)) # (1, we have
18g (R,))) & oir (d)
201
(67)
Before providing the proof of Lemma 2 we analyze an intermediate logistic map which, as will be shown later, closely resembles the dynamic of (.j (U.):
Te+1 (1 + na n7. 22, Zo =
(logistic map)
The following two lemmas characterize the dynamic single logistic map, as well as the dynamic of the ratio between two different logistic maps.
28

---

Published as a conference paper at ICLR 2023
Lemma 3 (Iteration complexity of logistic map). Suppose that < 0.1o. Then. for the logistic (log 4o /a) + log(Ao/E)) iterations.
map we have TT 2 a within T log(1+no)
Lemma (Separation between two logistic maps). Let 01 dg be such that 01 02 > and
F++1 = (1 + 1101 - -1/
To = a gt+1 = (1 + 1102 11/))
yo
Then, within T ( log 16m
iterations, we have
160 01-22
21-22
01 - s oT Sob
The proofs of Lemmas 3 and 4 are deferred to Appendix F.4 We are ready to provide the proof of Lemma
Proof of Lemma 2 Similar to the proof of Lemma T use an inductive argument. Suppose that we have (U,,) I (U,, X
to is when the second phase starts. According to Lemma [
r'a log(d))+o// S oir log(d)a
201
Therefore, the base case of our induction holds. Next, 201 oi
suppose that some time within the second phase, we have (U.) S 01
v(i,j) +
201 1oL_s 2-1
(1, Our goal is to show that [B:y(UHH E 010
To this goal, we consider two cases.
Case I: S r or j S r and (i,j) * (1, We have
[B,(U++) S (1 17 (o+ + o. 23x(U.) 23,(0 2n
Σ |Btx (U.)),((a) kfi.
+1²
Σ
(0, 3, Bi (U)
Bu (u((3)Ba (ox
kd
(a)
(1 + 1l
(1 -o 10,0, + 0 (axzi)) (B(((U)))
22 201-0
Here in (a) we used the assumption by Lemma 4 we have
[B,y(U,)) $ 01a
2 01a I5, v(i.j) # (1, Hence, 201
2r,
(U(+)) E 010
24
where 1 siirrolsii
(68)
Case II: il: [ r For (U,(U,) such that i,j 2 r + 1, its dynamic is characterized by
(B,+())) (1 23, (B (U)() + 2n Σ
18xx (U.)33x),x
+ 11 Σ
130h (U)
(69)
kd
< (1+ + 11 (da )) (B,)
0 ( log (+))
Hence, for t S nn, (Bry(U)) since assume $ (7) 01/ This completes our inductive proof for (B,y(U.) S
1 log ( we have (B,( 1 (1+1 no (da
)da=))
(B)((U)) S
29

---

Published as a conference paper at ICLR 2023
20g 204-0--0,
01a 201 M(i,j) + (1, in the second phase. Finally, we turn to (U,) One can write
(1((++1)= = (1+ 2n (a) - (hi(U))) - 21 Σ (U,(
571
+ n/
Σ (1:(U1)811((03)883(1)1 201 Σ ( +oS
j.k
(70)
(a)
[ - 2n (01 + 0..mo7 - (u1( 0 (da=z))) ))) B11(U,)
(1 + 2n (0.99950, + 0.5no7 -1n (Ui() (U,).
Here in (a) we used the fact that 13,, S oir' log(d)a
Sa v(i.) & (1,
In (b), we used the assumption that S d
201/// The above inequality together with Lemma entailss entails
that within [ log [1600cm O yo
log ( a iterations, we have (n,(U,) 2 0.990
This completes the proof of Lemma
2 and marks the end of Substage 1.
Next, we move on to Substage
Substage 2. In Substage 2, we show that the second component 322 (U,) converges to d2 within iterations while the other coefficients remain small. To this goal, we first study the one-step dynamic of (U;)::):
log [ α
322 (22(1)))) (1 + 2n (02 (2.())))) 322
21) Σ 82 (U,)
j+2
(71)
+ 2
Σ j.k
B2s (U) - 202 Σ /2 (U.) - -33322(U.)
Different from the dynamic of 31 (U.). not all the coefficients (U,) with 2 or j are smaller than (22(U:) at the beginning of Substage 2. In particular, the basis coefficient (U.))
may much larger than (2,) at the beginning of Substage To see this, note that, according
to Equation 68 we have 312(U,) S a 21-02 and (uz(U,) S a
Hence, may be possible to have 1B12(U:)) x (U,) >> (U.) Therefore, the term 2 in Equation 76 must
be handled with extra care. Note that if we can show 322 (U,) > (U,)). then 2n (U.) can
be combined with the first term in the RHS of Equation 76 and the argument made in Substage can be repeated to complete the proof of Substage 2. However, our provided bound in Equation 68
can only imply (1
(22( Therefore, that (U,) < 02 (U,) along the trajectory. Upon controlling (U,)) we can then show the convergence of (U,) similar to our analysis for B11(U) in Substage 1.
need to provide a tighter analysis show
To control the behavior of (32 (U.), study the ratio (t) (T,)
(,a We will show that w(t) < do
along the trajectory. To this goal, we will show that w(t) can only increase for O yes
log at ))
iterations. Therefore, its maximum along the solution trajectory happens = O no1 log o. )) Therefore, by bounding the maximum, we can show that w()) remains small throughout the solution trajectory.
30

---

Published as a conference paper at ICLR 2023
First, at the initial point, we have w(0) Σ 64 64 (d)r'a?, which satisfies our claim. We next provide an upper bound for w(t + 1) based on (t).. Note that
(1 + (01 + +02 (U)) + poly(a)))
(2,
w(t + 1)5
(1 (1+ ++ poly(a)))2
B22(U,)
(1+ π(α (2)11(0)))) + poly(o)) w(t)
(72)
1+ - -poly(a))
< (1 + (a - -2:1 - ((1)
< ((1 + no1)) (0)) 0.4r 122) w(t).
Due to the first inequality, wi( can be increasing only until 311 (U,) 2 3 + poly(a). On the other hand, due to the dynamic of B11 in Substage 1, we can show that 811 2 = poly(a) in at most
T O sorr log (a))
iterations. Therefore, w()) takes its maximum T O log 21
n
On the other hand, we know that (U, satisfies
Bn (U2+)) ((1)==(()))))))))))) (1o1)))) 21B811 (U,) + poly((a) Bn (U.).
(73)
Hence, we can bound (t) as
T T-1
((T) < II
((1 + n01)² 2n - 0.4 0 (0))
(=0
T
(1 + )² 2n (U.) 0.4n) 0102 (1((U+1)
w())
(1 + 1)11)2 12 21 (U, = pa)
Bni(U)
(74)
(=0
B11 (U++1) w(0)
0(1)01102)
(U,)
t=
< 256(1
(2 0202)) Bn (UT) (d).
Here in (a) we used the fact that a Z (?ar 01/7 Hence, we have
((noz)
(T)) $ (1 10---14p 1(1)2(100000000 0102)) (d) $
a
ldg)
(75)
01
Due to our assumption a S (K log2 (d)) tion 76 can be lower bounded as
ver
we conclude that w()) - 0..11o2. Therefore, equa-
32e (2+()) 2 (1 1.99n (o2 922 (U.))) (2,( 2m
Σ 82 (U,)
j>2
(76)
+ n?
Σ
20 Σ (U,) -U)
j.k
The rest of the proof is a line by line reconstruction of Substage 1 and hence omitted for brevity.
Substage 3 < k
Via identical argument Substage 2, we can show that for each Substage
3 < k < we have
0.990 S Buk (U,) S okk
(77)
within Tí = O
log (1
iterations. This completes the proof of the first statement of Proposi-
tion
3
To prove the second statement of Proposition 3 as well as Theorem we next control the residual terms. First, we consider the residual term (U,) where either S or S Note that B,,(U) (U,( and hence we can assume & r without loss of generality. We will show that Biy (U.) decreases linearly once the corresponding signal (U,) converges to the vicinity of
31

---

Published as a conference paper at ICLR 2023
J₁- To this goal, it suffices to control the largest component Smax (U1) max (3,z(U)) (U)) Without loss of generality, we assume that the index (i, attains the maximum at time t, i... 9max (U,) = (U.).. One can write
(U(+11
$ (1+ 10 (B)) +2 Σ (3)((
+1² Σ Bu (U (a, o,) T,) (u(( (78)
(a)
$ (1 70.9(;; + +,,)),,,(U (1 10.00,)11.
Here in (a) we used the fact that 0.990; S (U,) < 0g s r and the fact that Suax (U+) = [B,(U,)), Hence, conclude that Bmax (U++1) I (1 0.91)o,)@aan (Ut). Therefore, within additional O nor log (2)) iterations, have (B,((U)) S r' r'log(d)a² for all (i,5) such that i i,ri
The remaining residual terms, i.e., those coefficients (U,) for which i, r, can be bounded via the same approach in Case II of substage 1. In particular, we can show that |8ky (U²) & (8,y(U)) S log(d)) (o)(d) Vi,j j > r. For brevity, we omit this step. This completes the proof of the second statement of Proposition 3
Finally, to prove Theorem we show that once (B,( S r' lolo log(d)d (d)aa V(i,J) + (k, 1s
the signals (U.) will further converge to a + O (a?) within
log (+t) iterations. To
see this, we simplify the dynamic of A&x (U,) as
BEN (U++1) (1 + (ox (e.((U)))))))))))))))))))))))( BuR (U:) - 21/ Σ (U))
j#k
+ 2
Σ
Bky (UN((U))au(U)))B 20k Σ 3 (U:) + +o
(79)
jd
(1 (BRA (U:) - oz)))² Bax (U,) - O (ndra log (d)at)
which leads to
Tk - Ben (U++1)
= n) (2 n( - =Bx (ox (u,) + O (1 log (d)a*)
< (1 1.98 ( BRx((U)) + O (ndr ld)a"")
(80)
Hence,
within additional
0
1 log ( 10
iterations, we have l0x (Ru(UU+))
o (
L dr log log (d)aa) for every Sr.
In conclusion, we have
|U 1 = Σ (Bi(UT)) - -Bt)² < 1 1og (d)
(81)
within O (L log Zx iterations. This completes the proof of Theorem
nor
F.4 ANALYSIS OF THE LOGISTIC MAP
In this section, we provide the proofs of Lemmas 3 and
32

---

Published as a conference paper at ICLR 2023
UPPER BOUND OF ITERATION COMPLEXITY
Recall the logistic map
Tt+1 (1 + 1j0 -)0; 20 a.
(82)
Here the initial value satisfies 0 < a < a. Vaskevicius et al. (2019) provide both upper and lower bounds for that follows the above logistie map. However, their bounds are not directly applicable to our setting. Hence, we need to develop new proof for Lemma
Proof of Lemma
We divide the dynamic into two stages: (a) 20, and (b) 20.
We consider [log (1) substages, where in each Substage k. we have a S 2t S ae kt1 Suppose that t is the number of iterations in Substage k. One can write
Stage 1: 2t 0.
1 (1 + 110 )fax)ar
(83)
2 (1 + 1j0 naxe k )ar.
Hence, it suffices to find the smallest min such that
ae (1 + na n1e 2 kt1
(84)
Solving this inequality leads to
Iiin log(1 (1 + nj nack+1) +1)*
(85)
Based on the above equality, we provide an upper bound for min:
min
log(1 + 1jo)
log(1 yoe ++1 /11 + n)))
3V
1
log(1 no)
mactt1
1+no mnae
(86)
11
+ 1tno gactFT
log(1 + no))
(6) + nae=+1
log(1 + no))
Here in (a) we used the fact that log(1 + x) S T S Lo. Hence, we have le < 111 *11 Therefore, the total iteration complexity of Stage is [og[1+no]
Yo > -1 and in (b) we used the fact that
upper bounded by
+ ntt
[lo (2a/a)1+ log(1 n7))
log(4or/a) log(1 + njo)"
(87)
log(1 17o)
k=0
Stage 2: rz 2 10. In this stage, we rewrite the equation 82 as
--+++1 = (1 ( (a xz).
(88)
Via a similar trick, we can show that within additional Tn log(4r/s) log(1+7)))
iterations, we have
It F which be achieved in the total number of iterations T Ti log( (oa) (log(40/or log(4a/ss) iterations. This completes the proof of Lemma
T2
SEPARATION BETWEEN TWO INDEPENDENT SIGNALS
In this section, show that there is sharp separation between two logistic maps with signals 0 d2 provided that 0. * 02. In particular, suppose that 01 02 2 7 > and
F (1 1jō1
118,, )
3o 02
(89)
Vt+1 (1 + 1-02 my::
yo a.
33

---

Published as a conference paper at ICLR 2023
16e
Proof of Lemma By Lemma we have æT 2 01 - = within log(I+no)) log
iterations.
ca
Therefore, it suffeee to show that y/ remains small for T. To this goal, note that
1/++1 (1 + 1 11b)) (1 + 1022) (1 + (
(90)
Hence, we need to bound I = (1 + no2)" Taking logarithm of both sides, we have
log(T) (T) = 7 log(1 1 + 102)
16o log(1 + 702))
(91)
log
log(1 mon))
Now, we provide a lower bound for the ratio log(1 + 11 10 hts//log(11 + n02):
log(1 + ((11-02)
log(1 + n/o11))
1-1102
log(1 + 1702)
log(1 + 702)
gl., -02)
1+11d2
1+7002
2 1 +
1702
(92)
(a)
01
2 1
202
01 + 02
202
where (a) follows from the assumption 1) S
I. Therefore, we have
40.
20 202/(r1t0))
1 2
202
16o
exp exp
log
(93)
EQ
02
EQ
which implies that
[6o
202/(01+oz)
01-02 160
yr S
"1
(94)
8a
8
This completes the proof of Lemma
G PROOF FOR TENSOR DECOMPOSITION
In this section, 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 S i S r', we assume that sin (0), x.) S y for some small 7 We will show that us (t) aligns with along the whole optimization trajectory. To this goal, we define may (t) = (t) zy) for every < S and S d. Recall that a multi-index with length l. We define |A|a the number of times index k appears one of the elements of A Evidently, have 0 < AA < l. Based on these definitions, one can write
BA = ΣII uu 24) IAle 1=1==1
E1
A U
(95)
Now, it suffices to study the dynamic of (ty()). In particular, we will show that vi (t) remains small except for the top-r diagonal elements Ui (4), S j S r, which will approach of intuition more concrete, we divide the terms (tt into three parts:
1/1 To make this
signal terms defined as (t, 1 S.
diagonal residual terms defined as ;; (t), r S S d. and
off-diagonal residual terms defined as uzy(t) (t), Vi $j.
For instance, assume that (1, 1,2). Then, |A/1 2 and |A/s
34

---

1. Published as a conference paper at ICLR 2023
Moreover, define (t) m (uy as the maximum element of the off-diagonal residual terms at every iteration When there is no ambiguity, we will omit the dependence on iteration t. For example, we write Vi (t) and V = (t).. Similarly, when there is no ambiguity, we write (t) or BA in lieu of Bx
Our next lemma characterizes the relationship between AA and Uaj-
1-1 Then, Lemma 5, Suppose that max 10/ S + V. and V S of d
. For Br with A₁- (f.....5) we have
[ax, - u/il S r'V!.
(96)
For BA with least two different indices A, we have
13 S 2ro, V.
(97)
The proof of this lemma is deferred to Appendix G.1 be upper bounded by maxips |pai]. Next, we control Viy by providing both lower and upper bounds on its dynamics.
Lemma 5 reveals that the magnitude of BA can
Proposition (One-step dynamics for viy(!)). Suppose that we have V(1) S of a and (a(t) S ot" vi s sisda d. Moreover, suppose that the step-size satisfies nj S loi
2. Then.
For the signal term va 1is we have
va(l) + T (o: (() 22'--110 2d'- ((),2())) (1) id'no 7 v'(I).
(98)
For the diagonal residual term Uza vt)(t), + S d, we have
(t 1) (() nlu 1 (1) 2npld' 01 +
V2())
(99)
For the off off-diagonal term Vtt,, we have
v(t+ + < V(t) 3nlo,
(100)
The proof of this proposition is deferred to Appendix G.2 dynamics, we next provide bound on the growth rate of Cij.
Equipped with the above one-step
Proposition 8. Suppose that the initial point satisfies sin (0), S y and |,)( a// with a S a 7 S
in 1 Moreover, suppose that the step-size satisfies n E ioi Then, within nla, a iterations,
For the signal term va va(t), 10,115 i 11,1
r, we have
[ot,(t)) ai| s 8dd 2 1 21122
(101)
For the diagonal residual term vu((), 15isis d, we have
|oa(1") < !/1
(102)
For the off-diagonal term (t), we have
v(t)) < <211(((0) s 127)
(103)
The proof of this proposition is deferred to Appendix prove Theorem
G.3
With the above proposition, we are ready
35

---

Published as a conference paper at ICLR 2023
Proof of Thereom We have the following decomposition
|T T 1| Σ 3A)-
(104)
Hence, it suffices to bound each 1AA - 1. Combining Lemma 5 and Proposition g
we have for
every [A* 38 Sjsr
Lemma
(A,((") - -
(1g - u3| + (1))
Proposition
8d
oj T a ? 2 + 2dar
1-2
(105)
16d
43 o
a ?²
where in the last inequality, we used 2 a s 1, and S 1. For the remaining diagonal elements Bay: r 1 S j S d, we have
(1') Lemmn
(1)) (())+++V(())
(106)
Propocition
+ 2day
For the general BA with at least two different indices in the multi-index A, we have
Lemma
[1a(1))
2ro 1 V(t*)
(107)
Propositio
Aro a 117
Hence, we conclude
[T((') - I
16d o0,' 1-3
a211 +
Σ
(2'a + 2dwr!) + d' -rro," 1 allly
(108)
1=+1
I 8rdd y01 1/1
which completes the proof of the theorem.
G.1 PROOF OF LEMMA 5
Proof. We first analyze Br,- Note that
r'
B4,
(us 2,)²
2/g
= 1/1
+
Σ
Vi.
(109)
i=1
i
Hence,
BBg u//l =
Σ
<
(110)
ifj
where we used the definition of V. For general BA where there are at least two different elements in the multi-index A, we have
BAI
II
AA/e sk
1=1
AEA
|Als 1
53
(11I)
|Al y +
|Als y 1111
j=r+1
S Sr0 7 + (r' - r)
max
|0/ Vi
12r+1
< V,
36

---

Published as a conference paper at ICLR 2023
where in the last inequality, we used the assumption that V x o1 o d and max | x
01 V. This completes the proof.
G.2 PROOF PROPOSITION 7
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 + with superscript '+'. For instance, we write Uij Fiy (t) and Dij Uty(t + 1).
Recall that Bij zz), For simplicity, denote μ of Hence, by assumption, have Va S μr. v1 S is r. We first provide the exact dynamic of Økj in the following lemma.
Lemma 6. The one-step dynamic of Vij takes the following form
b Vif + nnLo, 1'
ryl kE(r"NE
Σ
(112)
17 Σ sE[[-1] A:
BA
LAlk
2--1
sk
kEA,Ky
Proof. Recall that L =
2 Ex (Bx - 03)² and BA =x=E D
IekA (u;, 2x) AAle
Moreover, we have
37 oj for 1 i S r, and 3 otherwise. We first calculate
Vu,
Σ |A|.
II KEA.
(ui, zx)
(eie, za) (M6-1 2
(113)
Σ |A|s
II KEA.
|Alk
-A
2
Hence, the partial derivative of L(I) with respect to u.
Vu,C C = ESA - - = EGA - 3N
Σ
|A/s
II
ik
IAle
25.
VA
SEA
A
Note that {Zsbjena are unit orthogonal vectors. Hence, we have
(Vu, Bu, ={ { |Als (IIEAA,Z,
M1 Us
|Ag-11
if jEN
(114)
if j $ A.
By the definition of Cij. its update rule can be written in the following way
(u: 23}
Uiy 17 V L, 2z)
=Gig + Σ (332 BBa) 8A. 2j)
A
(115)
Uij +
Σ
(0) - BANAA,
II
Uik
AAle
rij
|M1--1
KEA.
(c)
Uij + Σ Σ
IAL - Bx)
Vik
AAl
1ij
* A|
KEA,KAS
37

---

Published as a conference paper at ICLR 2023
Here in (a), used the update rule for u, In (b), we applied equation 114 to exclude those A without j. In (c). we simply rearranged the above equation according to the cardinality [Als- We further isolate the term that only has Day:
UG + + BA, 11
Σ
S
Σ
BA
II KEAKAS
AAlk Uik
-E(1-1
i Bi + tn( o - ml
Σ
(116)
REA,N
-
7 E 8 Σ sE|1-1 A=|A|
BA
II REAAF
Uik
|Alk
Here in (a), rearranged terms and isolated the term with |A|; = 1. Note that the remaining terms must satisfy 1 & [Al5 I 1 - 1, which indicates that there must be at least 2 different indexes in A which in turn implies 3 O. In (b), we used the definition of Bx,- This completes the proof of Lemma 6
Equipped with Lemma we are ready prove Proposition 7
Proof of Proposition 7 The proof is divided into three parts:
Signal Term: un (f),sis
We first consider the signal terms ();( S j S r. First, upon setting =j j in Lemma
we have
# 0i + np
0i u//
Σ
U/
1-1
1
Σ
BA
II KEA
IAIx
A:|A/==s
(117)
2 u4 pl (o. - 0/3 mV\)
1-1
S
Σ A:JA|=
BA
II
U
JAle
U
KEARYAi
Now we aim to control (A) = BA (IlREA
uike
[Ale)
2if for |A
= =SS fl...., 1}. We have
r'
(4) = Σ
II |N
II
JALe
111-11
Uik
U
1=1
AEA
KEAR
(0)
|Als vi "0 Dj
(118)
(6)
20- 1--11-122 r' max
AAl 2 --1-
5 1/
jti
(0) 21 -32 + max p 1
< 21 222 3/2 + 7 -1
In (a), we used the fact that |uy| < V, Vi # j. (b), we isolated the term with j = and bounded the remaining terms with their maximum value. In (c), we used the fact that V S μ..
After substituting equation T18 into equation
we have
+ 2 Ua + mp (a: 2/e - vV 26
s
(122
21-
+ 1' p1
A:AA==
(119)
1-1
2 Ui + nn (o; J/4 - v'v)
o/ -
- n (1
21-3
+
11
|- vV) Σ
-
(d
38

---

Published as a conference paper at ICLR 2023
Id'-1 Therefore, we have where Cf = (.) Note that Σ 1)
-
vtilt 1)
vn + mpl (0. 1/2 - rV')) 0/4 ld' ld' ' (0 21 vv2 + r' x VV)
Un ni (0; 0/3 - 2 1/211 ld "pr 1 l'
(120)
+ ml [o; - 1/ 2d 2d' - 222 1-1 - ld "p1p"pp m -v'
US nl (a; d/ 2d'-1 10 222222 o/1 ld not
id'no,
+ Vv
where in the last inequality, we used the fact that r' d.
Diagonal Residual Term: r.rr+1sis
In this consider the terms vaa with ++1 < sis S d, which is similar to the case <is Without loss of generality, we assume that U/j 2 0. The case Un S 0 can be argued in an identical fashion. By equation
I17 have
o = Un - ml
Σ Uhi
Ui
8
Σ
BN
II
Vik
Ala
s-1
A:|A/ams
NEA
(121)
S Use - m7l (U/a - r'V') w//
1)
Σ
Σ A:AA|==
BA
II
Uik
JAIs
AEA
For (A) = 3A (IlKEAAZe
Uik
IAlk)
uit we further have
(A)
= f=1
II JAla
II KEA
JAlk Uik
Sh
NEA
JAla ah
JA|e
Σ
LA II
II
LA
jh
MEW
ViEN
KEA.NN
(122)
II
2|A|k
28-1
1
|Als
2
V 22
Uak
max J
20
2 2- V
Therefore, we obtain
0/ S va - 171 (r/s - v'V') r// + mldd 1 1-lv
(123)
I 0/4 nlo 2nd p -1!
Off-diagonal Residual Term: V(()
Finally, we characterize the dynamic of (t)) = max |0,, (t)) To this goal, we first consider the dynamic of each Vi such that i + j. Without loss of generality, we assume that Vij 2 0. One can write
1 Vij + m/
0f
Σ
o/
1- -1-1
Σ =E[l-1]
Σ A:
BA
II
Vik
JAle
FE(r")
REA,KY
< V + nlln Vi-1 + t
17
Σ
BA
IAlE
(124)
<E(1
A:|A|s=s
REA.
< V + 2nloo -7 Σ
Σ
BA
II
U
JAlk
1-1
sE[--]
A:AA| =
kEA,KyS
39

---

Published as a conference paper at ICLR 2023
where in the last inequality the assumption V E o d Similar to the previous case, next bound (A) 3A (ITREARA
V/kk
|Ale
Vij"
First note that [Al, 1]. Hence, have
that
(A) = Σ II
II
Alk Likk
REA
KEA
2|A
II KEA
JAlx
JAle ik
h-1
max hig
ij
KEA.
kEA,
(125)
20
II
II
JAle
h-1
max Regi
Uhk
1/j
kEA,Ky
1
We further note that |A|a
|A|: S s. Hence, it can be lower bounded by
(A) 2 - av--s-s -r'p'-'v'
(126)
Pluggin our estimation for (A) into equation 124 we finally have
S V + 2njor, V + nlld r' 'v' S 3nlo,
(127)
The last inequality comes from the assumption that V < 1 dl
This completes the proof of Proposition
G.3 PROOF OF PROPOSITION
g
In this section, we provide the proof of Proposition We will show that the signal terms vx(t), 1 S i S quickly converge to o."", and the residual terms remain small. To this goal, we first study the dynamic of (t).
Iteration complexity of the off-diagonal residual term V(t).
To start with the proof, first study the time required for the off-diagonal term V go from (0)) to (0), i.e., Ti minezo (V(() 20). By Proposition 7 know
Vtt + 1) s ((t) + (y (1 + 3nfo, Vi vi-2(1) v()).
(128)
Hence, for & < Ti, we have
V(t)= II
3nfo V (s)) V(0) < + 6nlo, V (0)))))
V(O).
==0
Note that V(Ti) 2 2 Therefore,
(1 6nlo, 20)) Ti v(0) > 20)v
(129)
Solving the above inequality for Ti, obtain
log(22/1)
log(2)
T log + 6nlo, v7-) 6n 01 V V1-2
(130)
On the other hand, our initial point satisfies sin u, (0), 2;) S 7 and [0]] 1/1, Hence, we have (0) I al1lyy Substituting this into the above equation, we conclude that
log(2)
T N G
(131)
Note that Ti
log(2) 2
1-2
1. Hence, we have V(I*) < 201/v((0> {(2))14
6tll" o
40

---

Published as a conference paper at ICLR 2023
Iteration complexity of the diagonal residual term (a((),r (o(t),,+1110 +1siso d.
Next, show that va((), r + s i S d will remain small during <t S tt
8
2
First
nl.
by equation 123 in the proof of Lemma we have for every 0 S S
(t + 1) S va(() + 2nddd o v
(132)
Note that V()) & (0), S 1", which leads to
(t) - (0) + Anild -1 rtat
- va 4sjld' (ppld'ot--pp o at"
(133)
< a + +0 (d'ma
2a
for 0 s f*, Here we used the assumption that a S di and 7 S 1-2
Iteration complexity of the signal term Ue (t), Sisr.
As the last piece of the proof, we show that by iteration 1* will converge the eigenvalue o./². First, recall that
8 nl..
a the signal tai 1sis S
(a 1) (t)) + ryl (o: - (1) - 6 (())((()
v/ 21d -
(134)
( (t) + m/l 04 (1) 4''- up (()) 2112) 84 2 4ld nol'nno--/ee
where in the last inequality we used the fact that (tt) < (0) for S S In light of the above inequality, we characterize the convergence of dis using similar method as in (Ma & Fattahi 2022a). particular, we divide our analysis into two phases.
Phase 1. the first phase, we have Us I (0. First, since U. (0) 2 1/ aVVI-~2 we can easily conclude that (t + 1) 2 (t) by induction. Hence, we can simplify the dynamic as
(t + 1) 2 (1+ (0.9990 - 2/ Ju (()) (t) 2 (1 + 0.49n)o,, 2 (tt) vt(()..
(135)
Next, we further split the interval I = 0, 0.50 1/0 into N = o (log (aso)"/a) sub-intervals {Io. IN- where Ia = [2 u.. (0). 2 uy (0)) Let Ta collect the iterations that en spends in Ig. Accordingly, let |T = L be the number of iterations that u/e spends within I, First note that
(a(t)= 2 2 (0)) for every E Ta. Hence, we have
(1 0.49nlo 2 2112 03
(0)) 4 22
(136)
which implics
log(2)
t4 0. 0.49nlo 2()
2 (1-2))
(137)
(0))
By summing over k =0
N -...
1, we can upper bound the required number of iterations T3
log (2)
4
T3
< 1
(1-12))
S n << Ti,
(138)
49nlo, 20)
k=0
0=
where the last inequality is due to our assumption 7 S
1 r
kl
1/1 Phase In the second phase, we have U4 2 0.50
We further simplify equation 134 as
va + 1) 2 (t) Tl (o. sd o, =
a²n² v(a(t) '
(139)
2 u (t) + nl (ō: - v (t) 2/2
1-1 (t).
where denote o; 8d 0,, بي 2?11?2 Then, via similar trick, within additional Ta s 1-2
- iterations, have (1) > a; -d' 22/1,2 similar argument the upper bound shows (1) a, - 10 42 2/1,, which completes the proof.
nlo,
41

---

Published as a conference paper at ICLR 2023
H AUXILIARY LEMMAS
Lemma (Bernoulli inequality). For 0 S x <
1 r-1 and r > 1, we have
(1+x)* +1+ 1 (r - 1)
(140)
42


Reviews:
[{"review_id": "ofYci6ODvxj", "reviewer": "Reviewer_yaF6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper analyzes the training dynamics of gradient descent based algorithms using a particular orthonormal basis for the function space spanned by the (function) iterates during training. Convergence to the function in the model space that minimizes the expected l2 loss is shown to occur at a novel convergence rate under some assumptions.", "summary_of_the_review": "Sorry for a very lightweight review.", "strength_and_weaknesses": "I am unable to fully evaluate this current work due to time mismanagement from my side, sorry.\n\nHere are however a few remarks:\n - I would have appreciated (I am not an expert) a more comprehensive comparison to previous convergence rates (theorem 1 and proposition 2 e.g.)\n - in theorem 1, the required learning rate depends in $\\frac{1}{\\sqrt{d}\\log{d}}$. For problems in high dimension, it implies a very small learning rate. Doesn't it force the GD dynamics to the gradient flow regime? In the CNN examples, it is common practice to use a learning rate as large as possible: would you say that your analysis still applies?\n - in section 3.4 \"Surprisingly, A-CK [...] captures the underlining solution of different gradient-based algorithms\" -> can you elaborate on what makes A-CK so special that its SVD can be used in your framework, instead of any other choice of basis. How is it \"surprising\" ?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XyUs6Vh-ox", "reviewer": "Reviewer_5bKV", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the convergence behavior of SGD. Specifically, it says that there is a way to project the learning process onto a set of orthogonal basis functions so that the leading coefficients in the basis functions exhibit 'smoothly' convergence behavior. While not explicitly stated in the paper, I think that implies certain generalization properties of the SGD (still a bit confused see below). The authors then proceed to analyze three specific applications/models and demonstrate the power of their generic framework. \n\nI have some clarification questions: \n\nQ1. If I understand the result correctly, it is a convergence result of the optimization problem and it does not explicitly say anything about generalization errors. But because the coefficients of leading basis functions get stabilized over time, does that also imply some generalization result, and if so, is there an explicit way to characterize that. \nQ2. For application 2 (matrix completion), the result looks too strong (also related to my Q1). Specifically, this is only a convergence result but at the end it also seems to imply all local optimal are the same. Ge, Jin, and Zheng's result made a quite heavy effort in matrix manipulation to prove the no spurious local optimal result. I wonder how this result manages to circumvent the technical barriers there. Also, it seems these results usually require some incoherence assumptions. I cannot find the counterpart in this paper too. ", "summary_of_the_review": "1. it is a strong result and presents a new and interesting way to interpret the SGD.  \n2. In fact, it actually appears to be too strong so I need some clarification. ", "strength_and_weaknesses": "Strength: the authors find a quite interesting way to construct the orthogonal basis, and under some special circumstances (the applications), the basis can be explicitly written down. The analysis seems substantial. \n\nWeakness: it would be helpful if the authors can clarify my two questions above. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "MexVoOSXoUW", "reviewer": "Reviewer_iyzY", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors 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. 1) The authors prove that GD learns the coefficients of an appropriate function basis that forms the true model. 2) The authors 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. 3) The authors 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.\n", "summary_of_the_review": "The paper proposes analyzing the convergence of gradient descent via basis function decomposition, which seems interesting. However, I have a few questions on this paper, see details in \"Strength and Weaknesses\"", "strength_and_weaknesses": "Strength: \n\n1. The paper is clearly written and easy to read.\n\n2. The analysis seems new to me.\n\nWeakness and questions to the authors (I am not an expert in this area): \n\n1. For the motivating example, what is the behavior of SGD with momentum for training ResNet-18? I am asking as SGD with momentum is the baseline method for training ResNet-18.\n\n\n2. What are the major advantages of the proposed analysis of GD via basis function decomposition over the existing analysis of GD? Indeed, we also analyze the eigendecomposition of the GD dynamics for the classical analysis.\n\n\n3. In practice, we may not be able to get the exact gradient. Can authors comment on how to extend the analysis to stochastic gradient descent and related algorithms?\n\n\n4. One weak point is the analysis relies on Assumption 1, which is much stronger than the assumption used in the classical analysis of gradient descent.\n\n\n5. For neural networks, how to choose the orthogonal basis? Perhaps I missed something, but after reading the paper, I still do not know how to do this.\n\n\n6. In Theorem 1, can authors comment on the choice of the step size \\eta? \n\n\n7. In section 3.1, how can we assume the kernel function are orthonormal for the neural tangent kernel? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Hzbr5kKiSNm", "reviewer": "Reviewer_CW4U", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the training dynamics of gradient descent (GD) on Hilbert Space ($L^2$). A key idea from this paper is to utilize orthogonalities from the Hilbert space to decompose the GD dynamics into a union of (possibly countably infinite) per eigenmode dynamics. In the general and realistic setting, the per eigenmode dynamics is in-trackable as each mode interacts with other modes in a complex and structural manner. To simplify the analysis, the authors make a key assumption, among many others: (`GIA`=Gradient Independent Assumption) gradients of the loss with respect to modes are `mutual` independent (orthogonal) throughout the whole training trajectory. This assumption reduces the original complex system into a 1-dimensional problem in the sense that each eigenmode evolves independently. With an additional Polyak-Lojasiewicz type assumption (called Gradient dominant in the paper), convergence to the global minima can be proved.    \n\nThe GIA does not strictly hold in general (the only known example in the paper is the kernel regression). Nevertheless, the authors can apply the idea of eigendecomposition to symmetry matrix/tensor decomposition (with some additional assumptions, and finer analysis to control the coefficient gradients) and recover/improve convergence results from prior work. \n\nOverall, I think the decomposition of the dynamics of GD into eigenmodes is natural and interesting. However, the GIA is too strong and unnatural, which trivializes the rich and complex dynamics of GD into a union of independent 1D dynamics. ", "summary_of_the_review": "The paper has several interesting insights. However, the assumptions on Theorem 1, in particular, the GIA, are too strong and unnatural, which basically remove all interesting structures in GD dynamics. ", "strength_and_weaknesses": "#Strength\n- The paper proposes a natural framework to study gradient descent dynamics in Hilbert space. \n- The paper recovers existing convergence results for symmetry matrix/ tensor decomposition.  \n\n# Weaknesses\n\n- The GIA is too strong and almost trivializes the dynamics. My understanding is that the rich, structural, and coupled dynamics between eigenmodes are the most interesting part of GD + representation learning. Note that this framework does not even hold for symmetry matrix/tensor decomposition, which is much simpler than neural networks. Note that the authors argue that, empirically, the assumptions (GIA, Gradient dominant) approximately hold. However, I don't think such empirical results suffice to justify the assumptions in Theorem 1. For example $\\sum_{j\\neq i} \\langle \\nabla \\beta(\\theta_j) , \\nabla\\beta(\\theta_i) \\rangle$ can be large even each $\\langle \\nabla \\beta(\\theta_j) , \\nabla\\beta(\\theta_i) \\rangle$ is small.    \n\n- The framework does not yield new results or solve a new problem. Further work needs to be done to justify the usefulness of the framework. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster