Title: Implicit Bias of (Stochastic) Gradient Descent for Rank-1 Linear Neural Network
Venue: NeurIPS 2023
Accepted: yes
Keywords: implicit bias

gradient descent

stochastic gradient descent

linear networks
Abstract: Studying the implicit bias of gradient descent (GD) and stochastic gradient descent (SGD) is critical to unveil the underlying mechanism of deep learning. Unfortunately, even for standard linear networks in regression setting, a comprehensive characterization of the implicit bias is still an open problem. This paper proposes to investigate a new proxy model of standard linear network,  rank-1 linear network, where each weight matrix is parameterized as a rank-1 form. For over-parameterized regression problem, we  precisely analyze the implicit bias of GD and SGD---by identifying a “potential” function such that GD converges to its minimizer constrained by zero training error (i.e., interpolation solution), and further characterizing the role of the noise introduced by  SGD in perturbing the form of this potential. Our results explicitly connect the depth of the network and the initialization with the implicit bias of GD and SGD. Furthermore, we emphasize a new implicit bias of SGD jointly induced by stochasticity and over-parameterization, which can reduce the dependence of the SGD's solution on the initialization. Our findings regarding the implicit bias are different from that of a recently popular model, the diagonal linear network. We highlight that the induced bias of our rank-1 model is more consistent with standard linear network while the diagonal one is not. This suggests that the proposed rank-1 linear network might be a plausible proxy for standard linear net.

Implicit Bias of (Stochastic) Gradient Descent for Rank-1 Linear Neural Network

Abstract

Studying the implicit bias of gradient descent (GD) and stochastic gradient descent (SGD) is critical to unveil the underlying mechanism of deep learning. Unfortunately, even for standard linear networks in regression setting, a comprehensive characterization of the implicit bias is still an open problem. This paper proposes to investigate a new proxy model of standard linear network, rank-1 linear network, where each weight matrix is parameterized as a rank-1 form. For overparameterized regression problem, we precisely analyze the implicit bias of GD and SGD-by identifying a 'potential' function such that GD converges to its minimizer constrained by zero training error (i.e., interpolation solution), and further characterizing the role of the noise introduced by SGD in perturbing the form of this potential. Our results explicitly connect the depth of the network and the initialization with the implicit bias of GD and SGD. Furthermore, we emphasize a new implicit bias of SGD jointly induced by stochasticity and over-parameterization, which can reduce the dependence of the SGD's solution on the initialization. Our findings regarding the implicit bias are different from that of a recently popular model, the diagonal linear network. We highlight that the induced bias of our rank-1 model is more consistent with standard linear network while the diagonal one is not. This suggests that the proposed rank-1 linear network might be a plausible proxy for standard linear net.

1 Introduction

Gradient Descent (GD) and its stochastic variant, Stochastic Gradient Descent (SGD), are probably the most important optimization techniques in deep learning. To unveil the underlying mechanism of modern neural networks, it is highly fundamental to understand the thrilling and mysterious properties of GD and SGD. Some recent works [24, 18, 3, 22] have made significant efforts in this direction by exploring their implicit bias : among all global minimum, i.e., interpolation solutions, what particular ones will GD and SGD prefer without adding an explicit regularization? This question highlights the crucial roles of these algorithms in the generalization performance of the trained models.

Among the vast number of architectures in deep learning, the first object to investigate is the deep linear network, i.e., without any nonlinear activations,

where W k 's are weight matrices and x is the input data. The linear network can be obviously treated as an over-parameterized model of standard linear model, θ T = W L · · · W 1 . However, the introduction of the over-parameterization brings significant non-convexity, and thus complicates dynamics during learning. Until now, the direct analysis of implicit bias of GD and, particularly, SGD for standard deep linear networks with any depth and initialization on regression problems is still an open problem.

(a)

(b)

(c)

O

O

O

Q

O

O

O

O

O

Q

O

O

O

O

O

O

O

O

O

O

O

O

In order to study the implicit bias of linear networks in regression settings, a simplified version of standard linear network, diagonal linear network (i.e., each neuron is only connected with a single neuron of the next layer [26, 22]), was proposed as a proxy to underpin the bias of the standard one. With this simplified model, the solution selected by GD solves a constrained norm minimization problem that interpolates between the ℓ 1 and ℓ 2 norms up to the initialization scale [26]. For SGD, the solution is closer to that of the sparse regression when compared to GD with the same initialization [22].

However, in Section 3.2, our theoretical analysis shows that the induced implicit bias of diagonal linear networks is not consistent with standard linear network, at least for GD in regression settings. It is then natural to ask: are there any proxy architectures that could produce similar implicit bias to that of standard linear networks? If so, we still expect the new proxy model be amenable to tractable analysis and provide insights for investigating the implicit bias of SGD and a wider spectrum of architectures. This is the main goal of our work.

On the other hand, [11] showed that weight matrices of linear networks for linearly separable classification will become low-rank when trained with GD and logistic loss. [16] also observed such low-rank bias for deep matrix factorization. [8] showed that SGD and weight decay jointly induces a low-rank bias in the weight matrices when training a neural network. In fact, for linear networks with a single output trained with GD, all weight matrices will automatically become rank-1 and will maintain this property during training when the initialization is balanced [11].

Our contributions. Inspired by this low-rank bias and to avoid the complicated analysis of standard linear networks, we propose a novel model rank-1 deep linear network as a plausible proxy of standard linear networks, to reveal the implicit bias of GD and SGD on over-parameterized linear models for regression problems. A depthL rank-1 deep linear network f ( x ; u, v ) is defined by

where u k ∈ R d k +1 and v k ∈ R d k are vectors while u and v are denoted as collections of u k and v k , respectively. See Figure 1 for the difference between three types of linear networks. In this work we will show that the formulation above could offer us the possibility to understand the bias of GD and SGD for standard linear networks through the lens of results on rank-1 linear networks.

With the proposed model, this paper targets on precisely identifying the 'potential' function V ( θ ) such that GD converges to its minimizer constrained by zero training error (i.e., interpolation solution), and further characterizing the role of the noise induced by adding sampling stochasticity to GD, i.e., SGD, in perturbing the form of V ( θ ) . For the convenience of theoretical analysis, we focus on the continuous versions of GD and SGD, i.e., we study gradient flow (GF) and stochastic gradient flow (SGF). In particular, this paper establishes the following findings:

For GF, we show that the solution implicitly minimizes a potential function V ( θ ) that depends on the initialization and depth subject to f ( x ; u, v ) achieving zero training error (Theorem 1). The single layer case recovers the standard linear regression results, while a depth larger than one immediately changes the form of V ( θ ) , which clearly connects the implicit bias of GF with the model architecture. We emphasize that our results explicitly reveal how depth and initialization jointly influence the implicit bias of GF.

More importantly, in Theorem 2, by showing the similarity of the implicit bias of GF for standard and rank-1 linear networks, we conclude that rank-1 linear networks are standard linear networks with special initialization when trained with GF, highlighting that our rank-1 linear network is a plausible proxy of standard linear networks. This offers us the possibility to explore the implicit bias of SGD for standard linear networks by analyzing the rank-1 linear nets. On the other hand, the diagonal linear networks exhibit drastically different implicit bias of GF when compared to standard linear networks, e.g., GF prefers ℓ 2 minimum norm solution for standard linear networks when the initialization is small while, on the contrary, it prefers such solution when the initialization is large for diagonal linear networks.

For SGF, we show that the sampling noise brings an extra effect that depends on several hyper-parameters such as the learning rate, batch size and network depth compared to GF. When L > 1 , this extra effect 'alleviates' the influence of the initialization, i.e., the final solution reduces its dependence on the initialization. More intriguingly, when L = 1 , i.e., without over-parameterization, this effect brought by SGF immediately disappears. Thus this implicit bias is jointly induced by both model depth and stochasticity from SGD . To the best of our knowledge, we are the first to show such an implicit bias jointly affected by the two factors through the lens of the rank-1 networks, which is also empirically verified on standard linear and nonlinear networks. Our findings on SGF are summarized in Theorem 3.

Organization. In Section 2, we summarize the notations and the setup of our work. Section 3 and Section 4 present our main results where Section 3 focues on the implicit bias of GD while Section 4 is about SGD. Numerical experiments are presented in Section 5 and we conclude this paper in Section 6. Some technical details are deferred to Appendix.

Related Works

The study of implicit bias of GD has been pioneered by [24] on linearly separable classification problem and was generalized to GD for deep neural networks and different training strategies [11, 18, 5, 21, 15, 17]. Recently extensive works have made progress in this direction by focusing on the diagonal linear network model. Assuming the existence of a perfect solution in the regression setting, [26] showed that the solution selected by GF interpolates between ℓ 2 minimum norm solutions and ℓ 1 ones up to the initialization scale. Besides the full-batch gradient descent, [22] then analyzed SGF for diagonal linear networks and concluded that the sampling noise brought by SGF reduces the effective initialization scale when compared with GF, leading its solution to be closer to a sparse one. Aside from the sampling noise induced by SGF, [9, 23] also investigated the influence of the label noise in the diagonal linear networks setting. For GD and SGD with moderate learning rate, [7] studied their implicit bias for diagonal linear networks and revealed the corresponding influence of the finite learning rate.

This paper considers the rank-1 deep linear network to study the implicit bias of GD and SGD. We briefly introduce differences between the settings considered here and those in previous works. For GD on standard linear networks, [3] did not restrict the initialization for the 2-layer case while [28] required the initialization for standard linear networks to be nearly-zero without restricting the number of layers. [25] considered GD for shallow ReLU nets. In this work, we consider linear networks without requirements either on the scale of the initialization or the number of layers. For GD and SGD on the diagonal linear networks [3, 26, 22], we point out that the results are different with that for standard linear networks. Note that this is not to downgrade diagonal linear network, and instead the point is to show that different architectures could induce different implicit bias. The rank-1 linear networks considered here are close to the standard linear networks since it can be seen as standard linear networks with special initialization. Furthermore, the classification problem [24, 18, 5] and linear regression problem [1] have been comprehensively investigated and we focus on the over-parameterized regression. Finally, [4, 10, 19, 27, 9] focused on the flatness of the loss landscape and model parameters while our analysis is for the overall parametrization θ .

2 Preliminaries and Setup

Notations. Given a dataset of n samples { ( x i , y i ) } n i =1 , x i ∈ R d represents a d -dimensional data vector with scalar label y i ∈ R . Weuse X ∈ R n × d to denote the data matrix and use y = ( y 1 , . . . , y n )

to denote the collection of y i . ∥ · ∥ denotes the ℓ 2 -norm and ∥ · ∥ F is the Frobenius norm. ⟨ a, b ⟩ is the inner product. For any vector u and matrix A , we use u (0) and A (0) to denote their initialization. We let tr () be the trace operator. For the weight matrix W k ∈ R d k × d k +1 , W k ; ij is its i -th row j -th column element. For a parametric model f ( x ; θ ) = θ T x with parameters θ , we use L ( θ ) = 1 n ∑ n i ℓ i ( θ ) to represent the empirical loss where ℓ i ( θ ) is the loss function for the sample ( x i , y i ) .

Over-parameterized regression. We focus on the regression problem where n < d and assume the existence of solutions that perfectly fit the dataset, i.e., there exist interpolating parameters θ ∗ such that ∀ i ∈ { 1 , . . . , n } : ⟨ x i , θ ∗ ⟩ = y i . The quadratic loss ℓ i ( θ ) is used in our setting. The empirical loss is then L ( θ ) = 1 n ∑ n i =1 ℓ i ( θ ) = 1 n ∑ n i =1 ( y i -⟨ x i , θ ⟩ ) 2 .

