Title: ON THE DYNAMICS AND CONVERGENCE OF WEIGHT NORMALIZATION FOR TRAINING NEURAL NETWORKS Anonymous authors Paper under double-blind review

Abstract: We present a proof of convergence for ReLU networks trained with weight normalization. In the analysis, we consider over-parameterized 2-layer ReLU networks initialized at random and trained with batch gradient descent and a fixed step size. The proof builds on recent theoretical works that bound the trajectory of parameters from their initialization and monitor the network predictions via the evolution of a "neural tangent kernel" (Jacot et al. 2018) . We discover that training with weight normalization decomposes such a kernel via the so called "length-direction decoupling". This in turn leads to two convergence regimes and can rigorously explain the utility of WeightNorm. From the modified convergence we make a few curious observations including a natural form of "lazy training" where the direction of each weight vector remains stationary.

1. INTRODUCTION:

Dynamic normalization in neural networks is a re-parametrization procedure between the layers that improves stability during training and leads to faster convergence. This approach was popularized with the introduction of Batch Normalization (BatchNorm) in [20] and has led to a plethora of additional normalization methods, notably including Layer Normalization (LayerNorm) [6] and Weight Normalization (WeightNorm) [28] . WeightNorm was proposed as a method that emulates BatchNorm and benefits from similar stability and convergence properties. Moreover, WeightNorm has the advantage of not requiring a batch setting, therefore considerably reducing the computational overhead that is imposed by BatchNorm [16] . WeightNorm is widely used in training of neural networks and is the focus of this work.

Today, normalization methods are ubiquitous in the training of neural nets since in practice they significantly improve the convergence speed and stability in training. Despite the impressive empirical results and massive popularity of dynamic normalization methods, explaining their utility and proving that they converge when training with non-smooth, non-convex loss functions has remained an unsolved problem. In this paper we provide sufficient conditions on the data, initialization, and over-parameterization for dynamically normalized ReLU networks to converge to a global minimum of the loss function, and rigorously illustrate the utility of normalization methods.

Consider the class of 2-layer ReLU neural networks f :

R d → R parameterized by (W, c) ∈ R m×d × R m as f (x; W, c) = 1 √ m m k=1 c k σ(w k x). (1.1)

Here the activation function is the ReLU, σ(s) = max{s, 0} [26] , m denotes the width of the second layer, and f is normalized accordingly by a factor √ m. We investigate gradient descent training with WeightNorm for (1.1), which re-parameterizes the network in terms of (V, g, c) ∈ R m×d ×R m ×R m as

f (x; V, g, c) = 1 √ m m k=1 c k σ g k • v k x v k 2 . (1.2)

This gives a similar parameterization to [14] that study convergence of gradient optimization of convolutional filters on Gaussian data. We consider the regression task, optimizing with respect to the L 2 loss with random parameter initialization and focus on the over-parametrized regime, meaning that m > n, where n is the number of training samples.

The neural network function class (1.1) has been studied in many papers including [3, 15, 33, 36] along with other similar over-parameterized architectures [1, 14, 23] . An exuberant series of recent works prove that feed-forward ReLU networks converge to zero training error when trained with gradient descent from random initialization. Nonetheless, to the best of our knowledge, there are no proofs that ReLU networks trained with dynamic normalization on general data converge to a global minimum. This is in part because normalization methods completely change the optimization landscape during training. Here we show that neural networks of the form given above converge at linear rate when trained with gradient descent and WeightNorm. The analysis is based on the over-parameterization of the networks, which allows for guaranteed descent while the gradient is non-zero.

For regression training, a group of papers studied the trajectory of the networks' predictions and showed that they evolve via a "neural tangent kernel" (NTK) as introduced by Jacot et al. [21] . The latter paper studies neural network convergence in the continuous limit of infinite width overparameterization, while the works of [3, 15, 27, 33, 36] analyze the finite width setting. For finitewidth over-parameterized networks, the training evolution also exhibits a kernel that takes the form of a Gram matrix. In these works, the convergence rate is dictated by the least eigenvalue of the kernel. We build on this fact, and also on the general ideas of the proof of [15] and the refined work of [3] .

Compared with un-normalized training, we prove that normalized networks follow a modified kernel evolution that features a "length-direction" decomposition of the NTK. This leads to two convergence regimes in WeightNorm training and explains the utility of WeightNorm from the perspective of the NTK. In the settings considered, WeightNorm significantly reduces the amount of over-parameterization needed for provable convergence, as compared with un-normalized settings.

The decomposition of the NTK also connects to observations of [12] that discuss "lazy training" which refers to a training regime where the weights of the network stay close to their initialization (see Section 6) . Further, we present a more careful analysis that leads to improved over-parameterization bounds as compared with [15] .

In this work we rigorously analyze the dynamics of weight normalization training and its convergence from the perspective of the neural tangent kernel. We discover WeightNorm training has two regimes with distinct behaviors. The main contributions of this work are:

• We prove the first general convergence result for dynamically normalized 2-layer ReLU networks trained with gradient descent. Our formulation does not assume the existence of a teacher network and has mild assumptions on the training data.

• We explain the utility of normalization methods via a decomposition of the neural tangent kernel.

In the analysis we highlight two distinct convergence regimes and give a concrete example of "lazy training" for finite-step gradient descent.

• It is shown that finite-step gradient descent converges for all weight magnitudes at initialization and we significantly reduce the amount of over-parameterization required for provable convergence as compared with un-normalized training.

The paper is organized as follows. In Section 2 we provide background on WeightNorm and derive key evolution dynamics of training in Section 3. We present and discuss our main results, alongside with the idea of the proof, in Section 4. We discuss related work in Section 5, and offer a discussion of our results and their implications to dynamic normalization training and "lazy training" in Section 6. Proofs are presented in the Appendix.

2. WEIGHTNORM:

Here we give an overview of the WeightNorm procedure and review some known properties of normalization methods.

Notation We use lowercase, lowercase boldface, and uppercase boldface letters to denote scalars, vectors and matrices resp. We denote the Rademacher distribution as U {1, -1} and write N (µ, Σ) for a Gaussian with mean µ and covariance Σ. Training points are denoted by x 1 . . . x n ∈ R d and parameters of the first layer by v k ∈ R d . We use σ(x) := max{x, 0}, and write • 2 , • F for the spectral and Frobenius norms for matrices. λ min (A) is used to denote the minimum eigenvalue of a matrix A and •, • denotes the Euclidean inner product. For a vector v denote the 2 vector norm as v 2 and for a positive definite matrix S define the induced vector norm v S := √ v Sv. The projections of x onto u and u ⊥ are defined as x u := uu x u 2 2 , x u ⊥ := x Iuu u 2

2:

. Denote the indicator function of event A as 1 A and for a weight vector at time t, v k (t), and data point x i we denote

1 ik (t) := 1 {v k (t) xi≥ 0} .

WeightNorm procedure For a single neuron σ(w x), WeightNorm re-parametrizes the weight

w ∈ R d in terms of v ∈ R d , g ∈ R as w(v, g) = g • v v 2 , σ g • v x v 2 . (2.1)

This decouples the magnitude and direction of each weight vector (referred as the "length-direction" decomposition). In comparison, for BatchNorm each output w x is normalized according to the average statistics in a batch. We can draw the following analogy between WeightNorm and BatchNorm if the inputs x i are centered (Ex = 0) and the covariance matrix is known (Exx = S).

In this case, batch training with BatchNorm amounts to

EQUATION

From this prospective, WeightNorm is a special case of (2.2) with S = I [22, 28] .

Properties of WeightNorm:

We start by giving an overview of known properties of WeightNorm that will be used to derive the gradient flow dynamics of WeightNorm training.

For re-parameterization (2.1) of a network function f that is initially parameterized with a weight w, the gradient ∇ w f relates to the gradients ∇ v f, ∇ g f by the identities

∇ v f = g v 2 (∇ w f ) v ⊥ , ∇ g f = (∇ w f ) v .

This implies that ∇ v f • v = 0 for each input x and parameter v. For gradient flow, this orthogonality results in v(0) 2 = v(t) 2 for all t. For gradient descent (with step size η) the discretization in conjunction with orthogonality leads to increasing parameter magnitudes during training [4, 19, 28] , as illustrated in Figure 1 ,

v(s + 1) 2 2 = v(s) 2 2 + η 2 ∇ v f 2 2 ≥ v(s) 2 2 . (2.3) v k (0) dvk dt (0) v k (t) α α v k (0) -∇ vk L v k (s)

Figure 1 : WeightNorm updates for gradient flow and gradient descent. For gradient flow, the norm of the weights are preserved, i.e., v k (0) 2 = v k (t) 2 for all t > 0. For gradient descent, the norm of the weights v k (s) 2 is increasing with s.

Problem Setup:

We analyze (1.1) with WeightNorm training (1.2), so that

f (x; V, c, g) = 1 √ m m k=1 c k σ g k • v k x v k 2 .

We take an initialization in the spirit of [17, 28] :

v k (0) ∼ N (0, α 2 I), c k ∼ U {-1, 1}, and g k (0) = v k (0) 2 /α. (2.4)

Where α is the variance of v k at initialization. The initialization of g k (0) is therefore taken to be independent of α. We remark that the initialization (2.4) gives the same initial output distribution as in methods that study the un-normalized network class (1.1). The parameters of the network are optimized using the training data {(x 1 , y 1 ), . . . , (x n , y n )} with respect to the square loss

EQUATION

where

f = (f 1 , f 2 , . . . , f n ) = (f (x 1 ), f (x 2 ), . . . , f (x n ))

and y = (y 1 , y 2 , . . . , y n ) .

3. EVOLUTION DYNAMICS:

We present the gradient flow dynamics of training (2.5) to illuminate the modified dynamics of WeightNorm as compared with vanilla gradient descent. In Appendix C we tackle gradient descent training with WeightNorm where the predictions' evolution vector df dt is replaced by the finite difference f (s + 1) -f (s). For gradient flow, each parameter is updated in the negative direction of the partial derivative of the loss with respect to that parameter. The optimization dynamics give

dv k dt = - ∂L ∂v k , dg k dt = - ∂L ∂g k . (3.1)

We consider the case where we fix the top layer parameters c k during training. In the overparameterized regime, the dynamics of c k and g k turn out to be the same.

To quantify convergence, we monitor the time derivative of the i-th prediction, which is computed via the chain rule as

∂f i ∂t = m k=1 ∂f i ∂v k dv k dt + ∂f i ∂g k dg k dt .

Substituting (3.1) into the i-th prediction evolution and grouping terms yields

∂f i ∂t = - m k=1 ∂f i ∂v k ∂L ∂v k T i v - m k=1 ∂f i ∂g k ∂L ∂g k T i g . (3.2)

The gradients of f i and L with respect to v k are written explicitly as

∂f i ∂v k (t) = 1 √ m c k • g k (t) v k (t) 2 • x v k (t) ⊥ i 1 ik (t), ∂L ∂v k (t) = 1 √ m n i=1 (f i (t) -y i ) c k • g k (t) v k (t) 2 x v k (t) ⊥ i 1 ik (t). Thus T i v (t) in (3.

2) can be calculated as

T i v (t) = n j=1 (f j (t) -y j ) 1 m m k=1 1 jk (t)1 ik (t) c k • g k (t) v k (t) 2 2 x v k (t) ⊥ j , x v k (t) ⊥ i .

Defining the v-orthogonal Gram matrix V(t) as

EQUATION

we can write T i v as

T i v (t) = n j=1 V ij (t) α 2 (f j (t) -y j ).

Note that V(t) is the induced neural tangent kernel [21] for the parameters v of WeightNorm training. While it resembles the Gram matrix H(t) studied in [3] , here we obtain a matrix that is not piece-wise constant in v since the data points are projected onto the orthogonal component of v. We compute T i g in (3.2) analogously. The associated derivatives with respect to g k are

∂f i ∂g k (t) = 1 √ m c k v k (t) 2 σ(v k (t) x i ), ∂L ∂g k (t) = 1 √ m n j=1 (f j (t) -y j ) c k v k (t) 2 σ(v k (t) x j ),

and we obtain

T i g (t) = m k=1 1 m n j=1 (f j (t) -y j ) c k v k (t) 2 2 σ(v k (t) x j )σ(v k (t) x i ). Given that c 2 k = 1, define G(t) as G ij (t) = 1 m m k=1 σ(v k (t) x i )σ(v k (t) x j ) v k (t) 2 2 (3.4)

hence we can write

T i g (t) = n j=1 G ij (t)(f j (t) -y j ).

Combining T v and T g , the full evolution dynamics are given by

EQUATION

