Implicit Bias of (Stochastic) Gradient Descent for Rank-1 Linear Neural Network
Authors: Bochen Lyu
Zhanxing Zhu
Venue: NeurIPS 2023 poster (2023)
Keywords: implicit bias
gradient descent
stochastic gradient descent
linear networks
TL;DR: We investigate the implicit bias of GD and SGD for a novel linear network that can be seen as a proxy of standard 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
Bochen Lyu
Zhanxing Zhu Changping National Lab & Peking University Beijing, China zhanxing. zhu@pku edu.cn
DataCanvas Lab, DataCanvas
Beijing, China
bochen. wateenn1@@mmal...cmmmmmmmm com
Abstract
Studying the implicit bias of gradient descent (GD) and stochastic gradient descent (SGD) critical to unveil the underlying mechanism of deep learning. Unfortu- nately, even for standard linear networks in regression setting, comprehensive characterization the implicit bias is still an open problem. This paper pro- poses to investigate a new proxy model of standard linear network, rank- I linear network, where each weight matrix is parameterized as a rank-1 1 form. For over- parameterized regression problem,
precisely analyze the implicit bias of GD and SGD-by identifying a "potential" function such that GD converges to its min- imizer constrained by zero training error (i.e., interpolation solution), and further characterizing the role the noise introduced by SGD in perturbing the form this potential. Our results explicitly connect the depth of the network and the initial- ization with the implicit bias of GD and SGD. Furthermore, we emphasize a new implicit bias of SGD jointly induced by stochasticity and wver-parameterization. which can reduce the dependence of the SGD's 's solution on the initialization. Our findings regarding the implicit bias are different from that of 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 not. This suggests that the proposed rank-1 1 linear network might be plausible proxy for standard linear net.
Introduction
Gradient Descent (GD) and its stochastic variant, Stochastic Gradient Descent (SGD), 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 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 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,
F( W) WL Wi
where W₂ 's are weight matrices and is the input data. The linear network can be obviously treated as an over-parannn paraterrmeter model of standard linear model, 0" = Wz ... W1. However, the introduction of the over- over-paraaae parameterization brings significant non-convexity, and thus complicates dynamics during learning. Until now, the direct analysis of implicit bias GD and, particularly, SGD for standard deep linear networks with any depth and initialization regression problems is still open problem.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(a)
(b)
Figure 1: (a). Standard linear networks. (b). Rank-1 linear networks. One neuron in the middle
hidden layer is active and fully connected. (c). Diagonal linear networks. Every neuron is active but not fully connected.
In order to study the implicit bias of linear networks in regression settings, simplified version of standard linear network, diagonal linear network (i.e., each neuron is only connected with a single neuron the next layer [26] 222), was proposed as a proxy to underpin the bias of the standard one. With this simplified model, the solution selected by GD solves constrained norm minimization problem that interpolates between the & and lg norms up to the initialization scale (26). For SGD, the solution is closer 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 tractable analysis and provide insights for investigating the implicit bias of SGD and 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. [ showed that SGD and weight decay jointly induces a low-rank bias the weight matrices when training neural network. In fact, for linear networks with single output trained with GD, all weight matrices will automatically become rank-1 and will maintain this property during training when the initialization balanced 11
Our contributions. Inspired by this low-rank bias and to avoid the complicated analysis of standard linear networks, we propose novel model rank-1 deep linear network as a plausible proxy of standard linear networks, to reveal the implicit bias of GD and SGD over-vorannn wer-paramettrized linear models for regression problems. A depth-L L rank- deep linear network of (x; u, v) is defined by
u Wr. W. W = 0=x,, 8.1. W₁ = uwt Vk E (1,...,--1) b-1
where uy Rd and Us E dR are vectors while u and are denoted collections of us and Uk. respectively. See Figure]]oo 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 (0) 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(0). 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(z u, v) achieving zero training error (Theorem [T. The single layer case recovers the standard linear regression results, while a depth larger than one immediately changes the form of ((0), 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 importantl, 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, c.g., GF prefers lo 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 saper-parameters such as the learning rate, batch size and network depth compared to GF. When L > this extra effect "alleviates" the influence of the initialization, i.c., the final solution reduces its dependence on the initialization. More intriguingly, when L = 1, i.c., without over- over-parameiest 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
we summarize the notations and the setup of our work. Section and Section present our main results where Ixction[Jocuue on the implicit bias of GD while Section is about SGD. Numerical experiments are presented in Section S and we conclude this paper in Some technical details are deferred to Appendix.
Organization. In Section 2
Section 6
Related Works
The study of implicit bias of GD has been pioneered by (2 on linearly separable classification problem and was generalized to GD for deep neural networks and different training strategies
[1 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 & minimum norm solutions
and & 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 26 also investigated the influence of the label noise in the diagonal linear networks setting. For GD and SGD with moderate learning rate, M studied their implicit bias for diagonal linear networks and revealed the corresponding influence of the finite learning rate.
This paper considers the rank-1 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 the diagonal linear networks 3 26 22), we point out that the results 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 standard linear networks with special initialization. Furthermore, the classification problem [4 18 5 and linear regression problem 1 have been comprehensively investigated and we focus on the werrparaaeterizee regression. Finally, 1O 19 27 9 focused on the flatness of the loss landscape and model parameters while our analysis is for the overall parametrization 0.
2 Preliminaries and Setup
Notations. Given dataset of n samples (11, 1 2, E IR represents d-dimensional data vector with scalar label yı IR. We use X e Rgd to denote the data matrix and use y (y1...., ,Ba)
3

---

to denote the collection of V/. ||lll denotes the no-morm and || p is the Frobenius norm. (a. b) is the inner product. For any vector u and matrix A, we use u(0) and A()) to denote their initialization. We let () be the trace operator. For the weight matrix W4 IR xdv+1, Waii is its i-th row j-th column element. For parametric model f(a;; 0) 0 with parameters 0, use [()) : Σ (0) to represent the empirical loss where (0) is the loss function for the sample (vi, Bi)
Over-parameterizet regression. We focus on the regression problem where < and assume the existence of solutions that perfectly fit the dataset, i.e., there exist interpolating parameters 0 such that Vi E {l,..., n (zi, 0*) = y4. The quadratic loss (0) is used in our setting, The empirical loss is then (()) = Ei=1 4(0) = - E (ys ())
Rank-1 deep linear networks. In this paper we consider the rank-1 1 deep linear network f(;; u, v) (see Eq. (2). We use and v to denote the collections of ug 's and 's respectively. For convenience,
Pk: w W 1k+24 PL-1 1, and p_a WK-1- W2u1, P-1
such that the network can be written as f(z; v) - Pk 2 Through this paper, we treat the depth as a parae--er of the network and try to precisely characterize its role in the implicit bias.
u 1 = T x for l E fI....,IL L-1) when
Definition 1 (Balanced initialization for rank-1 linear networks). Given an L-layer rank-1 I linear 1}, the balanced initialization means that
network Eq. 2 for any & E {1......
(U++ (0), uA (0)) 6 (0) uu () (0)
1, [1)]1 [0)]| 11 (0)ll-
We add more discussions on the initialization in Appendix
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 the continuous version of GD for the rank- deep linear networks f(;; u, v) Eq. (2). Note that the model parameters are updated according to the gradient flow
duy V L()), duk+1 -V Vn,
dt
dt
for h E (1..... L- 1) and duz/dt = V., ((0), which clarifies that uy and Uk. rather than Wk. are the model parameters.
Previous work (26) showed that the solution selected by GF interpolates between the & minimum norm solution and &a one depending on the initialization scale for diagonal linear networks. In this section we examine whether this is the case for the rank- I deep linear network with depth L. For
convenience, we define
2L
2(L - 2(L-1)
Ω2
2L-1 AL
2
where L is the number of layers. Recall that 0(0) represents the initialization of 0, 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) 7 and if the gradient flow solution ((cc) satisfies that Xx( = y. then gradient flow converges to a minimizer of the potential function V(0):
((0)) arg min (()),
S.S. X0 y
0
where
((0)
V(0) =
00
(6)
|[0(0]]a=**

---

Figure 2: The plot of V(0) for differ- ent numbers of layers and initialization scales in two-dimensional parameter space. The initialization in the first col- umn 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 sim- ilar implicit bias for rank-1 linear ner- works of different layers. When we in- crease the initialization scale, the shape of the contour for linear regression po- tential function does not change, while for linear networks as shown in the last column, the contour of the potential func- tion gradually presents sharp angle as we increase the number of layers.
25
21
COT
-38
While previous works focused 2-layer standard linear networks 3 multiple-layer linear networks with nearly zero initialization (28). i.e., |0(0)]] Rs O, Theorem builds the potential function V(0) that depends both the initialization (()) and the depth of the network I. explicitly. We also plot the contours of (0) for different numbers of layers and scales of initialization in Fig. 2-dimensional space. Theorem [ crreorlyllleerl clearly reveals how parreetttttiiiioonnnn and the initialization of 0 guide GD to select different solutions.
Effects of depth. When L = 1, (0) becomes 1110 2 (0, ((0) which is the same as the potential of the least square ||0 - 8(0)]2 except for constant term. When L > 1, V( is no longer an Euclidean distance. The form of V (0) depends the depth and GD will favor different interpolating solutions. A particular interesting case is when L - 0o where we have Q - and Λy - 1. As result, V(0) becomes V(0) - ||0]| (0, ((0)/ /l][())))))) which reflects the difference between the norm of and the norm of its projection on the direction the initialization ((0). Therefore, the direction of the initialization (()) matters for the potential function (00) and the final solution ((o)) while they are rather less sensitive to the scale of 0(0). This may serve as a benefit of the ave-parramttrizzaionn 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 ((
Effects of initialization. When L is finite, inspect the effects of the initialization. V(0) 00 /R, [0(0) ((22-11) (0, ((0)/11((0)) The second term will be more important when the initialization scale |((0)] || getting larger thus initialization affects the implicit bias of GD more significantly. On the other hand, as ||O(0)] - 0, the second term vanishes. Thus we have following corollary:
rewrite
Corollary 1.1. Under conditions of Theorem 7 if we further assume that the initialization is in- finitesimal [[0(0)]] - 0 and the depth is finite, then the GF solution 0(00) is an Ez-norm minimization solution:
((oo) min ||@||
s.t. X8
We now summarize the effects of initialization on the training regime as follows. (i). According to (6). 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 I 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: (fi). For vanishing initialization, according to Corollary 1.1 an mnomr minimization predictor which can not be captured by the NTK kernel returned. Therefore, as the initialization becomes smaller, we "escape" from the lazy regime and falls into the "Anti-NTK" regime defined
5

---

0...7
n-.-
Wa- .......
Wa-tne... 19
W (8)
We-10e..
Wean(()
W (asherltt
....
...
Wa-saa(0
...
(Runnit
Wasnintt
Figure 3: GD maintains the shape of rank-1 initialization for the weight matrices (the upper panel), but destroys that of the diagonal initialization for standard linear networks (the lower panel), as formally described by Proposition T
by 3] Furthermore, based on Corollary I.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. Sshows 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 f2-norm minimization solution for both 2- layer standard linear networks and rank-1 linear networks with small initialization drives to further investigate the comparison and similarities between different network architectures. For diagonal linear networks, GD prefers f no-ormm 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 fo-norm minimization solution when the initialization 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(r; W) = W WL-1 W1x with any depth L. which corresponds to the parameterization 8d = w W ... W1. with balanced initialization (Definition 2).
Theorem 2 (Implicit bias of GF for standard linear networks with any initialization). For an I-layer standard linear network, if the initialization is balanced across layers, i.e., W (0)W = & (0) for all layers, and if the gradient flow solution 6( ((00) satisfies that X6((o) = y. then gradient flow converges to minimizer of the potential function V(0):
W& WT
8(oo) arg min Vsta(0),
s.t. X6 y,
where
L (fta(f) = L
|10| +
8
((0)
[1((0)] '
Although we state before that the parameters of a rank-1 deep linear network are uy 's and v 's rather than Ws,s the form of the potential for rank-1 linear networks Eq. (6) is very similar to that of standard linear networks Eq. 7 ((0) for an L-layer rank- linear network is the same as Vara( for a (2L- 1) -lyee standard linear network. In the following. we explain the reason behind this fact and conclude that rank-1 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 useful proposition regarding special kind of initialization for standard linear networks.
Proposition (Effects of diagonal and rank-1 initialization for standard linear networks). Given an L-layer standard linear network f(; W) ww W W1 where W₁ E Rdsi1xds for k E (1,..., .L,, if the weights are initialized such that only one column of W. is non-zero when is even and only one row of V is non-zero otherwise, i.e., for an integer p
Vk 2p + 1 E (I,....AA L} Weay (0) = 0 i & cu. where Ck E (I,....dp d
vk- = 2p E (1....,AI L} W Kag (0) = 0 if j * Ck.

---

then for any > 0:
vk
L} (() = 0 ifi cui vk 2p [1....,L} Wkits(t) O if + ck.
Furthermore, if the weights initialized the diagonal shape, i.e, vk fi,...,4 1), Wkas (0) = 0 if + j. then GD does not maintain the diagonal shape of weight matrices.
Remark. A (2L 1) l-layer standard linear network initialized as in Proposition @with -1 for all layers has the same formulation as our I-layer rank-1 architecture. According to Proposition TT 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 (2L llayer 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 in the diagonal net, GD will not maintain this property (the lower panel of Fig. B i.e., the diagonal structure cannot be seen as 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, 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 I 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 wer-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
(t + = =u ? (.,,, U (t + 1) uh+ (1)
Σ UK (0) IEB.
(8)
iEB,
for i E f1...., L- 1} where denotes the iteration step, 17 is the learning rate, is the batch-size, and B, consists of points randomly sampled from the uniform distribution 2111,
The continuous modelling techniques for SGD have been widely applied in recent works [ 9 23 22 to study the dynamics of SGD, In our setting, recalling the definition of Pk and p- in Eq. 3 the continuous counterpart of SGD, SGF, is given by the following set of stochastic differential equations (SDE):
Continuous Modelling of SGD.
2 oi x TPpp- + dt 2/ nb
nc (pap (pkp- u) ddW,
dug
(9)
n
2
nC
dus41
T x
rppp p dl + 2/ nb (php- x" aw.
(10)
n
where r (f( u. v) y/ f u, yr,) E I2" is the residual, and W, is standard Brownian motion in R For the parameterization of A (Eq. (2). we now aim to characterize the implicit bias of SGD by showing the existence of a function V(0) 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 1 linear networks). For the rank-1 linear network Eq. (2) that is trained with SGF (Eq. (9) and (IO) in the over prretaaranttintttionnn regression, the initialization balanced across layers, then the dynamics of gives us
dV dV.VS VoLdi + 21
X dwM nb
(11)
7

---

6(0)
x(0)
Origin
SGF
Origin
Toastry d
GF
off
SGP
f5r
PGr Ogur
.000)
SSF
2 =
(a) The implicit bias jointly induced by stochasticity and architectures "alleviates" the effects of initialization
(b) trajectory of a starting from the origin converging to a(x))
Figure 4: (a) The blue trajectory is for gradient flow (GF) that converges to fGr (blue square). The green trajectory that converges Ssr (green square) is for stochastic gradient flow (SGF). ((0) denotes the initialization of and 022 (red square) is the én-norm minimum solution. Black dot marks the origin. Compared GF. (SG is closer 0% since its dependence Note that when ((0) is nearly zero (the left figure), both OüF and OSG are close 0ū, (Corollary I.I]. (b) The red arrow denotes the direction of (i(o)) while the black one denotes that of 0(0), which, as shown in the figure, is particularly important to the direction of (<(o). See more experiments in Appendix A.3
the initialization is reduced.
where
1 1011 0 (
2ALN [((0]]
(((s)) (PL(0)))( 0(s)ds (12) 1|0(s)
VS (0, = 12
nb
with P1 (0) I 00 /|001 being the orthogonal projection operator of 0. Furthermore, if the SGF solution 8( (c) satisfies that X0 = then the model parameter 0 converges to a minimizer of V (0, 20) as -
((x)) arg min v 00), S.L. XO X0= =y
One can immediately notice that the R.H.S of Eq. (II includes noise term, which can interpreted as that the model parameter follows a mirror flow with a noise term. Furthermore, Eq. (2)) which characterizes the optimization geometry, depends on time explicitly. This partly reflects stochastic nature the optimization trajectory is not deterministic. Furthermore, as in [7]. it is also possible to generalize the current vanishing learning rate results the moderate learning rate analysis by deriving stochastic mirror descent recursion with time varying potentials.
Implicit bias jointly induced by stochasticity and architectures. The sampling noise of SGD introduces extra term (the last term of Eq. (22) that depends several parameters such the learning rate and the data matrix when compared with that GD (Eq. (6)) Interestingly, this extra term only exists when the model is simultaneously over-pane peer-aaamtteriee and trained with the existence of the sampling noise. If the model is not over-paraman peeetttramttri,z,d i.e., standard linear regression when == 1, then År. and this term disappears. On the other hand, if there is sampling noise this term also vanishes. Indeed, the I. = 1 case has been widely studied by recent works [. and turns that GD and SGD share similar implicit bias when there is weer-parameterizaaionn Theorem confirms this phenomenon and, more importantly and intriguingly. reveals a new connection between the rechiecctureiinduee vver-paramterrization and the implicit bias of the optimization algorithms.
"Alleviating" the effects of initialization. When I. - co, the last term of Eq. 122 brought by the sampling noise can seen as "alleviating" the influence of the initialization ((0) such that the training dynamics reduces its dependence on the initialization (Fig. 4a)). 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 + oo, i.c., infinitely deep linear network, we have dL - 1, therefore # ds. In a d-dimensional space, the integral in Eq. 12 becomes 2n 16 L(0)tr (P) (0)XTX) ToT
let a(t) H (()) ( X) 0/100|ds be the position of particle A that starts from the origin. Then the integral of Eq. amounts to the final position a<(2)) of A, which moves

---

SOD
920
(a)
(b)
(c)
(d)
Figure 5: For different ||((0]]]: (a) ((O(o0), 0%2) for rank-1 linear networks. (b) (0x),),, 0*) for
rank-1 linear networks. (c) (oo), 07, for standard linear networks. (d) ((0())), 0*) for standard linear networks.
along the direction of {(t)] with speed proportional to L (+(t)) tr (P1 (0(t)) xT)) Since || - ||11 < tr (P1 x"x) I || where | is finite, it is highly likely that the direction of a((o)) heavily depends on 0(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 result, in Eq. 112 the effect coming from the initialization 0z ()) is alleviated by the extra term -0" ((o) induced by the sampling noise and over-parameterization In this sense, V (0, oo) is closer to a simple f2-norm ||O]| when compared with GD, i.e. the SGD solution is more likely be an lg minimum norm solution when compared to GD. We also present the trajectory of a in Fig. 4bb] The trajectory is obtained by first training rank-1 linear network using SGD for 5000 iterations and computing (t)) at every step followed by projection of the trajectory (22 into a 2-dimensional space according to the method described in [13] It can be seen that the direction of the final position ((0)) is close to that of ((0) ((0(20), ((0) /( ([o(()))) ((())))))))))))) 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 be closer to that of sparse regression when compared with GD. Our Theorem does not show such bias that the solution interpolates between the & --rmm and é --normm 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 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 D01,, 02) 0 821 /||22 to measure the relative difference between 0₃ and 82. For a linear network and the parameterization 6" = w W Wi, we use ((co) to denote the solution returned by the optimization algorithms (GD SGD) and 6(0) to denote the corresponding initialization. 8" is the ground truth solution, and 8%, is the l minimum norm solution (Corollary 1.1. Details of the experiments and more numerical experiments are deferred Appendix A. and we focus on the results here.
Rank-1 linear networks. As shown in Fig. 5(a) for all different initialization scales [[((0]], [|- D(8(20), (t,) is smaller if 0(oc) returned by SGD, i.e., the SGD solution closer to the & minimum norm solution 87 when compared to the GD solution. Similarly. Fig. ((b) shows benefit of such alleviating dependence of initialization bias of SGD: its solution is closer to the ground truth solution 0* when compared to the GD solution for different |[O(O)]]. Furthermore, in ((a) we show (0((),, (((()) along training, which further reveals that the final solution returned by SGD closer to 07, Note that when the initialization scales are small, both GD solution and SGD solution

---

are close to 22g These experiments well support our theoretical findings. Besides, we present additional experiments to compare diagonal and rank- linear networks in Appendix A.2
Lene Helworts
Non- Insa Nittoos
SGD (8 0000
CD (0. 0006)
5GD (0 22:37
5D (2 03000071
00-30
GD (3 (0.06)
30 (3.12431
09:20
SGD 10 0676)
5CD (0 124134
03 (0 2004)
50 (2 00-133111
5GD (0 2019)
SGD (0 33333
68
17
17
15
16
(a)
(b)
(c)
Figure 6: Numbers in the bracket denote the scale of the initialization. (a). D( 07 along training for rank-1 linear networks. (b). D(f(t), 0ty) along training for standard linear networks. (c). Test error along training for non-linear networks.
Standard linear networks and non-linear networks. Proposition | states similarity between rank- I 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) [(d] (oo(),), ). fryy) and (0(o)), 8*) are also smaller if we run SGD, which supports the generalization of conclusion of rank-1 linear network standard linear networks. D(O(t), 02 along training plotted in Fig. ((b)trrtee further supports this phenomenon. For non-linear networks (Fig. [6(c), we report the test
error on newly sampled test for both GD and SGD along training, where SGD solutions have smaller test eтror when compared to GD solutions.
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 stan- dard linear networks better than diagonal linear networks. We further reveal the joint role of over-paree 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 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 in current approach. Furthermore, exact characterization of the stochastic integral is absent.
Acknowledgments
B. Lyu thanks Di Wang for useful discussion. Z. Zhu supported by Beijing Nova Program (No. 202072) from Beijing Municipal Science Technology Commission.
References
[1] Alnur Ali, Edgar Dobriban, and Ryan Tibshirani. The implicit regularization of stochastic gradient flow for least squares. International Conference on Machine Learning, 2020.
[2] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit accelearation by overparameterization. In arXiv: 1802.06509 2018.
[3] 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 NeurlPS, 2021.
[4] Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and langevin processes, In International Conference on Machine Learning, 2020.
10

---

[5] Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. Conference on Learning Theory, 2020.
[6] Simon S. Du, Wei Hu, and Jason D. Lee. Algorithmic regularization in learning deep homoge- neous models: Layers are automatically balanced. In arXiv: 1806,00900, 2018.
[7] Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s)gd over diagonal lin- ear networks: Implicit regularisation, large stepsizes and edge of stability. In arXiv:2302.08868, 2023.
[8] 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.05799, 2022.
[9] 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.
[10] 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.
[11] Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019.
[12] Kei Kobayashi. Stochastic calculus for a time-changed semimartingale and the associated stochastic differential equations. In arXiv:0906.5388 2010.
[13] Hao Li, Zheng Xu, Gavin Taylor. Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In NeurlPS, 2018.
[14] 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.
[15] 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.
[16] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank Icarning. In International Conference Learning Representations, 2021.
[17] Bochen Lyu and Zhanxing Zhu. Implicit bias of adversarial training for deep neural networks. In International Conference on Learning Representations, 2022.
[18] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. International Conference on Learning Representations, 2020.
[19] Stephan Mandt, Matthew D. Hoffman, and David M. Blei. A variational analysis of stochastic gradient algorithms. In International Conference on Machine Learning, 2016.
[20] Kenneth S. Miller. On the inverse of the of matrices. In Mathematics Magazine, 1981.
[21] 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.
[22] Scott Pesme, Loucas Pillaud-Vivier and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. In NeurlPS, 2021.
[23] 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.
[24] 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.
11

---

[25] Francis Williams, Matthew Trager, Claudio Silva, Daniele Panozzo, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. In NewrIPS, 2019.
[26]
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.
[27] Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In NeurlPS, 2018.
[28] 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.
12

---

Appendix
Organization of Appendix. In Appendix we present details of the numerical experiments and additional numerical experiments. We give more discussions of our work in Appendix B Appendix provides missing technical details of Section while Appendix provides those Section
Details of Numerical Experiments and Additional Experiments
we present the details of the numerical experiments. We compare the rank-1 I linear networks with the diagonal linear networks empirically in [ Finally, in A.3 we conduct additional experi- ments to further verify the "alleviting" effect of the SGD sampling noise mentioned in Theorem In
In A.
A.4
we conduct experiments when balanced initialization condition not satisfied.
Data. We conduct wer-parameterized regression with different linear networks. For the dataset 100 and 2; (0, T). For
{(r,, (e), where x; Rdd and y/ R, we set n i E {1,... (I,...,n), yi is generated by y/ 0 where 0 Rd i.e., 0" the ground truth solution. We let 20 components of (! be informative.
A.1 Details of Numerical Experiments in Section 5
Large initialization, ||@(0)]] = 124.77
GD
SGD
102
0,0(=)),0;)
10a
10
10
10a
10a
10a
iterations
Figure 7: ((((t), 07,) along training for rank- 1 linear networks when the initialization is extremely large. SGD solution is closer to 0% when compared to the GD solution for rank-1 linear networks.
now present the details of numerical experiments conducted Section
l minimum norm solution. To get the & minimum norm solution 0% we train single layer linear network with zero initialization using GD for 20000 iterations, since GD will return an &z minimum norm solution solution in this case according to Corollary I.I and (26
Rank-1 linear networks. For the rank-1 linear network f fa; u, v) 1 WL- Wi where Wi u.eu [ddxxd+1) let and the batch size is if we run SGD. We construct different rank-1 linear networks as follows: for randomly sampled
and Vk (I...., L} d4 100. The learning rate is 10-
IR 100 we let the initialization 6, (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 0% and GD and SGD solutions, respectively, of each trained rank-l networks with different initialization scales.
13

---

2. In Fig. 5(b) we measure the distances between 6 and GD and SGD solutions, respectively. for each trained rank-l networks with different initialization scales.
3. Fig. 6(a) plots the distances between 07 and the model parameters 0 along training when the inttalizttion scales are different for both GD and SGD. The numbers in the bracket denote ||](0)]|-
4. Fig. is about the distances between 07 and the model parameters 0 along training for both GD and SGD when |[0(0) is extremely large.
Standard linear networks. For the standard linear network fa; W' = wf,W Wiz where Wa E Rdexde+1, we let L = 4 and Ak E (1, ds 100. The learning rate is 10-3 and the batch size is 4 we run SGD. Other settings are similar to that of rank-1 linear networks.
1. In Fig. 5(c)
we plot the distances between 07 and GD and SGD solutions, respectively. Similar to the case of rank-1 linear networks, for all initialization scales, ((0())), 07 is smaller if the network is trained with SGD when compared to GD.
2. Fig. 5(d) presents the results of the distances between H" and GD and SGD solutions.
3. Fig. 6(b) plots the distances between 0; and the model parameters along training when the initialization scales [[((0]] are different for both GD and SGD. The numbers in the bracket denote [I0(0)]]-
Non-linear networks. For the non-linear network f(x; = (W where W& E 1d.xde11, we let L and Vk {1, (1.2.3) d+ 100. The learning rate is 10 and the batch size is 4 if we run SGD. We use the ReLU activation axx) = ReLU(xz). We use the same dataset as in the experiments of rank-1 I linear networks. Since the non-linear networks do not have the overall parameterization of 0 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 by norm of this vector as the scale of the initialization of a non-linear network, i.e., we use V(Ex, I|W (O)ll?) l|?=) as the initialization scale where ||. || is the Frobenius norm. Due to the same reason, we can not measure quantities such as D(0, 0%, therefore, we report the test error of the model in newly sampled test set instead. Fig. plc)s 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
3.2 3.2 indicate that diagonal linear networks exhibit different implicit bias in compari- son with rank-1 and standard linear networks, e.g., both rank-1 and standard linear networks prefer & 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.
Results in Section
In particular, we use the same settings as in the experiments for rank-1 linear networks in A. while only change the model to diagonal linear networks. As in previous works 3 parameterization of diagonal linear network is
22
the re-
0 W4 C W4 - w_ O w-,
where W4 E R 100 w- E R 100 and @ is the elementwise product. = (1, 1) E R00 we set the initialization as
w4 (0) Cc, (- = Cc,
where C' is positive constant measuring the initialization scale. For each diagonal linear network with different C, we run GD for 20000 iterations and calculate D( )c), and D(o(co), 8*). The results are plotted in Fig. 8 where, for convenience of comparing the implicit bias of GD for rank-1 linear networks with that diagonal linear networks, we also plot the results of rank-1 linear networks of Fig. 5[a) and Fig. 5bb) in Fig. 8(a)and Fig. ((b) respectively.
As shown in Fig. ((a) as the initialization scale ([0(()) for rank-1 I linear networks, for diagonal linear networks) increases, D(0(co), 07,) decreases for rank-1 linear linear networks trained with
14

---

0.5
0.4
((B(a). 8.)
0.3
GD
SGD
0.2
GD, diagonal
10-3
10-2
10-1
109
Initialization Scale
(a)
0.85
GD
0.80
SGD
0.75
GD, diagonal
D( )
0.70
0.65
0.55
0.50
10-3
10
10-1 10a
Initialization Scale
(b)
Figure 8: For different initialization scale ([0(()] for rank-1 linear networks and C for diagonal linear networks): (a) D(6(2o), 0%22) for rank-1 1 linear nets and diagonal linear networks (the green solid line). (b) ((o(co), 0*) for rank-1 I linear nets and diagonal linear nets (the green solid line).
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- I linear networks (also standard linear networks).
A.3
Additional Experiments for the "Alleviating" Effect in Theorem
Recall the form of V (0, in Theorem
3
0 (()
2
(' (0(s)) tr (Pi (0(s))x²X)
VS t)
||0| Rz +
(ss)ss,
22
[[[(0]]|>
nb
[0(s)
we
(t) 0F (L))(0)
(13)
[0(0)
2 nb
L (((s)) (P (0(s))X*XX
qe (t)
((s)ds.
(14)
[0(s)] 2->r
To quantitatively measure the "alleviating" effect the SGD sampling noise, we train rank-1 linear networks with different initialization scales using SGD. The batch size is 4 and the learning rate
15

---

2.5
Pa. 0.2849
Pe - Qa. 0.2849
2.0
Pa. 0.5135
Pe - qa. 0.5135
1.5
Pa. 0.8696
Pe ao. 0.8696
0.5
0.0
10a
10
10a
10a
10
iterations
Figure 9: SGD sampling noise alleviates the dependence on the initialization. Numbers after the comma denote the initialization scales. The solid lines are for pe (Eq. (13 and the dotted lines are for po qo (Eq. 14
is 5x 10-5, For the rank-1 linear network (æ; u, v) W Wi-1 Wix where Wa uaa Rd xd+1,, we let = and vk (1...., L.) dg = 100. We calculate both po(t) and go (t) along training, where (t) measures the alleviating effect of the SGD sampling noise and their difference pe((c)= (xc)) 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, qo. is equivalent to make the dependence of S on the initialization closer to (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 (()) since the loss decays along training. To further support this argument, here we present the training loss when we perform the experiments of Fig. /4b) in Fig. [0] It can be seen that, for a random initialization, the loss, the magnitude of the speed of a, 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 (co).
A.4 Additional Experiments for Biased Initialization
In this section, provide additional experiments to show that conclusion still holds when removing the balanced initialization condition (Definition IT.
To make the initialization unbalanced, we add small perturbation to the balanced initialization. Specifically, we define
1
|| |u
|| |
A
2L
as the scale of the perturbation to the balanced initialization (larger A implies that the initialization is more unbalanced). All the other experiment details are kept unchanged as in Section 5 As shown in Fig. and Fig. 12 we still observe similar phenomenons as in the case of balanced initialization, e.g., SGD solutions are closer to the f2-norm minimization solution compared to GD, when a small perturbation added to the balanced initialization. Thus the implicit bias not unique the balanced initialization. In particular:
16

---

10
10a
10a
L(
10
10-3
10-5
10-7
10
10a
102
10a
iterations
Figure 10: The empirical loss [(0) along training for Fig. 4(b)
We report ((O(o)), 07,) for both GD and SGD for different levels of perturbation A (denoted in the title of each figure) in Fig. [T] the last figure, we fix the initialization scale and report ( (o)). 07 of both GD and SGD for different A. It can be seen that, without the balanced initialization, GD and SGD still prefer l f--norm minimization solution 0% for small initialization, while the SGD solution is closer to 0r due to its initialization reduction effect.
We report ((t()), 0%2)) during optimization for both GD and SGD for different levels of perturbation A and the same scale of initialization ([){()) 0.8696) in Fig. 12 which further clearly reveals that there are still similar phenomenons when A A as in the case when the initialization balanced.
B More Discussions
Our work proposes the rank-1 1 linear network which is 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 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 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 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 1 linear networks are consistent with those of standard linear networks. We also reveal the key role of the veerpparameterization in characterizing the implicit bias of SGD, namely that it will only be different with that of GD for wver-parameterization model.
The inconsistency between the implicit bias of GD and SGD for diagonal linear networks and rank-1 linear networks leads us suggest intriguing questions for future work such 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 parametriza- tion 0. This is different with another line of recent work I0
19 27 9 which focused on the
17

---

balanced init
balanced init
SGD, nan-balanced init
SGD, nan-balanced init
init
SGD, balanced init
SGD, balanced init
D101=1.911
Diet+1,8
6.2
2.2
63-
5.3-
5.0
3.0
10
[16101]]
[10101]
Am0.0943
init
SGD, nan balanced init
GD. balanced init
SGD, balanced Init
0.25
D181=),01
6.5
0:11
GD
6.2
0.14
30
30
309
3,000
0.003
0.010
3.535
0.003
0.025
[10100]]
Figure 11: (00(o)), 0%2) ) for different |[((0]] when the initialization is unbalanced (A o 0, larger A means the initialization is more unbalanced). We use solid lines for the results of GD and dashed lines for SGD. For results under the balanced initialization, we use blue lines; for the results when small perturbation is added the balanced initialization, i.e., A A O, we use red lines.
10a
GD, balanced init
SGD, balanced init
GD, 4 0.0111
D(o(t), )
SGD, A 0.0111
GD, A = 0.0491
SGD, A 0.0491
10a
10
10
10
10*
iterations
Figure 12: ((O(t), 07,) along training for rank-1 linear networks when the initialization is unbalanced (A # 0).
18

---

dllua-1ll flatness of the loss landscape by directly analyzing the independent model parameters, i.c., and for the rank-1 linear networks, which is also crucial to fully understand the learning dynamics.
is worth to mention that removing the balanced initialization (Definition 1 also 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 0, which 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 that the induced mirror flow potential should be composed of two parts: the original potential presented in Section 3 and perturbation due the imbalance of the initialization to it. This implies that the fn-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 much more complex time varying mirror flow potential. We believe that the exact theoretical characterization the implicit bias of SGD without the balanced initialization is 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 0 that is necessary for our analysis. Moreover, the exact characterization of the stochastic integral also absent in the current work, while expect that the integral term L())tr (P)) x) T the property of the training dynamics.
in Eq. [12) might have close relation with
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 for rank-1 linear networks and Section C.3 proves Theorem for standard linear networks.
For a rank-1 linear network f(z;
w W Wir where W₁ tx for any k (1...., L 1}. recalling the definition Eq. (3) and that the network f(x; u, v) can be written as f(z;u, v) p UpP-kUf =uf For convenience, we let & W[WL- W221.
C.1 Balanced initialization
For rank-1 linear network Eq. 2 dynamics of gradient flow is given by
duy 2 dt
PkP of x" ruutt,
(15)
n
du 2
n PkP v T rug.
(16)
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 (x; u, v) trained with gradient flow, we have
dll dllux]1 dlleall
V b 13.
(17)
dt. dt dt
iee, layer norms grow the same rate. Furthermore, if E (1...., L. ||ux(0)] || (0)] = [[)], we have
d d( ux)
(VR, ux-1)
Vk E (2, ........ L 1-1
(18)
dt
19

---

Proof. Using Eq.
16). we have
c uuell
dup
T
2
2 [of x
Uk
PkP 20 x2 ruk+1uu
(19)
dt
dt
11
11
1
dI 0x+11
du
2 PMP- u1j X rug Uk+1
2 cuj x"
Un
(20)
2
dl.
dl
Therefore, both
dilue and
mell do not depend on lt and are same then Eq. 17 follows.
dt
We now discuss Eq. (8)). Since assume vk (1,..., T 1| (0)l and Eq. (17 implies that for any t > O:
ux+ (0)
(0)
[|)]2 I| (0)1l² 110 (t)l2 [10)+1 [t)x(+)]12 [1xx (0)1)
(21)
we have
[t)] (1)+1(((111 ||x(1))1 [[x
(22)
To show Eq.
18
we note that
(Uk+1, up)
dux++ dl
Uk + T duk
uk+1 di
dl
2
P X r(1 2 (1|2x12 2 + 11x+zl1?)
n
2 Eof
| 12 || x
4 Eof
||
x
(23)
(Uk+11 ux)
(UK+1, uk)
where we use (T) in the second equality and the third equality is because & = PPP-X (0x+1, ux). As result, the above equation implies that
1 d( (UK-1, ux))2
4 -o x rrlulu
(24)
2 dt
which does not depend on k, and Eq. 18 follows.
To simplify the analysis, in Theorem D we have required the balanced initialization across layers (Definition D.. Recall that the balanced initialization defined as
Definition 1 (Balanced initialization for rank-1 linear networks). Given an L-layer rank-1 linear network Eg. for any le E {1....,L 1), the balanced initialization means that
(uk+1 (0), us (0))
10 (1) = 1,
(25)
[0)+((0]] [0)] 11 (0)
(26)
Eq. (25) states that uy of the k-th layer is aligned with Uk+1 of the (k + 1)-th layer in direction while Eq. (26) means they have the same magnitudes as (0). The balanced initialization has been
suggested by several previous works 3
2 6 28
for standard linear networks defined as follows.
Definition (Balanced initialization for standard linear networks). Given an L-layer standard linear network f (a;; W) w.. W1-1 Wiæ, for any & E {L....,Le 1}. the balanced initialization means that
WO) +1 (0) Wo(O)WZ (0)
for any E (1).....)).
This directly means that Wi++ (0) and (o) share same singular values and W right singular vector aligns with the left singular vector of (O) In our case, such reasoning gives us
Uk++ 1 (0)
Wy (0)
(27)
Ui 1 [0)]
1| (0)l
1k (0)
| (0) 24 (0)
(28)
where Eq. similar to Eq. 25 which shows that (0) aligns with 6k+ 1 (0), and we adapt the
condition 28
Eq. 226 for rank-1 linear net since Up and uy are the independent model parameters.
20

---

A nice property of GD is that the balanced property across layers will be maintained during training (t((t) W (t) for 0 and E {I...., L}. which can be showed by taking derivative with respect to time on W A+11 (t+(())++1(0 and W (t). For the rank-1 linear network case, according to Lemma 1 and Eq. (25). (va+1 (1), (x)) 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.
6 IT 21 i.c., +++((t) W++
C.2 Proof of Theorem
In this section, we prove Theorem |
Basically the idea is show the existence of potential function V(O) such that the model parameter 0 follows a mirror descent with respect to V(0):
((oo) = arg min V(0) s.t. X0
(29)
This method is called infinitesimal mirror descent (IMD) approach and can be found in, e.g., B] Note that, to apply this method, the dynamics be strict. For example, given the linear model f(: 0) = 0 x for 0 I both parameterization of 0 with standard linear networks, i.e. 0" = W Wr.11... W1, and a much simpler one 0 = cv for e E IR and e E IRd 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 1 linear networks is different with 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.
should satisfy certain condition, which might
We first present a useful Lemma in (20
Lemma 2. If H has rank I and is invertible, then
1 G G
(G HJ =
(30)
1+9
where g tr (HG-1).
We now prove Theorem
Proof. Recall that R with T4 (fi v.), we let
L-1
Σ
ok:
(31)
k=1
Φy = P.PP2
|u 11 + |1
(32)
The key step is to derive the dynamics of the overall parameter 0 which can be done by noting that
d0
dg
dup
1 dt
+
dt
dt
where, according to Lemma
T
dg
1-1
d d(u2+1, uk)
= Σ PkP-K k
dt
dt
k=1
2 2, Xvi Σ (1uo=x(lull@++vvve +1
1
2 OoT T
(33)
dvi
2 x
(34)
dt
d0
2
= (5" I + f2lne) XTr.
(35)
dt
21

---

Eq. 35
can be rewritten
(5 + Do)ot) +0100
-1 -1 d0
__xxx x
(36)
dt
According to Lemma
note that
$1102 |?
tr
5²
5²
the inverse appeared in Eq. (36) is
of x 4 v. 11
000
(5 I + P (oloooo
I
(37)
<²
E²
11101e2
is now left for us to express e² and & in terms of In the following, we assume the balanced initialization in Theorem [ and apply Lemma
where ç is given by 1. <². This can be done by noting that ||01| = 1/2 112
e² II ux)+
(38)
Note that (vk-1, uz)² grow at the same rate for different & according to Lemma and (Ue-1; ux) are the same at initialization for different due to our assumption, we have (UK-1, ux)² u1) and g² (u2, u1) 2(L-1) Note that
1 (12, w1)²
4 50 x
4 cot X2 111111 1 u
r1
u1
(39)
2
n
n
we have (12, 41) - () =
| |ae (0) Since we have (v2, u1)
= ||1l|* = || Il at initialization according to our assumption, e² can be finally written as
<² || (()1))
(40)
As a result,
e|| = ||ا 1|ا
||ا || = |0
5² ||0| 41
(41)
2. <b/ By taking some simple algebra, we have
+ 2(L - 1)5²
c
1+1|0 1|2
(42)
|6
2
2(L 1)
Now Eq. 36 becomes
20-11) 21-11
00
d0
2 XTr.
|
I
(43)
181 2(L-1) 1)
||0|12
dl.
n
These conditions are now sufficient for us to find the form of the potential V()). Suppose that (0) can be written
V = (I + TT
(44)
for some vector h and satisfies the following relation:
V2V(0)= v3V(())
-00
|11|
1
(45)
0012
2(L 1)
then Eq. 43
gives us
2 _T
dt (VoV(0))
(46)
22

---

and the integration relation
VoV(0)- - VoV())
Σ
Ti fz(T)dd
(47)
==1
where we let f = -2 /// Requiring VoV(0)lo= = and denoting A; = -oo (T(T) gives us the condition at t oc:
VA (0)(0))=)00>>> = Σ x4Ar
(48)
i=1
Eq. 48 coincides with the KKT stationary condition of the optimization problem 29 we can prove the theorem by deriving the explicit form of V(0).
Therefore,
Solving V(0).
According to Eq.
44
derive the following relation:
0
0gV(0)
r-
hT
(49)
|0|
BBT
88T
(Ø)
+
V'1
|01
1
11011
|
V
08T
||0||
(50)
|
|
Comparing this with Eq.
45
we conclude that
T"
4"
1
(51)
V
|11||
2[1-13
which, by noting that ihen In (((101))
('( can be solved as follows
1119111))
V (11011)
1
(((1O1)
|01| 2L
In ((100)))
In 1011
2L
2L
1
((10)
|10||
(52)
2L
Furthermore, since
2E-1
|0
Eq. 445 is automatically satisfied when @1)) has the form of Eq. (5I). It now left for us to get the form of h, which can be done by noting that
(oV(0(())
0(0)
[|((0)]
h 0.
[0(0)]
Thus the final form of ((0) is
1
0(0)
V(0) = Z2
|10| 92
0"
(53)
||((0]]
This completes the proof.
C.2.1 Remove the assumption of convergence the interpolation solution
The assumption that Xo(oc) Theorem can be removed the dimension the span of X larger than the number of samples n, i.e., when dim (span(X²) nu This be proved as follows.
Since max(dim(span(X" n,, this condition is in fact dim(span (x)
23

---

Proposition 2. For the over-parameterized regression of rank-1 linear networks Eq. 2 and the dataset {(1, Bi) where x; E IR and n if
dim (span (x")) 2 n
then the gradient flow solution ((o) satisfies that
(0() = y-
Proof. To prove this proposition, we study the dynamics of the loss function L= E, r?. where Ti = (0, v.)) yt, that given by
dC 0L do
dt
0A dt
25
1 [1+(LL- 1) A0T
TT X
x
1010
n
n
45a
Xr + +2(L 11) T- x X00 X
n²
11012
462
|xTr||2 + 2(T 1) (0 x)
n
|1011
S 0,
where we have the equality in the last line when r r= = (O.... 0) E IRd, i.e., X0= = y. or X which is not possible since have assumed that dim(span(Xx)) 2 Therefore, L(((t)) keeps decreasing until (0) = y. i.c., until GD finds the interpolation solution. Noting that ming ( complete the proof.
C.3 Proof of Theorem
In this section we prove Theorem The techniques are similar those C.2 and we still need a time wrapping technique introduced in [3] to derive the form of Vita((), since the condition for applying the IMD approach violated in this case. Recall that, for the standard linear networks fa; W) wfW,, WL-1 Wi = 05x æ where W; E Rd xdx+1, our purpose is to find a potential function Vsta(() such that the gradient flow solution ((oo) satisfies that
((co) arg min Vsta(f), s.t.X0 =y
(54)
Since we assume the balanced initialization (Definition 2 then according to 11 2, the norms of all layers grow the same rate and are the same for any O:
WT+ (4)
W&(OWW (1).
Furthermore, following the procedure of (2). we obtain that the dynamics of 0
d0
||0|| 2(-11)
I I+ (L - 1)
00
x r.
(55)
dt
|0
We now present the proof.
Proof. Eq.
55 can be written
T0
de
|0||
I (L
1)
-T
|001
dt
where, according to Lemma 2 the inverse in above equation is given by
000
00T
I 1+ (L 1)
I
(I
|
(L- (1-1)
||01
1 tr
1011
L
T0T
(56)
10122
24

---

As result, we have
||@|| حجيدة
1
00
d0
-
(57)
L
0
dt
In this case, still assume that (0) has the following form
(it = (( + 3
(58)
for some constant vector 3 E Rd. Then following a similar procedure as in C.2 we have that
0
0e sta((())) = Vied
+ h²
(59)
|0||
V |00
(1
V atd
0g
(0))
100|
(60)
V
[
In the IMD approach, (f) should satisfy that
di (8g Vora(0))
-
(61)
which requires that ViVaaa |0|| " (1 41-
4-1 88 de But this is not possible. Therefore,
|
we multiply a time re-scale factor ( as long as g()) is positive, to both sides of Eq. (57) and
only require that (0) satisfies the above relation under the new time scale R - such that r' g(0):
0001100
s
201 L2
L L
00
do
=(0)x"
(62)
di
Then the limit point at t = oo in Eq. (61) is also visited at the point r r= Jo g( in Eq. 62
We now solve the explicit form of Vtta(O that satisfies Eq. 662) By comparing Fq. (60) and the right hand side of Eq. (57). we obtain that the following relation should be satisfied:
V ||@]|
V
00T
000)]101 20-- (1
L 00
7
|00
(63)
|0112
11012
std
This implies that need:
the terms in the bracket on both sides should match:
|
std
L
L-1
std
std
std
L
In Vild
In |10|| +C
= Veid
||0| + +1
(64)
L
L +
for some constant C and C", where we can simply choose C'
the terms outside the bracket both sides should also match:
std
| 20-4)
g()) Vo |10|
(65)
|||||
To obtain the form of the constant vector 3, we note that (0(O) = 0. which immediately gives us
[e(o]] 8(0) [((O]]]
3
3 =((0)](((0)]
(66)
Combining all these terms, have the final form of (o..(():
Vata(0) L
L
|0||| + 0(0)) -
(67)
As in C.2 denoting X; -100 ri(s)ds and noting that Veee (0(0)) = give us
0gV.ta
Σ zzdi,
which is exactly the KKT stationary condition of the optimization problem 54
25

---

Convergence to the interpolation solution. Similar to C.2.1 we can also show the convergence to the interpolation solution when dim (span(_X")) 2 n by deriving the dynamics of L:
dc 0C d0
dt
00 dt
2 2|0||
A0
Tx X I + (L 1)
X²r
11012
n
2|0||
||r T x 1 (L- (0 (0"x""))
n
Since we assume that until GD finds the interpolation solution.
2 n, df/dtt < until X0(t) y, i.e., C keeps decreasing
C.4 Proof of Proposition
I
In this section, prove Proposition[] by analyzing the gradient of model parameters. It is helpful to recall that for a matrix A, we use A:y to denote its i-th row j-th column element. For weight matrices, e.g. Wi. we use Anyy to denote its i-th row j-th column element.
Proof. For a standard linear network that has the initialization of Proposition I and L is an odd number, i.e., for integer p
Vk = 2p + E Vk: 2p E (1...., L} Waa (0) 0f + ck.
Wayy (0) = 0 i & cx. where Ck E (1,...dx-i)
note that the L 1-th layer satisfies that W WT--i Wi.-1zjjaa where 651 1 if = otherwise 5st O. write the networks at t Oas
[x; W) ΣΣ
Whs Wr.iWr. Wn- (Wi.2....Wzzj
3
Then the gradient w.r.t WT-1 at is
(ww,,-(0))
T (WL-2 -Wix))
(68)
p=1
(Vw,-, 0 CL-11
(69)
This means that the parameter W 13y will be updated only when j ct-1- As result, the initialization shape for WL-1 will be maintained, i.c., only the non zero column of W at t = 0 will be updated and all other elements of WL-1 will be zero for any > 0 since the corresponding gradients vanish. Similarly, for WL-1, we can write the network at t = 0 as:
f(x:W)= Σ> [w). Wo.WWL))),&0cccm 1e-0ssS----22111 WL- (WL-3.. Wix),, -...)
then
((w,, =
and the initialization shape for WL-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
Weat = 0 if Hj, Vk E (I....,I.- 1}.
For any & E (I,..., 1}. the gradient w.r.t Wt at is
(Vw, ((0),, = Σ r (wf ... W ... ...Waa)
where, clearly, there is not any constraint on which elements of (w W Rd and (Wk-1*" Wix) _d are zeros if we do not require that many diagonal elements of the initialization
26

---

matrices are zeros, which clearly violates the diagonal initialization requirements. Thus we can not conclude that
(V (()), = 0 i
i.e, the off-diagonal elements of W will also be updated to be non-zeros. Thus the diagonal initialization of weight matrices not be maintained. The conclusions for other layers be easily derived by following similar arguments.
D Proofs for Section
In this section, present the technical details for Section In particular, introduce our modelling details of the SGD dynamics in D.I. derive the SDE of the model parameter 0 in D.2 and, finally, prove Theorem
D.3
D.1 SDE modelling details
Similar to the case GD, to derive the implicit bias of SGD for rank-1 linear networks, we need to first derive the dynamies of the overall model parameter 0. For this purpose, the structure of the noise of SGD is crucial. For convenience, we first discuss basic SGD where only one data randomly sampled at each step, while the generalization to batch-size SGD is straightforward.
Structure of the noise. We present the details for uk, and the case for uk+1 similar. Recall that the empirical loss is - = E, & / is the learning rate and that the network can be written as f (a; u, v) = P&P->
T 1++ Uy e we start with the SGD update equation for where we let je denote the index of the sampled data the t-th step:
jl ux (t) - 2c
0L
Obie
(t + 1) = (t) - Ouk
+
Ouu
Our
Ouk
201 of x2xx"Pepe rptp-kuk+1
(t)
n
+ 2n
T
uk+1
n
2 of x" TPEP-ATE+1 kV 10 T x Zie
(1)
(70)
where let ej. be the basis vector in R" such that the je-th element is 1 while all other elements 0 and
2 In,-
2, = E,.1r, 3o Cov [2,i]
(71)
As result of this, the noise of SGD for un now
4n
((u((t)
(pap- n
x U
(72)
Continuous Modelling of SGD.
The continuous modelling techniques for SGD have been widely
applied in recent works 9 23 22 study the dynamics of SGD. In our setting, the continuous counterpart of SGD established as follows. First, the discrete SGD updating equations Eq. (8) can be equivalently written as
(t 1) (t) v/ us ((0) + V uv ((0) - V (0))
Uk+1 ( + v (1) o, V/ (() + 7) V V [()) Vo (0)
for ls E (1...., L 1} where both Vu, V L()) (0) and V L(A) V V Lice) (f) are zero-mean noises with covariance matrices in Rd xd, E(uR) 4C(pkP=*) n
bi 7 T X0104410
4[(ppp-b)
E(1+11)=
01
x Nuyukut
27

---

with and p-k defined in Eq. (3. Second, as we identify the noise covariance, letting 1j 0 then we obtain the continuous counterpart of the discrete SGD that of stochastic differential equations (SDE):
2 _u X TPAP- + 2V (PuP- dWW
nf
duy
(73)
71
2 -u x
1
dus+1
Tphp Phpp-11 sudd 2 (pap- ) x aW,
(74)
n
where we let (f(1; u, v) yi..... f u, v) yn)² E Rn be the residuals and W/ is standard Brownian motion in R Similarly, recalling the definition of 1 Wi.1 ... W2 the SDE of 01 is
dup 2 ex rdt + 2/ n
n 2/22 exTaw aW
(75)
Generalization to batch-SGD. When (a positive constant) data points B, are sampled in each iteration of SGD, i.e., batch-SGD, we change the SGD update equation as follows (taking eg as an example)
(t 1) ux
1
Σ
Oli.
Ouk
J.EB.
This only changes the noise [t)x(t)) to batch version
(t)) - (t)),
which only affects the noise part of the SDE uy and leads to become a batch version SDE:
du: 2 xT www."PpPPea +P + 2
nc nb
(PoP-)) X aaw.
This is equivalent to re-scale the learning rate 1 to
5'
and leaving other parts unchanged. Thus the generalization batch-SGD is strighhforward--iinpy replacing all v) with 176.
D.2 The SDE of 0
In this section, we carefully derive the continuous dynamics of SGD for the parameterization of our rank-1 I linear networks. We first discuss the balanced initialization condition.
Balanced initialization. Similar the for GD (Definition 2
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 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.
(uk+1 (t), un (t))²
D1+1 (t)]
and
1 (t)] (t)l || (+)
for > 0 and is E (1...., 1). during the training rank-1 1 linear networks.
With the equations for uı and 14+11 now derive the SDE of 0 summarized in the following lemma.
Mooe details of this modelling technique be found in 14
28

---

Lemma 3. For an L-layer rank-1 linear network Eq. the stochastic gradient flow of
if we assumee balanced initialization, then
25a
nC
d0
(T rdt + 25²
(1(0)XTaw: aw,
S
2L 3 00"
X X0dt
n||0| म..
2 |
where H()) I 2(L - 1)000)010000 /10
Proof. According the Ito's Lemma, we have
dl = d(fvi) = dgvi + Edv + dÉdv1.
(76)
Thus obtain the SDE of 0, we need to analyze every term of the above equation. We first give de.
The form of de. we = the uh and a/ = 1|1x+1112 + || 12, we first characterize the SDE of wk. According the Ito's calculus, obtain that
dwr d(u ux) ()))+
dode + uktz dus + duk dun-
(77)
where, by applying Eq. (73) and Eq. 77 and noting that (dW))² dt,
n oi T TP + +2 n
n PUK- x dWW,
Amc
(Ppp X Xeidt.
Combining the above two terms gives us the SDE wi
nc dwy TT x rPaP-kth + 4n(pkp-))2 wkt x X dt + 2/ PkP-L x dW,
Since & d([ wk). its SDE can be done by repeatedly applying the Ito's Lemma and the SDE of wi:
dg = d(
Wkb
k=1
E
dw
Σ
E
-dwaddw
(78)
Wk
2
WkWWk
k.k=
1,tfk
For convenience, we first define several helper notations:
L-1 1
$
k=1
1=1
1-1
$
Σ
wRwh.
29

---

Now plugging the form of day, into
gives us the first term of de:
2 of X
dt 4n x Xu
Σ 3 dt
n
nC g²
(of (uT'Tumii)
+ 2
2
4n x2 X
7 p2 + 2 52 (of aw)), (79)
oi x
2 Φ dt +
n
n
n
and applying again (dW))² dt and the form of dwa gives each term of the second sum of dg:
E
4 v vre of x" Xvidt.
WwWw mwwww
Summing all V and
obtain the SDE of E:
2 of x
An X x
Xui
4n S x
Xu1
de
rE
22 +
D3 dt
72
n
n
9 <2 (uf) 2/
2
of xf
X An of
dt + 2/ nnc ( + 4>))
22 (vf x" aw,) (80)
n The SDE of U1 is much simpler. To get this, start with the SGD update equation for ey:
n
n
dc
dc
0b
vi + (t) 11 Our
+
Ovi
Oe
21 ex² + 21
1 Σ
ut)
Exrri Tj. E:
n
211 exx + 2n Zir:
(t)
(81)
which implies that the noise covariance in this case is
An C X X. n
(())
Then using similar approach as that of uk, we get the SDE of 0n
2 ExT rdt 2v n
d
<x. aW-.
(82)
Now it is sufficient for us to derive the form of d0.
The form of dA. Combined with the SDE of &, we now have
d0 = d(Evr) = Edvi + vide + dEde
(83)
For the + term, as we already have the form of dv, in Eq. 82) and de in Eq. them into and obtain that:
80
we simply plug
2 22x rdt + +2
nc n (xF + tu [ 2 of x2 =
4n
of x X
( + D3)
dt
n
n
n
I 2 1 x awW aw
+ 2/
2
(52 + 2 )prad XTrdi + An x Xvı
( + xz)
n
(521 + 2 ( xTaW. 7
+ +2
(84)
30

---

For the term, we only need to consider the dWW, terms of Eq. (2) and 80
C D D1X" Xeidi.
Combining the above two equations gives the final SDE of di:
d0 - - - 111+ x" rdt
4nf
[ 1 ( + Da)olof] x" Xuidt
+ +2 nf (6I+ ddW
(85)
On the other hand, recall that our assumptions regarding the initial conditions of f(z; u, v) in the Lemma (Definition D and following similar techniques as in the case for GD, we have that for any
t 0:
wk: ||x( = 11xx (1)1 [1)) (1)1
2. Ak w = (t), (t)) = 1| (t)ll* (t)l
Plugging these terms back to the definitions of 11. 422 and $3. we obtain that
2 = =2||v/ and 5² = |101| |101 4
(86)
in the case for GD and
2(L - 1)
2(L 1)5²
(87)
|0i1|
11012
L (4 -1
42 =
(88)
|
1101
4 2(L - 1)(1 -2) 2) 2T - 1)(L 2)
(89)
||22/1
|10|
where we use that ||0| v/ |E][|=]||- Therefore, the final form of do is now
25a
nC (10)X aW
dl = (1)0) (110)xxrrdi rdt + 2
8nC(L 1) 1) n
2L
00T
I
X Xedt
1|10 -
|
where
000
H(0) = 1 2(L
1)
11011
D.3 Proof for Theorem
In this section, we determine the form (0) such that 0 follows a stochastic mirror flow
UC dt 2 9c n
d0g d (0.
(xT))).
(90)
which then proves the claims Theorem 3 For this purpose, we start with manipulating the SDE of 0 derived in Lemma 3 Note that Eq. (85) can be written as
(6 + ( -1 d0 4nf Pdt 2
x rdt + 2/ nC xT kdW,
(91)
where
==
(5 + ( 42) xx Xen.
31

---

To solve the inverse appeared in the above equation, we apply Lemma 2 and noting that
tr (522/u = 21 0111011?;
then
1
2 C
5²
+ 11 11 2
g
1
D1v10'
I
+ 1 ||
1
1
00
1
g2
|1012
where we use that 0 = {Ui in the last equality, which enables us to simplify P:
p (52 I 2 p
[ I + ( + Ds)uret x X
00
22 + 23
Φ1
1
00
x
1 +
8²
11012
c²
110112
4 + 43
|1012
$ + 43
Φ
007
x x0
|1012
82
1101
11
$1
i2
D3
41
00
X XO.
4
0012
2
Thus Eq. (91) for the overall SDE of 0 now becomes
00T
4nf
OC dt + +2 (xFaw.).
d8
Pdt 00
(92)
e2
||0|2
n
The balanced initialization gives us Eq.
87
88
and
and recall that
dL = 2(L 1)
2L 1
thus we can further rewrite
1
1
A0T
00T
1
C2
101
22
1 + 2(L-1)
11012
10
R8T
I d
(93)
x2
|1@112
and
00T
p
1
[E2 ( 40)00"' xT X0
x0
10012
E I +
( 21 -12)
A0T X
2
10112
2(L 1)5a
Et
+ ((L + 4(L 1)(L 21) -Ax) -2(L 1)Ac)
|1l| 000)
x" x0
~2
1
22
3(L
1) 00
2(L - 1)
x XO
1
2T.
1
2(T.
1)g²
(
000
X X0.
(94)
2110112
32

---

Thus, Eq. 91 now
can be rewritten as
1
00T
S (L
1)
00T
I
dB
x Xtdt
1012
|1011
2 11012
0C
+ 2V 7 (xTd)))
a.w)).
(95)
00
n
Finding the form of (0). We now proceed to find the form of (0). For convenience, apply time re-scaling technique such that dr 15 dt (dW, / I [gW according to (12]) the above equation. For convenience, we still use to represent the time after re-scaling. Then the above equation becomes
Ar.
00 11013
SnC (L d0
1)
(1-2
00
X x0dtt
11012
1011
n
dL dt
NEEC (x* aw)) 22
(96)
00
Recall that (0) E Rd is the initialization of 0, since (0) for SGD should have similar form as that for GD, we borrow from the GD results and first define a constant vector y E Rd
7 ||((0)] 3t (0))
Suppose now that V has the following form:
2L
($
|0 |1
+ + g(() for g E Bd,
(97)
2T.
by similar techniques as the case of GD, we obtain the first and second derivatives of V* w.r.t 0:
V& (0,t) t)
||0||
21-1
+ ? gt).
00
@g,8(0,t)
A1. 1011
2(6-1)
||0|
There exists a function of 6, G(Ø) whose exact form can be derived but not necessary for us, corresponding to our V t) defined above, such that dS,, t) can be written as
0²
d8g sa 0, (a, V ( (0,1) ()) d0 a. (3oVS(0) (7,((00))) ()) dt 0000
d.v
(0, ()d)²²
8
= Ou (8, V (0, 1)) d0 + a (a.v + G(0)
0 0 (0, t)
dt
0000
Now suppose that we can choose partieular g(t), thus a particular G(0), such that the following relation is satisfied:
i 0 S (0,1000)) - G(O)
On (0,1))
0000
SnL (I--10 -1)
(I
00T
x X0,
(98)
11012
|1011
then we can rewrite the SDE of 0g $0, (0,t) as follows:
02
d (0) de (80 t)) do + a (8. dt G(0)
0g S (0, t)dt
000a
100
SnC (L
1)g²
000
1 > 1012
d0
x" X
(99)
n 11012
||0|
'Note that when L = 1 we do not have the second dt term in the LHS of the equation, this term is brought by adding layers the model.
33

---

we have that the d00, appeared in which will give us the desired "mirror flow" equation:
nEEC x" dww
d8g -VV + 2V
It is now suffice for us to find the particular g that makes Eq. (98) satisfied. For convenience, recall that for vector we use ay to denote its f--th component, we define the following helper notations:
nf
= 2
E R,
(100)
00
B
2(T. 1)
X" E R
(101)
1012
d0,, O,dl + Σ BuidWa
(102)
(0,
D,(0,
0r 7a + gu((), D(O,0) ER,
(103)
00
J
0p = if p- = a otherwise ópo 0,
(104)
where the exact form of O can be obtained from Lemma but is not necessary for us. In the following, we use D represent ((0,t) t) for convenience. For our purpose of choosing the particular g()., dD(6, t) dog 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
ODD
dD = 0, Ddi + 3g Dd0 +
Σ
d d0,.
2 00 00.
P.a
where the last term is crucial. From the SDE of (Lemma can be written as
d8 d0pd0,
k² Σ Bio B WW
122 Σ Bio Bp
On the other hand, according to
00
0g = V8 (0)
T As.
a0
||012
we have that, with explicit subseripts of vectors,
@D θ
0,0m
[101|~>e d7
(8up d
00,00
00,
1001
2
=>c11101 -v -2 -8/pp0
A 00,
(1101-Aa-
20,000
-1c11011 -A -28
Ar (Az AL
+ 2 + 60w00)
(pp (Ae + 2) 0 1101)a+
A1
|10||a (p 5p00p)
(105)
1
Note that the above expression implies that expression of d0 dl0, derived above, we have that
82 is in fact rank-3 tensor. Combined with the 3030
K2XE
OpDa 1012
ΣΣ
0 + 8p +0u
(ppr (Az. 2)
Bop
(106)
2110
p.O
H
Poo
where E Rdxdxd is a rank-3 tensor. To get the exact form of we need have the exact forms of
34

---

1. E.o H Bip Note that HBB E R and there are two different terms, both of which will induce same vector. In particular,
Σ po BuyDBBBBy Bir P.
(BB00 una 0 B.B.. = ))"" = (BB"
Σ P
Thus have
Σ
H uo Bos Bip = 2(B" B0)
Rio5
where
00
x X I 2(L -1) 00
BB 0
I 2+2(L
1011 -1010
0
1101
10T
[+2(LL 1) 11012
X x (0 + 2(L 1)0)
00T)
= (2L - 1) [+ 2(L - 1)
[
x X0.
(107)
1012
2. Enog Pe BB B..
ByB
Using the matrix notation, we can easily find that
P.a.a
where, recall the definition of in Eq. IOT
00
I -- (Ar. 2) 00
(I+2(L=
00
B'PBB= = X
+ 2 1) |0|
|0||3
X
1013
00
X
T (2(1- 1) (Ae 2) 2(T. 1)(A2 + 2))
|0112
+ 2(L - 1) 10
X
x
|001
000
10
X
2(2L 1)
+
2(T
1)
xT
||011
||011
00
X I (8L² - 10L + 4)
X
(108)
0012
Taking the trace of the above equation gives us the second term of
0
tr (r ur xx=x) (SI.2 100 +
X X0.
|
Now removing the μ subscript of the derived Ep.o.s H8 Bri Bip we have the final form of by combining it with the result of tr (BPP
Bip and recovering the matrix notation,
- 1) (1 (T+2(1.-) 2(L 2(2L 2(21-
00
x" X0
x0x0
2
1011
122 2r
Otr (x"x) - (8I22 100 + 4) 00T
XT X8
2110||2
|1012
1.2
10 tr (x"x)e 4(L 4(L-1) 1) [ 2L
A
00
x X0
dt.
(109)
2
|101
|102||
Thus the SDE of D now becomes
dD = -0 Ddo Dd0
2T -1)
L
100
+ 0.D
22 tr (XTX)
1
xT x
dt
|101
2L - 1 |
35

---

To find the particular g(t) such that Eq. (98) is satisfied, we need to require that
Ar 2 X
1)
L
00
0,D 2
X X0
|011Ar+2
2L 1
|
SnL (T.
1)
(I
00
x² X0.
n
||012
|1012
Recall that D = |101] + g(t), we have
0, ((),
which, noting that k² 45-n//nn n,, further gives us that the following relation should be satisfied
g'(t) n Ar+ )xtr(x"x) 0 4(L - 1) 2L - 1 ||02|
L 00T
X x
8nL (L - 11)) (1
AAT
X XO.
1011
2 1101
As result of this, can, noting that ç² ||0 from ||0||2 =/ 11? give the required g(t) that makes Eq. (98) satisfied by solving the following equation:
g((t) 2 (XTx) x) 6100 A
0
n 8n 1) +
L
1)
x"x0
1
2T. -
21 d
tr (11-
00T
x X
0.
(110)
|0i1
n
Now let the orthogonal projection operator of be P1 (0)
00 Tap.
then we can solve g(() as
s(t) = 2Am L [(0)[10| ÀL 2 (PL X) ds.
(111)
With this partieular g(), then give (0,tt)
1 1 1 + rr
22
[ [(0)]10]] 20tr (PI ds
8(0)
Vs (0,
(112)
Ω2
n
11((0)
that satisfies the relation
d8g (0,vv(0) -oldt + +21 nL5 xaW
Moreover, a particular interesting case is that, as + cc.
21
ds 0(0)
lim V(O,t) 1+00
||l|| + gr n J L()) (0)
(113)
||||
[0(00]
36


Reviews:
[{"review_id": "JXpnfgtGdX", "reviewer": "Reviewer_dEGz", "summary": "This 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", "questions": "- 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?", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "### 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.", "weaknesses": "### 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed comments and valuable suggestions. Below we address your concerns.\n- - -\n### Weaknesses Part\n1. **On the balanced initialization:** We agree with the reviewer that the balanced initialization is special. To address this concern, we first present the empirical observations of the influence of adding a small perturbation to the balanced initialization, then theoretically discuss the main difficulty and effects of removing this assumption. \n\n\n   From the empirical aspect, we provide here a series of numerical experiments to show the influence of removing the balanced initialization. Specifically, we define \n    $$\n        \\Delta = \\frac{1}{2L - 1}\\sum_{k = 1}^{L - 1} \\frac{| ||v_{k + 1}||^2 - ||u_k||^2|}{||u_k||^2}\n    $$\n    as the scale of the perturbation to the balanced initialization. All the other experiment details are kept unchanged as in the Numerical Experiments section of our paper. We still observe similar phenomenons as in the case of the balanced initialization, e.g., SGD solutions are closer to the $\\ell_2$-norm minimization solution compared to GD, when a small perturbation is added. Thus the implicit bias is not unique to the balanced initialization. Please see the figures in the attached PDF for the experimental results. In particular: \n\n      - We report $D(\\theta(\\infty), \\theta_{\\ell_2}^*)$ for both GD and SGD for different levels of perturbation $\\Delta$ (denoted in the title of each figure) in the first four figures. In the first three figures, we use solid lines for the results of GD and dashed lines for SGD. For results under the balanced initialization, we use blue lines; for the results when a small perturbation is added to the balanced initialization, i.e., $\\Delta \\neq 0$, we use red lines. In the last figure of the first four figures, we fix the initialization scale and report $D((\\theta(\\infty), \\theta_{\\ell_2}^*)$ of both GD and SGD for different $\\Delta$. It can be seen that, without the balanced initialization, GD and SGD still prefer $\\ell_2$-norm minimization solution $\\theta_{\\ell_2}^*$ for small initialization, while the SGD solution is closer to $\\theta_{\\ell_2}^*$ due to its initialization reduction effect. \n\n    - We report $D(\\theta(t), \\theta_{\\ell_2}^*)$ during optimization for both GD and SGD for different levels of perturbation $\\Delta$ and the same scale of initialization ($||\\theta(0)|| = 0.8696$) in the last figure. This figure further clearly reveals that there are still similar phenomenons when $\\Delta\\neq 0$ as in the case when the initialization is balanced.  \n\n   From the theoretical aspect, balanced initialization enables us to derive the exact dynamics of the overall parameter $\\theta$, which is necessary to precisely characterize the implicit bias of GD/SGD. In particular, \n    we need to solve a polynomial equation of the first layer weight $||v_1||$ of the form\n    $$\n        ||v_1||^2\\left[\\prod_k^{L - 1}(||v_1||^2 + \\delta_k)(||v_1||^2 + \\gamma_k) \\right]- ||\\theta||^2 = 0,\n    $$\n    where $\\delta_k = ||v_{k + 1}(0)||^2 - ||v_1(0)||^2$ and $\\gamma_k = ||u_{k}(0)||^2 - ||v_1(0)||^2$ are constants that measure the imbalance of the initialization. The main difficulty of removing the balanced initialization lies in that, according to the theory of polynomial equations, the above equation does not admit explicit solutions for general $L$ when $\\delta_k\\neq 0 $ and $\\gamma_k \\neq 0$ (for $L = 2$ we can explicitly solve the above equation and remove the balanced initialization assumption). Thus it is difficult to discuss arbitrary initialization without the balanced initialization assumption. \n    \n    We may only asymptotically solve the above equation when $\\delta_k$ and $\\gamma_k$ are small perturbations, in the sense that $||v_1|| = ||\\theta||^{1 / (2L - 1)} + \\epsilon g(||\\theta||, L)$ where $\\epsilon$ is a small constant determined by $\\delta_k$ and $\\gamma_k$ and $g$ is a function of $||\\theta||$ and $L$. In this way, 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 the paper and a perturbation due to the imbalance of the initialization to it. This implies that the $\\ell_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.\n\n    Thanks again for the question on the balanced initialization, and we will include the above discussion in the revision.\n2. **On Proposition 1:** We agree with the reviewer that only assuming rank$(W_i)$=1 does not guarantee rank$(W_i)$=1 during training, and that having zero entries indeed is critical in proving the theorem. Proposition 1 aims to show that rank-1 linear networks can be  cast as standard linear networks with special initialization, and that diagonal linear networks should be treated as a different architecture from standard linear network.\n- - -\n### Questions Part\n1. **On the small perturbation to the balanced initialization:** When a small perturbation is added to the balanced initialization, similar phenomenons still exist as in the case of balanced initialization. Please see our above response to Weakness 1 and the attached empirical results.\n2. **On the convergence rates:** Since the corresponding optimization problem is a non-convex one and that the over-parameterization equivalently adds momentum to the dynamics of $\\theta$, we believe that the convergence rate needs a thorough and fine analysis, which will become more complex when considering SGD. Since our focus is on the implicit bias aspect, we will leave such analysis for future works."}, {"review_id": "Zbl7i2CtE4", "reviewer": "Reviewer_B6bN", "summary": "The 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.", "questions": "Please see the weaknesses.", "limitations": "This paper has no negative societal impacts. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "-- 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", "weaknesses": "-- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review and questions. Below we answer your questions.\n- - -\n### Weaknesses Part\n1. **On the broader implications of the results:** Since the rank-1 linear networks share similar implicit bias of gradient descent with standard linear networks, we expect the initialization reduction effect characterized by Theorem 3 for rank-1 linear networks can also be generalized to standard linear and nonlinear networks. Indeed, we empirically verify such generalization in Fig. 5(c), Fig. 5(d), and Fig. 6(b), where, for standard linear networks, the solutions selected by SGD is closer to the $\\ell_2$-norm solution than that selected by GD, and SGD exhibits a lower test loss. \n    \n    Furthermore, we also empirically verify this phenomenon for non-linear networks in Fig. 6(c), where SGD solutions enjoy lower test losses than GD solutions for all the different initialization scales. These results are consistent with that for rank-1 linear networks and imply that SGD is less sensitive to the initialization scales than GD. \n2. **On the experimental details:** The dataset is a synthetic over-parameterizated regression dataset. Due to the lack of space, we defer all the experimental details to Appendix. Please refer to *Appendix A.1* of the supplementary materials for these details, where we specify all the dataset details and optimization hyper-parameters. \n3. **Missing related works:** Thanks for suggesting this related work, which further supports our motivation of investigating rank-1 linear networks. We will include the discussion of it in the revision."}, {"review_id": "INO9pfqajg", "reviewer": "Reviewer_hVR8", "summary": "The 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.", "questions": "3) 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.", "limitations": "No such limitations.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Studying 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.", "weaknesses": "1) 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable suggestions and appreciation of our work. Below we address your concerns. \n- - -\n### Weaknesses Part\n1. **On the motivation of rank-1 linear networks:** We would like to first clarify that we do not interpret the rank-1 linear network as an overall better proxy than diagonal linear networks (DLNs). Instead, we claim that the rank-1 linear network is a better proxy for *standard linear networks* from the perspective of implicit bias of gradient flow. This is because they share similar bias for both large and small initialization and rank-1 linear networks are in fact standard linear networks with special initialization (Proposition 1). These are *differences* between rank-1 linear networks and DLNs, *not shortcomings* of DLNs, which reveal that the implicit bias is affected by the difference between the architectures of rank-1 and diagonal linear networks. And we agree with the reviewer that the transition between rich and lazy regime of DLNs is meaningful and significant. We will make this statement clearer in the revision as suggested by the reviewer. \n    \n    In this sense, we believe that the rank-1 linear network can be an alternative prototype model as DLNs for the theoretical understanding of deep learning models, e.g., it allows us to  characterize the joint effect induced by the architecture and sampling noise of SGF. \n2. **Regarding the misleading usages of GD and SGD:** \n    - **On the suggestion (i).** Thanks for pointing this out. We agree with the reviewer that the using of GD and SGD, instead of GF and SGF, is misleading, and that GD and SGD might be more complicated than GF and SGF. We will fix this in the revision by suitably replacing GD and SGD with GF and SGF.\n    - **On the suggestion (ii).** Thanks for suggesting this nice related work, which reveals the difference between (S)GD and the infinitesimal step-size version, (S)GF. The observation that moderate step-size induces an equivalent initialization reduction effect (thus the solution is closer to the $\\ell_1$-norm solution when compared to the infinitesimal step-size version) is highly enlightening. We will include the discussion of this related work in the revision.\n- - -\n### Questions Part\n3. **Rich or lazy regimes:** Imposing rank-1 matrices does not forbid the influence of scale of the initialization. For rank-1 linear networks with finite $L$, when the initialization is large, the lazy regime is also reached. As the initialization becomes smaller, we \"escape\" from the lazy regime and enter in a non-NTK regime. Below we present a detailed discussion.\n   - **(i) Large initialization.** We first discuss the large initialization case. According to Woodworth et al., COLT2020, 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. In particular, let $x$ and $x'$ be two arbitrary inputs and $\\theta(0)$ be the initialization, we can derive the NTK kernel $K(x, x')$ for the $L$-layer rank-1 linear networks:\n    $$\n        K(x, x') \\propto \\left< x, \\left(I + 2 (L - 1)\\frac{\\theta(0)\\theta^T(0)}{||\\theta(0)||^2_2}\\right)x'\\right>,\n    $$\n    which is fixed during training for large initialization. Please see Fig. 2 for how increasing the initialization scale and the number of layer affect the shape of the corresponding mirror flow potential.\n   - **(ii) Small initialization.** For small initialization, according to Corollary 1.1, an $\\ell_2$-norm minimization predictor which can not be captured by the NTK is returned. Therefore, as the initialization becomes smaller, we \"escape\" from the lazy regime.\n   - **(iii) Comparison to the rich regime.** Since the solution of rank-1 linear network is still captured by a kernel for small initialization regime, we may not claim that this regime is the same as the rich regime in DLNs. Instead, we should interpret it as a non-NTK regime.\n\n    Finally, $L \\to \\infty$ (for any finite scale of the initialization) in Theorem 1 has an effect of aligning the solution (if it is an interpolation solution) to the initialization. This is because, as $L \\to \\infty$, the corresponding mirror flow potential is composed of two parts, the $\\ell_2$ norm plus a term that is minimized when the solution aligns with the initialization.\n4. **On the Analysis of non-infinitesimal step size:** As motivated by Even et al., 2023, we need to derive the Stochastic Mirror Descent recursion with time varying potentials to generalize to the moderate step-size analysis. Since the proof technique of Even et al., 2023 does not seem to restricted to DLNs, we believe such generalization is possible.\n5. **On the overall comment:** We thank the reviewer again for the valuable comments, and we will modify the paper accordingly as suggested by the comments in the revision."}, {"review_id": "c8y2q2yYLT", "reviewer": "Reviewer_zKfq", "summary": "The 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).", "questions": "If 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.", "limitations": "The 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].", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "There 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.", "weaknesses": "I 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments and interesting questions. Below we address your concerns.\n- - -\n### Weaknesses Part\n1. **Merit of studying single-output linear networks.** We acknowledge that single-output linear networks are a special type of linear networks in the setting with rotational invariant initialization. However, we argue that single-output networks are representative prototypes for comparing the implicit of S(GD) for different architectures with balanced initialization. All the previous works [3,23,19] (citations in the paper) on diagonal linear networks studied the single-output case, as our rank-1 network has done. One of main contributions in the paper is to show that even in single-output scenario, different architectures exhibit very different implicit bias, i.e. diagonal linear networks implicit minimize $\\ell_1$-norm while rank-1  and standard linear networks minimize $\\ell_2$ norm up to their depth.\n2. **Proxy of linear networks.** Through our analysis Theorem 1 and 2, rank-1 and standard linear network show the same implicit bias with GD, while diagonal linear network does not. Thus, *in term of implicit bias*, rank-1 net is a better proxy of standard linear network with sufficient tractability for analysis. This is not to downgrade diagonal linear network, and instead the point is to show that different architectures could induce different implicit bias. \n\n   Through this proxy, **another major contribution of our work is the implicit bias of SGD on rank-1 net to shed light on standard linear and nonlinear nets**. We characterize the joint effect induced by the model architecture, initialization, and stochastic sampling noise (Theorem 3). This implicit bias of SGD for rank-1 net is tractable for theoretical analysis while standard linear net is not.  We expect that the results for rank-1 linear networks can also be generalized to standard linear and non-linear networks. Indeed, we empirically verify such generalization in Fig. 5(c), Fig. 5(d), and Fig. 6(b), where, for standard linear networks, the solutions selected by SGD is closer to the $\\ell_2$-norm solution than that selected by GD, and SGD exhibits a lower test loss. Furthermore, we also empirically verify this phenomenon for non-linear networks in Fig. 6(c), where SGD solutions enjoy lower test losses than GD solutions for all the different initialization scales. These results are consistent with the rank-1 linear networks and imply that SGD is less sensitive to the initialization scales than GD.\n3. **Motivation of studying rank-1 network:** Complete and rigorous characterization of implicit bias of neural network is still an open problem in the deep learning community. Typically there exists a tradeoff between analytical tractability and generality when conducting theoretical underpinning. We admit that diagonal and rank-1 linear nets are special cases due to their analytical tractability, but still can shed light on standard linear and nonlinear networks, as demonstrated empirically.\n4. **Comparison to [13]:** First, [13] investigated matrix factorization, i.e. regression with multiple-outputs,  showed  the learned matrix has a low-rank bias. On the other hand, both of our model and diagonal linear networks studied the overparameterized single-output linear regression. Thus, the setting is quite different. Second, it still remains as an open question that whether it is appropriate to compare the similarity of sparsity induced by diagonal linear net and matrix factorization, since the setting is different. \n\n    Our investigation on the implicit bias of (S)GD is designed in a prudent manner that we studied rank-1 model in the exactly same setting as diagonal linear network has done to compare the difference of their implicit bias.  \n- - -\n### Questions Part\nSince the rank-1 linear network and standard linear network show the same implicit bias with GD, we expect a similar generalization for the SGD case. From the empirical aspect, our numerical experiments show similar phenomenons for standard linear networks as in the case for rank-1 linear networks, which supports such generalization. Please see Fig. 5(c), Fig. 5(d), and Fig. 6(b) for the empirical observations.\n\nFrom the theoretical aspect, the analysis of SGD for rank-1 linear network is tractable, while the precise analysis for standard linear network (general single-output linear network) is still an open problem, which is one of reasons why we introduce the rank-1 linear network.\n- - -\n### Limitations Part\nPlease see our response to the Weaknesses 4 for the comparison of our work and [13]. "}]

Author rebuttal:
We present the figures for the supplementary experiments in the attached PDF.

Meta-review: This paper proposes to consider a theoretical model *rank-1 linear neural network* and analyzed the implicit bias of GD and SGD when the network converges to the interpolating solution. It is shown that GD and SGD minimize corresponding potential functions and the proposed frame-work is more consistent result to the standard linear network than well known simpler models such as diagonal linear network.   
The idea is interesting and the theoretical founding of this paper sheds new light to the community. The numerical experiments well demonstrate the effectiveness of the approach.  
In summary, this is a good paper. I would like to recommend acceptance.
Decision: Accept (poster)