Rank-1 deep linear networks. In this paper we consider the rank-1 deep linear network f ( x ; u, v ) (see Eq. (2)). We use u and v to denote the collections of u k 's and v k 's, respectively. For convenience,

such that the network can be written as f ( x ; u, v ) = ρ k v T k +1 u k ρ -k v T 1 x for k ∈ { 1 , . . . , L -1 } when L ≥ 2 . Through this paper, we treat the depth L as a hyper-parameter of the network and try to precisely characterize its role in the implicit bias.

Definition 1 (Balanced initialization for rank-1 linear networks) . Given an L -layer rank-1 linear network Eq. (2) , for any k ∈ { 1 , . . . , L -1 } , the balanced initialization means that

We add more discussions on the initialization in Appendix C.1.

3 Equivalence Between Implicit Bias of GD for Standard and Rank-1 Nets

3.1 Implicit bias of GF for rank-1 linear networks

In this section, we characterize the implicit bias of the continuous version of GD for the rank-1 deep linear networks f ( x ; u, v ) Eq. (2). Note that the model parameters are updated according to the gradient flow

for k ∈ { 1 , . . . , L -1 } and dv 1 /dt = -∇ v 1 L ( θ ) , which clarifies that u k and v k , rather than W k , are the model parameters.

Previous work [26] showed that the solution selected by GF interpolates between the ℓ 1 minimum norm solution and ℓ 2 one depending on the initialization scale for diagonal linear networks. In this section we examine whether this is the case for the rank-1 deep linear network with depth L . For convenience, we define

where L is the number of layers. Recall that θ (0) represents the initialization of θ , we now state the main theorem regarding the implicit bias of GF for rank-1 linear networks below:

Theorem 1 (Implicit bias of GF for rank-1 linear networks) . For a rank-1 linear network, if the initialization is balanced across layers (Definition 1) and if the gradient flow solution θ ( ∞ ) satisfies that Xθ ( ∞ ) = y , then gradient flow converges to a minimizer of the potential function V ( θ ) :

where

5.0

2.5

0.0

2.5

4

2

0

2

4

Layers = 1

Scale = 1

0

4

8

12

16

16

16

16

20

20

20

20

5.0

2.5

0.0

2.5

4

2

0

2

4

Scale = 6

0

4

8

12

12

16

16

16

16

20

20

20

24

5.0

2.5

0.0

2.5

4

2

0

2

4

Scale = 11

0

5

10

10

15

20

20

20

25

5.0

2.5

0.0

2.5

4

2

0

2

4

Layers = 5

1

2

3

4

5

6

6

6

6

7

7

7

7

5.0

2.5

0.0

2.5

4

2

0

2

4

1.5

3.0

4.5

4.5

6.0

7.5

7.5

9.0

5.0

2.5

0.0

2.5

4

2

0

2

4

0.0

2.5

5.0

7.5

10.0

12.5

5.0

2.5

0.0

2.5

4

2

0

2

4

Layers = 100

1

2

3

4

5

5

6

6

6

6

5.0

2.5

0.0

2.5

4

2

0

2

4

1.5

3.0

4.5

6.0

6.0

7.5

5.0

2.5

0.0

2.5

4

2

0

2

4

0.0

2.5

5.0

7.5

10.0

12.5

Figure 2: The plot of V ( θ ) for different numbers of layers and initialization scales in a two-dimensional parameter space. The initialization in the first column is [0.08, 0.06] and is multiplied by different scale factors for the other two columns. In the first column where the initialization is nearly zero, GD has similar implicit bias for rank-1 linear networks of different layers. When we increase the initialization scale, the shape of the contour for linear regression potential function does not change, while for linear networks as shown in the last column, the contour of the potential function gradually presents a sharp angle as we increase the number of layers.

While previous works focused on 2-layer standard linear networks [3] or multiple-layer linear networks with nearly zero initialization [28], i.e., ∥ θ (0) ∥ ≈ 0 , Theorem 1 builds the potential function V ( θ ) that depends on both the initialization θ (0) and the depth of the network L explicitly . We also plot the contours of V ( θ ) for different numbers of layers and scales of initialization in Fig. 2 in a 2-dimensional space. Theorem 1 clearly reveals how over-parameterization and the initialization of θ guide GD to select different solutions.

Effects of depth. When L = 1 , V ( θ ) becomes 1 2 ∥ θ ∥ 2 -⟨ θ, θ (0) ⟩ which is the same as the potential of the least square ∥ θ -θ (0) ∥ 2 except for a constant term. When L > 1 , V ( θ ) is no longer an Euclidean distance. The form of V ( θ ) depends on the depth and GD will favor different interpolating solutions. A particular interesting case is when L → ∞ where we have Ω L → 1 and λ L → 1 . As a result, V ( θ ) becomes V ( θ ) →∥ θ ∥ - ⟨ θ, θ (0) / ∥ θ (0) ∥⟩ , which reflects the difference between the norm of θ and the norm of its projection on the direction of the initialization θ (0) . Therefore, the direction of the initialization θ (0) matters for the potential function V ( θ ) and the final solution θ ( ∞ ) while they are rather less sensitive to the scale of θ (0) . This may serve as a benefit of the over-parameterization in the sense that it makes the network more stable to different scales of the initialization. As a comparison, both direction and scale of θ (0) are crucial for the least square potential function ∥ θ -θ (0) ∥ 2 .

Effects of initialization. When L is finite, to inspect the effects of the initialization, we can rewrite V ( θ ) = ∥ θ ∥ Ω L / Ω L -∥ θ (0) ∥ 1 / (2 L -1) ⟨ θ, θ (0) / ∥ θ (0) ∥⟩ . The second term will be more important when the initialization scale ∥ θ (0) ∥ is getting larger thus the initialization affects the implicit bias of GD more significantly. On the other hand, as ∥ θ (0) ∥ → 0 , the second term vanishes. Thus we have the following corollary:

Corollary 1.1. Under conditions of Theorem 1, if we further assume that the initialization is infinitesimal ∥ θ (0) ∥ → 0 and the depth is finite, then the GF solution θ ( ∞ ) is an ℓ 2 -norm minimization solution:

We now summarize the effects of initialization on the training regime as follows. (i). According to [26], for any D -homogeneous model ( D is a positive integer), the lazy regime (or NTK regime) is reached for large initialization. Since the rank-1 linear network is a homogeneous model, the NTK regime is reached for large initialization and the implicit bias is given by the RKHS norm predictor accordingly; (ii). For vanishing initialization, according to Corollary 1.1, an ℓ 2 -norm minimization predictor which can not be captured by the NTK kernel is returned. Therefore, as the initialization becomes smaller, we 'escape' from the lazy regime and falls into the 'Anti-NTK' regime defined

by [3]. Furthermore, based on Corollary 1.1, it is now worth to mention that, assuming vanishing initialization, linear regression (i.e., 1-layer linear networks), 2-layer standard linear networks [3] and rank-1 linear networks with finite depth share a similar implicit bias if trained with GD. And the first column of Fig. 2 shows that rank-1 linear networks with different number of layers exhibit potential contours with a similar shape.

3.2 Comparison between different architectures

The aforementioned phenomenon that GD will return an ℓ 2 -norm minimization solution for both 2layer standard linear networks and rank-1 linear networks with small initialization drives us to further investigate the comparison and similarities between different network architectures. For diagonal linear networks, GD prefers ℓ 1 -norm minimization solution with nearly-zero initialization, which is drastically different from the case for rank-1 and standard linear networks. On the contrary, it prefers ℓ 2 -norm minimization solution when the initialization is sufficiently large [26]. This inconsistency begs us to ask the question will standard linear networks and rank-1 linear networks share a similar implicit bias of GD when the initialization is large? To answer this question, in the following, we first analyze the implicit bias of GD for standard linear networks f ( x ; W ) = w T L W L -1 · · · W 1 x with any depth L , which corresponds to the parameterization θ T = w T L W L -1 · · · W 1 , with balanced initialization (Definition 2).

Theorem 2 (Implicit bias of GF for standard linear networks with any initialization) . For an L -layer standard linear network, if the initialization is balanced across layers, i.e., W T k +1 (0) W k +1 (0) = W k (0) W T k (0) for all layers, and if the gradient flow solution θ ( ∞ ) satisfies that Xθ ( ∞ ) = y , then gradient flow converges to a minimizer of the potential function V ( θ ) :

where

Although we state before that the parameters of a rank-1 deep linear network are u k 's and v k 's rather than W k 's, the form of the potential for rank-1 linear networks Eq. (6) is very similar to that of standard linear networks Eq. (7): V ( θ ) for an L -layer rank-1 linear network is the same as V std ( θ ) for a (2 L -1) -layer standard linear network. In the following, we explain the reason behind this fact and conclude that rank-1 linear network can be seen as a qualified proxy for the standard linear network when studying the implicit bias of GD, and potentially SGD.

Rank-1 linear networks are standard linear networks with special initialization. We first present a useful proposition regarding a special kind of initialization for standard linear networks.

Proposition 1 (Effects of diagonal and rank-1 initialization for standard linear networks) . Given an L -layer standard linear network f ( x ; W ) = w T L W L -1 · · · W 1 x where W k ∈ R d k +1 × d k for k ∈ { 1 , . . . , L } , if the weights are initialized such that only one column of W k is non-zero when k is even and only one row of W k is non-zero otherwise, i.e., for an integer p

̸

̸

then for any t > 0 :

̸

̸

Furthermore, if the weights are initialized as the diagonal shape, i.e., ∀ k ∈ { 1 , . . . , L -1 } , W k ; ij (0) = 0 if i = j , then GD does not maintain the diagonal shape of weight matrices.

Remark. A (2 L -1) -layer standard linear network initialized as in Proposition 1 with c k = 1 for all layers has the same formulation as our L -layer rank-1 architecture. According to Proposition 1, the special structure of the initialization of such standard linear network will be maintained if we run GD, which suggests that it will always have the same formulation as our L -layer rank-1 linear network, see the upper panel of Fig. 3. Therefore, an L -layer rank-1 linear network can be seen as a (2 L -1) -layer standard one with special initialization, and they exhibit similar implicit bias of GD. On the other hand, if the standard linear network is initialized as the diagonal shape as in the diagonal net, GD will not maintain this property (the lower panel of Fig. 3), i.e., the diagonal structure cannot be seen as a standard one. In this sense, the diagonal network exhibits special implicit bias particularly due to its structure, while the rank-1 linear network is a more qualified proxy of standard linear networks.

4 Implicit bias of SGD for Rank-1 Linear Networks

Recently, [23, 22] developed the characterization of the implicit bias of GD with label noise and SGD in diagonal linear networks. However, as mentioned earlier, its diagonal structure is special in the sense that the corresponding results can not be generalized to the case for standard linear networks. On the contrary, our results in Section 3 reveal that rank-1 linear networks can be seen as standard linear networks with special initialization. In this section, to take a step forward towards understanding the implicit bias of SGD for standard linear networks, we explore the continuous part of SGD, stochastic gradient flow (SGF), for the rank-1 linear networks in the over-parameterized regression setting. We begin with the introduction of the definition of SGD and our modelling techniques.

SGD. Unlike the full-batch GD where the parameters are updated according to Eq. (4), the SGD dynamics is