Denote Λ(t) := V(t) α 2 + G(t) and write df dt = -Λ(t)(f (t) -y). We note that V(0), G(0), defined in (3.3), (3.4) , are independent of α: Observation 1 (α independence). For initialization (2.4) and α > 0 the Gram matrices V(0), G(0) are independent of α.

This fact is proved in Appendix A. When training the neural network in (1.1) without WeightNorm (see [3, 15, 36] ) , the corresponding neural tangent kernel

H(t) is defined by ∂fi ∂t = m k=1 ∂fi ∂w k dw k dt = - m k=1 ∂fi ∂w k ∂L ∂w k = - n j=1 H ij (t)(f j -y j )

and takes the form

EQUATION

The analysis presented above shows that vanilla and WeightNorm gradient descent are related as follows.

Proposition 1. Define V(0), G(0), and H(0) as in (3.3), (3.4) , and (3.6) respectively. then for all α > 0,

V(0) + G(0) = H(0).

Thus, for α = 1,

∂f ∂t = -Λ(0)(f (0) -y) = -H(0)(f (0) -y).

That is, WeightNorm decomposes the NTK in each layer into a length and a direction component. We refer to this as the "length-direction decoupling" of the NTK, in analogy to (2.1). From the proposition, normalized and un-normalized training kernels initially coincide if α = 1. The utility of normalization methods can be attributed to the modified NTK Λ(t) that occurs when the WeightNorm coefficient, α, deviates from 1. For α 1 the kernel Λ(t) is dominated by G(t), and for α 1 the kernel Λ(t) is dominated by V(t). We elaborate onthe details of this in the next section. In practice, by (2.3) as training progresses, v k 2 grow monotonically, leading to larger α and a transition from the V-dominated to the G-dominated regime. In our analysis we will study the two regimes α > 1 and α < 1 in turn.

4. MAIN CONVERGENCE THEORY:

In this section we discuss our convergence theory and main results. From the continuous flow (3.5), we observe that the convergence behavior is described by V(t) and G(t). The matrices V(t) and G(t) are positive semi-definite since they can be shown to be covariance matrices. This implies that the least eigenvalue of the evolution matrix Λ(t) = 1 α 2 V(t) + G(t) is bounded below by the least eigenvalue of each kernel matrix, λ min (Λ(t)) ≥ max{λ min (V(t))/α 2 , λ min (G(t))}.

For finite-step gradient descent, a discrete analog of evolution (3.5) holds. However, the discrete case requires additional care in ensuring dominance of the driving gradient terms. For gradient flow, it is relatively easy to see linear convergence is attained by relating the rate of change of the loss to the magnitude of the loss. Suppose that for all t ≥ 0,

λ min Λ(t) ≥ ω/2, with ω > 0. (4.1)

Then the change in the regression loss is written as

d dt f (t) -y 2 2 = 2(f (t) -y) df (t) dt = -2(f (t) -y) Λ(t)(f (t) -y) (4.1)

≤ -ω f (t) -y 2 2 . Integrating this time derivative and using the initial conditions yields

f (t) -y 2 2 ≤ exp(-ωt) f (0) -y 2 2

, which gives linear convergence. The focus of our proof is therefore showing that (4.1) holds throughout training.

By Observation 1 we have that V and G are independent of the WeightNorm coefficient α (α only appears in the 1/α 2 scaling of Λ). This suggests that the kernel Λ(t) = 1 α 2 V(t) + G(t) can be split into two regimes: When α < 1 the kernel is dominated by the first term 1 α 2 V, and when α > 1 the kernel is dominated by the second term G. We divide our convergence result based on these two regimes.

In each regime, (4.1) holds if the corresponding dominant kernel, V(t) or G(t), maintains a positive least eigenvalue. Having a least eigenvalue that is bounded from 0 gives a convex-like property that allows us to prove convergence. To ensure that condition (4.1) is satisfied, for each regime we show that the corresponding dominant kernel is "anchored" (remains close) to an auxiliary Gram matrix which we define in the following for V and G.

Define the auxiliary v-orthogonal and v-aligned:

EQUATION

EQUATION

For now, assume that V ∞ and G ∞ are positive definite with a least eigenvalue bounded below by ω (we give a proof sketch below). In the convergence proof we will utilize over-parameterization to ensure that V(t), G(t) concentrate to their auxiliary versions so that they are also positive definite with a least eigenvalue that is greater than ω/2. The precise formulations are presented in Lemmas B.4 and B.5 that are relegated to Appendix B.

To prove our convergence results we make the assumption that the x i s have bounded norm and are not parallel. Assumption 1 (Normalized non-parallel data). The data points (x 1 , y 1 ), . . . , (x n , y n ) satisfy x i 2 ≤ 1 and for each index pair i = j,

x i = β • x j for all β ∈ R \ {0}.

In order to simplify the presentation of our results, we assume that the input dimension d is not too small, whereby d ≥ 50 suffices. This is not essential for the proof. Specific details are provided in Appendix A. Assumption 2. For data points x i ∈ R d assume that d ≥ 50.

Both assumptions can be easily satisfied by pre-processing, e.g., normalizing and shifting the data, and adding zero coordinates if needed.

Given Assumption 1, V ∞ , G ∞ are shown to be positive definite. Lemma 4.1. Fix training data points {(x 1 , y 1 ), . . . , (x n , y n )} satisfying Assumption 1. Then the v-orthogonal and v-aligned Gram matrices V ∞ and G ∞ , defined as in (4.2) and (4.3), are strictly positive definite. We denote the least eigenvalues λ min (V ∞ ) =: λ 0 , λ min (G ∞ ) =: µ 0 .

Proof sketch Here we sketch the proof of Lemma 4.1. The main idea, is the same as [15] , is to regard the auxiliary matrices V ∞ , G ∞ as the covariance matrices of linearly independent operators. For each data point x i , define φ i (v) := x v ⊥ i 1 {x i v≥0} . The Gram matrix V ∞ is the covariance matrix of {φ i } i=1:n taken over R d with the measure N (0, α 2 I). Hence showing that V ∞ is strictly positive definite is equivalent to showing that {φ i } i=1,...n are linearly independent. Unlike [15] , the functionals under consideration are not piecewise constant so a different construction is used to prove independence. Analogously, a new set of operators, θ i (v) := σ(x v i ), is constructed for G ∞ . Interestingly, each φ i corresponds to dθi dv . The full proof is presented in Appendix D. As already observed from evolution (3.5), different magnitudes of α lead to two distinct regimes that are discussed below. We present the main results for each regime.

V-DOMINATED CONVERGENCE:

For α < 1 convergence is dominated by V(t) and λ min (Λ(t)) ≥ 1 α 2 λ min (V(t)). We present the convergence theorem for the V-dominated regime here. Theorem 4.1 (V-dominated convergence). Suppose a neural network of the form (1.2) is initialized as in (2.4) with α ≤ 1 and that Assumptions 1,2 hold. In addition, suppose the neural network is trained via the regression loss (2.5) with targets y satisfying y ∞ = O(1). If m = Ω n 4 log(n/δ)/λ 4 0 , then with probability 1 -δ, 1. For iterations s = 0, 1, . . . , K, the evolution matrix Λ(s) satisfies λ min (Λ(s)) ≥ λ0 2α 2 .

2.. WeightNorm training with gradient descent of step:

-size η = O α 2

V ∞ 2 converges linearly as

f (s) -y 2 2 ≤ 1 - ηλ 0 2α 2 s f (0) -y 2 2 .

The proof of Theorem 4.1 is presented in Appendix C. We will provide a sketch below. We make the following observations about our V-dominated convergence result.

The required over-parameterization m is independent of α. Further, the dependence of m on the failure probability is log(1/δ). This improves previous results that require polynomial dependence of order δ 3 . Additionally, we reduce the dependence on the sample size from n 6 (as appears in [3] ) to n 4 log(n).

In Theorem 4.1, smaller α leads to faster convergence, since the convergence is dictated by λ 0 /α 2 . Nonetheless, smaller α is also at the cost of smaller allowed step-sizes, since η = O(α 2 / V ∞ 2 ). The trade-off between step-size and convergence speed is typical. For example, this is implied in Chizat et al. [12] , where nonetheless the authors point out that for gradient flow training, the increased convergence rate is not balanced by a limitation on the step-size. The works [4, 19, 32] define an effective step-size (adaptive step-size) η = η/α 2 to avoid the dependence of η on α.

G-DOMINATED CONVERGENCE:

For α > 1 our convergence result for the class (1.2) is based on monitoring the least eigenvalue of G(t). Unlike V-dominated convergence, α does not affect the convergence speed in this regime. Theorem 4.2 (G-dominated convergence). Suppose a network of the form (1.2) is initialized as in (2.4) with α ≥ 1 and that Assumptions 1, 2 hold. In addition, suppose the neural network is trained via the regression loss (2.5) with targets y satisfying y

∞ = O(1). If m = Ω max n 4 log(n/δ)/α 4 µ 4 0 , n 2 log(n/δ)/µ 2 0

, then with probability 1 -δ, 1. For iterations s = 0, 1, . . . , K, the evolution matrix Λ(s) satisfies λ min (Λ(s)) ≥ µ0 2 .

2.. WeightNorm training with gradient descent of step-size:

η = O 1 Λ(t)

converges linearly as

f (s) -y 2 2 ≤ 1 - ηµ 0 2 s f (0) -y 2 2 .

We make the following observations about our G-dominated convergence result, and provide a proof sketch further below.

Theorem 4.2 holds for α ≥ 1 so long as m = Ω max n 4 log(n/δ)/µ 4 0 α 4 , n 2 log(n/δ)/µ 2 0 . Taking α = n/µ 0 gives an optimal required over-parameterization of order m = Ω n 2 log(n/δ)/µ 2 0 . This significantly improves on previous results [15] for un-normalized training that have dependencies of order 4 in the least eigenvalue, cubic dependence in 1/δ, and n 6 dependence in the number of samples n. In contrast to V-dominated convergence, here the rate of convergence µ 0 is independent of α but the over-parameterization m is α-dependent. We elaborate on this curious behavior in the next sections.

Proof sketch of main results:

The proof of Theorems 4.1 and 4.2 is inspired by a series of works including [3, 13, 15, 33, 36] . The proof has the following steps: (I) We show that at initialization V(0), G(0) can be viewed as empirical estimates of averaged data-dependent kernels V ∞ , G ∞ that are strictly positive definite under Assumption 1. (II) For each regime, we prove that the corresponding kernel remains positive definite if v k (t) and g k (t) remain near initialization for each 1 ≤ k ≤ m. (III) Given a uniformly positive definite evolution matrix Λ(t) and sufficient overparameterization we show that each neuron, v k (t), g k (t) remains close to its initialization. The full proof is presented in Appendix B for gradient flow and Appendix C for finite-step gradient descent.

While the spirit of the proof is familiar from other works, the different convergence regimes explain the utility of WeightNorm. We elaborate further on this in the discussion.

5. RELATED WORK:

Dynamic normalization theory A number of recent works attempt to explain the dynamics and utility of dynamic normalization methods. The original works [20, 28] of BatchNorm and Weight-Norm resp. suggest that dynamic normalization methods improve training by fixing the intermediate layers' output distributions. The works of Bjorck et al. [8] and Santurkar et al. [29] argue that BatchNorm may improve optimization by improving smoothness of the Hessian of the loss, therefore allowing for larger step-sizes with reduced instability. Hoffer et al. [18] showed that the effective step-size in BatchNorm is divided by the magnitude of the weights, this followed the work of WNgrad [32] that introduces an adaptive step-size algorithm based on this fact. Following the intuition of WNGrad, Arora et al. [4] proved that for smooth loss and network functions, the diminishing "effective step-size" of normalization methods lead to convergence with optimal convergence rate for properly initialized step-sizes. The work of [22] explains the accelerated convergence of BatchNorm from a "length-direction decoupling" perspective. The authors along with [9] analyze the linear least squares regime, with [22] presenting a bisection method for finding the optimal weights. Robustness and regularization of Batch Normalization is investigated in [25] and improved generalization is analyzed empirically. Shortly after the original work of WeightNorm, [35] showed that for a single precptron WeightNorm may speed-up training and emphasized the importance of the norm of the initial weights. Additional stability properties were studied by [34] via mean-field analysis. The authors show that gradient instability is inevitable even with BatchNorm as the number of layers increases; this is in agreement with [7] for networks with residual connections. The work of [5] suggests initialization strategies for WeightNorm and derives lower bounds on the width to guarantee same order gradients across the layers.

Over-parametrized neural networks:

There has been a significant recent literature studying the convergence of un-normalized over-parametrized neural networks. The analysis of the majority of the works relies on the width of the layers. These include 2-layer networks trained with Gaussian inputs and outputs from a teacher network [30] , [24] and [14] (with WeightNorm). Assumptions on the data distribution are relaxed in [15] and the works that followed [3, 33, 36] . Our work is inspired by the mechanism presented in this chain of works. Wu et al. [33] extend convergence results to adaptive step-size methods and propose AdaLoss. Recently, the global convergence of over-parameterized neural networks was also extended to deep architectures [2, 13, 37, 38] . In the over-parameterized regimes, Arora et al. [3] develop generalization properties for the networks of the form (1.1). In addition, in the context of generalization, Allen-Zhu et al. [1] illustrates good generalization for deep neural networks trained with gradient descent. Cao and Gu [10] and [11] derive generalization error bound of gradient descent and stochastic gradient descent for learning over-parameterization deep ReLU neural networks.

6. DISCUSSION:

In this section we interpret our main results and make a connection of the convergence theory with "lazy training".

Chizat et al. [12] analyze recent works of over-parameterized convergence based on the NTK and note that re-scaling the network's outputs during gradient flow training leads to fast convergence while at the same time the weights remain close to their initialization. This is interpreted as a "linear solution", meaning that the direction of the weights does not change at training. The authors also refer to this type of convergence as "lazy training".

In our G-dominated convergence, we observe a new type of "lazy training" that is different from the one presented in [12] . We refer to our G regime as "lazy training" since the directions of the weights remain fairly stationary. Nonetheless, in G-dominated "lazy training", the magnitudes of the weights (g k ) change. Moreover, we observe this phenomenon not only in the continuous flow setting but also in the finite step gradient descent setting. Unlike [12] where it is argued that neural networks do not necessarily follow the "lazy training" regime, we believe that G-dominated convergence actually is common but that it emerges at later stages in training, after the weights have adopted their directions in the V-dominated regime and improves stability.

Overall training with WeightNorm leads to a gradual transition from V-dominated to the G-dominated regimes. We recall that since the weights grow (see (2.3)), the WeightNorm coefficient α increases during training progressively. The direction of the weights changes rapidly at the earlier stages of training when α is small, and G-dominated convergence ensues as α grows, leading to improved stability. For α > 1 this allows us to ease the requirements made on the over-parameterization (less is sufficient) and step size (a bigger step size is possible, of order η = O(1/ G(t) 2 )). From this perspective, the utility of WeightNorm is attributed to the allowed larger step-sizes and increased stability at the later stages of training (G-dominated), all while maintaining fast convergence at the beginning of training (V-dominated convergence).

Dynamic normalization is the most common optimization set-up of current deep learning models, yet understanding the convergence of such optimization methods is still an open problem. In this work we present a proof giving sufficient conditions for convergence of dynamically normalized 2-layer ReLU networks trained with gradient descent. To the best of our knowledge this is the first proof showcasing convergence of gradient descent training of neural networks with dynamic normalization and general data, where the objective function is non-smooth and non-convex. The notion of "length-direction decoupling" is clarified by the neural tangent kernel Λ(t) that naturally separates in our analysis into "length", G(t), and "direction", V(t)/α 2 , components. For α = 1 the decomposition initially matches un-normalized training. Yet we discover that in general, normalized training with gradient descent leads to 2 regimes dominated by different pieces of the neural tangent kernel. We note that training typically commences in the V-dominated regime and transitions into the G-dominated regime as training proceeds and the magnitude of the weights grow. Our improved analysis is able to reduce the amount of over-parameterization that was needed in previous convergence works in the un-normalized setting and in the G-dominated regime, we prove convergence with a significantly lower amount of over-parameterization as compared with un-normalized training.

APPENDIX:

We present the detailed proofs of the main results of the paper below. The appendix is organized as follows. We provide proofs to the simple propositions regarding the NTK presented in the paper in Appendix A, and prove the main results for V-dominated and G-dominated convergence in the settings of gradient flow and gradient descent in Appendices B,C. The proofs for gradient flow and gradient descent share the same main idea, yet the proof for gradient descent has a considerate number of additional technicalities. In Appendices D and E we prove the lemmas used in the analysis of Appendices B and C respectively.

A WEIGHTNORM DYNAMICS PROOFS:

In this section we provide proofs for Proposition 1, which describes the relation between vanilla and WeightNorm NTKs and Observation 1 of the paper.

Proof of Proposition 1::

We would like to show that V(0) + G(0) = H(0). For each entry, consider

(V(0) + G(0)) ij = 1 m m k=1 x v k (0) ⊥ i , x j v k (0) ⊥ 1 ik (0)1 jk (0) + 1 m m k=1 x v k (0) i , x j v k (0) 1 ik (0)1 jk (0).

Note that

x i , x j = x v k (0) i + x v k (0) ⊥ i , x v k (0) j + x v k (0) ⊥ j = x v k (0) ⊥ i , x j v k (0) ⊥ + x v k (0) i , x j v k (0) .

This gives

(V(0) + G(0)) ij = 1 m m k=1 x i , x j 1 ik (0)1 jk (0) = H ij (0)

which proves the claim.

Proof of Observation 1::

We show that the initialization of the network is independent of α. Take α, β > 0, and for each k,

initialize v α k , v β k as v α k (0) ∼ N (0, α 2 I), v β k (0) ∼ N (0, β 2 I). Then v α k (0) v α k (0) 2 ∼ v β k (0) v β k (0) 2 ∼ Unif(S d-1 ) (in distribution).

Hence the distribution of each neuron σ v k (0) v k (0) 2 at initialization is independent of α. Next for g k (0), we note that

EQUATION

g α k (0)v α k (0) v α k (0) 2 ∼ g β k (0)v β k (0) v β k (0) 2 ∼ N (0, I),

for all α, β. This shows that the network initialization is independent of α and is equivalent to the initialization of the un-normalized setting. Similarly, inspecting the terms in the summands of V(0), G(0) shows that they are also independent of α. For

V ij (0) = 1 m m k=1 1 ik (0)1 jk (0) αc k • g k (0) v k (0) 2 2 x v k (0) ⊥ i , x v k (0) ⊥ j the terms 1 ik (0), x v k (0) ⊥ i

are independent of scale, and the fraction in the summand is identically 1. G(0) defined as

G ij (0) = 1 m m k=1 1 ik (0)1 jk (0) x v k (0) i , x v k (0) j

is also invariant of scale since the projection onto a vector direction v k (0) is independent of scale.

B CONVERGENCE PROOF FOR GRADIENT FLOW:

In this section we derive the convergence results for gradient flow.

The main results are analogous to Theorems 4. 

f (t) -y 2 2 ≤ exp(-λ 0 t/α 2 ) f (0) -y 2 2 .

This theorem is analogous to Theorem 4.1 but since here, the settings are of gradient flow there is no mention of the step-size. It is worth noting that smaller α leads to faster convergence and appears to not affect the other hypotheses of the flow theorem. This "un-interuptted" fast convergence behavior does not extend to finite-step gradient descent where the increased convergence rate is balanced by decreasing the allowed step-size.

The second main result for gradient flow is for G-dominated convergence. Theorem B.2 (G-dominated convergence). Suppose a network from the class (1.2) is initialized as in (2.4) with α > 1 and that assumptions 1, 2 hold. In addition, suppose the neural network is trained on the regression loss (2.5) with target y satisfying y ∞ = O(1). Then if m = Ω max n 4 log(n/δ)/α 4 µ 4 0 , n 2 log(n/δ)/µ 2 0 , WeightNorm training with gradient flow converges at a linear rate, with probability 1 -δ, as

f (t) -y 2 2 ≤ exp(-µ 0 t) f (0) -y 2 2 .

B.1 PROOF SKETCH:

To prove the results above we follow the steps introduced in the proof sketch of Section 4. The main idea of the proofs for V and G dominated convergence are analogous and a lot of the proofs are based of Du et al. [15] . We show that in each regime, we attain linear convergence by proving that the least eigenvalue of the evolution matrix Λ(t) is strictly positive. For the V-dominated regime we lower bound the least eigenvalue of Λ(t) as λ min (Λ(t)) ≥ λ min (V(t))/α 2 and in the G-dominated regime we lower bound the least eigenvalue as λ min (Λ(t)) ≥ λ min (G(t)).

The main part of the proof is showing that λ min (V(t)), λ min (G(t)) stay uniformly positive. We use several lemmas to show this claim.

In each regime, we first show that at initialization the kernel under consideration, V(0) or G(0), has a positive least eigenvalue. This is shown via concentration to an an auxiliary kernel (Lemmas B.1, B.2), and showing that the auxiliary kernel is also strictly positive definite (Lemma 4.1). Lemma B.1. Let V(0) and V ∞ be defined as in (3.3) and (4.2), assume the network width m satisfies m = Ω n 2 log(n/δ)

λ 2 0

. Then with probability 1 -δ,

V(0) -V ∞ 2 ≤ λ 0 4 .

Lemma B.2. Let G(0) and G ∞ be defined as in (3.4 ) and (4.3), assume m satisfies m = Ω n 2 log(n/δ)

µ 2 0

. Then with probability 1 -δ,

G(0) -G ∞ 2 ≤ µ 0 4 .

After showing that V(0), G(0) have a positive least-eigenvalue we show that V(t), G(t) maintain this positive least eigenvalue during training. This part of the proof depends on the over-parameterization of the networks. The main idea is showing that if the individual parameters v k (t), g k (t) do not change too much during training, then V(t), G(t) remain close enough to V(0), G(0) so that they are still uniformly strictly positive definite. We prove the results for V(t) and G(t) separately since each regime imposes different restrictions on the trajectory of the parameters.

For now, in Lemmas B.3, B.4, B.5, we make assumptions on the parameters of the network not changing "too much"; later we show that this holds and is the result of over-parameterization. Specifically, over-parameterization ensures that the parameters stay at a small maximum distance from their initialization.

V-dominated convergence To prove the least eigenvalue condition on V(t), we introduce the surrogate Gram matrix V(t) defined entry-wise as

Vij (t) = 1 m m k=1 x v k (t) ⊥ i , x v k (t) ⊥ j 1 ik (t)1 jk (t). (B.1)

This definition aligns with V(t) if we replace the scaling term αc k g k (t)

v k (t) 2

2 in each term in the sum

V ij (t) by 1.

To monitor 

V(t) -V(0) we consider V(t) -V(0) and V(t) -V(t) in Lemmas B.

V(t) -V(0) 2 ≤ λ 0 8 . Lemma B.4. Define R g = λ 0 48n(m/δ) 1/d , R v = αλ 0 96n(m/δ) 1/d . (B.2)

Suppose the conditions of Lemma B.3 hold, and that

v k (t)-v k (0) 2 ≤ R v , g k (t)-g k (0) 2 ≤ R g for all 1 ≤ k ≤ m. Then with probability 1 -δ, V(t) -V(0) 2 ≤ λ 0 4 .

G-dominated convergence:

We ensure that G(t) stays uniformly positive definite if the following hold. Lemma B.5. Given v 1 (0), . . . , v k (0) generated i.i.d. as in (2.4), suppose that for each k,

v k (t) - v k (0) 2 ≤ √ 2παµ0

8n(m/δ) 1/d =: Rv , then with probability 1 -δ,

G(t) -G(0) 2 ≤ µ 0 4 .

After deriving sufficient conditions to maintain a positive least eigenvalue at training, we restate the discussion of linear convergence from Section 4 formally.

Lemma B.6. Consider the linear evolution df dt = -G(t)

+ V(t) α 2 (f (t) -y) from (3.5). Suppose that λ min G(t) + V(t) α 2 ≥ ω 2 for all times 0 ≤ t ≤ T . Then f (t) -y 2 2 ≤ exp(-ωt) f (0) -y 2 2

for all times 0 ≤ t ≤ T .

Using the linear convergence result of Lemma B.6, we can now bound the trajectory of the parameters from their initialization.

Lemma B.7. Suppose that for all 0 ≤ t ≤ T , λ min G(t)

+ 1 α 2 V(t) ≥ ω 2 and |g k (t) -g k (0)| ≤ R g ≤ 1/(m/δ) 1/d . Then with probability 1 -δ over the initialization v k (t) -v k (0) 2 ≤ 4 √ n f (0) -y 2 αω √ m =: R v (B.3)

for each k and all times 0 ≤ t ≤ T .

Lemma B.8. Suppose that for all 0 ≤ t ≤ T , λ min G(t)

+ 1 α 2 V(t) ≥ ω 2 .

Then with probability 1 -δ over the initialization

|g k (t) -g k (0)| ≤ 4 √ n f (0) -y 2 √ mω =: R g

for each k and all times 0 ≤ t ≤ T.

The distance of the parameters from initialization depends on the convergence rate (which depends on λ min (Λ(t))) and the width of the network m. We therefore are able to find sufficiently large m for which the maximum parameter trajectories are not too large so that we have that the least eigenvalue of Λ(t) is bounded from 0; this proves the main claim.