for k ∈ { 1 , . . . , L -1 } where t denotes the iteration step, η is the learning rate, b is the batch-size, and B t consists of b points randomly sampled from the uniform distribution U [1 , n ] .

Continuous Modelling of SGD. The continuous modelling techniques for SGD have been widely applied in recent works [1, 9, 23, 22] to study the dynamics of SGD. In our setting, recalling the definition of ρ k and ρ -k in Eq. (3), the continuous counterpart of SGD, SGF, is given by the following set of stochastic differential equations (SDE):

where r = ( f ( x 1 ; u, v ) -y 1 , . . . , f ( x n ; u, v ) -y n ) T ∈ R n is the residual, and W t is a standard Brownian motion in R n . For the parameterization of θ (Eq. (2)), we now aim to characterize the implicit bias of SGD by showing the existence of a function V ( θ ) such that the solution of the stochastic dynamics converges to its minimizer under the constraint of zero-training loss as follows.

Theorem 3 (Implicit bias of SGF for rank-1 linear networks) . For the rank-1 linear network Eq. (2) that is trained with SGF (Eq. (9) and (10) ) in the over-parameterization regression, if the initialization is balanced across layers, then the dynamics of θ gives us

̸

(a) The implicit bias jointly induced by stochasticity and architectures 'alleviates" the effects of initialization

12

10

8

6

4

2

0

4

3

2

1

0

Trajectory of

(0)

(

)

(b) The trajectory of α starting from the origin converging to α ( ∞ )

where

with P ⊥ ( θ ) = I -θθ T / ∥ θ ∥ 2 being the orthogonal projection operator of θ . Furthermore, if the SGF solution θ ( ∞ ) satisfies that Xθ = y , then the model parameter θ converges to a minimizer of V S ( θ, ∞ ) as t →∞ :

One can immediately notice that the R.H.S of Eq. (11) includes a noise term, which can be interpreted as that the model parameter θ follows a mirror flow with a noise term. Furthermore, Eq. (12), which characterizes the optimization geometry, depends on time explicitly. This partly reflects its stochastic nature-the optimization trajectory is not deterministic. Furthermore, as in [7], it is also possible to generalize the current vanishing learning rate results to the moderate learning rate analysis by deriving a stochastic mirror descent recursion with time varying potentials.

Implicit bias jointly induced by stochasticity and architectures. The sampling noise of SGD introduces an extra term (the last term of Eq. (12)) that depends on several parameters such as the learning rate and the data matrix X when compared with that of GD (Eq. (6)). Interestingly, this extra term only exists when the model is simultaneously over-parameterized and trained with the existence of the sampling noise. If the model is not over-parameterized, i.e., standard linear regression when L = 1 , then λ L = 0 and this term disappears. On the other hand, if there is no sampling noise this term also vanishes. Indeed, the L = 1 case has been widely studied by recent works [1], and it turns out that GD and SGD share similar implicit bias when there is no over-parameterization. Theorem 3 confirms this phenomenon and, more importantly and intriguingly, reveals a new connection between the architecture-induced over-parameterization and the implicit bias of the optimization algorithms.

'Alleviating' the effects of initialization. When L →∞ , the last term of Eq. (12) brought by the sampling noise can be seen as 'alleviating" the influence of the initialization θ (0) such that the training dynamics reduces its dependence on the initialization (Fig. 4(a)). Due to the difficulty of explicitly solving the stochastic integral in Eq. (12), we give here a qualitative interpretation of this alleviating effect. As L →∞ , i.e., infinitely deep linear network, we have λ L → 1 , therefore the integral in Eq. (12) becomes -2 η nb ∫ ∞ 0 L ( θ ) tr ( P ⊥ ( θ ) X T X ) θ ∥ θ ∥ ds. In a d -dimensional space, we let α ( t ) := ∫ t 0 L ( θ ) tr ( P ⊥ ( θ ) X T X ) θ/ ∥ θ ∥ ds be the position of a particle A that starts from the origin. Then the integral of Eq. (12) amounts to the final position α ( ∞ ) of A , which moves

10

3

10

2

10

1

10 0

||

(0)||

0.0

0.1

0.2

0.3

0.4

0.5

D

(

(

),

*

2

)

SGD

(a)

10

3

10

2

10

1

10 0

||

(0)||

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

D

(

(

),

*

)

SGD

(b)

10

1

10 0

||

(0)||

0.1

0.2

0.3

0.4

0.5

D

(

(

),

*

2

)

Standard linear networks

SGD

(c)

10

1

10 0

||

(0)||

0.500

0.525

0.550

0.575

0.600

0.625

0.650

D

(

(

),

*

)

Standard linear network

SGD

(d)

along the direction of θ ( t ) / ∥ θ ( t ) ∥ with speed proportional to L ( θ ( t )) tr ( P ⊥ ( θ ( t )) X T X ) . Since ∥ X ∥ 2 F - ∥ X ∥ 2 2 ≤ tr ( P ⊥ ( θ ( t )) X T X ) ≤ ∥ X ∥ 2 F where ∥ X ∥ 2 F is finite, it is highly likely that the direction of α ( ∞ ) heavily depends on θ (0) . This is because the velocity magnitude, which is proportional to the value of empirical loss L ( θ ( t )) , is large when t = 0 and quickly shrinks along its trajectory. As a result, in Eq. (12), the effect coming from the initialization θ T θ (0) / ∥ θ (0) ∥ is alleviated by the extra term -θ T α ( ∞ ) induced by the sampling noise and over-parameterization. In this sense, V S ( θ, ∞ ) is closer to a simple ℓ 2 -norm ∥ θ ∥ when compared with GD, i.e., the SGD solution is more likely to be an ℓ 2 minimum norm solution when compared to GD. We also present the trajectory of α in Fig. 4(b). The trajectory is obtained by first training a rank-1 linear network using SGD for 5000 iterations and computing α ( t ) at every step followed by a projection of the trajectory { α ( t ) } 5000 t =1 into a 2-dimensional space according to the method described in [13]. It can be seen that the direction of the final position α ( ∞ ) is close to that of θ (0) ( ⟨ α ( ∞ ) , θ (0) ⟩ / ( ∥ α ( ∞ ) ∥ θ (0) ∥ ) = 0 . 9314 ). We also conduct additional experiments in Appendix A.3 to further verify our finding.

Comparison with diagonal linear networks. In [22], the authors showed that the extra effects of the sampling noise of SGD is equivalent to multiplying a shrinking coefficient which depends on the training dynamics to the initialization scale. This effect leads the solution of SGD to be closer to that of sparse regression when compared with GD. Our Theorem 3 does not show such bias that the solution interpolates between the ℓ 2 -norm and ℓ 1 -norm minimization solution. Therefore the conclusion of implicit bias of SGD for diagonal linear networks can not be directly applied to rank-1 linear networks and standard linear networks. This reveals an important finding that the implicit bias of GD or SGD is strongly tied with network architecture. On the other hand, similar to the case of diagonal linear networks, Theorem 3 also reveals an initialization cancellation effect induced by the SGD sampling noise. It is interesting for future work to explore whether this effect can be generalized to other architectures and can be seen as a special benefit brought by SGD.

5 Numerical Experiments

In this section, we consider the over-parameterized regression problem with rank-1 linear, standard linear, and non-linear networks for different initialization scales to verify our theoretical claims. We define the distance D ( θ 1 , θ 2 ) = ∥ θ 1 -θ 2 ∥ 2 / ∥ θ 2 ∥ 2 to measure the relative difference between θ 1 and θ 2 . For a linear network and the parameterization θ T = w T L W L -1 · · · W 1 , we use θ ( ∞ ) to denote the solution returned by the optimization algorithms (GD or SGD) and θ (0) to denote the corresponding initialization. θ ∗ is the ground truth solution, and θ ∗ ℓ 2 is the ℓ 2 minimum norm solution (Corollary 1.1). Details of the experiments and more numerical experiments are deferred to Appendix A.1 and we focus on the results here.

Rank-1 linear networks. As shown in Fig. 5(a), for all different initialization scales ∥ θ (0) ∥ , D ( θ ( ∞ ) , θ ∗ ℓ 2 ) is smaller if θ ( ∞ ) is returned by SGD, i.e., the SGD solution is closer to the ℓ 2 minimum norm solution θ ∗ ℓ 2 when compared to the GD solution. Similarly, Fig. 5(b) shows a benefit of such alleviating dependence of initialization bias of SGD: its solution is closer to the ground truth solution θ ∗ when compared to the GD solution for different ∥ θ (0) ∥ . Furthermore, in Fig. 6(a), we show D ( θ ( t ) , θ ∗ ℓ 2 ) along training, which further reveals that the final solution returned by SGD is closer to θ ∗ ℓ 2 . Note that when the initialization scales are small, both GD solution and SGD solution

are close to θ ℓ ∗ 2 . These experiments well support our theoretical findings. Besides, we present additional experiments to compare diagonal and rank-1 linear networks in Appendix A.2.

10 0

10 1

10 2

10 3

10 4

iterations

10

1

10 0

D

(

(

),

*

2

)

GD (0.0086)

SGD (0.0086)

GD (0.0676)

SGD (0.0676)

GD (0.2849)

SGD (0.2849)

(a)

10 0

10 1

10 2

10 3

10 4

iterations

10

1

10 0

D

(

(

),

*

2

)

Standard Linear Networks

GD (0.0207)

SGD (0.0207)

GD (0.1048)

SGD (0.1048)

GD (0.3311)

SGD (0.3311)

(b)

10 0

10 1

10 2

10 3

10 4

iterations

10 2

6×10 1

7×10 1

8×10 1

9×10 1

Non-linear Networks

GD (8.1744)

SGD (8.1744)

GD (14.4536)

SGD (14.4536)

GD (19.4559)

SGD (19.4559)

(c)

Standard linear networks and non-linear networks. Proposition 1 states similarity between rank1 and standard linear nets, therefore they have similar results. We conduct numerical experiments for standard linear networks as in the case of rank-1 linear networks and show in Fig. 5(c) and 5(d) that D ( θ ( ∞ ) , θ ∗ ℓ 2 ) and D ( θ ( ∞ ) , θ ∗ ) are also smaller if we run SGD, which supports the generalization of conclusion of rank-1 linear network to standard linear networks. D ( θ ( t ) , θ ∗ ℓ 2 ) along training plotted in Fig. 6(b) further supports this phenomenon. For non-linear networks (Fig. 6(c)), we report the test error on a newly sampled test set for both GD and SGD along training, where SGD solutions have smaller test error when compared to GD solutions.

6 Discussion & Conclusion

Our work proposes the rank-1 linear network that is a plausible proxy of standard linear networks. We analyze the implicit bias of GD and SGD for this new net and find that it approximates standard linear networks better than diagonal linear networks. We further reveal the joint role of over-parameterization and stochasticity in characterizing the implicit bias of SGD. Similar to the diagonal linear networks, our results also reveal an initialization alleviating effect of SGD sampling noise, suggesting a future direction that investigates whether such effect is general across different architectures. See Appendix B for more discussions.

Limitation. We do not generalize the analysis to non-linear networks due to its lack of the form of θ as in our current approach. Furthermore, an exact characterization of the stochastic integral is absent.

Acknowledgments

B. Lyu thanks Di Wang for useful discussion. Z. Zhu is supported by Beijing Nova Program (No. 202072) from Beijing Municipal Science Technology Commission.

References