Before proving the main results in the case of gradient flow, we use two more technical lemmas. Lemma B.9. Suppose that the network is initialized as (2.4) and that y ∈ R n has bounded entries |y i | ≤ M . Then f (0) -y 2 ≤ C n log(n/δ) for some absolute constant C > 0. Lemma B.10 (Failure over initialization). Suppose v 1 (0), . . . , v k (0) are initialized i.i.d. as in (2.4) with input dimension d. Then with probability 1 -δ,

max k∈[m] 1 v k (0) 2 ≤ (m/δ) α 1/d .

In addition by (2.3), for all t ≥ 0, with probability 1 -δ,

max k∈[m] 1 v k (t) 2 ≤ (m/δ) α 1/d .

Remark (Assumption 2). Predominately, machine learning applications reside in the high dimensional regime with d ≥ 50. Typically d 50. This therefore leads to an expression (m/δ) 1/d that is essentially constant. For example, if d = 50, for max k∈[m] 1 v k (0) 2 ≥ 10, one would need m/δ ≥ 10 80 (the tail of χ 2 d also has a factor of (d/2)! • 2 d/2 which makes the assumption even milder). The term (m/δ) 1/d therefore may be taken as a constant for practicality,

max k∈[m] 1 v k (0) 2 ≤ C α .

While we make Assumption 2 when presenting our final bounds, for transparency we do not use Assumption 2 during our analysis and apply it only when we present the final over-parameterization results to avoid the overly messy bound. Without the assumption the theory still holds yet the over-parameterization bound worsens by a power 1 + 1/(d -1). This is since the existing bounds can be modified, replacing m with m 1-1 d .

Proof of Theorem B.1::

Note that m = Ω n 4 log(n/δ)/λ 4 0 implies that Lemma B.1 holds. Further since the gradient flow updates are continous (since the time derivative with respect to each parameter may be bounded), there exist a small time T for which the parameters are bounded from initialization for 0 ≤ t ≤ T and for all 1 ≤ k ≤ m,

v k (t) -v k (0) 2 ≤ R v , |g k (t) -g k (0)| ≤ R g .

The bounded trajectory above, along with the over-parameterization, ensures that with probability 1 -δ over the initialization that Lemma B.4 holds and that

λ min G(t) + 1 α 2 V(t) ≥ λ min (V(t))/α 2 ≥ λ 0 2α 2

for times 0 ≤ t ≤ T . The condition on the eigenvalue of the evolution matrix along with the bounded trajectory of g k implies that Lemmas B.7, B.8 hold for at least times 0 ≤ t < T . Define T 0 to be the first failure point of either Lemmas B.8, B.7, we have that T 0 ≥ T > 0. For 0 ≤ t < T 0 , by substituting m = Ω n 4 log(n/δ)/λ 4 0 and using the bound on f (0) -y 2 of Lemma B.9, a direct calculation utilizing Lemma B.7 shows that

v k (t) -v k (0) 2 B.7 ≤ α √ n f (0) -y 2 √ mλ 0 ≤ R v .

Similarly by Lemma B.8 m ensures that

|g k (t) -g k (0)| B.8 ≤ α 2 √ n f (0) -y 2 √ mλ 0 ≤ R g .

The over-parameterization of m implies that the parameter trajectories stay close enough to initialization to satisfy the hypotheses of Lemmas B.3, B.4 and that λ min (Λ(t)) ≥ λ min (V(t))/α 2 ≥ λ0 2α 2 at time T 0 . This implies that Lemmas B.7, B.8 both hold at time T 0 which contradicts the definition of T 0 . Therefore we conclude that Lemmas B.7, B.8 hold for t > 0 which implies that Lemma B.4 holds for t > 0 hence Lemma B.6 guarantees linear convergence.

Here we consider the case where the convergence is dominated by G. This occurs when α > 1.

Proof of Theorem B.2::

Similarly to Theorem B.1, we note that m = Ω max n 4 log(n/δ)/α 4 µ 4 0 , n 2 log(n/δ)/µ 2 0 implies that Lemma B.2 holds. Further since the gradient flow updates are continous (since the time derivative with respect to each parameter may be bounded), there exist small time T for which the parameters are bounded from initialization for 0 ≤ t ≤ T and for all

1 ≤ k ≤ m v k (t) -v k (0) 2 ≤ Rv .

The bounded trajectory above, along with the over-parameterization, ensures that with probability 1 -δ over the initialization that Lemma B.5 holds and that

λ min G(t) + 1 α 2 V(t) ≥ λ min (G(t))/α 2 ≥ µ 0 2

for times 0 ≤ t ≤ T . The condition on the eigenvalue of the evolution matrix along with the bounded trajectory of g k implies that Lemmas B.7, B.8 hold for at least times 0 ≤ t < T . Define T 0 to be the first failure point of either Lemmas B.8, B.7, we have that T 0 ≥ T > 0. For 0 ≤ t < T 0 , by substituting m = Ω max n 4 log(n/δ)/α 4 µ 4 0 , n 2 log(n/δ)/µ 2 0 and using the bound on f (0) -y 2 of Lemma B.9, a direct calculation utilizing Lemma B.7 shows that

v k (t) -v k (0) 2 B.7 ≤ 4 √ n f (0) -y 2 αµ 0 √ m B.9 ≤ Cn log(n/δ) αµ 0 √ m ≤ Rv .

Similarly by Lemma B.8 m ensures that

|g k (t) -g k (0)| B.8 ≤ √ n f (0) -y 2 √ mµ 0 ≤ R g .

The over-parameterization of m implies that the parameter trajectories stay close enough to initialization to satisfy the hypotheses of Lemmas B.5 and that λ min (Λ(t)) ≥ λ min (G(t))/α 2 ≥ µ 0 /2 at time T 0 . This implies that Lemmas B.7, B.8 both hold at time T 0 which contradicts the definition of T 0 . Therefore we conclude that Lemmas B.7, B.8 hold for t > 0 which implies that Lemma B.5 holds for t > 0 hence Lemma B.6 guarantees linear convergence with rate µ 0 /2.

Note that if α is large, the required complexity on m is reduced. Taking α = Ω( n/µ 0 ) gives the improved bound

m = Ω n 2 log (n/δ) µ 2 0 .

C FINITE STEP-SIZE TRAINING:

The general technique of proof for gradient flow extends to finite-step gradient descent. Nonethless, proving convergence for WeightNorm gradient descent exhibits additional complexities arising from the discrete updates and joint training with the new parameterization (1.2). We first introduce some needed notation.

Define S i (R) as the set of indices k ∈ [m] corresponding to neurons that are close to the activity boundary of ReLU at initialization for a data point x i ,

S i (R) := {k ∈ [m] : ∃ v with v -v k (0) 2 ≤ R and 1 ik (0) = 1{v x i ≥ 0}}.

We decreases linearly for some convergence rate ω during gradient descent training for all iteration steps s = 0, 1, . . . , K with step-size η as

f (s) -y 2 2 ≤ (1 -ηω 2 ) s f (0) -y 2 2 .

Then for each k we have

|g k (s) -g k (0)| ≤ 4 √ n f (0) -y 2 √ mω

for iterations s = 0, 1, . . . , K + 1. Lemma C.3. Under the assumptions of Lemma C.2, suppose in addition that |g k (s) -g k (0)| ≤ 1/(m/δ) 1/d for all iterations steps s = 0, 1, . . . K . Then for each k,

v k (s) -v k (0) 2 ≤ 8 √ n f (0) -y 2 α √ mω for s = 0, 1, . . . , K + 1.

To prove linear rate of convergence we analyze the s + 1 iterate error f (s + 1) -y 2 relative to that of the s iterate, f (s) -y 2 . Consider the network's coordinate-wise difference in output between iterations, f i (s + 1) -f i (s), writing this explicitly based on gradient descent updates yields

f i (s + 1) -f i (s) = 1 √ m m k=1 c k g k (s + 1) v k (s + 1) 2 σ(v k (s + 1) x i ) - c k g k (s) v k (s) 2 σ(v k (s) x i ). (C.1)

We now decompose the summand in (C.1) looking at the updates in each layer,

f i (s + 1) -f i (s) = a i (s) + b i (s) with a i (s) = 1 √ m m k=1 c k g k (s + 1) v k (s + 1) 2 σ(v k (s) x i ) - c k g k (s) v k (s) 2 σ(v k (s) x i ), b i (s) = 1 √ m m k=1 c k g k (s + 1) v k (s + 1) 2 σ(v k (s + 1) x i ) -σ(v k (s) x i ) .

Further, each layer summand is then subdivided into a primary term and a residual. a i (s), corresponding to the difference in the first layer

c k g k (s+1) v k (s+1) 2 -c k g k (s) v k (s) 2

, is subdivided into a I i (s) and a II i (s) as follows:

a I i (s) = 1 √ m m k=1 c k g k (s + 1) v k (s) 2 - c k g k (s) v k (s) 2 σ(v k (s) x i ), (C.2)

a II i (s) = 1 √ m m k=1 c k g k (s + 1) v k (s + 1) 2 - c k g k (s + 1) v k (s) 2 σ(v k (s) x i ). (C.3) b i (s)

is sub-divided based on the indices in the set S i that monitor the changes of the rectifiers. For now, S i = S i (R) with R to be set later in the proof. b i (s) is partitioned to summands in the set S i and the complement set,

b I i (s) = 1 √ m k ∈Si c k g k (s + 1) v k (s + 1) 2 σ(v k (s + 1) x i ) -σ(v k (s) x i ) , b II i (s) = 1 √ m k∈Si c k g k (s + 1) v k (s + 1) 2 σ(v k (s + 1) x i ) -σ(v k (s) x i ) .

With this sub-division in mind, the terms corresponding to convergence are a I (s), b I (s) whereas a II (s), b II (s) are residuals that are the result of discretization. We define the primary and residual vectors p(s), r(s) as

p(s) = a I (s) + b I (s) η , r(s) = a II + b II (s) η . (C.4)

If the residual r(s) is sufficiently small and p(s) may be written as p(s) = -Λ(s)(f (s) -y) for some iteration dependent evolution matrix Λ(s) that has

λ min (Λ(s)) = ω/2 (C.5)

for ω > 0 then the neural network (1.2) converges linearly when trained with WeightNorm gradient descent of step size η. We formalize the condition on r(s) below and later derive the conditions on the over-parameterization (m) ensuring that r(s) is sufficiently small. Property 1. Given a network from the class (1.2) initialized as in (2.4) and trained with gradient descent of step-size η, define the residual r(s) as in (C.4) and take ω as in (C.5). We specify the "residual condition" at iteration s as

r(s) 2 ≤ cω f (s) -y 2

for a sufficiently small constant c > 0 independent of the data or initialization.

Here we present Theorem C.1 which is the backbone of Theorems 4.1 and 4.2.

Theorem C.1. Suppose a network from the class (1.2) is trained via WeightNorm gradient descent with an evolution matrix Λ(s) as in (C.5) satisfying λ min (Λ(s)) ≥ ω/2 for s = 0, 1, . . . K. In addition if the data meets assumptions 1, 2, the step-size η of gradient descent satisfies η ≤

1 3 Λ(s) 2

and that the residual r(s) defined in (C.4) satisfies Property 1 for s = 0, 1, . . . , K then we have that

f (s) -y 2 2 ≤ 1 - ηω 2 s f (0) -y 2 2

for s = 0, 1, . . . , K.

Proof of Theorem C.1::

This proof provides the foundation for the main theorems. In the proof we also derive key bounds to be used in Theorems 4.1, 4.2. We use the decomposition we described above and consider again the difference between consecutive terms f (s + 1) -f (s),

f i (s + 1) -f i (s) = 1 √ m m k=1 c k g k (s + 1) v k (s + 1) 2 σ(v k (s + 1) x i ) - c k g k (s) v k (s) 2 σ(v k (s) x i ). (C.6)

Following the decompostion introduced in (C.2), a I i (s) is re-written in terms of G(s),

a I i (s) = 1 √ m m k=1 c k v k (s) 2 -η ∂L(s) ∂g k σ(v k (s) x i ) = - η m m k=1 c k v k (s) 2 n j=1 (f j (s) -y j ) c k v k (s) 2 σ(v k (s)x j )σ(v k (s)x i ) = -η n j=1 (f j (s) -y j ) 1 m m k=1 (c k ) 2 σ v k (s) x i v k (s) 2 σ v k (s) x j v k (s) 2 = -η n j=1 (f j (s) -y j )G ij (s),

where the equality holds by the gradient update rule g k (s + 1) = g k (s) -η∇ g k L(s). In this proof we also derive bounds on the residual terms of the decomposition which we will aid us in the proofs of Theorems 4.1, 4.2. a I i (s) is the primary term of a i (s), now we bound the residual term a II i (s). Recall a II i (s) is written as

a II i (s) = 1 √ m m k=1 c k g k (s + 1) v k (s + 1) 2 - c k g k (s + 1) v k (s) 2 σ(v k (s) x i ),

which corresponds to the difference in the normalization in the second layer. Since ∇ v k L(s) is orthogonal to v k (s) we have that

c k g k (s + 1) 1 v k (s + 1) 2 - 1 v k (s) 2 σ(v k (s) x i ) = c k g k (s + 1) 1 v k (s) 2 2 + η 2 ∇ v k L(s) 2 2 - 1 v k (s) 2 σ(v k (s) x i ) = -c k g k (s + 1)η 2 ∇ v k L(s) 2 2 v k (s + 1) 2 v k (s) 2 ( v k (s) 2 + v k (s + 1) 2 ) σ(v k (s) x i ) ≤ -c k g k (s + 1)η 2 ∇ v k L(s) 2 2 2 v k (s) 2 v k (s + 1) 2 σ v k (s) x i v k (s) 2 ,

where the first equality above is by completing the square, and the inequality is due to the increasing magnitudes of v k (s) 2 .

Since

0 ≤ σ v k (s) xi v k (s) 2

≤ 1, the above can be bounded as

|a II i (s)| ≤ 1 √ m m k=1 g k (s + 1)η 2 ∇ v k L(s) 2 2 2 v k (s) 2 v k (s + 1) 2 ≤ 1 √ m m k=1 η 2 1 + R g (m/δ) 1/d 3 n f (s) -y 2 2 (m/δ) 1/d α 4 m = η 2 n 1 + R g (m/δ) 1/d 3 f (s) -y 2 2 (m/δ) 1/d α 4 √ m . (C.7)

The second inequality is the result of applying the bound in equation (E.1) on the gradient norm ∇ v k L(s) 2 and using Lemma B.10.

Next we analyze b i (s) and sub-divide it based on the sign changes of the rectifiers. Define the set S i := S i (R) as in Lemma C.1 with R taken to be such that v k (s + 1) -v k (0) 2 ≤ R for all k. Take b II i (s) as the sub-sum of b i (s) with indices k from the set S i . b I i (s) corresponds to the sub-sum with the remaining indices. By the definition of S i , for k ∈ S i we have that 1 ik (s + 1) = 1 ik (s). This enables us to factor 1 ik (s) and represent b I i (s) as a Gram matrix similar to V(s) with a correction term from the missing indices in S i .

b I i (s) = - 1 √ m k ∈Si c k g k (s + 1) v k (s + 1) 2 η ∇ v k L(s), x i 1 ik (s) = - η m k ∈Si c k g k (s + 1) v k (s + 1) 2 c k g k (s) v k (s) 2 n j=1 (f j (s) -y j )1 ik (s)1 jk (s) x v k (s) ⊥ j , x i . Note that x v k (s) ⊥ j , x i = x v k (s) ⊥ j , x v k (s) ⊥ i therefore, b I i (s) = - η m k ∈Si c k g k (s + 1) v k (s + 1) 2 c k g k (s) v k (s) 2 n j=1 (f j (s) -y j )1 ik (s)1 jk (s) x v k (s) ⊥ j , x v k (s) ⊥ i . Define Ṽ(s) as Ṽij (s) = 1 m m k=1 αc k g k (s + 1) v k (s + 1) 2 αc k g k (s) v k (s) 2 1 jk (s)1 ik (s) x v k (s) ⊥ i , x v k (s) ⊥ j .

This matrix is identical to V(s) except for a modified scaling term

c 2 k g k (s+1)g k (s) v k (s) 2 v k (s+1) 2 .

We note however that

min c k g k (s) v k (s) 2 2 , c k g k (s + 1) v k (s + 1) 2 2 ≤ c k g k (s) v k (s) 2 c k g k (s + 1) v k (s + 1) 2 ≤ max c k g k (s) v k (s) 2 2 , c k g k (s + 1) v k (s + 1) 2 2

because g k (s), c 2 k are positive. Hence the matrix Ṽ(s) satisfies the hypothesis of Lemma B.4 entirely. We write b I i (s) as

b I i (s) = -η/α 2 n j=1 (f j (s) -y j )( Ṽij (s) -Ṽ⊥ ij (s)),

where we have defined

Ṽ⊥ ij (s) = 1 m k∈Si αc k g k (s) v k (s) 2 αc k g k (s + 1) v k (s + 1) 2 1 ik (s)1 jk (s) x v k (s) ⊥ i , x v k (s) ⊥ j . (C.8)

We then bound the magnitude of each entry Ṽ⊥ ij (s):

Ṽ⊥ ij (s) = 1 m k∈Si αc k g k (s) v k (s) 2 αc k g k (s + 1) v k (s + 1) 2 1 ik (s)1 jk (s) x v k (s) ⊥ i , x v k (s) ⊥ j ≤ (1 + R g (m/δ) 1/d ) 2 |S i | m . (C.9)

Lastly we bound the size of the residual term b II i (s),

|b II i (s)| = - 1 √ m k∈Si c k g k (s + 1) v k (s + 1) 2 σ(v k (s + 1) x i ) -σ(v k (s) x i ) ≤ g k (s + 1)η|S i | • ∇ v k L(s) 2 √ m v k (s + 1) 2 ≤ η|S i |(1 + (m/δ) 1/d R g ) ∇ v k L(s) 2 α √ m .

Where we used the Lipschitz continuity of σ in the first bound, and took R g > 0 that satisfies |g k (s + 1) -g k (0)| ≤ R g in the second inequality. Applying the bound (E.1),

|b II i (s)| ≤ η|S i | √ n(1 + R g (m/δ) 1/d ) 2 f (s) -y 2 α 2 m . (C.10)

The sum f (s + 1) -f (s) = a I (s) + a II (s) + b I (s) + b II (s) is separated into the primary term ηp(s) = a I (s) + b I (s) and the residual term ηr(s) = a II (s) + b II (s) which is a result of the discretization. With this, the evolution matrix Λ(s) in (C.5) is re-defined as

Λ(s) := G(s) + Ṽ(s) -Ṽ⊥ (s) α 2 and f (s + 1) -f (s) = -ηΛ(s)(f (s) -y) + ηr(s). Now we compare f (s + 1) -y 2 2 with f (s) -y 2 2 , f (s + 1) -y 2 2 = f (s + 1) -f (s) + f (s) -y 2 2 = f (s) -y 2 2 + 2 f (s + 1) -f (s), f (s) -y + f (s + 1) -f (s), f (s + 1) -f (s) . Substituting f (s + 1) -f (s) = a I (s) + b I (s) + a II (s) + b II (s) = -ηΛ(s)(f (s) -y) + ηr(s) we obtain f (s + 1) -y 2 2 = f (s) -y 2 2 + 2(-ηΛ(s)(f (s) -y) + ηr(s)) (f (s) -y) + η 2 (Λ(s)(f (s) -y) -r(s)) (Λ(s)(f (s) -y) -r(s)) ≤ f (s) -y 2 2 + (f (s) -y) (-ηΛ(s) + η 2 Λ 2 (s))(f (s) -y) + ηr(s) (I -ηΛ(s))(f (s) -y) + η 2 r(s) 2 2 .

Now as λ min (Λ(s)) ≥ ω/2 and η = 1 3 Λ(s) 2 , we have that

(f (s) -y) (-ηΛ(s) + η 2 Λ 2 (s))(f (s) -y) = -η(f (s) -y) (I -ηΛ(s))Λ(s)(f (s) -y) ≤ - 3ηω 8 f (s) -y 2 2 .

Next we analyze the rest of the terms and group them as q(s), q(s) := ηr(s

) (I -ηΛ(s))(f (s) -y) + η 2 r(s) 2 2 ≤ η r(s) 2 f (s) -y 2 + η 2 r(s) 2 2 .

By Property 1 we have

q(s) ≤ ηcω f (s) -y 2 2 (1 + ηcω) ≤ 2cηω f (s) -y 2 2 ,

so that q(s) ≤ c ηω f (s) -y 2 2 , for c sufficiently small. Substituting, the difference f (s + 1) -y is bounded as

f (s + 1) -y 2 2 ≤ f (s) -y 2 2 -ηω(1 -η Λ(s) 2 ) f (s) -y 2 2 + c ηω f (s) -y 2 2 ≤ (1 -ηω(1 -η Λ(s) 2 ) + c ηω) f (s) -y 2 2 ≤ (1 -ηω/2) f (s) -y 2 2

, for well chosen absolute constant c. Hence for each s = 0, 1, . . . , K,

f (s + 1) -y 2 2 ≤ (1 -ηω/2) f (s) -y 2 2

so the prediction error converges linearly.

In what comes next we prove the necessary conditions for Property 1, and define the appropriate ω for the V and G dominated regimes, in order to show λ min (Λ(s)) ≥ ω/2.

Proof of Theorem 4.1::

To prove convergence we would like to apply Theorem C.1 with ω/2 = λ0 2α 2 . To do so we need to show that m = Ω n 4 log(n/δ)/λ 4 0 guarantees that Property 1 holds and that λ min (Λ(s)) ≥ λ 0 /2α 2 . For finite-step gradient training, take

R v = αλ 0 192n(m/δ) 1/d , R g = λ 0 96n(m/δ) 1/d . (C.11)

Note the residual r(s) and the other terms b I (s), b II (s) depend on the sets S i that we define here using R v . We make the assumption that

v k (s) -v k (0) 2 ≤ R v and |g k (s) -g k (0)| ≤ R g for all k

and that s = 0, 1, . . . K + 1, this guarantees that b I (s) and Λ(s) are well defined. Applying Lemmas B.1, B.4 with R v , R g defined above, we have that λ min ( Ṽ(s)) ≥ 5λ0 8 . Then the least eigenvalue of the evolution matrix Λ(s) is bounded below

λ min (Λ(s)) = λ min G(s) + Ṽ(s) -Ṽ⊥ (s) α 2 ≥ λ min Ṽ(s) -Ṽ⊥ (s) α 2 = λ min ( Ṽ(s) -Ṽ⊥ (s)) α 2 ≥ 5λ 0 8α 2 - Ṽ⊥ (s) 2 α 2 .

The first inequality holds since G(s) 0 and the last inequality is since λ min ( Ṽ(s)) ≥ 5λ0 8 . To show λ min (Λ(s)) ≥ λ0 2α 2 we bound Ṽ⊥ (s) 2 ≤ λ0 8 . By (C.9), we have

| Ṽ⊥ ij (s)| ≤ (1 + R g (m/δ) 1/d )|S i | m ≤ (1 + R g (m/δ) 1/d ) √ 2 Rv √ πα + 16 log(n/δ) 3m .

Substituting R v , R g and m, a direct calculation shows that

| Ṽ⊥ ij (s)| ≤ λ 0 8n , which yields Ṽ⊥ (s) 2 ≤ Ṽ⊥ (s) F ≤ λ 0 8 .

Hence λ min (Λ(s)) ≥ λ0 2α 2 for iterations s = 0, 1, . . . K.

We proceed by showing the residual r(s) satisfies property 1. Recall r(s) is written as

r(s) = a II (s) η + b II (s) η .

and Property 1 states that r(s) 2 ≤ cηλ0 α 2 f (s) -y 2 for sufficiently small absolute constant c < 1. This is equivalent to showing that both a II (s), b II (s) satisfy

a II (s) 2 ≤ cηλ 0 α 2 f (s) -y 2 , (C.12) b II (s) 2 ≤ cηλ 0 α 2 f (s) -y 2 . (C.13)

We consider each term at turn. By (C.10),

b II (s) 2 ≤ √ n max i b II i (s) ≤ max i ηn(1 + R g (m/δ) 1/d ) 2 |S i | f (s) -y 2 α 2 m ≤ CmR v ηn f (s) -y 2 α 2 m ≤ λ 0 η f (s) -y 2 α 2 • nCR v .

In the above we used the values of R v , R g defined in (C.11) and applied Lemma C.1 in the third inequality. Taking m = Ω n 4 log(n/δ)/λ 4 0 with large enough constant yields

b II (s) 2 ≤ cλ 0 η f (s) -y 2 α 2 .

Next we analogously bound a II (s) via the bound (C.7),

a II (s) 2 ≤ √ n max i a II i (s) ≤ η 2 n 3/2 1 + R g (m/δ) 1/d 3 f (s) -y 2 2 (m/δ) 1/d α 4 √ m ≤ ηλ 0 f (s) -y 2 α 2 • η 1 + R g (m/δ) 1/d 3 n 3/2 f (s) -y 2 (m/δ) 1/d λ 0 α 2 √ m ≤ ηλ 0 f (s) -y 2 α 2 • η α 2 • Cn 2 log(n/δ) λ 0 √ m ≤ cηω f (s) -y 2 .