Alnur Ali, Edgar Dobriban, and Ryan Tibshirani. The implicit regularization of stochastic gradient flow for least squares. In International Conference on Machine Learning , 2020.

Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit accelearation by overparameterization. In arXiv: 1802.06509 , 2018.

Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: beyond infinitesimal mirror descent. In NeurIPS , 2021.

Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and langevin processes. In International Conference on Machine Learning , 2020.

Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory , 2020.

Simon S. Du, Wei Hu, and Jason D. Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In arXiv: 1806.00900 , 2018.

Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s)gd over diagonal linear networks: Implicit regularisation, large stepsizes and edge of stability. In arXiv:2302.08982 , 2023.

Tomer Galanti, Zachary S. Siegel, Aparna Gupte, and Tomaso Poggio. Sgd and weight decay provably induce a low-rank bias in neural networks. In arXiv:2206.05794 , 2022.

Jeff Z. HaoChen, Colin Wei, Jason D. Lee, and Tengyu Ma. Understanding the implicit bias of the noise covariance. In Conference on Learning Theory , 2020.

Stanislaw Jastrzebski, Zac Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Amos Storkey, and Yoshua Bengio. Three factors influencing minima in sgd. In International Conference on Learning Representations , 2018.

Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations , 2019.

Kei Kobayashi. Stochastic calculus for a time-changed semimartingale and the associated stochastic differential equations. In arXiv:0906.5385 , 2010.

Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In NeurIPS , 2018.

Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research , 2019.

Yan Li, Etan X. Fang, Huan Xu, and Tuo Zhao. Implicit bias of gradient based adversarial training on separable data. In International Conference on Learning Representations , 2020.

Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations , 2021.

Bochen Lyu and Zhanxing Zhu. Implicit bias of adversarial training for deep neural networks. In International Conference on Learning Representations , 2022.

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations , 2020.

Stephan Mandt, Matthew D. Hoffman, and David M. Blei. A variational analysis of stochastic gradient algorithms. In International Conference on Machine Learning , 2016.

Kenneth S. Miller. On the inverse of the sum of matrices. In Mathematics Magazine , 1981.

Mor Shpigel Nacson, Suriya Gunasekar, Jason D. Lee, Nathan Srebro, and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models. In International Conference on Machine Learning , 2019.

Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. In NeurIPS , 2021.

Loucas Pillaud-Vivien, Julien Reygner, and Nicolas Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Conference on Learning Theory , 2020.

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on spearable data. In Journal of Machine Learning Research , 2018.

Francis Williams, Matthew Trager, Claudio Silva, Daniele Panozzo, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. In NeurIPS , 2019.

Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory , 2020.

Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In NeurIPS , 2018.

Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. In International Conference on Learning Representations , 2021.

Appendix

Organization of Appendix. In Appendix A we present details of the numerical experiments and additional numerical experiments. We give more discussions of our work in Appendix B. Appendix C provides missing technical details of Section 3 while Appendix D provides those of Section 4.

A Details of Numerical Experiments and Additional Experiments

In A.1, we present the details of the numerical experiments. We compare the rank-1 linear networks with the diagonal linear networks empirically in A.2. Finally, in A.3, we conduct additional experiments to further verify the 'alleviting' effect of the SGD sampling noise mentioned in Theorem 3. In A.4, we conduct experiments when the balanced initialization condition is not satisfied.

Data. We conduct over-parameterized regression with different linear networks. For the dataset { ( x i , y i ) } n i =1 where x i ∈ R d and y i ∈ R , we set n = 40 , d = 100 and x i ∼ N (0 , I ) . For i ∈ { 1 , . . . , n } , y i is generated by y i = θ ∗ T x i where θ ∗ ∈ R d , i.e., θ ∗ is the ground truth solution. We let 20 components of θ ∗ be informative.

A.1 Details of Numerical Experiments in Section 5

10 0

10 1

10 2

10 3

10 4

iterations

10 1

10 2

D

(

(

),

*

2

)

Large initialization, ||

(0)|| = 124.77

SGD

We now present the details of numerical experiments conducted in Section 5.

ℓ 2 minimum norm solution. To get the ℓ 2 minimum norm solution θ ∗ ℓ 2 , we train a single layer linear network with zero initialization using GD for 20000 iterations, since GD will return an ℓ 2 minimum norm solution solution in this case according to Corollary 1.1 and [26].

Rank-1 linear networks. For the rank-1 linear network f ( x ; u, v ) = w T L W L -1 · · · W 1 x where W k = u k v T k ∈ R d k × d k +1 , we let L = 3 and ∀ k ∈ { 1 , . . . , L } : d k = 100 . The learning rate is 10 -3 and the batch size is 4 if we run SGD. We construct different rank-1 linear networks as follows: for a randomly sampled ˜ θ ∈ R 100 , we let the initialization θ i (0) of the i -th rank-1 networks have the same direction as ˜ θ but with different scales. We then train each rank-1 linear network with GD and SGD, respectively, for 20000 iterations. In particular:

Fig. 5(a) presents the results of the distances between θ ∗ ℓ 2 and GD and SGD solutions, respectively, of each trained rank-1 networks with different initialization scales.

In Fig. 5(b), we measure the distances between θ ∗ and GD and SGD solutions, respectively, for each trained rank-1 networks with different initialization scales.

Fig. 6(a) plots the distances between θ ∗ ℓ 2 and the model parameters θ along training when the initialization scales are different for both GD and SGD. The numbers in the bracket denote ∥ θ (0) ∥ .

Fig. 7 is about the distances between θ ∗ ℓ 2 and the model parameters θ along training for both GD and SGD when ∥ θ (0) ∥ is extremely large.

Standard linear networks. For the standard linear network f ( x ; W ) = w T L W L -1 · · · W 1 x where W k ∈ R d k × d k +1 , we let L = 4 and ∀ k ∈ { 1 , 2 , 3 } : d k = 100 . The learning rate is 10 -3 and the batch size is 4 if we run SGD. Other settings are similar to that of rank-1 linear networks.

In Fig. 5(c), we plot the distances between θ ∗ ℓ 2 and GD and SGD solutions, respectively. Similar to the case of rank-1 linear networks, for all initialization scales, D ( θ ( ∞ ) , θ ∗ ℓ 2 ) is smaller if the network is trained with SGD when compared to GD.

Fig. 5(d) presents the results of the distances between θ ∗ and GD and SGD solutions.

Fig. 6(b) plots the distances between θ ∗ ℓ 2 and the model parameters θ along training when the initialization scales ∥ θ (0) ∥ are different for both GD and SGD. The numbers in the bracket denote ∥ θ (0) ∥ .

Non-linear networks. For the non-linear network f ( x ; W ) = w T L σ ( W L -1 · · · σ ( W 1 x )) where W k ∈ R d k × d k +1 , we let L = 4 and ∀ k ∈ { 1 , 2 , 3 } : d k = 100 . The learning rate is 10 -3 and the batch size is 4 if we run SGD. We use the ReLU activation σ ( x ) = ReLU( x ) . We use the same dataset as in the experiments of rank-1 linear networks. Since the non-linear networks do not have the overall parameterization of θ as in the linear networks case, to measure the initialization scale, we first straight all weight matrices to vectors and stack them to get a single vector, then we calculate the ℓ 2 norm of this vector as the scale of the initialization of a non-linear network, i.e., we use √ ( ∑ k ∥ W k (0) ∥ 2 F ) as the initialization scale where ∥ · ∥ F is the Frobenius norm. Due to the same reason, we can not measure quantities such as D ( θ, θ ∗ ℓ 2 ) , therefore, we report the test error of the model in a newly sampled test set instead. Fig. 6(c) plots the test error of the model along training when the initialization scales are different for both GD and SGD. The numbers in the bracket denote initialization scales.

A.2 Additional Experiments of Comparison with Diagonal Nets

Results in Section 3.2 indicate that diagonal linear networks exhibit different implicit bias in comparison with rank-1 and standard linear networks, e.g., both rank-1 and standard linear networks prefer ℓ 2 minimum norm solution for GD when the initialization is nearly-zero while, on the contrary, diagonal linear networks prefer such solution when the initialization is sufficiently large. In this section, we empirically compare the implicit bias for rank-1 linear networks and diagonal linear networks to show this phenomenon.

In particular, we use the same settings as in the experiments for rank-1 linear networks in A.1 while only change the model to diagonal linear networks. As in previous works [3, 22], the reparameterization of diagonal linear network is

where w + ∈ R 100 , w -∈ R 100 and ⊙ is the elementwise product. Let e = (1 , · · · , 1) T ∈ R 100 , we set the initialization as

where C is a positive constant measuring the initialization scale. For each diagonal linear network with different C , we run GD for 20000 iterations and calculate D ( θ ( ∞ ) , θ ∗ ℓ 2 ) and D ( θ ( ∞ ) , θ ∗ ) . The results are plotted in Fig. 8, where, for convenience of comparing the implicit bias of GD for rank-1 linear networks with that of diagonal linear networks, we also plot the results of rank-1 linear networks of Fig. 5(a) and Fig. 5(b) in Fig. 8(a) and Fig. 8(b), respectively.

As shown in Fig. 8(a), as the initialization scale ( ∥ θ (0) ∥ for rank-1 linear networks, C for diagonal linear networks) increases, D ( θ ( ∞ ) , θ ∗ ℓ 2 ) decreases for rank-1 linear linear networks trained with

10

3

10

2

10

1

10 0

Initialization Scale

0.0

0.1

0.2

0.3

0.4

0.5

D

(

(

),

*

2

)

SGD

GD, diagonal

(a)

10

3

10

2

10

1

10 0

Initialization Scale

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

D

(

(

),

*

)

SGD

GD, diagonal

(b)

both GD and SGD, while it increases for diagonal linear networks trained with GD. This indicates the drastic difference between the implicit bias of GD exhibited by diagonal linear networks and rank-1 linear networks (also standard linear networks).

A.3 Additional Experiments for the 'Alleviating' Effect in Theorem 3

Recall the form of V S ( θ, t ) in Theorem 3

we let

To quantitatively measure the 'alleviating' effect of the SGD sampling noise, we train 3 rank-1 linear networks with different initialization scales using SGD. The batch size is 4 and the learning rate

10 0

10 1

10 2

10 3

10 4

iterations

0.0

0.5

1.0

1.5

2.0

2.5

p

, 0.2849

p

q

, 0.2849

p

, 0.5135

p

q

, 0.5135

p

, 0.8696

p

q

, 0.8696

is 5 × 10 -5 . For the rank-1 linear network f ( x ; u, v ) = w T L W L -1 · · · W 1 x where W k = u k v T k ∈ R d k × d k +1 , we let L = 3 and ∀ k ∈ { 1 , . . . , L } : d k = 100 . We calculate both p θ ( t ) and q θ ( t ) along training, where q θ ( t ) measures the alleviating effect of the SGD sampling noise and their difference p θ ( ∞ ) -q θ ( ∞ ) is the alleviated initialization dependence of the SGD solution compared to GD solution.

As shown in Fig. 9, the effect coming from the SGD sampling noise, q θ , is equivalent to make the dependence of V S on the initialization closer to 0 (after about 1000 iterations, every dotted line is closer to the x-axis compared to the corresponding solid line with the same color), thus it controls the dependence of the SGD solution on the initialization. This phenomenon further verifies our claims.

A.3.1 Training loss for Fig. 4(b)

Fig. 4(b) indicates that the final direction of the integral term in Eq. (12) highly depends on the initialization θ (0) since the loss decays along training. To further support this argument, here we present the training loss when we perform the experiments of Fig. 4(b) in Fig. 10. It can be seen that, for a random initialization, the loss, the magnitude of the speed of α , has a high value at the start of the training, and decays very quickly, which explains why the direction of θ (0) is crucial to that of α ( ∞ ) .