In the above we applied Lemma B.9 once again. The last inequality holds since m = Ω(n 4 log(n/δ)/λ 4 0 ) and η = O α 2

V(s) 2 , hence r(s) satisfies Property 1. Now since Theorem C.1 holds with ω = λ 0 /α 2 we have that the maximum parameter trajectories are bounded as

v k (s) -v k (0) 2 ≤ R v and g k (s) -g k (0) ≤ R g for all k and every iteration s = 0, 1, . . . , K + 1 via Lemmas C.2, C.3.

To finish the proof, we apply the same contradiction argument as in Theorems B.1, B.2, taking the first iteration s = K 0 where one of Lemmas C.2, C.3 does not hold. We note that K 0 > 0 and by the definition of K 0 , for s = 0, 1, . . . , K 0 -1 the Lemmas C.2, C.3 hold which implies that by the argument above we reach linear convergence in iteration s = K 0 . This contradicts one of Lemmas C.2, C.3 which gives the desired contradiction, so we conclude that we have linear convergence with rate λ 0 /2α 2 for all iterations.

Proof of Theorem 4.2::

For G-dominated convergence, we follow the same steps as in the proof of Theorem 4.1. We redefine the trajectory constants for the finite step case

Rv := √ 2παµ 0 64n(m/δ) 1/d , R g := µ 0 48n(m/δ) 1/d .

To use Theorem C.1 we need to show that m = Ω n 4 log(n/δ)/α 4 µ 4 0 guarantees Property 1, and that λ min (Λ(s)) ≥ µ 0 /2. We again note that the residual r(s) and b I (s), b II (s) depend on the sets S i that we define here using Rv above as S i := S i ( Rv ).

We start by showing the property on the least eigenvalue. We make the assumption that we have linear convergence with ω/2 = µ 0 /2 and step-size η for iterations s = 0, . . . K so that Lemmas C.2, C.3 hold. Via an analogous analysis of the continous case we reach that m = Ω n 4 log(n/δ)/µ 4 0 α 4 implies

v k (s) -v k (0) 2 ≤ 16α √ n f (0) -y 2 α √ mµ 0 ≤ Rv , |g k (s) -g k (0)| ≤ 8 √ n f (0) -y 2 √ mµ 0 ≤ R g .

for s = 0, . . . K + 1 by Lemmas C.2, C.3 and that Λ(s), b I (s) are well defined. Using the bounds on the parameter trajectories, Lemma B.5 and Rv defined above yield λ min (G(s)) ≥ 5µ0 8 . The least eigenvalue of the evolution matrix Λ(s) is bounded below as

λ min (Λ(s)) = λ min G(s) + Ṽ(s) -Ṽ⊥ (s) α 2 ≥ λ min (G(s)) -Ṽ⊥ (s) 2

since Ṽ(s) 0 and α ≥ 1. We bound the spectral norm of Ṽ⊥ (s) 2 , for each entry i, j we have by (C.9) that

| Ṽ⊥ ij (s)| ≤ (1 + R g (m/δ) 1/d )|S i | m ≤ (1 + R g (m/δ) 1/d ) √ 2 Rv √ πα + 16 log(n/δ) 3m ≤ 8 Rv √ 2πα ≤ µ 0 8n .

where in the above inequalities we used our bounds on Rv , R g and m. Then the spectral norm is bounded as

Ṽ⊥ (s) 2 ≤ Ṽ⊥ (s) F ≤ µ 0 /8.

Hence we have that λ min (Λ(s)) ≥ µ 0 /2 for s = 0, 1, . . . K.

Next we show the residual r(s) satisfies Property 1. Recall r(s) is written as

r(s) = a II (s) η + b II (s) η .

Property 1 states the condition r(s) 2 ≤ cωη f (s) -y 2 for sufficiently small c < 1 with ω = µ 0 . This is equivalent to showing that both a II (s), b II (s) satisfy that

a II (s) 2 ≤ cηµ 0 f (s) -y 2 , (C.14) b II (s) 2 ≤ cηµ 0 f (s) -y 2 , (C.15)

for sufficiently small absolute constant c. For b II (s) we have that (C.10) gives

b II (s) 2 ≤ √ n max i b II i (s) ≤ max i η(1 + R g (m/δ) 1/d ) 2 |S i |n f (s) -y 2 α 2 m .

Next applying Lemmas C.1 and B.9 in turn yields

≤ Cm Rv ηn f (s) -y 2 α 2 m ≤ ηµ 0 f (s) -y 2 Rv nα 2

Substituting m = Ω n 4 log(n/δ)/µ 4 0 α 4 for a large enough constant and R v we get

b II (s) 2 ≤ cηµ 0 f (s) -y 2 .

Analogously we bound a II (s) 2 using (C.7),

a II (s) 2 ≤ √ n max i a II i (s) ≤ η 2 n 3/2 1 + R g (m/δ) 1/d 3 f (s) -y 2 2 (m/δ) 1/d α 4 √ m ≤ ηµ 0 f (s) -y 2 • η 1 + R g (m/δ) 1/d 3 n 3/2 f (s) -y 2 (m/δ) 1/d µ 0 α 4 √ m ≤ ηµ 0 f (s) -y 2 • η α 2 • Cn 2 log(n/δ) α 2 µ 2 0 √ m ≤ cηµ 0 f (s) -y 2 .

Where we have used Lemma B.9 in the third inequality and substituted m = Ω(n 4 log(n/δ)/α 4 µ 4 0 ) noting that η = O 1 Λ(s) 2 and that α ≥ 1 in the last inequality. Therefore we have that r(s) satisfies Property 1 so that Theorem C.1 holds. By the same contradiction argument as in Theorem 4.1 we have that this holds for all iterations.

D ADDITIONAL TECHNICAL LEMMAS AND PROOFS OF LEMMAS FROM APPENDIX B:

Proof of Lemma 4.1:

We prove Lemma 4.1 for V ∞ , G ∞ separately. V ∞ can be viewed as the covariance matrix of the functionals φ i defined as

φ i (v) = x i I - vv v 2 2 1{v x i ≥ 0} (D.1)

over the Hilbert space V of L 2 (N (0, α 2 I)) of functionals. Under this formulation, if φ 1 , φ 2 , . . . , φ n are linearly independent, then V ∞ is strictly positive definite. Thus, to show that V ∞ is strictly positive definite is equivalent to showing that

c 1 φ 1 + c 2 φ 2 + • • • + c n φ n = 0 in V (D.2)

implies c i = 0 for each i. The φ i s are piece-wise continuous functionals, and equality in V is equivalent to

c 1 φ 1 + c 2 φ 2 + • • • + c n φ n = 0 almost everywhere.

For the sake of contradiction, assume that there exist c 1 , . . . , c n that are not identically 0, satisfying (D.2). As c i are not identically 0, there exists an i such that c i = 0. We show this leads to a contradiction by constructing a non-zero measure region such that the linear combination i c i φ i is non-zero.

Denote the orthogonal subspace to x i as D i := {v ∈ R d : v x i = 0}. By Assumption 1,

D i ⊆ j =i D j

This holds since D i is a (d -1)-dimensional space which may not be written as the finite union of subspaces D i ∩ D j of dimension d -2 (since x i and x j are not parallel). Thus, take z ∈ D i \ j =i D j . Since j =i D j is closed in R d , there exists an R > 0 such that

B(z, 4R) ∩ j =i D j = ∅.

Next take y ∈ ∂B(z, 3R) ∩ D i (where ∂ denotes the boundary) on the smaller disk of radius 3R so that it satisfies y 2 = max y ∈∂B(z,3R)∩Di y 2 . Now for any r ≤ R, the ball B(y, r) is such that for all points v ∈ B(y, r) we have v x ⊥ i 2 ≥ 2R and v xi 2 ≤ R. Then for any r ≤ R, the points v ∈ B(y, r) ⊂ B(z, 4R) satisfy that

x v ⊥ i 2 ≥ x i 2 - x i • v v 2 ≥ x i 2 1 - R 2R ≥ x i 2 2 .

Next we decompose the chosen ball B(y, r) = B + (r) ∨ B -(r) to the areas where the ReLU function at the point x i is active and inactive

B + (r) = B(y, r) ∩ {x i v ≥ 0}, B -(r) = B(y, r) ∩ {x i v < 0}.

Note that φ i has a discontinuity on D i and is continuous within each region B + (r) and B -(r). Moreover, for j = i, φ j is continuous on the entire region of B(y, r) since B(y, r) ⊂ B(z, 4R) ⊂ D c j . Since we have that φ j is continuous in the region, the Lebesgue differentiation theorem implies that for r → 0, φ i satisfies on B + (r), B -(r):

lim r→0 1 µ(B + (r)) B + (r) φ i = x y ⊥ i = 0, lim r→0 1 µ(B -(r)) B -(r) φ i = 0.

For j = i φ j is continuous on the entire ball B(y, r) hence the Lebesgue differentiation theorem also gives

lim r→0 1 µ(B + (r)) B + (r) φ i = φ j (y), lim r→0 1 µ(B -(r)) B -(r)

φ i = φ j (y).

We integrate c 1 φ 1 + . . . c n φ n over B -(r) and B + (r) separately and subtract the integrals. By the assumption, c 1 φ 1 + • • • + c n φ n = 0 almost everywhere so each integral evaluates to 0 and the difference is also 0,

0 = 1 µ(B + (r)) B + (r) c 1 φ 1 + • • • + c n φ n - 1 µ(B -(r)) B -(r) c 1 φ 1 + • • • + c n φ n . (D.3)

By the continuity of φ j for j = i taking r → 0 we have that

1 µ(B + (r)) lim r→0 B + (r) φ j - 1 µ(B -(r)) B -(r) φ j = φ j (y) -φ j (y) = 0.

For φ i the functionals evaluate differently. For B -(r) we have that

1 µ(B -(r)) lim r→0 B -(r) φ i = 1 µ(B -(r)) lim r→0 B -(r) 0 = 0,

while the integral over the positive side, B + (r) is equal to

1 µ(B + (r)) B + (r) φ i (z)dz = 1 µ(B + (r)) B + (r) x z ⊥ i dz = x y ⊥ i .

By construction, x y ⊥ i 2 > R and is non-zero so we conclude that for (D.3) to hold we must have c i = 0. This gives the desired contradiction and implies that φ 1 , . . . φ n are independent and V ∞ is positive definite with λ min (V ∞ ) = λ 0 .

Next we consider G ∞ and again frame the problem in the context of the covariance matrix of functionals. Define

θ i (v) := σ v x i v 2

for v = 0. Now the statement of the theorem is equivalent to showing that the covariance matrix of {θ i } does not have zero-eigenvalues, that is, the functionals θ i s are linearly independent. For the sake of contradiction assume ∃ c 1 , . . . , c n such that

c 1 θ 1 + c 2 θ 2 + • • • + c n θ n = 0 in V (equivalent to a.e).

Via the same contradiction argument we show that c i = 0 for all i. Unlike φ i defined in (D.1), each θ i is continuous and non-negative so equality "a.e" is strengthened to "for all v",

c 1 θ 1 + c 2 θ 2 + • • • + c n θ n = 0.

Equality everywhere requires that the derivatives of the function are equal to 0 almost everywhere.

Computing derivatives with respect to v yields

c 1 x v ⊥ 1 1{v x 1 ≥ 0} + c 2 x v ⊥ 2 1{v x 2 ≥ 0} + • • • + c n x v ⊥ n 1{v x n ≥ 0} = 0. Which coincide with c 1 φ 1 + • • • + c n φ n

By the first part of the proof, the linear combination c 1 φ 1 + • • • + c n φ n is non-zero around a ball of positive measure unless c i = 0 for all i. This contradicts the assumption that the derivative is 0 almost everywhere; therefore G ∞ is strictly positive definite with λ min (G ∞ ) =: µ 0 > 0.

We briefly derive an inequality for the sum of indicator functions for events that are bounded by the sum of indicator functions of independent events. This enables us to develop more refined concentration than in Du et al. [15] Proof of Lemma D.1:

Bound S as S = 1 m m k=1 1 A k ≤ 1 m m k=1 1 B k .

We apply Bernstein's concentration inequality to reach the bound. Denote

X k = 1 B k m and S = m k=1 X k . Then Var(X k ) ≤ EX 2 k = (1/m) 2 P(X k ) + 0 ≤ p m 2 , E S = E m k=1 X k ≤ p.

Applying Bernstein's inequality yields

P( S -E S ≥ t) ≤ exp -t 2 /2 m k=1 EX 2 k + t 3m .

Fix δ and take the smallest t such that P( S -E S ≥ t) ≤ δ. Denote t = r • E S, either P( S -E S ≥ E S) ≤ δ, or t = rE S corresponds to r ≥ 1. Note that t = rE S ≤ rp. In the latter case, the bound is written as We prove the claim by applying concentration on each entry of the difference matrix. Each entry V ij (0) is written as

P( S -E S ≥ rp) ≤ exp -(pr)

V ij (0) = 1 m m k=1 x v k (0) ⊥ i , x v k (0) ⊥ j αc k • g k v k 2 2 1 ik (0)1 jk (0). At initialization g k (0) = v k (0) 2 /α, c 2 k = 1 so V ij (0) simplifies to V ij (0) = 1 m m k=1 x v k (0) ⊥ i , x v k (0) ⊥ j 1 ik (0)1 jk (0).

Since the weights v k (0) are initialized independently for each entry we have

E v V ij (0) = V ∞ ij .

We measure the deviation V(0) -V ∞ via concentration. Each term in the sum

1 m m j=1 x v k (0) ⊥ i , x v k (0) ⊥ j 1 ik (0)1 jk (0) is independent and bounded, -1 ≤ x v k (0) ⊥ i , x v k (0) ⊥ j 1 ik (0)1 jk (0) ≤ 1.

Applying Hoeffding's inequality to each entry yields that with probability 1 -δ/n 2 , for all i, j,

|V ij (0) -V ∞ ij | ≤ 2 log(n 2 /δ) √ m .

Taking a union bound over all entries, with probability 1 -δ,

|V ij (0) -V ∞ ij | ≤ 4 log(n/δ) √ m .

Bounding the spectral norm, with probability 1 -δ,

V(0) -V ∞ 2 2 ≤ V(0) -V ∞ 2 F ≤ i,j |V ij (0) -V ∞ ij | 2 ≤ 16n 2 log(n/δ) m . Taking m = Ω n 2 log(n/δ) λ 2 0 therefore guarantees V(0) -V ∞ 2 ≤ λ 0 4 .

Proof of Lemma B.2::

This is completely analogous to B.1. Recall G(0) is defined as,

G ij (0) = 1 m m k=1 x v k (0) i , x v k (0) j c 2 k 1 ik (0)1 jk (0)

with c 2 k = 1 and v k (0) ∼ N (0, α 2 I) are initialized i.i.d. Since each term is bounded like B.1 the same analysis gives

G ij (0) -G ∞ ij 2 2 ≤ 16n 2 log(n/δ) m . Taking m = Ω n 2 log(n/δ) µ 2 0 therefore guarantees, G(0) -G ∞ 2 ≤ µ 0 4 .

Proof of Lemma B.3::

For a given R, define the event of a possible sign change of neuron k at point x i as

A i,k (R) = {∃v : v -v k (0) 2 ≤ R, and 1{v k (0) x i ≥ 0} = 1{v x i ≥ 0}} A i,k (R) occurs exactly when |v k (0) x i | ≤ R, since x i 2 = 1

and the perturbation may be taken in the direction of -x i . To bound the probability A i,k (R) we consider the probability of the event

P(A i,k (R)) = P(|v k (0) x i | < R) = P(|z| < R).

Here, z ∼ N (0, α 2 ) since the product v k (0) x i follows a centered normal distribution. The norm of x i 2 = 1 which implies that z computes to a standard deviation α. Via estimates on the normal distribution, the probability on the event is bounded like

P(A i,k (R)) ≤ 2R α √ 2π .

We use the estimate for P(A i,k (R)) to bound the difference between the surrogate Gram matrix and the Gram matrix at initialization V(0). Recall the surrogate V(t) is defined as,

Vij (t) = 1 m m k=1 x v k (t) ⊥ i , x v k (t) ⊥ k 1 ik (t)1 jk (t).

Thus for entry i, j we have

| Vij (t) -V ij (0)| = 1 m m k=1 x v k (t) ⊥ i , x v k (t) ⊥ j 1 ik (t)1 jk (t) -x v k (0) ⊥ i , x v k (0) ⊥ j 1 ik (0)1 jk (0)

This sum is decomposed into the difference between the inner product and the difference in the rectifier patterns terms respectively:

x v k (t) ⊥ i , x v k (t) ⊥ j -x v k (0) ⊥ i , x v k (0) ⊥ j , 1 ik (t)1 jk (t) -1 ik (0)1 jk (0) . Define Y k ij = x v k (t) ⊥ i , x v k (t) ⊥ j -x v k (0) ⊥ i , x v k (0) ⊥ j 1 ik (t)1 jk (t) , Z k ij = x v k (0) ⊥ i , x v k (0) ⊥ j 1 ik (t)1 jk (t) -1 ik (0)1 jk (0) . Then | Vij (t) -V ij (0)| = 1 m m k=1 Y k ij + Z k ij ≤ 1 m m k=1 Y k ij + 1 m m k=1 Z k ij . To bound | 1 m m k=1 Y k ij | we bound each |Y k ij | as follows. |Y k ij | = x v k (t) ⊥ i , x v k (t) ⊥ j -x v k (0) ⊥ i , x v k (0) ⊥ j 1 ik (t)1 jk (t) ≤ x v k (t) ⊥ i , x v k (t) ⊥ j -x v k (0) ⊥ i , x v k (0) ⊥ j = x i , x j -x v k (t) i , x v k (t) j + x v k (0) i , x v k (0) j -x i , x j = x i v k (t) v k (t) 2 • v k (t) v k (t) 2 , x j v k (t) v k (t) 2 • v k (t) v k (t) 2 -x v k (0) i , x v k (0) j = x i v k (t) v k (t) 2 • x j v k (t) v k (t) 2 -x v k (0) i , x v k (0) j = x i v k (0) v k (0) 2 • x j v k (0) v k (0) 2 + x i v k (t) v k (t) 2 - v k (0) v k (0) 2 • x j v k (t) v k (t) 2 + x j v k (t) v k (t) 2 - v k (0) v k (0) 2 • x i v k (0) v k (0) 2 -x v k (0) i , x v k (0) j ≤ x i v k (t) v k (t) 2 - v k (0) v k (0) 2 • x j v k (t) v k (t) 2 + x i v k (t) v k (t) 2 - v k (0) v k (0) 2 • x j v k (t) v k (t) 2 ≤ 2 v k (t) v k (t) 2 - v k (0) v k (0) 2 2 .

Therefore, we have

1 m m k=1 Y k ij ≤ 2 m m k=1 v k (t) v k (t) 2 - v k (0) v k (0) 2 2 ≤ 4R v (2m/δ) 1/d α ≤ 8R v (m/δ) 1/d α ,

where the first inequality follows from Lemma B.10. Note that the inequality holds with high probability 1 -δ/2 for all i, j.

For the second sum,

| 1 m m k=1 Z k ij | ≤ 1 m m k=1 1 A ik (R) + 1 m m k=1 1 A jk (R) so we apply Lemma D.1 to get, with probability 1 -δ/2n 2 1 m m k=1 Z k ij ≤ 2R v α √ 2π 2 + 2 √ 2πα log (2n 2 /δ) 3mR v ≤ 8R v α √ 2π since m satisfies m = Ω (m/δ) 1/d n 2 log(n/δ) αλ0 . Combining the two sums for Y k ij and Z k ij , with probability 1 -δ 2n 2 , | Vij (t) -V ij (0)| ≤ 8R v α √ 2π + 8R v (m/δ) 1/d α ≤ 12R v (m/δ) 1/d α .

Taking a union bound, with probability 1 -δ/2,

V(t) -V(0) F = i,j | Vij (t) -V ij (0)| 2 ≤ 12nR v (m/δ) 1/d α .

Bounding the spectral norm by the Frobenous norm,

V(t) -V(0) 2 ≤ 12nR v (m/δ) 1/d α .

Taking R v = αλ0 96n(m/δ) 1/d gives the desired bound.

V(t) -V(0) 2 ≤ λ 0 8 .

Proof of Lemma B.4::

To bound V(t) -V(0) 2 we now consider V(t) -V(t) 2 . The entries of V ij (t) are given as

V ij (t) = 1 m m k=1 x v k (t) ⊥ i , x v k (t) ⊥ j 1 ik (t)1 jk (t) αc k • g k v k (0) 2 2 .

The surrogate V(t) is defined as

Vij (t) = 1 m m k=1 x v k (t) ⊥ i , x v k (t) ⊥ j 1 ik (t)1 jk (t).

The only difference is in the second layer terms. The difference between each entry is written as

|V ij (t) -Vij (t)| = 1 m m k=1 x v k (t) ⊥ i , x v k (t) ⊥ j 1 ik (t)1 jk (t) αc k • g k v k (t) 2 2 -1 ≤ max 1≤k≤m α 2 g k (t) 2 v k (t) 2 2 -1 . Write 1 = α 2 g 2 k (0) v k (0) 2 2 , since v k (t) 2 is increasing in t according to (2.3) α 2 g k (t) 2 v k (t) 2 2 -1 = α 2 g k (t) 2 v k (t) 2 2 - α 2 g k (0) 2 v k (0) 2 2 ≤ 3R g (m/δ) 1/d + 3R v (m/δ) 1/d /α.

The above inequality is shown by considering different cases for the sign of the difference g k (t)g k (0). Now

α 2 g k (t) 2 v k (t) 2 2 - α 2 g k (0) 2 v k (0) 2 2 = αg k (t) v k (t) 2 + αg k (0) v k (0) 2 αg k (t) v k (t) 2 - αg k (0) v k (0) 2 ≤ αg k (0) + αR g v k (0) 2 + αg k (0) v k (0) 2 αg k (t) v k (t) 2 - αg k (0) v k (0) 2 ≤ (2 + R g (m/δ) 1/d ) αg k (t) v k (t) 2 - αg k (0) v k (0) 2 ≤ (2 + R g (m/δ) 1/d ) max α(g k (0) + R g ) v k (0) 2 - αg k (0) v k (0) 2 , α(g k (0) -R g ) v k (0) 2 + R v - αg k (0) v k (0) 2 ≤ (2 + R g (m/δ) 1/d ) max R g (m/δ) 1/d , R g (m/δ) 1/d + R v (m/δ) 1/d /α ≤ 3R g (m/δ) 1/d + 3R v (m/δ) 1/d /α,

where the second inequality holds due to Lemma B.10 with probability 1 -δ over the initialization.

Hence:

V(t) -V(t) 2 ≤ V(t) -V(t) F = i,j | Vij (t) -V ij (t)| 2 ≤ 3nR g (m/δ) 1/d + 3nR v (m/δ) 1/d /α. Substituting R v , R g gives V(t) -V(t) 2 ≤ λ 0 8 .

Now we use Lemma B.3 to get that with probability 1 -δ

EQUATION

combining we get with probability 1 -δ

V(t) -V(0) 2 ≤ λ 0 4 .

We note that the source for all the high probability uncertainty 1 -δ all arise from initialization and the application of Lemma B.10.

Proof of Lemma B.5::

To prove the claim we consider each entry i, j of G(t) -G(0). We have,

|G ij (t) -G ij (0)| = 1 m m k=1 σ v k (t) x i v k (t) 2 σ v k (t) x j v k (t) 2 -σ v k (0) x i v k (0) 2 σ v k (0) x j v k (0) 2 ≤ 1 m m k=1 σ v k (t) x i v k (t) 2 σ v k (t) x j v k (t) 2 -σ v k (t) x i v k (t) 2 σ v k (0) x j v k (0) 2 + 1 m m k=1 σ v k (t) x i v k (t) 2 σ v k (0) x j v k (0) 2 -σ v k (0) x i v k (0) 2 σ v k (0) x j v k (0) 2 ≤ 2 v k (t) v k (t) 2 - v k (0) v k (0) 2 2 ≤ 2 Rv (m/δ) 1/d α .

In the last inequality we used the fact,

v k (0) v k (0) 2 - v k (t) v k (t) 2 2 ≤ v k (t) -v k (0) 2 v k (0) 2 ≤ (m/δ) 1/d α v k (t) -v k (0) 2 ,

where the first inequality uses that v k (0) 2 ≤ v k (t) 2 and is intuitive from a geometrical standpoint. Hence, 

G(t) -G(0) 2 ≤ G(t) -G(0) F = i,j |G ij (t) -G ij (0)| 2 ≤ 2n Rv (m/δ)

(y i -f i (t)) 1 √ m c k g k (t) 1 v k (t) 2 x v ⊥ i 1 ik (t) 2 ≤ 1 √ m n i=1 |y i -f i (t)| c k g k (t) v k (t) 2 .

Now using (2.3) and the initialization v k (0) = αg k (0), we bound c k g k (t) v k (t) 2 ,

c k g k (t) v k (t) 2 ≤ c k g k (0) + R g v k (0) 2 ≤ 1 α 1 + αR g / v k (0) 2 .