A.4 Additional Experiments for Biased Initialization

In this section, we provide additional experiments to show that our conclusion still holds when removing the balanced initialization condition (Definition 1).

To make the initialization unbalanced, we add a small perturbation to the balanced initialization. Specifically, we define

as the scale of the perturbation to the balanced initialization (larger ∆ implies that the initialization is more unbalanced). All the other experiment details are kept unchanged as in Section 5. As shown in Fig. 11 and Fig. 12, we still observe similar phenomenons as in the case of the balanced initialization, e.g., SGD solutions are closer to the ℓ 2 -norm minimization solution compared to GD, when a small perturbation is added to the balanced initialization. Thus the implicit bias is not unique to the balanced initialization. In particular:

10 0

10 1

10 2

10 3

iterations

10

7

10

5

10

3

10

1

10 1

10 3

10 5

(

)

Wereport D ( θ ( ∞ ) , θ ∗ ℓ 2 ) for both GD and SGD for different levels of perturbation ∆ (denoted in the title of each figure) in Fig. 11. In the last figure, we fix the initialization scale and report D (( θ ( ∞ ) , θ ∗ ℓ 2 ) of both GD and SGD for different ∆ . It can be seen that, without the balanced initialization, GD and SGD still prefer ℓ 2 -norm minimization solution θ ∗ ℓ 2 for small initialization, while the SGD solution is closer to θ ∗ ℓ 2 due to its initialization reduction effect.

̸

We report D ( θ ( t ) , θ ∗ ℓ 2 ) during optimization for both GD and SGD for different levels of perturbation ∆ and the same scale of initialization ( ∥ θ (0) ∥ = 0 . 8696 ) in Fig. 12, which further clearly reveals that there are still similar phenomenons when ∆ = 0 as in the case when the initialization is balanced.

B More Discussions

Our work proposes the rank-1 linear network which is a plausible proxy of standard linear networks with some neurons fully connected with neurons in its last and next layers. By showing that the proposed rank-1 linear networks are standard linear networks with special initialization, our conclusions may be generalized to standard linear networks. In comparison, the diagonal linear network, a special kind of linear networks that receives a lot of attention recently, does not have fully connected neurons. Furthermore, we find that the implicit bias of both GD and SGD for diagonal linear networks are not consistent with ours. The diagonal linear networks also exhibit drastically different implicit bias of GD when compared to standard linear networks, while the conclusions for rank-1 linear networks are consistent with those of standard linear networks. We also reveal the key role of the over-parameterization in characterizing the implicit bias of SGD, namely that it will only be different with that of GD for over-parameterization model.

The inconsistency between the implicit bias of GD and SGD for diagonal linear networks and rank-1 linear networks leads us to suggest intriguing questions for future work such as what about other architectures and is there any unified analytical approach for studying implicit bias of GD and SGD for different architectures? And it is interesting to reveal whether the 'alleviating' initialization effect of the SGD sampling noise is general accross different architectures.

We precisely characterize the implicit bias of both GD and SGD for rank-1 linear networks, where the dependence on the initialization and depth is explicit and clear. In this sense, we take a step forward in the direction of characterizing the implicit bias of optimization algorithms.

Finally, our analysis characterizes the implicit bias of SGD through analyzing the overall parametrization θ . This is different with another line of recent work [4, 10, 19, 27, 9] which focused on the

10

3

10

2

10

1

10 0

||

(0)||

0.0

0.1

0.2

0.3

0.4

0.5

D

(

(

),

*

2

)

=0.0047

GD, non-balanced init

SGD, non-balanced init

GD, balanced init

SGD, balanced init

10

3

10

2

10

1

10 0

||

(0)||

0.0

0.1

0.2

0.3

0.4

0.5

D

(

(

),

*

2

)

=0.0111

GD, non-balanced init

SGD, non-balanced init

GD, balanced init

SGD, balanced init

10

3

10

2

10

1

10 0

||

(0)||

0.0

0.1

0.2

0.3

0.4

0.5

D

(

(

),

*

2

)

=0.0943

GD, non-balanced init

SGD, non-balanced init

GD, balanced init

SGD, balanced init

0.000

0.005

0.010

0.015

0.020

0.025

0.12

0.14

0.16

0.18

0.20

0.22

D

(

(

),

*

2

)

SGD

̸

̸

10 0

10 1

10 2

10 3

10 4

iterations

10 0

D

(

(

t

),

*

2

)

GD, balanced init

SGD, balanced init

GD,

=0.0111

SGD,

=0.0111

GD,

=0.0491

SGD,

=0.0491

̸

flatness of the loss landscape by directly analyzing the independent model parameters, i.e., u and v for the rank-1 linear networks, which is also crucial to fully understand the learning dynamics.

It is worth to mention that removing the balanced initialization (Definition 1) is also a promising direction. As verified by the numerical experiments in Appendix A.4, similar phenomena exist for unbalanced initialization. From the theoretical aspect, balanced initialization enables us to derive the exact dynamics of the overall parameter θ , which is necessary to precisely characterize the implicit bias of GD/SGD. And it is difficult to discuss arbitrary initialization without the balanced initialization assumption. The effect of removing the balanced initialization is that the induced mirror flow potential should be composed of two parts: the original potential presented in Section 3 and a perturbation due to the imbalance of the initialization to it. This implies that the ℓ 2 -norm solution is still returned for small initialization. On the other hand, the case for SGD is much more complicated: the Brownian motion term of the corresponding SDE will also be affected by the imbalance of the initialization, which in turn induces a much more complex time varying mirror flow potential. We believe that the exact theoretical characterization of the implicit bias of SGD without the balanced initialization is a valuable future direction.

There are also some limitations in the current work. For example, although the conclusions of numerical experiments conducted on non-linear networks resemble that of rank-1 and standard linear networks, we can not directly generalize the current theoretical analysis to non-linear neural networks, which normally do not have overall parametrization vectors as θ that is necessary for our analysis. Moreover, the exact characterization of the stochastic integral is also absent in the current work, while we expect that the integral term L ( θ ) tr ( P ⊥ ( θ ) X T X ) θ ∥ θ ∥ in Eq. (12) might have close relation with the property of the training dynamics.

C Proofs for Section 3

In this section, we present the technical details of Section 3. In particular, Section C.1 discusses the balanced initialization, Section C.2 proves Theorem 1 for rank-1 linear networks and Section C.3 proves Theorem 2 for standard linear networks.

For a rank-1 linear network f ( x ; u, v ) = w T L W L -1 · · · W 1 x where W k = u k v T k for any k ∈ { 1 , . . . , L -1 } , recalling the definition Eq. (3) and that the network f ( x ; u, v ) can be written as f ( x ; u, v ) = ρ k v T k +1 u k ρ -k v T 1 x . For convenience, we let ξ = w T L W L -1 · · · W 2 u 1 .

C.1 Balanced initialization

For a rank-1 linear network Eq. (2), dynamics of gradient flow is given by

Based on this set of dynamics, we first discuss he following useful lemma that characterizes the dynamics of norms of model parameters:

Lemma 1. For f ( x ; u, v ) trained with gradient flow, we have

i.e., layer norms grow at the same rate. Furthermore, if ∀ k ∈ { 1 , . . . , L -1 } : ∥ u k (0) ∥ = ∥ v k +1 (0) ∥ = ∥ v k (0) ∥ , we have

Proof. Using Eq.(16), we have

Therefore, both d ∥ u k ∥ 2 dt and d ∥ v k +1 ∥ 2 dt do not depend on k and are same then Eq. (17) follows.

We now discuss Eq. (18). Since we assume ∀ k ∈ { 1 , . . . , L -1 } : ∥ u k (0) ∥ = ∥ v k +1 (0) ∥ = ∥ v k (0) ∥ and Eq. (17) implies that for any t > 0 :

we have

To show Eq. (18), we note that

where we use Eq.(16) in the second equality and the third equality is because ξ = ρ k ρ -k ⟨ v k +1 , u k ⟩ . As a result, the above equation implies that

which does not depend on k , and Eq. (18) follows.

To simplify the analysis, in Theorem 1, we have required the balanced initialization across layers (Definition 1). Recall that the balanced initialization is defined as

Definition 1 (Balanced initialization for rank-1 linear networks). Given an L -layer rank-1 linear network Eq. (2) , for any k ∈ { 1 , . . . , L -1 } , the balanced initialization means that

Eq. (25) states that u k of the k -th layer is aligned with v k +1 of the ( k + 1) -th layer in direction while Eq. (26) means they have the same magnitudes as v 1 (0) . The balanced initialization has been suggested by several previous works [3, 2, 6, 28] for standard linear networks defined as follows.

Definition 2 (Balanced initialization for standard linear networks) . Given an L -layer standard linear network f ( x ; W ) = w T L W L -1 · · · W 1 x , for any k ∈ { 1 , . . . , L -1 } , the balanced initialization means that

for any k ∈ { 1 , . . . , L } .

This directly means that W k +1 (0) and W k (0) share same singular values and W k +1 's right singular vector aligns with the left singular vector of W k (0) . In our case, such reasoning gives us

where Eq. (27) is similar to Eq. (25), which shows that u k (0) aligns with v k +1 (0) , and we adapt the condition (28) to Eq. (26) for rank-1 linear net since v k and u k are the independent model parameters.

A nice property of GD is that the balanced property across layers will be maintained during training [6, 11, 2], i.e., W T k +1 ( t ) W k +1 ( t ) = W k ( t ) W T k ( t ) for t > 0 and k ∈ { 1 , . . . , L } , which can be showed by taking derivative with respect to time on W T k +1 ( t ) W k +1 ( t ) and W k ( t ) W T k ( t ) . For the rank-1 linear network case, according to Lemma 1 and Eq. (25), ⟨ v k +1 ( t ) , u k ( t ) ⟩ are the same for all k ∈ { 1 , . . . , L -1 } , thus GD also maintains the balanced property for rank-1 deep linear networks. Although the balanced initialization conditions are slightly strict, it can be approximately accurate if the initialization scale is not large, which is rather common in practice. Under such initialization conditions, we are able to precisely characterize the implicit bias of GD and focus more on the effects coming from the overall initialization of model parameters, rather than the difference between layers.

C.2 Proof of Theorem 1

In this section, we prove Theorem 1. Basically the idea is to show the existence of a potential function V ( θ ) such that the model parameter θ follows a mirror descent with respect to V ( θ ) :

This method is called infinitesimal mirror descent (IMD) approach and can be found in, e.g., [3]. Note that, to apply this method, the dynamics of θ should satisfy certain condition, which might be strict. For example, given the linear model f ( x ; θ ) = θ T x for θ ∈ R d , both parameterization of θ with standard linear networks, i.e. θ T = w T L W L -1 · · · W 1 , and a much simpler one θ = cv for c ∈ R and v ∈ R d do not satisfy the condition of applying the IMD approach, while the rank-1 linear networks satisfy the condition, which implies that the parameterization of rank-1 linear networks is different with a scalar times a vector. Furthermore, the above example also confirms our motivation of studying rank-1 linear networks as a proxy of standard linear networks, especially considering that the implicit bias of SGD for rank-1 linear networks is more amenable.

We first present a useful Lemma in [20]:

Lemma 2. If H has rank 1 and G is invertible, then

where g = tr ( HG -1 ) .

We now prove Theorem 1.

Proof. Recall that r ∈ R n with r i = ( f i -y i ) , we let

The key step is to derive the dynamics of the overall parameter θ which can be done by noting that

where, according to Lemma 1,

Eq. (35) can be rewritten as

As a result,

ξ 6 / Φ . By taking some simple algebra, we have

Now Eq. (36) becomes

These conditions are now sufficient for us to find the form of the potential V ( θ ) . Suppose that V ( θ ) can be written as

for some vector h and satisfies the following relation:

then Eq. (43) gives us

According to Lemma 2, note that the inverse appeared in Eq. (36) is

It is now left for us to express ξ 2 and Φ in terms of θ . In the following, we assume the balanced initialization in Theorem 1 and apply Lemma 1.

ξ 2 . This can be done by noting that ∥ θ ∥ 2 = ξ 2 ∥ v 1 ∥ 2 , where ξ 2 is given by

Note that ⟨ v k +1 , u k ⟩ 2 grow at the same rate for different k according to Lemma 1 and ⟨ v k +1 , u k ⟩ 2 are the same at initialization for different k due to our assumption, we have ⟨ v k +1 , u k ⟩ 2 = ⟨ v 2 , u 1 ⟩ 2 and ξ 2 = ⟨ v 2 , u 1 ⟩ 2( L -1) . Note that

we have ⟨ v 2 , u 1 ⟩ 2 - ⟨ v 2 (0) , u 1 (0) ⟩ 2 = ∥ v 1 ∥ 4 - ∥ v 1 (0) ∥ 4 . Since we have ⟨ v 2 , u 1 ⟩ 2 = ∥ u 1 ∥ 4 = ∥ v 1 ∥ 4 at initialization according to our assumption, ξ 2 can be finally written as

and the integration relation

where we let ˜ r = -2 | ξ | r/n . Requiring ∇ θ V ( θ ) | θ = θ (0) = 0 and denoting λ i = ∫ ∞ 0 ˜ r i ( τ ) dτ gives us the condition at t = ∞ :

Eq. (48) coincides with the KKT stationary condition of the optimization problem (29). Therefore, we can prove the theorem by deriving the explicit form of V ( θ ) .

Solving V ( θ ) . According to Eq. (44), we can derive the following relation:

Comparing this with Eq. (45), we conclude that

which, by noting that ∂ ∥ θ ∥ ln ˆ V ′ ( ∥ θ ∥ ) = ˆ V ′′ ( ∥ θ ∥ ) ˆ V ′ ( ∥ θ ∥ ) , can be solved as follows

Eq. (45) is automatically satisfied when ˆ V ( ∥ θ ∥ ) has the form of Eq. (51). It is now left for us to get the form of h , which can be done by noting that

Furthermore, since

Thus the final form of V ( θ ) is

This completes the proof.

C.2.1 Remove the assumption of convergence to the interpolation solution

The assumption that Xθ ( ∞ ) = y in Theorem 1 can be removed if the dimension of the span of X T is larger than the number of samples n , i.e., when dim ( span ( X T ) ) ≥ n . 1 This can be proved as follows.

Proposition 2. For the over-parameterized regression of rank-1 linear networks Eq. (2) and the dataset { ( x i , y i ) } n i =1 where x i ∈ R d and n < d , if

then the gradient flow solution θ ( ∞ ) satisfies that

Proof. To prove this proposition, we study the dynamics of the loss function L = 1 n ∑ i r 2 i , where r i = ⟨ θ, x i ⟩ -y i , that is given by

where we have the equality in the last line when r = (0 , . . . , 0) T ∈ R d , i.e., Xθ = y , or X T r = 0 , which is not possible since we have assumed that dim ( span ( X T )) ≥ n . Therefore, L ( θ ( t )) keeps decreasing until Xθ ( t ) = y , i.e., until GD finds the interpolation solution. Noting that min θ L ( θ ) = 0 , we complete the proof.

C.3 Proof of Theorem 2

In this section we prove Theorem 2. The techniques are similar to those in C.2, and we still need a time wrapping technique introduced in [3] to derive the form of V std ( θ ) , since the condition for applying the IMD approach is violated in this case. Recall that, for the standard linear networks f ( x ; W ) = w T L W L -1 · · · W 1 x = θ T x where W k ∈ R d k × d k +1 , our purpose is to find a potential function V std ( θ ) such that the gradient flow solution θ ( ∞ ) satisfies that

Since we assume the balanced initialization (Definition 2), then according to [11, 6, 2], the norms of all layers grow at the same rate and are the same for any t > 0 :

Furthermore, following the procedure of [2], we obtain that the dynamics of θ is

We now present the proof.

Proof. Eq. (55) can be written as

where, according to Lemma 2, the inverse in above equation is given by

As a result, we have

In this case, we still assume that V std ( θ ) has the following form

for some constant vector β ∈ R d . Then following a similar procedure as in C.2, we have that

In the IMD approach, V std ( θ ) should satisfy that

which requires that ∂ 2 θ V std ( θ ) = ∥ θ ∥ 2( L -1) L ( I -L -1 L θθ T ∥ θ ∥ 2 ) dθ dt . But this is not possible. Therefore, we multiply a time re-scale factor g ( θ ) , as long as g ( θ ) is positive, to both sides of Eq. (57) and only require that V std ( θ ) satisfies the above relation under the new time scale τ : R → R such that τ ′ = g ( θ ) :

Then the limit point at t = ∞ in Eq. (61) is also visited at the point τ = ∫ ∞ 0 g ( θ ( s )) ds in Eq. (62). We now solve the explicit form of V std ( θ ) that satisfies Eq. (62). By comparing Eq. (60) and the right hand side of Eq. (57), we obtain that the following relation should be satisfied:

This implies that we need:

the terms in the bracket on both sides should match:

for some constant C and C ′ , where we can simply choose C ′ = 1 ;

the terms outside the bracket on both sides should also match:

To obtain the form of the constant vector β , we note that ∂ θ V ( θ (0)) = 0 , which immediately gives us

Combining all these terms, we have the final form of V std ( θ ) :

As in C.2, denoting λ i = ∫ ∞ 0 r i ( s ) ds and noting that ∂ θ V std ( θ (0)) = 0 give us

which is exactly the KKT stationary condition of the optimization problem (54).

Convergence to the interpolation solution. Similar to C.2.1, we can also show the convergence to the interpolation solution when dim ( span ( X T ) ) ≥ n by deriving the dynamics of L :

Since we assume that dim ( span ( X T ) ) ≥ n , d L /dt < 0 until Xθ ( t ) = y , i.e., L keeps decreasing until GD finds the interpolation solution.

C.4 Proof of Proposition 1

In this section, we prove Proposition 1 by analyzing the gradient of model parameters. It is helpful to recall that for a matrix A , we use A ij to denote its i -th row j -th column element. For weight matrices, e.g., W k , we use W k ; ij to denote its i -th row j -th column element.

Proof. For a standard linear network that has the initialization of Proposition 1 and L is an odd number, i.e., for an integer p

̸

̸

note that the L -1 -th layer satisfies that W L -1; ij = W L -1; ij δ jc L -1 where δ jl = 1 if j = l otherwise δ jl = 0 , we can write the networks at t = 0 as

Then the gradient w.r.t W L -1 at t = 0 is

̸

This means that the parameter W L -1; ij will be updated only when j = c L -1 . As a result, the initialization shape for W L -1 will be maintained, i.e., only the non-zero column of W L -1 at t = 0 will be updated and all other elements of W L -1 will be zero for any t > 0 since the corresponding gradients vanish. Similarly, for W L -1 , we can write the network at t = 0 as:

then

̸

and the initialization shape for W L -2 will also be maintained. Following a similar procedure, we conclude that all the initialization shapes will be maintained for any t > 0 .

We now consider the diagonal initialization for weight matrices, namely that

̸

For any k ∈ { 1 , . . . , L -1 } , the gradient w.r.t W k at t = 0 is

where, clearly, there is not any constraint on which elements of ( w T L · · · W k +1 ) T ∈ R d and ( W k -1 · · · W 1 x ) ∈ R d are zeros if we do not require that many diagonal elements of the initialization

matrices are zeros, which clearly violates the diagonal initialization requirements. Thus we can not conclude that

̸

i.e., the off-diagonal elements of W k will also be updated to be non-zeros. Thus the diagonal initialization of weight matrices can not be maintained. The conclusions for other layers can be easily derived by following similar arguments.

D Proofs for Section 4

In this section, we present the technical details for Section 4. In particular, we introduce our modelling details of the SGD dynamics in D.1, derive the SDE of the model parameter θ in D.2, and, finally, prove Theorem 3 in D.3.

D.1 SDE modelling details

Similar to the case of GD, to derive the implicit bias of SGD for rank-1 linear networks, we need to first derive the dynamics of the overall model parameter θ . For this purpose, the structure of the noise of SGD is crucial. For convenience, we first discuss the basic SGD where only one data is randomly sampled at each step, while the generalization to batch-size SGD is straightforward.

Structure of the noise. We present the details for u k , and the case for v k +1 is similar. Recall that the empirical loss is L = ∑ i ℓ i / n, η is the learning rate and that the network can be written as f ( x ; u, v ) = ρ k ρ -k v T k +1 u k v T 1 x , we start with the SGD update equation for u k where we let j t denote the index of the sampled data at the t -th step:

where we let → e j t be the basis vector in R n such that the j t -th element is 1 while all other elements are 0 and

As a result of this, the noise of SGD for u k is now

Continuous Modelling of SGD. The continuous modelling techniques for SGD have been widely applied in recent works [1, 9, 23, 22] to study the dynamics of SGD. In our setting, the continuous counterpart of SGD is established as follows. First, the discrete SGD updating equations Eq. (8) can be equivalently written as

for k ∈ { 1 , . . . , L -1 } where both ∇ u k L ( θ ) - ∇ u k ℓ i ( t ) ( θ ) and ∇ v k +1 L ( θ ) - ∇ v k +1 ℓ i ( t ) ( θ ) are zero-mean noises with covariance matrices in R d k × d k

with ρ k and ρ -k defined in Eq. (3). Second, as we identify the noise covariance, letting η → 0 2 , then we obtain the continuous counterpart of the discrete SGD that is a set of stochastic differential equations (SDE):

where we let r = ( f ( x 1 ; u, v ) -y 1 , . . . , f ( x n ; u, v ) -y n ) T ∈ R n be the residuals and W t is a standard Brownian motion in R n . Similarly, recalling the definition of ξ = w T L W L -1 · · · W 2 u 1 , the SDE of v 1 is

Generalization to batch-SGD. When b (a positive constant) data points B t are sampled in each iteration of SGD, i.e., batch-SGD, we can change the SGD update equation as follows (taking u k as an example)

This only changes the noise Σ( u k ( t )) to a batch version

which only affects the noise part of the SDE of u k and leads it to become a batch version SDE:

This is equivalent to re-scale the learning rate η to

and leaving other parts unchanged. Thus the generalization to batch-SGD is straightforward-simply replacing all η with η b .

D.2 The SDE of θ

In this section, we carefully derive the continuous dynamics of SGD for the parameterization of our rank-1 linear networks. We first discuss the balanced initialization condition.

Balanced initialization. Similar to the case for GD (Definition 2), we also assume the balanced initialization Eq. (26) across layers. Although the dynamics of SGD is different with that of GD, it still applies the gradient to update the parameters at every step that will maintain the balanced property thus the dynamics of SGD will also maintain the balanced property, i.e.,

for t > 0 and k ∈ { 1 , . . . , L -1 } , during the training of rank-1 linear networks.

With the equations for u k and v k +1 , we now derive the SDE of θ summarized in the following lemma.

and

Lemma 3. For an L -layer rank-1 linear network Eq. (2) , if we assume balanced initialization, then the stochastic gradient flow of θ is

where H ( θ ) = I +2( L -1) θθ T / ∥ θ ∥ 2 .

Proof. According to the Ito's Lemma, we have

Thus to obtain the SDE of θ , we need to analyze every term of the above equation. We first give dξ .

The form of dξ . Let ω k = v T k +1 u k and ψ k = ∥ v k +1 ∥ 2 + ∥ u k ∥ 2 , we first characterize the SDE of ω k . According to the Ito's calculus, we obtain that

where, by applying Eq. (73) and Eq. (74) and noting that ( d W t ) 2 = dt ,

Combining the above two terms gives us the SDE of ω k

Since ξ = d ( ∏ L -1 k =1 ω k ) , its SDE can be done by repeatedly applying the Ito's Lemma and the SDE of ω k :

̸

For convenience, we first define several helper notations:

̸

Now plugging the form of dω k into ♠ gives us the first term of dξ :

and applying again ( d W t ) 2 = dt and the form of dω k gives us each term of the second sum of dξ :

Summing all ♡ and ♠ , we obtain the SDE of ξ :

The SDE of v 1 is much simpler. To get this, we start with the SGD update equation for v 1 :

which implies that the noise covariance in this case is

Then using a similar approach as that of u k , we get the SDE of v 1

Now it is sufficient for us to derive the form of dθ .

The form of dθ . Combined with the SDE of ξ , we now have

For the ♣ term, as we already have the form of dv 1 in Eq. (82) and dξ in Eq. (80), we simply plug them into ♣ and obtain that:

For the ♠ term, we only need to consider the d W t terms of Eq. (82) and (80):

Combining the above two equations gives us the final SDE of dθ :

On the other hand, recall that our assumptions regarding the initial conditions of f ( x ; u, v ) in the Lemma (Definition 1) and following similar techniques as in the case for GD, we have that for any t > 0 :

Plugging these terms back to the definitions of Φ 1 , Φ 2 and Φ 3 , we obtain that

where we use that ∥ θ ∥ = ∥ ξv 1 ∥ = | ξ |∥ v 1 ∥ . Therefore, the final form of dθ is now

where

D.3 Proof for Theorem 3

In this section, we determine the form of V S ( θ ) such that θ follows a stochastic mirror flow

which then proves the claims of Theorem 3. For this purpose, we start with manipulating the SDE of θ derived in Lemma 3. Note that Eq. (85) can be written as where

as in the case for GD and

To solve the inverse appeared in the above equation, we apply Lemma 2 and noting that

then

where we use that θ = ξv 1 in the last equality, which enables us to simplify P :

Thus Eq. (91) for the overall SDE of θ now becomes

The balanced initialization gives us Eq. (87), (88), and (89), and recall that

thus we can further rewrite

and

Thus, Eq. (91) now 3 can be rewritten as

Finding the form of V S ( θ ) . We now proceed to find the form of V S ( θ ) . For convenience, we apply a time re-scaling technique such that dτ = | ξ | dt ( d W τ = √ | ξ | d W t according to [12]) to the above equation. For convenience, we still use t to represent the time after re-scaling. Then the above equation becomes

Recall that θ (0) ∈ R d is the initialization of θ , since V S ( θ ) for SGD should have similar form as that for GD, we borrow from the GD results and first define a constant vector γ ∈ R d

Suppose now that V S has the following form:

by similar techniques as in the case of GD, we obtain the first and second derivatives of V S ( θ ) w.r.t θ :

There exists a function of θ , G ( θ ) whose exact form can be derived but not necessary for us, corresponding to our V S ( θ, t ) defined above, such that d∂ θ V S ( θ, t ) can be written as

Now suppose that we can choose a particular g ( t ) , thus a particular G ( θ ) , such that the following relation is satisfied:

then we can rewrite the SDE of ∂ θ V S ( θ, t ) as follows:

which will give us the desired 'mirror flow' equation:

It is now suffice for us to find the particular g ( t ) that makes Eq. (98) satisfied. For convenience, recall that for a vector a we use a µ to denote its µ -th component, we define the following helper notations:

where the exact form of O can be obtained from Lemma 3 but is not necessary for us. In the following, we use D to represent D ( θ, t ) for convenience. For our purpose of choosing the particular g ( t ) , dD ( θ, t ) = d∂ θ V S ( θ, t ) should match the R.H.S of Eq. (99). To make this relation clear, according to the Ito's Lemma, we obtain that

where the last term is crucial. From the SDE of θ (Lemma 3), we have that the dθ ρ dθ σ appeared in ♠ can be written as

On the other hand, according to

we have that, with explicit subscripts of vectors,

Note that the above expression implies that ∂ 2 D ∂θ∂θ is in fact a rank-3 tensor. Combined with the expression of dθ ρ dθ σ derived above, we have that

where H ∈ R d × d × d is a rank-3 tensor. To get the exact form of ♠ , we need to have the exact forms of

1. ∑ ρ,σ,i H ρ µσ B σi B iρ . Note that HBB ∈ R d and there are two different terms, both of which will induce a same vector. In particular,

Thus we have where

2. ∑ ρ,σ,i P σ ρ B i σ B ρ i . Using the matrix notation, we can easily find that

where, recall the definition of B in Eq. (101),

Taking the trace of the above equation gives us the second term of ♠ :

Now removing the µ subscript of the derived ∑ ρ,σ,i H ρ σ,µ B σi B iρ and recovering the matrix notation, we have the final form of ♠ by combining it with the result of tr ( B T PB )

Thus the SDE of D now becomes

To find the particular g ( t ) such that Eq. (98) is satisfied, we need to require that

Recall that D = ∥ θ ∥ -λ L θ + γ + g ( t ) , we have

which, noting that k 2 = 4 ξ 4 η L /n , further gives us that the following relation should be satisfied

As a result of this, we can, noting that ξ 2 = ∥ θ ∥ 2 λ L from ∥ θ ∥ 2 = ξ 2 ∥ v 1 ∥ 2 , give the required g ( t ) that makes Eq. (98) satisfied by solving the following equation:

Now let the orthogonal projection operator of θ be P ⊥ ( θ ) = I -θθ T ∥ θ ∥ 2 , then we can solve g ( t ) as

With this particular g ( t ) , we can then give V S ( θ, t )

that satisfies the relation

Moreover, a particular interesting case is that, as L →∞ ,

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper describes the implicit bias of GD and (a continuous approximation of) SGD in balanced rank-1 linear networks, which are linear networks where every second layer has only 1 neuron and there is only one output neuron. This bias takes the form of a potential equal to one term that scales with the parameter norm, a second term that describes the influence of the initialization and a third term that describes the effect of SGD and scales with the learning rate.\n\nIt is then argued that this model is a better approximation of general linear networks than previous work which describes the implicit bias of diagonal linear networks (showing a sparsity bias for small initializations that vanishes for large initializations).\n\nSTRENGTHS:\nThere are few models where the implicit bias of GD/SGD can be described with an explicit formula, it is nice to have an explicit formula. The mathematical results appear to be all correct.\n\nThe result are discussed in details, though I do not agree with some of the main conclusions made from the results.\n\nWEAKNESSES:\nI strongly disagree with the general point that the result presented in this paper are a better description of the implicit bias of general linear networks, outside of the case where there is only one output, and I think that the one output case is not very interesting as can be seen from a simple symmetry argument.\n\nLet me first describe why single output linear networks are not very interesting. In this setting the only thing that matters is the value of $\\theta$ orthogonal to the data span. Linear networks are invariant under orthogonal transformations of the inputs, thus if we rotate the first layer weights around the span of $X$ the final learned $\\theta$ will be rotated by the same rotation. Thus assuming a rotation invariant initialization of the weights (which is the case for Gaussian initialization and approximately true for other traditional initializations), we know that the distribution of the learned $\\theta$ is invariant under rotation, the expected $\\theta$ will thus minimize the $L_2$ norm, and the only interesting thing is the amount of noise. This paper answers this second question (the amount of noise and also how it depends on the initialization $\\theta_0$), but randomness w.r.t. the initialization of the parameters is always detrimental (the expected risk equals the risk of the expected predictor plus the variance). There is thus only disadvantage to using a deep linear network instead of a traditional linear model ($L=1$).\n\nI therefore think that linear networks are only interesting when there are multiple outputs, in which case they exhibit a low-rank bias (see citation [13] from the paper) which cannot be observed when there is a single output. And I would argue that diagonal linear networks are closer to these multiple outputs linear network than the model presented in this paper. Indeed they both feature a similar transition between a NTK regime without sparsity for large initializations, and for small initialization a dynamic that jumps from saddle to saddle [Jacot et al., Saddle-to-Saddle Dynamics in Deep Linear Networks: Small Initialization Training, Symmetry, and Sparsity], leading to some form of sparsity. Of course the resulting sparsities are different, with a sparsity of the entries of the learned vector for diagonal networks and a sparsity of rank for linear networks.\n\nThe authors motivate the study of rank-1 linear networks for the insight they give into models we care about: general linear networks or more importantly nonlinear networks. But there already exists (incomplete) description of the implicit bias of linear networks [13], of shallow nonlinear networks [Bach, Breaking the Curse of Dimensionality\nwith Convex Neural Networks; Abbé et al., The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks], and even for deep nonlinear networks [Jacot, Implicit Bias of Large Depth Networks: a Notion of Rank for Nonlinear Functions]. If the authors think that their approach better describes these bias or could lead to better descriptions, they should motivate this at least a little bit.\n\nFurthermore, the fact that in these settings the implicit bias cannot be expressed explicitely suggests that the models such as diagonal linear networks and rank-1 linear networks where such explicit formula exist are exceptions, not the rule. I personnaly think that the techniques used to derive such explicit formulas have little hope of extending to these more complex cases, and some alternative techniques (as cited above) have already shown much more potential.\n\nTo summarize I think that the results, while correct and precise, only describe the implicit bias of linear networks in a very specific and not particularly interesting case. And I do not think the results of this paper are very representative of the more interesting cases, since they fail to capture any form of sparsity, even though different forms of sparsity have been observed and seem to play avery important role in these more interesting settings.\n\nQUESTIONS:\nIf I understand correctly Theorem 1 and 2 state the same thing for rank-1 linear networks and general single output linear networks, so only the SGD result (Theorem 3) is restricted rank-1 linear networks. Do you expect a similar direct generalization for the SGD case? \n\nSince you assume balancedness anyways, the Theorem 1 of [Arora et al., On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization] shows that the widths of the network does not matter for the dynamics of $\\theta$ as long as they are larger than 1. Since rank-1 networks are linear networks where every second width is 1, this suggests that everything should extend directly to general single output linear networks. In that case, I think it would be clearer to state everything for single output linear networks, this way there is no need to introduce a whole new model of rank-1 linear networks and argue that they behave similarly.\n\nLIMITATIONS:\nThe authors only cite [13] but do not compare the results, nor explain why they observe something very different (no sparsity). They should explain that the reason they cannot observe rank sparsity is because they only have 1 output. Instead some of the conclusions from the paper could be interpreted as saying that linear networks have no sparsity bias (when they emphasize that linear networks minimize the $L_2$ norm instead of the $L_1$ norm) which we know is not true when there are multiple outputs [13].", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduce a new simple model of neural networks: the *rank-1 linear neural network*. This model consists of $L$ layers of linear networks, each of the form $u_kv_k^\\top$ (a rank one matrix parameterized by two vectors).\nThe paper motivates this model by shortcomings of other simplified models (such as diagonal linear networks) and by the natural low rank structure of the solution found by GD over linear networs.\nAssuming convergence of the gradient flow (and stochastic gradient flow respectively) to an interpolator, the solution found is entirely described by an implicit regularization problem that depends on the initialization and the depth of the network (and on the noise and the trajectory of the iterates respectively). These results are then discussed to derive some insights on the role of noise, depth and initialization.\n\nSTRENGTHS:\nStudying simple models of neural networks is a very interesting question, that has been extensively studied in the previous couple of years. In this line of work, the present submission takes an orthogonal step and introduces a new model, which I find very interesting.\nThen, the authors try to relate their results to the roles of important parameters of (stochastic) gradient flow: noise, depth and initialization.\n\n\nThe analysis is classical: deriving a (stochastic) mirror flow (with varying potentials) evolution for the model $\\theta$, in order to directly have the potential that the model minimizes. This approach and the whole paper are rigorous, state their results in a clear and concise way, with clear notation, leading to a paper that is easy to follow. I think that the writing is here a strength of the paper.\n\nWEAKNESSES:\n1) I am not totally convinced by the motivation of rank-1 linear networks. This model is motivated by noting that linear NNs are biased towards low-rank structure (the conjecture is that they minimize some kind of weighted nuclear norm, for small initializations). But this is the case for small initializations, where they are in the so-called ‘‘rich regime''. Thus, does imposing rank 1 + arbitrary scale really model better real NNs and linear NNs and their implicit biases than diagonal linear nets ?\nThen, I kind of disagree with the noted shortcomings of diagonal linear networks: DLNs make appear a transition between rich and kernel regimes (Woodworth et al. Colt2020): for small initializations, sparse vectors are promoted ($\\ell^1$ regularization), while for large initializations, they are in an $\\ell^2$ regularization, which amounts to the lazy training regime. These phenomenons (rich and lazy regimes) appear in linear networks, but also in more complex architectures.\n\n2) Throughout all the papers (and in particular in the title!), gradient flow and stochastic gradient flow (the processes studied in the paper) are replaced by GD and SGD, which is misleading ! We indeed expect GD and SGD to be much more complicated than GF and SGF to study.\nRelated to this, I think that (i) GD and SGD should be replaced by GF and SGF when appropriate, and (ii) the authors should mention the recent work \n*S) GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stability*, Even et al. 2023, that studies GD and SGD (with non infinitesimal learning rates) on DLNs, when discussing the implicit bias of GD and SGD over DLNs in related works.\nIndeed, this work derives the implicit bias for GD and SGD, and show that the nature of the two is totally different from GF in the rich regime.\n\nQUESTIONS:\n3) About the results of the implicit bias, I think the influence of the different parameters should be related to phenomenons that arise in more complex networks. Can the influence of scale be related to any rich or lazy regimes ? Or does imposing rank-1 matrices forbids this ?Also, for $L\\to\\infty$ in theorem 1, does this amount to aligning the solution to the initialization (for any scale) ?\n\n4) Can the tools used by *S) GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stability*, Even et al. 2023, to derive implicit biases with non-infinitesimal stepsizes in the rank-1 linear model ? If this does not seem too hard, I think this could be a good remark to write (just as a remark, I do not suggest the authors do the whole analysis agai of course).\n\nOverall, I think the paper is well written. The maths and the results are easy to follow, and the ideas are novel and interesting. However, I have several concerns regarding the relevance of the proposed model. I think it is a bit overstated (see point 1. above). If the authors modify their paper in a revised version in order to include suggested comments, I think this paper should be accepted, for the originality of the proposed model. I don't think (I am quite sure in fact, but maybe the authors can change my mind about it) that this model is more relevant than DLNs as proxy for more complex architectures, so I think the model could be just motivated as a step towards the implicit bias of linear networks, and not as better proxy than DLNs.\n\nLIMITATIONS:\nNo such limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates the implicit biases of gradient descent (GD) and stochastic gradient descent (SGD) when training multilayer linear networks. It presents an analysis of specific potential functions that are minimized during the training process to achieve a perfect fit of the labels. Notably, the potential function's form is contingent on the depth of the network, providing insights into the impact of network depth on the implicit biases induced by GD/SGD.\n\nSTRENGTHS:\n-- The paper provides concrete results characterizing the actual norms that are being implicitly minimized by GD when training multilayer linear models. I think it is a valuable question to tackle as we lack an understanding of what are the actual implicit biases of GD and to understand what are the complexity measures that are being minimized by our optimizers.\n-- The paper is well-written and discusses the literature quite well. It discussed the relevant papers that it is based on and its position against them.\n\nWEAKNESSES:\n-- The paper focuses on linear neural networks with a rank-1 constraint for each layer and assumes convergence to a perfect fit. While this specific setting is analyzed, it would be valuable to discuss the broader implications of the results. Considering that practical neural networks typically employ non-linear activations and do not necessarily have the rank-1 constraint, it would be interesting to explore if the findings can shed light on these more realistic scenarios.\n-- The experimental details in the paper require clarification. It is unclear what specific settings were used, such as the optimization hyperparameters and the dataset employed (is it a real-world dataset or a synthetic one).\n-- The paper is missing some relevant related work. For example, in https://arxiv.org/abs/2206.05794 they showed that when training a neural network with SGD and weight decay provably induces a low-rank bias in the weight matrices.\n\nQUESTIONS:\nPlease see the weaknesses.\n\nLIMITATIONS:\nThis paper has no negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes the rank-1 linear network, a linear neural network where the even layers are restricted to row vectors and the odd layers are restricted to column vectors, as a new proxy of the deep neural networks.\nIn the first part, they gave two reasons why rank-1 linear networks are (said to be) close to linear networks:\n(i) They analyze the deterministic gradient flow and find that the minimum L2-norm solution is preferred when the initialization scale is small, which is similar to linear neural networks and dissimilar to diagonal linear networks.\n(ii) Also, they show that when a linear neural network is initialized to have a rank-k structure (k rows are non-zero in even layers and k-columns are non-zero in odd layers), this structure is kept during update with gradient flow, while diagonal structure is not kept.\nIn the second part, they analyzed the SDE update of rank-1 linear networks and give the explicit form of dynamics and the asymptotic solution.\nThe form includes a term depends on both stochasticity and depth, and implies that the effect of initialization gets smaller as the depth goes to infinity.\n\nSTRENGTHS:\n### The paper is well structured and provides convincing motivation for the setting, although the topic is a little niche.\n\n(I think) the biggest result of this paper is the analysis of SGD solutions of rank-1 linear networks, and main critics would be about the validity of the setting and technical contributions.\nThat is, a rank-1 linear network is a special case of linear neural networks, and thus it is a \"proxy of proxy\" of practical deep neural networks, and the proofs are not so technically insightful because it somewhat heavily depends on the special initialization.\nHowever, before considering the SGD dynamic, they spare a certain space on motivating the analysis of rank-1 linear networks by analyzing the GD dynamic and thoroughly comparing with linear neural networks than diagonal linear networks.\nThe discussion following each of theoretical results, which compares rank-1 linear networks with other architectures, is interesting, and this paper constructs an acceptable story on why this network architecture is important as a collection of theorems.\n\nThe \"motivation\" part mainly consists of two theorems.\n\n- First, they derives the explicit form of the minimizer, which is attained as a result of deterministic gradient flow.\nThe form consists of the depth, L2-norm, and the inner-product with the initialization.\nBy changing the depth, different properties of the solution can be observed.\nWhen the depth is finite, the solution minimizes a sum of L2-norm and a term reflecting initialization (among the zero-loss solutions).\nAs the network gets deeper, the effect of the latter gradually vanishes.\nOne of the important implications of this result is that the minimum L2-norm solution is preferred when the initialization scale is small, which is similar to 2-layer linear neural networks while diagonal linear networks are known to do so when the scale is large.\n\n- Second, they show that rank-k property can be preserved in standard linear neural networks during optimization. They consider the setting when a linear neural network is initialized so that k rows are non-zero in even layers and k-columns are non-zero in odd layers, and is trained with deterministic gradient flow. They show that this structure is preserved during optimization, while if the linear network is initialized so that only diagonal elements are non-zero, the GF does not keep this structure.\n\nIn these senses, they argue that rank-1 restriction should be an acceptable proxy of linear neural networks, (compared to diagonal linear networks).\n\n### It is an interesting result that the dependency on the scale of initialization can be mitigated by the depth and stochasticity. \n\nAs a result of the analysis of SDE dynamics of rank-1 linear networks, they obtain a transformation of the SDE and the explicit form of the asymptotic solution by letting the time go to infinity, under assuming the special initialization.\nThere exists a term that vanishes when either the depth is one or stochasticity is removed.\nWe can interpret this term (not by a rigorous mathematical proofs and with some leaps, though, ) as an implicit bias mitigating the effects from the initialization direction and making the solution closer to the minimum L2-norm solution compared to the GD under the same depth.\n\nOverall, in both GD and SGD, they showed that letting the depth go to infinity decreases the effects of initialization, and in SGD, with some intuitive discussion, they showed that this effect is strengthened by stochasticity.\n Although this is only a possible \"proxy of proxy\" of DNNs in practice, but this paper at least poses an interesting question of whether this proposed benefit of depth/stochasticity holds in practical DNNs. \n\n### Sufficient experiments as a theory paper.\n\nAlthough I personally want to see more experimental results for non-linear DNNs, I think their main theorems are sufficiently verified with the experiments.\n\nOverall, this paper is well-written. Although this is not a paradigm changing paper, I think that accepting this paper will attract ML theory community and encourage reconsideration of proxy-models.\n\nWEAKNESSES:\n### The balanced initialization is somewhat a strong assumption.\n\nThis paper assumes a special initialization of the network.\nI know that this kind of assumptions are popular among the analysis of standard linear networks.\nHowever, because the paper is proposing a new architecture as a proxy of standard linear networks and an alternative of diagonal linear network, and the main contribution is the new implicit bias jointly induced by stochasticity and depth, I think whether the implicit bias is unique to this initialization should be addressed.\nI would like to know more experimental results about whether this initialization is necessary, and more theoretical discussion on this point.\n\n### In Proposition 1., it seems that sparseness is more important than being rank-1. Thus I doubt that real-world NNs have this rank-1 structure.\n\nIt seems that even if rank(W_i) = 1 at initialization in Proposition 1., this does not guarantee that rank(W_i)=1 during optimization. Rather, having zero entries seems critical in proving the theorem. I think it is possible that matrices in NNs becomes rank(W_i)=1 during optimization, by using L2-reguralization and learning single-index models, for example. However, in most cases, it seems that real NNs cannot become this \"rank-1 structure\". Thus I am not sure whether this theorem is so meaningful.\n\nQUESTIONS:\n- When a small perturbation is added to the balanced initialization, then does the discrepancy between the balanced initialization and the perturbed one will grow during optimization? Are there theoretical or empirical observations? \n\n- Can you derive the convergence rate guarantees?\n\nLIMITATIONS:\nThe main limitations of this papers are that the architecture is a \"proxy of proxy\" of practical DNNs and that a special initialization is requited.\nFor the first, the author clearly remarks the limitation about extension to nonlinear settings, and I indeed think that this limitation is unavoidable. \nFor the latter. I hope the authors address the above question during the discussion period.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}