By Lemma B.10, we have that with probability 1 -δ over the initialization,

α/ v k (0) 2 ≤ C(m/δ) 1/d .

Hence αR g / v k (0) 2 ≤ 1. This fact bounds c k g k (t) v k (t) 2 with probability 1 -δ for each k,

c k g k (t) v k (t) 2 ≤ 2/α.

Substituting the bound,

d dt v k (t) 2 ≤ 2 α √ m n i=1 |f i (t) -y i | ≤ 2 √ n α √ m f (t) -y 2 ≤ 2 √ n α √ m exp(-ωt/2) f (0) -y 2 .

Thus, integrating and applying Jensen's inequality,

v k (t) -v k (0) 2 ≤ s 0 dv k (s) dt 2 ds ≤ 4 √ n f (0) -y 2 αω √ m .

Note that the condition |g k (t) -g k (0)| ≤ R g is stronger than needed and merely assuring that |g k (t) -g k (0)| ≤ 1/(m/δ) 1/d suffices.

small x so we have that with probability 1 -δ m , that v k (0) 2 ≥ α(m/δ) 

E PROOFS OF LEMMAS FROM APPENDIX C:

Proof of Lemma C.1: Fix R, without the loss of generality we write S i for S i (R). For each k, v k (0) is initialized independently via ∼ N (0, α 2 I), and for a given k, the event 1 ik (0) = 1{v x i ≥ 0} corresponds to |v k (0) x i | ≤ R. Since x i 2 = 1, v k (0) x i ∼ N (0, α 2 ). Denoting the event that an index k ∈ S i as A i,k , we have

P(A i,k ) ≤ 2R α √ 2π .

Next the cardinality of S i is written as

|S i | = m k=1 1 A i,k .

Applying Lemma D. Taking a union bound, with probability 1 -δ, for all i we have that

|S i | ≤ 2mR α √ 2π + 16 log(n/δ) 3 .

Proof of Lemma C.2::

To show this we bound the difference g k (s) -g k (0) as the sum of the iteration updates. Each update is written as

∂L(s) ∂g k = 1 √ m n i=1 (f i (s) -y i ) c k v k (s) 2 σ(v k (s) x i ) . As c k σ v k (s) xi v k (s) 2 ≤ 1, ∂L(s) ∂g k ≤ 1 √ m n i |f i (s) -y i | ≤ √ n √ m f (s) -y 2 .

By the assumption in the statement of the lemma,

∂L(s) ∂g k ≤ √ n(1 -ηω 2 ) s/2 f (0) -y 2 √ m .

Hence bounding the difference by the sum of the gradient updates:

|g k (K + 1) -g k (0)| ≤ η K s=0 ∂L(s) ∂g k ≤ 4η √ n f (0) -y 2 √ m K s=0 (1 - ηω 2 ) s/2 .

The last term yields a geometric series that is bounded as

1 1 -1 -ηω 2 ≤ 4 ηω ,

Hence

|g k (K + 1) -g k (0)| ≤ 4 √ n f (0) -y 2 ω √ m .

Proof of Lemma C.3:

To show this we write v k (s) as the sum of gradient updates and the initial weight v k (0). Consider the norm of the gradient of the loss with respect to v k ,

∇ v k L(s) 2 = 1 √ m n i=1 (f i (s) -y i ) c k g k (s) v k (s) 2 1 ik (s)x v k (s) ⊥ i 2

.

Since v k (s) 2 ≥ v k (0) 2 ≥ α(δ/m) 1/d with probability 1 -δ over the initialization, applying Cauchy Schwartz's inequality gives

∇ v k L(s) 2 ≤ (1 + R g (m/δ) 1/d ) √ n f (s) -y 2 α √ m . (E.1)

By the assumption on f (s) -y 2 this gives

∇ v k L(s) 2 ≤ 2 √ n(1 -ηω 2 ) s/2 f (0) -y 2 α √ m .

Hence bounding the parameter trajectory by the sum of the gradient updates:

v k (K + 1) -v k (0) 2 ≤ η K s=0 ∇ v k L(s) 2 ≤ 2 √ n f (0) -y 2 α √ m K s=1 1 - ηω 2 s/2

yields a geometric series. Now the series is bounded as 

:



Figures and Tables:

Rectifier sign-changes). Suppose v 1 (0), . . . , v k (0) are sampled i.i.d. as (2.4). In addition assume we have m = Ω (m/δ) 1/d n log(n/δ) αλ0 and v k (t) -v k (0) 2 ≤ αλ0 96n(m/δ) 1/d =: R v . Then with probability 1 -δ,

for monitoring the orthogonal and aligned Gram matrices during training. Lemma D.1. Let A 1 , . . . , A m be a sequence of events and suppose that A k ⊆ B k with B 1 , . . . , B m mutually independent. Further assume that for each k, P(B k ) ≤ p, and define S = 1 m m k=1 1 A k . Then with probability 1 -δ, S satisfies S ≤ p 2 + 8 log(1/δ) 3mp .

Figure 2: Training dynamics of networks from class (1.2) on digits 4,9 of MNIST

1, 4.2 but by considering gradient flow instead of gradient descent the proofs are simplified. In Appendix C we prove the main results from Section 4 (Theorem 4.1, 4.2) for finite step gradient descent.We state our convergence results for gradient flow.

Table content: <table/>

Now that we have established bounds on V(t), G(t) given that the parameters stay near initialization, we show that the evolution converges in that case:Proof of Lemma B.6: Consider the squared norm of the predictions f (t) -y 2 2 . Taking the derivative of the loss with respect to time, Since we assume that λ min G(t) + V(t)For now, assuming the linear convergence derived in Lemma B.6, we bound the distance of the parameters from initialization. Later we combine the bound on the parameters and Lemmas B.4, B.5 bounding the least eigenvalue of Λ(t), to derive a condition on the over-parameterization m and ensure convergence from random initialization.Proof of Lemma B.7: Denote f (x i ) at time t as f i (t). Since x

Table content: <table><tr><td/><td/><td>α 2</td><td colspan="2">≥ ω 2 , the derivative of the squared norm is bounded as</td></tr><tr><td/><td>d dt</td><td colspan="3">f (t) -y 2 2 ≤ -ω f (t) -y 2 2 .</td></tr><tr><td colspan="3">Applying an integrating factor yields</td><td/></tr><tr><td/><td/><td colspan="3">f (t) -y 2 2 exp(ωt) ≤ C.</td></tr><tr><td colspan="3">Substituting the initial conditions, we get</td><td/></tr><tr><td/><td colspan="4">f (t) -y 2 2 ≤ exp(-ωt) f (0) -y 2 2 .</td></tr><tr><td/><td/><td/><td/><td>1/d</td></tr><tr><td/><td/><td/><td/><td>α</td><td>√</td><td>2π</td><td>.</td></tr><tr><td>Taking Rv =</td><td colspan="4">√ 8n(m/δ) 1/d gives the desired bound. Therefore, with probability 1 -δ, 2παµ0</td></tr><tr><td/><td/><td colspan="2">G(t) -G(0) 2 ≤</td><td>µ 0 4</td><td>.</td></tr></table>

1 d where d is the input dimension. Applying a union bound, with probability 1 -δ, for all 1 ≤ k ≤ m, Now by (2.3) for t ≥ 0, v k (t) 2 ≥ v k (0) 2 so

Table content: <table><tr><td colspan="3">1 v k (0) 2</td><td>≤</td><td colspan="2">m/δ α</td><td>1/d</td><td>.</td></tr><tr><td>1 v k (t) 2</td><td>≤</td><td colspan="3">1 v k (0) 2</td><td>≤</td><td colspan="2">m/δ α</td><td>1/d</td><td>.</td></tr></table>

1, with probability 1 -δ/n,

Table content: <table><tr><td>|S i | ≤</td><td>2mR α √ 2π</td><td>+</td><td>16 log(n/δ) 3</td><td>.</td></tr></table>

Peer review (iclr2020):
*Summary of the contributions:* This paper deals with convergence, of the hidden layer of a 2-layers Relu Network trained with Weight Normalization which consists in decoupling the direction and the magnitude for the pre-activation layers. The authors show, under mild assumption, the linear convergence (with high probability) of the training loss in the over-parameterized setting. Conceptually the result says that using random feature selection (the last layer is fixed, initialized with Rademacher random variables), with a large enough network, the hidden weights do not need to move much from their initialization to get arbitrarily small training loss. Thus, the analysis of the whole dynamics is reduced to a perturbation analysis of the dynamics around the initialization. (I would like to emphasize that even though the final argument ends up to be a ‘simple’ strong convexity convergence proof, I acknowledge that the proof of such result is far from being straightforward.) The contributions of this paper are the following: in the setting introduced in [15] (the ref numbers correspond to the ones in the paper) the authors prove the linear convergence of the training loss for a large enough output layer of a 2-layer Relu network with Weight Normalization. Regarding to [15] the difference is that the authors study Weight Normalization. They are able to get better bounds on the over parametrization, and provide an analysis that highlights two training regimes: a fast one (V-dominated) and a stable one (G-dominated) that could explain the benefit of weight normalization. *Decision:* I am leaning for a weak reject for the following reasons: 1. The main issue to me is that this work builds up on the work by [15] without addressing the issues risen by [12]. One important takeaway from [12] that the ‘lazy training’ regime studied might not be a desired behavior and might not happen in practice. Thus, I think that extending the work form [15] to WeightNormalization would be interesting if the phenomenons described by the theory are backed-up by some practical evidence. For instance : “we believe that G-dominated convergence actually is common but that it emerges at later stages in training, after the weights have adopted their directions in the V-dominated regime and improves stability.“ is something that the authors could check in practice. 2. I found the discussion regarding the phase transition a bit unclear. The magnitude of the weights v_k is treated as it was alpha the coefficient used to parametrized their initialization. Then since the phase transition depends on the magnitude of alpha, it is argued that the model naturally transitions as |v_k| increases (which is the case because of the discrete gradient updates). However, to me to make that phase transition argument you would need to show that the scaling of the matrices at time actually depends on the norm of the current iterate v_k(t) 3. Finally, regarding the proof of Theorem 4.1 and 4.2, I found it way more clear in [12]. There is a key subtle bootstrap argument that is not explicitly stated in the proof of Theorem B.1 and B.2: The condition R’ < R (stated in Lemma 3.4 in [12]) is the key argument that leads to a sufficient bound on the overparametrization. I put some remark regarding the clarity of the proof in the minor comment section. I think it would improve the quality of the work but did not have a major impact on my grade. 4. The presentation is slightly misleading. The fact that the last layer is not optimized is slightly mentioned once and the author make it appear just like is was done without any loss of generality “We consider the case where the ck’s are fixed. This does not alter the capacity of the network since the gks are trainable (ck and gk are interchangeable analytically).”. Moreover, the omission that Theorem 4.1 and 4.2 are only true with high probability makes even that fact even more unclear. I think that this aspect (last layer is fixed) should be clearly emphasized in the revision. *Questions:* “we believe that G-dominated convergence is common”, can you provides evidence that G-dominated and V-dominated stages occurs the way you describe it in the discussion ? Same for the claim “The direction of the weights changes rapidly at the earlier stages of training when ? is small, and G-dominated convergence ensues as ? grows, leading to improved stability.” Particularly the link between your theory, the norm of v_k and alpha. (c.f point 2. of the decision section) I think that it is important to emphasize that in your work c_k are randomly sampled and then fixed (it is much more clear in [15], but I do not think that the reader should know the related work to realize that). You seem to argue that it is without loss of generality when you say “We consider the case where the ck’s are fixed. This does not alter the capacity of the network since the gks are trainable (ck and gk are interchangeable analytically).” but this statement seems wrong since there is a relu activation between c_k and g_k (thus you cannot interchange negative c_k with negative g_k). Can you comment this ? Can you clarify the bootstrap argument mentioned in section 3. of the decision section ? More precisely, can you justify the sentence “this contradicts one of Lemmas B.7, B.8 at time T_0.” in the proof of Theorem B.1 ? *Minor comments: (do not need to be addressed in the rebuttal)* - In the proof of Theorems B.1 and B.2 the statement “clearly T_0>0” should justified. This statement is true but a continuity argument that should be stated. However, the dependence in t could be easily avoided using the proof technique from [12] (Lemma 3.2) where the result is actually stated for any weight close enough to the initialization. - The title of [12] has changed and one author is now missing. (since you cite the NeurIPS version you should then cite the right title with all the authors) - I would not say that “smaller alpha leads to faster convergence” since the optimal step size depend on alpha. === After rebuttal === I would like to thank the author for their detailed answer. I still think that experimental evidence that the phenomenon described by the theory actually happens in practice should be provided. I maintain my score.