On the spectral bias of two-layer linear networks
Authors: Aditya Vardhan Varre
Maria-Luiza Vladarean
Loucas Pillaud-Vivien
Nicolas Flammarion
Venue: NeurIPS 2023 poster (2023)
Keywords: linear networks
spectral bias
low rank
singular values
mirror flow
TL;DR: We study the impact of initialization scale on the singular values of the hidden layer of a two layer fully connected linear network.

Abstract: This paper studies the behaviour of two-layer fully connected networks with linear activations trained with gradient flow on the square loss. We show how the optimization process carries an implicit bias on the parameters that depends on the scale of its initialization.  The main result of the paper is a variational characterization of the loss minimizers retrieved by the gradient flow for a specific initialization shape. This characterization reveals that, in the small scale initialization regime, the linear neural network's hidden layer is biased toward having a low-rank structure. To complement our results, we showcase a hidden mirror flow that tracks the dynamics of the singular values of the weights matrices and describe their time evolution. We support our findings with numerical experiments illustrating the phenomena.

On the spectral bias of two-layer linear networks
Aditya Varre EPFL
Maria-Luiza Vladarean
EPFL
maria-luiza maria-luiza.vaadaeeaaaaaaa wwwdarranneppfl.ccl
Loucas Pillaud-Vivien
Nicolas Flammarion
Courant Institute of Mathematics, NYU / Flatiron Institute 1pi aelllaudviiieeellliiroonisiiiiie org
EPFL
nicolas www.maaiommeaeell.l...mm
Abstract
This paper studies the behaviour of two-layer fully connected networks with linear activations trained with gradient flow the square loss. We show how the opti- mization process carries an implicit bias on the parameters that depends on the scale of its initialization. The main result of the paper a variational characteriza- tion of the loss minimizers retrieved by the gradient flow for a specific initialization shape. This characterization reveals that, the small scale initialization regime, the linear neural network's 's hidden layer is biased toward having a low-rank struc- ture. To complement our results, we showcase hidden mirror flow that tracks the dynamics of the singular values of the weights matrices and describe their time evolution. We support our findings with numerical experiments illustrating the phenomena.
Introduction
The most forceful driver of advancements in the field of Machine Learning over the past decades has been the success of deep neural networks. Amongst the striking qualities of these models is the fact that, despite being heavily overparametrized, their optimization consistently yields minima with good generalization properties. A beckoning research direction is thus to unravel the process through which neural networks learn internal representations for a given task [Bengio et al., 2013). 3). Understanding such phenomena crucial for lowering the interpretability barrier of these models and developing a principled approach to their training and deployment in practice.
Recent experimental evidence identified one of the likely paths towards achieving these goals the study the inherent regularization properties (or implicit biases) of training algorithms [Neyshabur et al., 2014, Zhang al., 2016]. These observations laid the foundation for new line of work [see, e.g.. Vardi, 2022] whose driving question which minimum, amongst the many, awaits at the tail end of optimization.
One of the determining factors for the implicit bias of gradient methods is the initialization scale, which controls their operational regime as shown by empirical studies [Chizat et al., 2019]. More precisely, gradient descent with a low-scale initialization capable of learning rich feature represen- tations from the data. Strikingly, despite overparameterization,, the hidden-layer neurons align in the direction of the features [Chizat et al., 2019, Atanasov et al., 2022] and learning representations reflects in the low-rank structure of the hidden layers. Our work aims precisely explain this phenomenon and quantify the impact the initialization scale on feature learning.
Unfortunately, studying such phenomena for the types of neural networks used in practice is mathe- matically challenging at present due to the con-linearity of their activations. Their less expressive linear counterparts, however, are more tractable and represent a good proxy due to their nonvco
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

loss landscape and non-linear learning dynamics. Consequently, the study of deep linear networks has received significant amounts of attention over the past years, and spans several important directions, including convergence [Arora et al., 2019a, 2018, Min et al., 2021], learning dynamics [Saxe et al., 2014, Braun et al., 2022] and the implicit bias of optimization algorithms [Azulay et al., 2021]. This work complements these previous approaches by mathematically describing the properties of their parameters at convergence, highlighting the implicit bias phenomenon, and further analyzing the evolution of weight matrices throughout the optimization process.
Specifically, this paper studies overparameterized vector regression problems on two-layer fully- connected linear neural networks. We show the following results when the network is trained with gradient flow (GF).
(i) In Theorem 3.1, prove that the zero-loss solutions retrieved by the gradient flow are the minimizers of a potential that depends on the initialization scale. Additionally, we provide explicit expressions for the singular values of the hidden layer weights, also as a function of the initialization scale. These characterizations reveal how low-magnitude initialization induces low-rank structure of the hidden layer.
(ii) In Theorem 3.2, we show that gradient flow on the parameters induces a mirror flow on the singular values. In the specific case of scalar regression, we show that the gradient flow on the weights is equivalent give the geometrical structure of the training dynamics of linear neural networks.
mirror flow on the linear predictor. These characterizations
(iii) In Proposition 1, design simple process to analytically describe how stochastic noise in the training algorithm can likewise induce low-rank structures in the weights regardless of the initialization scale.
We proceed by presenting related work in Section 1.1, formalizing the problem setup and assumptions in Section 2, stating and discussing our results in Section 3, and finally, we provide supporting numerical evidence in Section 4.
1.1 Related Work
The first pillar of our work addresses the implicit bias of GF and its stochastic variant in regression problems. One of the hallmarks of bias in this setting the impact of initialization scale: large initial weights induce a learning regime in which the parameters travel short distance to convergence and feature learning fails to happen (lazy regime). while small initialization effects polar opposite behaviour of the system (rich regime) Chizat et al. [2019], Woodworth et al. [2020]. Training and generalization in the lazy regime are well-studied [Jacot et al., 2018, Du et al., 2019a, Arora et al., 2019c, Soltanolkotabi et al., 2017]. however this scenario fails to capture the observed behaviour of neural networks practice [Ghorbani et al., 2019]. While the rich regime more faithfully approximates the feature learning abilities of these models, it is comparatively more challenging to analyze and few results are known. Amongst them are those concerning diagonal linear networks, where preference towards sparse representations is shown [Woodworth et al., 2020], and restricted setting of the matrix factorization problem, where the implicit bias leads to low-rank representations [Gunasekar et al., 2017, Arora et al., 2019b, et al., 2018]. We similarly study the rich representation learning regime and provide initialization scale-dependent statements on implicit bias for two-layer fully connected linear networks.
Most theoretical results on the implicit bias of GF in overparametrized models rely on the identification of related mirror flow in a reparametrized space [Gunasekar et al., 2018). Diagonal linear networks are amenable to this technique and therefore well-studied [Woodworth et al., 2020, Pesme et al., 2021]. For linear fully-connected networks, however, the existence of mirror flow is not always guaranteed [Li et al., 2022]. To partly alleviate this issue, Azulay et al. [2021] introduce a nonlinear lime-rescaling technique and show that for scalar least-squares regression on two-layer fully connected network with zero-balee balance initialization, the implicit bias selects low εg-norm predictors. We prove similar result under imbalanced initialization controlled by scale parameter, and characterize the weight matrices independently at convergence, thus presenting a nigher-resolution view of the problem.
Other works on linear networks include [Min et al., 2021] where convergence is studied in the presence of weight imbalance and implicit bias results are provided in the functional space; and [Yun

---

et al., 2021] where tensor networks are studied with the goal of unifying the implicit bias results for linear parameterization. In the case of linear networks, Yun et al. [2021] further show an implicit bias towards minimum E2 linear predictors for vanishingly small initializations. For classification in the case of linear networks, and Telgarsky [2019] show that the weights grow to infinity and the layers of the deep linear network align during the course of optimization. Timor et al. [2023] show a similar phenomenon happens for two-layer ReLU networks.
The second pillar concerns the learning dynamics of linear neural networks. The two-layer case optimized with GF on the square loss has been studied by Fukumizu [1998]. Saxe et al. [2014, 2019]. Braun et al. [2022]. The common setup of these works is that of zero- brronbalance initialization and whitened data. First, Saxe et al. [2014, 2019] provide expressions for the temporal evolution of singular values of the predictor by assuming decoupled dynamics and specific data-dependent initialization of the weights. This latter condition is alleviated by the approach of Fukumizu [1998] and Braun et al. [2022], Tarmoun et al. [2021] who solve a matrix Ricatti equation yielding solutions for the weight dynamies in the case where the network initialization has full rank. Finally, Gidel et al. [2019] loosen the whitened data assumption through perturbation analysis and provide the time-evolution of singular values of the weight matrices. Our work removes the requirement of zero-balanced initialization and full-rank network initialization, and gives formulas for the weights' evolution as a function of the initialization scale. We further provide mirror flows on the weights' singular values and show that components are learned in hierarchical manner for the case of whitened data.
Related work is further addressed in the following sections, as part of the discussion of results.
2 Preliminaries and problem setup
Notation. Time-dependent variables are written in bold fonts: we drop the / in A(1) and simply denote it as A. The time derivative of such variables is denoted _ A(1) as Á.
Vector Regression. The set-up is that of standard vector regression problems with inputs (f1,---.,77) ,Tn) (Rd)" and outputs (y1.-.. y.. E (RR+)" in the so-called overparametrized regime where d 2 n. Regarding the output dimension, the reader may keep in mind throughout the article that k < d, though the analysis holds for any k d pair. In order learn the input/output rule, we
minimize the square loss over class of parametric models H = (-)) : Rd + IR 0 E IRP) which we specify in the next paragraph. The train loss therefore can be written as
1 n
(0)
Σ 2n
(y4 - fe())))
(2.1)
Parameterization with Linear Network. We consider the parametric model of two-layer linear neural networks of width / E N this corresponds to the parametrization 0 (W1, W2). W1 E Rexl W2 E Riik and fx(x) expressivity, it is equivalent to the linear class of predictors given by fz) 3 with 3 = W, W.. We henceforth use the symbol B to denote the associated linear predictor of the network. An important consequence of this reparametrization is that the prediction function So is positively homogeneous of degree in 0: VA E R, it holds that fxo = fo. as it is the case for two-layer ReL.U networks. This property has important consequences in the loss landscape through which 0 goes.
W.J W,Tx: x. The model linear in the input x. and in terms of
Train loss. Assume momentarily that k = and denote ((x) = WT x E R'. It is then clear that the predictor rewrites as fe(x) = (d(()) W2). For this reason, call the hidden layer Wi the feature layer and the last layer W2 the weight matrix. We study the overparametrized setting where > d. Letting X := [x In. and Y = [V .y....., J..]. the loss becomes
C CW1, W.)
2N XWWW, - 2
(2.2)
For brevity, we ignore the N in (2.2) by implicitly rescaling the data (X,Y) X/UN).
3

---

Interpolators. Note that when Y E span(X] (X) and X is non-degenerate (which occurs with prob- ability one if e.g.. X. Y are Gaussian and d 2 n). there always exists a solution which attains zero loss, i.e., 3 E Rdxk such that X = Y. We emphasize the fact that there are two lev- els of overparametrization here: on one hand, when d n, the set of zero loss linear predictors Ip == {3 E Rd** X Y) is typically an affine set of dimension (d - k: On the other hand, since we also reparametrize / as a linear network of width > d, the manifold of interpolators in the reparametrized space of 0, defined by Ia == {0- (W1, W2) W, Iz}, J, is of dimension [(d+k) k) - nk. A natural question, therefore, is to which of these interpolators 0" Io does a given optimization algorithm converge. This concept is referred to as the implicit bias of an algorithm. The aim of this work to study that of gradient flow.
Gradient Flow. The dynamics induced in parameter space by running gradient flow (2.2) is given by
0 0= -
(2.3)
We wish to describe the implicit regularization properties of this continuous-time process, which the vanishing stepsize limit of (stochastic) gradient descent. While the latter algorithms incur additional regularization properties from using non-zero stepsizes [Keskar et al., 2017], the study of GF is an important stepping stone to understanding the implicit bias of gradient-based methods in practice. In terms of W1, W₂ the dynamics translates to
Wi x (Y x2, W. W₂ = WT
(2.4a)
W2)
(2.4b)
We emphasize a crucial point: even if the function 3 + || 2 is convex, its reparametrization in terms of W,. W2 is not. Non-convexity and non-linearity makes the analysis challenging and priori it is not even clear whether the time evolution of B can be expressed as closed system.
Initialization. One of our primary objects of study is the impact of initialization on the behaviour of GF. We describe here our initialization choice, to which we henceforth refer as Ly.
(a) Orthogonal feature layer: We initialize the inner layer such that the rows of W are orthogonal and scale with parameter 7 > 0. Mathematically, this translates to Wi /2>P for P E Idx in the Stiefel manifold Va(RR) == {P E Rdx such that PPT Id}. Initializing with an orthogonal matrix is studied by Pennington [2020], however from an optimization perspective. Note that when I is very large, this setting approximates the real-world scenario of initializing the hidden neurons with i.i.d. Gaussian vectors in TR! which are known to be almost orthogonal.
al. [2018], Hu al.
(b) Zero weight layer: In order to remove any initialization bias from the linear layer, we initialize at (2) O. This can be seen as the limiting case of initializing the weight layer with very small relative scale 7 < 7
As already mentioned Section 1.1, existing studies on linear networks assume "zero-balance initialization", namely that WTT (0) = W2, (0) [Saxe et al., 2014, Arora et al., 20190 9b, Azulay et al., 2021]. This condition introduces the invariant W,TW, W W [Du et al., 2018]. 8]. which holds for all times 2 O. This balancedness can be seen as degeneracy assumption on the flow, since it implies that W1 has at most rank k during the entire process, irrespective of the scale of initialization y. In contrast, we show that depending on y the feature layer W1 is biased (or not) toward a low-rank predictor, thus unveiling truly rich representation learning regime.
3 Main result: implicit bias and dynamics description
3.1 Implicit bias on the parameters
Non-convex gradient flows are generally not guaranteed to reach global minimizers of the objective and even when they do, such results are difficult to formally prove. Moreover, the existence of many zero-loss solutions with different generalization properties raises the question of which interpolating

---

network is yielded by training. An elegant answer to such questions is to express the resulting predictor as the optimum among all the possible interpolators of some new, a priori unspecified cost. In addition to the descriptive power of such variational formulations, they express a form of capacity control over the estimator which can be further used to describe its generalization abilities [Bartlett et al., 2020]. The following theorem adds to this series of works, by precisely deriving such characterization for GF in the setting of linear networks.
Theorem 3.1. Let (W1, W2) be the process that follows the GF equations (2.4a)-(2.4b, initialized according condition Iy, for some y > O. Then
(i) The parameters converge to global optimum of the loss
lim (Wi(t), Wz())) (WP W E Io. t1+00
(i) The linear predictor B converges to the minimum E2-norm interpolator
| def
lim 3(1) argmin
3..
X
(iii) We have the following variational characterization of the limiting parameters
(WI., W
argmin XW
W2 ||
W 1| 7 log (det (Wiw;))) (3.1)
Interpretation of the theorem. The theorem is divided into three parts which state that (i) the matrices converge to a zero loss solution, which a priori non-trivial since the loss non-convex; (ii) among all the interpolators in Zz. 3 converges to the minimum f--norm interpolator for all 7 > 0; and (iii) among all the interpolators in Ig. (W1, W2) converge the ones that minimize y-dependent potential. To fully capture the richness of this result, we observe that in the limit of
7 - problem (3.1) informally translates [Attouch, 1996]
lim (W W E argmin 7-0
W|l + Will
This is equivalent, in the space of linear predictors ,3 the minimum nuclear norm solution 3 argminx ||3]|. (which is also the minimum fn-norm interpolator for the problem we study). We informally derived this interpretation by taking first the limit + oo and only after 7 - 0. The theorem naturally does not hold if the two limits are reversed, since = 0 places the initialization at a saddle point of the loss, which is stationary point of the flow.
With increasing y, we move towards solutions with a large log det (W; WTT,, which is a smooth approximation of the rank [Fazel et al., 2003]. Intuitively, this means that solutions with increasing rank are preferred as grows. This scale-induced implicit bias is reminiscent of the rich and lazy regimes [Chizat et al., 2019, Woodworth al., 2020], albeit visible in the space of representations rather than in that of predictors. As such, our result for linear networks is akin to Woodworth et al. [2020]'s, s, which characterizes the rich and lazy regimes for simpler diagonal linear networks.
Comparison with works on implicit bias G. Azulay et al. [2021], Min al. [2021] also study the implicit bias phenomenon in linear networks, however, these results only address the structure of the final predictor B and not that of the factorized problem (W1, W2). As shown in Theorem 1. these works fall short of unveiling all the nuances of the implicit regularization induced by GF in the case of linear networks.
To give a precise example, consider the simplest case of scalar regression (k = for which both Theorem 3.1 and [Azulay et al., 2021] show that (3 is biased towards low-&z interpolators. This view is not complete, since there exist many pairs (W1. W2) such that W, W₂ 3.. Theorem 3.1 goes one step further and provides variational characterization of (W,, W2) at convergence. Moreover, it shows that when 7 0 all columns of W, align in the direction of B, thus creating rank one hidden layer. This is an example of rich representation learning, where W, is learning the only feature needed to make prediction.
5

---

Implicit bias of the singular values. Proceeding with the description of the spectral bias, we provide a characterization of Wi 's limiting singular values which highlights their dependence on fy- Corollary 3.1. [Singular values at the limit] Using the same quantities in Theorem 3.1, denote (ar (WF))......... Fa (a(W))) and (an (3.). (..........(......
the singular values of WT,, G, are
1/2 : ffor 1 2orlsis sissk k.
o: (WF (o.((.)) + +22 +
o, (WT) = (27) 1/2 for k <isd.
Discussion. Similar expressions for the singular values of W2 can be derived and are deferred to the Appendix B due to lack of space. In the vanishing initialization limit, only the first & singular values are activated and W resembles a rank k matrix. Conversely, for large 7 all the singular values are approximately equal in scale and W1 resembles an isotropic full-rank matrix.
To ease interpretation, we again focus on the case of scalar regression. Corollary 3.1 shows that, for small y, only one singular value grows while the others remain small constants. More precisely, when 3. []), the training model is approximately rank one, with only one spiked singular value. Conversely, when y ~ 22 [.]] the low-rank structure disappears. This result perfectly captures the rich learning regime at low initialization where the hidden layer learns the defining feature of the problem, whereas in the lazy regime (large ?) the singular values of the matrix hardly move and no structure is present in W.
Our analysis removes the assumptions of balanced/spectral initialization and whitened data of previous works studying the evolution of singular values [Saxe et al., 2014, Gidel et al., 2019], thus allowing us to reveal the dependence the scale of initialization.
3.2 Description of the dynamics
So far we have described the structure of the parameters of the neural network at convergence. Here, we show that the dynamics of the singular values of [ enjoy very particular property: it satisfies a mirror flow [Alvarez et al., 2004] with mirror potential that can be written explicitly.
Theorem 3.2. (Dynamics of the flow] With the same notations in Theorem 3.1,
(a) Mirror on singular values: The singular values of B, denoted by follow the mirror flow
u (Da) Cd,
where the potential writes ,, (Dg) D (Dap/1) D 72)
(b) Mirror on B3. If k: = 1, the dynamics of B3 can be characterized as mirror flow
1/2
dV bhy (B) --)=-1p -
V113I +72
VL dl,
(3.3)
where the potential writes v.y (B) = [VI13]] + 72 +
+7)/² 22 [VIIS] + +22 +
+7]
3/2
1/2
Mirror on Singular Values. For vector regression, GF on the parameters induces continuous-time mirror descent (which we also refer as mirror flow) with the hyperbolic entropy function [Ghai et al., 2020]. This extends Arora et al. [2019b]'s characterization of the evolution of singular values when the initialization is balanced. Note that our result does not fully characterize the evolution of the system, since the characterization of the singular vectors is absent. Still, some interesting comments can be made. In the rich regime in which ? + 0, the hyperbolic entropy Vr ~ (- | Thus, informally, for small y the gradient flow parameters is approximately equivalent to mirror flow on the nuclear norm. This is reminiscent of the case of diagonal linear networks where such an equivalence is proven rigorously [Pesme and Flammarion, 2023] and is known to lead to an incremental maddle-to-saddl dynamics [Li et al., 2021, Jacot et al., 2022, Berthier, 2022].
In the case of whitened data, i.e., XT X 7, we show that the singular vectors of (3 are stationary (see Appendix C. for details). Therefore, the mirror flow on the singular values characterizes the entire provide exact expression for the evolution of the singular values
system. In this case,

---

(Appendix C.12) by solving a matrix Ricatti equation [Bittanti et al., 1991]. In the limit of y - 0. we can show that, beyond the case of balanced initialization [Saxe et al., 2014, Gidel et al., 2019], the singular values are learned in a hierarchical manner. When y + 0 and with appropriately rescaled time, the limiting trajectory for the jth singular value di,8 can be seen as the jump process
0,3 In ())² -. 20,,0.
1
1
where 0s,8, is the 8th singular value of B Each singular value is activated at time - In (n) (20,, Therefore, we observe an incremental learning process, where the activation begins with the largest singular value and proceeds accordingly.
Mirror descent for scalar regression. The result (b) states that the GF on the parameters (W1, W2) implies mirror flow on the predictor ,3 with the potential 4/8. To be more precise, the evolution governed by mirror flow with the time scaled as function of ||@|]. This technique of time-warping was proposed Azulay et al. [2021] for the case of a linear network with a single neuron (1 = I) with balanced initializations. In contrast, with specific initialization shape, show the existence of a mirror flow for an arbitrary number neurons and unbalanced initialization any scale, The existence of mirror flow is surprising since the reparametrization defining linear networks is not commutative in general [Li et al., 2022]. However, due to the specific initialization we use, this problem can be circumvented by preserving certain commutative properties.
The equivalence with mirror descent enables us show that (B(t)) (11yt) (see Appendix C.8), thus providing a convergence rate for the training loss independent of the conditioning of data, in contrast to Min et al. [2021], Du al. [2019h]. Note that with decreasing initialization scale 7, the convergence speed diminishes, while according to the results in Theorems 3.1 better implicit bias achieved. This suggests the existence of a trade-off between optimization and implicit bias already observed in several works [Woodworth et al., 2020], where achieving better quality solutions linked to slower optimization. In contrast to this behaviour, for the case of balanced initialization [Braun et al., 2022] emphasizes a decoupling between the learning speed and the quality of solutions. Conversely, we stress that in the general setting (e.g., under imbalance) such a decoupling is absent.
3.3 Sketch of the proofs
In this section, we give short description of the proofs of the main results from the previous sections. The common theme of the following intermediate results is to identify natural invariants of the dynamics, which can be leveraged to understand the hidden mirror structure of the flow.
Lemma 3.1. Consider the dynamics of the gradient flow (2.4) initialized at (W, (0), W2 = (V27P,0). Let Zı := WTPT, Ze == PW and the residual R == X XZ, Z2), then the evolution of (Z1. Z2) is governed by the following ODE
==RZZ
i= 2 R.
(3.4)
Furthermore, the dynamics of gradient flow (2.4) is equivalent to (3.4), iee. (Wi(t), Wz(t)) (L), PT (t)) at any time l.
Lemma 3.1 derives an equivalent dynamics to equations (2.4). It shows that weights W W₂ always stay in the column span of the initialization P. thus restricting their evolution a subspace. Going forward, we derive the invariants of the dynamics (3.4).
Lemma 3.2. For the projected matrices given in (3.4), we have the following invariant,
Z Z1 Z,Z, = 2-I.
This invariant ensures that Z Z1 and ZzZ, commute which a crucial ingredient in the proofs of Theorems 3.1, 3.2. Now, we derive the evolution of a == Z Zz, which turns out to be the central quantity enabling our result. The lemma below describes certain properties of the evolution of a.
Lemma 3.3. Let
Z Zzz we have the following time evolution of parameters:
à = R aR a. and 3= (1 -aaa" T)
7

---

An outline of the proof of Theorem 3.1. With an ansatz on the potential that it is decomposable in terms of Z,, Zz. we derive KKT conditions for the constrained optimization problem
argmin 1/1 (1) U2 (Zz).
X2,
Using Lemmas 3.3, we show that a stays in span(X). We use the isotropic property of the imbalance from Lemma 3.2 to find appropriate functions 1h, /22 and finally prove Theorem 3.1. The proofs for theorem 3.1, 3.2 and corollary 3.1 can be found in Appendix B.
4 Further thoughts and perspectives
The previous section provided deep-dive into the dynamics of the gradient flow, which we comple- ment here with few steps in the direction of understanding the dynamics with stochastic gradients. We investigate stochastic gradient descent (SGD) by studying its simpler counterpart, label noise gradient descent (LNGD) Blanc al. [2020].
4.1
The role of noise
was observed that the noise in stochastic gradient descent has a parameter-depenndnt shape that induces certain regularization properties [HaoChen et al., 2021, Blanc et al., 2020]. Here, we study the properties of the noise shape induced the case of parameterization with linear neural networks.
Inspired from the analysis of HaoChen et al. [2021] and the large noise regime described by Pillaud- Vivien et al. [2022] in the context of diagonal neural networks, we design a process driven purely by noise and which carries the same geometric properties as SGD's noise. We consider the scalar (k: 1) regression problem with I d and, through an abuse of notation, denote W1 W, and W₂ a. The noise-driven process which we consider is:
dW = (dB:) a
da Wa=
(4.1)
where B, is a d-dimensional Brownian motion. Details on how this SDE captures the noise of SGD
deferred Appendix D. We show that, similarly to the rich regime of the gradient flow (7 (7- 0), this noise also carries a rich spectral bias but for any initialization. Indeed, we have the following result on the SDE dynamics. The proof can be found in Appendix Proposition 4.1. The dynamics (4.1) has the following convergence properties
(a) Variance explosion. The variance of the norms of W. explode, ie.,
lim 500
([)) + 00
lim t00
[Ia(0]]] + 20.
(b) Scale divergence. For d for any a O we have that,
lim t1000
E [+]] + 00
lim 1000
E [||((t]] [|a())]]]] - cc.
where = et ('a(s)ds is the exponential moving average of
(c) Alignment spectral bias. Denote the 8th row of W w, Using [w,, a] def w.a -aw.
lim E [|],, a]]] + 0.
For any two vectors u. v. [u, denotes the commutator of the vectors: remark that if [u, o] = 0. then u, are aligned, i.c, u cV, for some scalar e E I. First, notice that for d = 1, the SDE in fact corresponds to the geometric Brownian motion with no drift and the dynamics collapses to zero [Oksendal, 2013]. For dimension d > 2, the proposition states that the system diverges and the weights grow towards infinity. However, despite the fact that the norm grows, the commutator [w; al goes to zero, indicating that all the rows of w, align towards a. Overall, similarly the gradient flow in the rich regime, this induces low rank structure in W. This phenomenon can be further seen through the evolution of singular values, where the top singular value W grows unboundedly, whereas the remaining singular values decay to 0 as depicted in Figure 3a in the Appendix. This sheds some light on how SGD induces a particular parameter-des dependedd noise which implicitly biases the solutions towards having low-rank structure of the hidden layer [Andriushchenko et al., 2022].
8

---

10
can be found in the Singilar Vallna
20-11
Singslan
10-1
20
10a
10'
10
10"
10'
10F
107
10'
10a
10P
10
10
Itersttons
Eeerttons
(00 01)
GD
LNGD
(b)
Figure 1: (a) Vector regression with orthogonal initialization and scales 7 1, 10- and Gaussian initialization with entries from N (0, 0.01) (b) Scalar regression with Gradient Descent (GD) and Label Noise Gradient Descent (LNGD).
Intricate dynamics in presence of drift. Proposition 4.1 focuses on the process that is solely driven by noise. However, in general SGD also encompasses a drift term which corresponds to the dynamics studied in Section 3. The continuous-time SDE describing the process is
dW V (W, a) dt + (dB,) a
da = = (W, a) + Sw JBd
where indicates the scale of the noise. The presence of drift quickly complicates the analysis, but intuitively, the noise simplifies the model by inducing a rank reduction, whereas the drift terms prevent the weights from growing unbounded. This noise-driven mechanism relaxes the role of initialization. Empirically this is illustrated in Figure 1b. Gradient descent already exhibits a regularization effect as it increases only one singular value while keeping the others constant. However, gradient descent with label noise [Blanc et al., 2020] enhances this regularization effect by decaying the singular values and promoting low-rank representations. As a result, even for larger initialization scales, we observe the presence of low-rank structures in the hidden layer, unlike in gradient descent. The precise characterization of the this phenomenon is left for future research.
Experiments. We consider regression problem synthetic data, with n 5 samples of Gaussian data in IR (d id=11 10) and the labels in R (k: = 3) generated by a ground truth 3., E Rdd We consider network with width I 200. In Figure la, show the evolution of the top-4 singular values of the hidden layer W1. We use orthogonal initialization for the network with the two scales of initialization
1, 10 Note that, depicted by Corollary 3.1, for the smaller scale only the first & = 3 singular values are significant in comparison to the remaining d - k singular values. This shows that the matrix is approximately rank and the neurons align along three directions. In contrast, for the larger scale y = 1, the final weight matrix has rank d. To complement this, we also consider a Gaussian initialization with variance 0.01 specifically, we initialize the inner layer with d = 10 Gaussian random vectors in Re As described when >> d, the initialization is close to the orthogonal initialization. Hence, in this case, can see that only k singular values grow and the final model has an approximately rank k hidden layer. In figure 1b, we depict the time evolution of singular values for GD and LNGD on a scalar regression problem with orthogonal data in R (n, d = 5) and network with / = 200. Further details on Appendix.
Extension to non-linear activations. Huh et al. [2023], Andriushchenko et al. [2022] empirically demonstrate a low-rank phenomenon through extensive experiments on deep networks with non-linear activations. However, a comprehensive theoretical comprehension of this behavior remains elusive, despite some efforts addressing these issues [Boursier et al., 2022]. To show that our analysis extends beyond linear activations, we present toy experiment for ReLU networks (see Figure 2 and further details in Figure 5). Consider a scalar regression problem in a ReLU reacher-student setup. We generate a training set of size 10 sampled from random Gaussian distribution in R8 The training data (2,, 10 E R x IR is generated by a teacher ReL.U network with 2 neurons (wo, wi ), i.e.,
y/1 (w) xo) + aio æ1),

---

where is the ReLU non-linearity. We train a student network with 20 hidden neurons. Note that there are two relevant directions wg. w for the student network to learn, therefore we expect the hidden layer to represent these two directions (i.e., rank 2 hidden layer, and singular value decomposition with two non-zero singular values). This property is empirically verified in Figure 2. We plot the time evolution of singular values and when initialized at low-scale the network converges to an approximately rank-2 matrix. When initialized at a larger scale, the network weight matrix is high rank and the neurons do not learn the teacher directions.
Perspectives. Learning representations which can be transferred to downstream tasks is a key attribute for the success of deep learning [Bengio et al., 2013, LeCun et al., 2015]. In this work, we present an archetypal problem where for the same predictor in functional space, there exist multiple representations in parameter space, some of which can exhibit a rich structure. This scenario presents case for going beyond the characterization of implicit bias in the functional space [Parhi and Nowak, 2022] and fur-
Singular Values
a0ma 10-3
ther studying the implicit bias in the parameter space. Such characterizations facilitate the identification of crucial ingredients in training algorithms that enable effective feature learning.
100
Iterations
10-1
10
Figure 2:
The time evolution of singular val- ues of the hidden layer weights of 2-layer ReLU network when trained with gradient flow initialized with Gaussian random vari- ables with different scales. We consider a scalar regression problem in a teacher-student setup.
Limitations and Future Work. This paper tack- les the phenomenon of implicit bias, with the aim of furthering the understanding of how neural networks learn in practice. Unfortunately, practical models are highly nonlinear due their activations and rely on various heuristics to achieve state-of-the-ae perfor- mance, thus being difficult to grasp mathematically. This work therefore studies the simplified setting of two-layer linear neural networks. In terms of the assumptions we make, the orthogonality of initialization is only approximately faithful to practical settings where small random weights are used. Nevertheless, we are confident that this requirement can be loosened through a perturbation analysis in the vein Gidel et al. [2019]. Finally, our dynamical description of the system is yet to be completed in the vector regression case with non-whitened data. A careful set of assumptions is necessary here, and hopefully ones that are weaker than the restricted isometry property used in related works [Li et al., 2018]. Finally, we only partially describe the dynamics in the presence of stochastic noise and giving a full characterization remains desired objective of future investigations. Further discussion these aspects is presented in Appendix C.1.
Acknowledgments and Disclosure of Funding
AV is supported by Swiss data science fellowship.
References
F. Alvarez, J. Bolte, and O. Brahic. Hessian riemannian gradient flows in convex programming. SIAM journal on control and optimization, 43(2):477-501, 2004.
M. Andriushchenko, A. Varre, I.. Pillaud- and N. Flammarion. Sgd with large step sizes learns sparse features. arXiv preprint arXXiv:21000053377 2022.
S. Arora, N. Cohen, and E. Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In Proceedings of the 35th International Conference on Machine Learning, 2018.
S. Arora, N. Cohen, N. Golowich, and W. Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 201 9a. URL. httpss//openrevee/ wttp://eeeeettotiddddtdtrroooooe
10

---

net/ S. Arora, N. Cohen, W. Hu, and Y. Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019b.
Arora, S. Du, W. Hu, Li, and R. Wang. Fine-grained analysis optimization and generaliza- tion for overparameterized two-layer neural networks. In International Conference Machine Learning, pages 922-332. PMLR, 2019c.
A. Atanasov, B. Bordelon, and C. Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2022. URL attps://openrevee net/ cnnt/formmidd net/Eorrumiee--llfagaaaoooo
H. Attouch. Viscosity solutions of minimization problems. SIAM Journal on Optimization, 6(3): 769-806, 1996.
S. Azulay, E. Moroshko, M. S. Nacson, B. E. Woodworth, N. Srebro, A. Globerson, and D. Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference on Machine Learning, ICML 2021, 2021.
P. L.. Bartlett, P. M. Long, G. Lugosi, and A. Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 2020.
Y. Bengio, A. Courville, and P. Vincent. Representation learning: review and new perspectives. IEEE Trans. Pattern Anal. Mach. Intell., 2013.
R. Berthier. Incremental learning in diagonal linear networks. arXiv preprint vrXiv:2208.14673 2022.
S. Bittanti, A. J. Laub, and J. C. Willems. The riccati equation. 1991.
G. Blanc, N. Gupta, G. Valiant, and P. Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Conference on Learning Theorry COLT 2020, Proceedings of Machine Learning Research. PMLR, 2020.
E. Boursier, L. Pillaud-Vivien, and N. Flammarion. Gradient flow dynamics of shallow reLU networks for square loss and orthogonal inputs. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL htttttt::/ppereee/eeee forum?id=In -117cc-1UU
L. Braun, C. C. J. Dominé, J. E. Fitzgerald, and A. M. Saxe. Exact learning dynamics of deep linear networks with prior knowledge. A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https::///ppsssee Gorum?id=1Jx2vng-
net/
L. Chizat, E. Oyallon, and F. Bach. On Lazy Training in Differentiable Programming. Curran Associates Inc., Red Hook, NY, USA, 2019.
Du, Lee, H. Li, Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pages 1675-1685. PMLR, 2019a.
S. S. Du, W. Hu, and D. Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In Advances in Neural Information Processing Systems, 2018 88
S. S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over- parameterized neural networks. International Conference on Learning Representations, 2019b.
M. Fazel, H. Hindi, and Boyd. Log-det heuristic for matrix rank minimization with applications to hankel and euclidean distance matrices. In Proceedings of the 2003 American Control Conference, 2003., volume 3, pages 2156-2162 vol.3, 2003. doi: 10. 10.1109/ACC.2203. 220....3...000.3333...
K. Fukumizu. Effect of batch learning in multilayer neural networks. Gen, 1(04):1E-03, 1998.
U. Ghai, E. Hazan, and Y. Singer. Exponentiated gradient meets gradient descent. In Proceedings of the 31st International Conference on Algorithmic Learning Theory, volume 117 Proceedings of Machine Learning Research, 2020.

---

B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Limitations of lazy training of two-layers neural network. Advances in Neural Information Processing Systems. 32, 2019.
G. Gidel, F. Bach, and S. Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019.
S. Gunasekar, B. E. Woodworth, S. Bhojanapalli, B. Neyshabur, and N. Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
Gunasekar, Lee, D. Soudry, and N. Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pages 1832-1841. PMLR, 2018.
HaoChen, C. Wei, J. Lee, and T. Ma. Shape matters: Understanding the implicit bias of the noise covariance. In M. Belkin and S. Kpotufe, editors, Conference on Learning Theory, COLT 2021, 15-19 August 2021, Boulder, Colorado, USA, 2021.
W. Hu, L. Xiao, and J. Pennington. Provable benefit of orthogonal initialization in optimizing deep linear networks. In International Conference on Learning Representations, 2020.
M. Huh, H. Mobahi, R. Zhang, B. Cheung, P. Agrawal, and P. Isola. The low-rank simplicity bias in deep networks. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL httpp://openrrevew ottp://viieeeeeeooooooooaooaaaa
A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
A. Jacot, F. Ged, B. Şimşek, C. Hongler, and F. Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity, 2022.
Z. Ji and M. Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019. URL HJf1g30qKX.
N. S. Keskar, D. Mudigere, Nocedal, M. Smelyanskiy, and P. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. URL https: /ttp:///ppppeiieeeeeee wwwpeeiwww...eeeeeeeeeeeeeeeeeee
J. Lasserre. A trace inequality for matrix product. IEEE Transactions on Automatic Control, 1995.
Y. LeCun, Y. Bengio, and G. Hinton. Deep learning. nature, 521((553):436 2015.
Y. Li, T. Ma, and H. Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pages 2-47. PMLR, 2018.
Z. Li, Y. Luo, and K. Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2021.
Z. Li, T. Wang, J. D. Lee, and S. Arora. Implicit bias of gradient descent on reparametrized models: On equivalence to mirror descent. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https http///pppeeeeee.e/ee/e/ee net/ 20V..
H. Min, S. Tarmoun, R. Vidal, and E. Mallada. On the explicit role of initialization the convergence and implicit bias of overparametrized linear networks. In Proceedings of the 38th International Conference on Machine Learning, 2021.
B. Neyshabur, Tomioka, and N. Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv: 141.v:1412.61144 2014.
B. Oksendal. Stochastic differential equations: an introduction with applications. Springer Science
& Business Media, 2013.
12

---

R. Parhi and R. Nowak. What kinds of functions do deep neural networks learn? insights from variational spline theory. SIAM Journal on Mathematics of Data Science, 2022.
Pennington, S. Schoenholz, and S. Ganguli. The emergence of spectral universality in deep networks. In International Conference on Artificial Intelligence and Statistics, pages 1924-1932. PMLR, 2018.
S. Pesme and N. Flammarion. Saddle-to-saddle dynamics in diagonal linear networks, 2023.
S. Pesme, I., Pillaud- Vivien, and N. Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34: 29218-29230 2021.
L. Pillaud-Vivien, J. Reygner, and N. Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Conference on Learning Theory, pages 2127-2159. PMLR, 2022.
A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. 2nd International Conference on Learning Representations, ICLR 2014, 2014.
A. M. Saxe, J. L. McClelland, and S. Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 155623):1153 2019.
M. Soltanolkotabi, A. Javanmard, and J. Lee. Theoretical insights into the optimization landscape of over- peerpparamttriized shallow neural networks. IEEE Transactions on Information Theory, 65: 749,769, 2017.
Tarmoun, G. Franca, B. D. Haeffele, and R. Vidal. Understanding the dynamics of gradient flow in overparameterized linear models. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10153-10161 PMLR, 18 24 Jul 2021.
N. Timor, G. Vardi, and O. Shamir. Implicit regularization towards rank minimization in relu networks. In Proceedings of The 34th International Conference on Algorithuic Learning Theory, Proceedings of Machine Learning Research, 2023.
G. Vardi. On the implicit bias in deep-learning algorithms. arXiv preprint arXiv:2208. 12591, 2022.
B. Woodworth, Gunasekar.
Lee, Moroshko, P. Savarese, Golan, D. Soudry, and N. Srebro. Kernel and rich regimes in overparametrized models. In J. Abernethy and S. Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673. PMLR, 09-11 Jul 2020.
C. Yun, S. Krishnan, and H. Mobahi. A unifying view implicit bias in training linear neural networks. In International Conference on Learning Representations, 2021. URL https:// openreview www.eeww.ee/foooe maetfformmiin d=
C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint Moiv:1611.0035 2016.
13

---

Organization
The appendix is organized as follows,
section A, we present the experiment details.
In section B, we present the proof of Theorems 3.1, 3.2 and Corollary 3.1.
section C contains the proofs of the supporting lemmas.
In the final section D, we discuss our choice on the noise model and present the proof of Proposition 4.1.
For the results referenced in the main section, the convergence rate of mirror flow be found at C.8, the time evolution of singular values and their limiting jump process is available at C.12, the stationarity of singular vectors in the orthogonal case at C.I1, the discussion on the noise model at D.
Notations,
For matrices of appropriate dimensions, we use [A, B3 to denote AB BA.
A Experiment Details
Experiments. We discretize the SDE (4.1) with a step-size 1/V7. We simulate three parallel runs and track the evolution of singular values and the evolution of alignment using the commutator of the row w;, a, i.e., [w, a] w aw ). We consider the SDE for dimemsion d' 2. The evolution is initialized at (0)) = 1 and a()) = O. As seen in Figure 3a, the noise shape inherently induces a low-rank structure where intensifies singular value and significantly diminishes the other singular value. As predicted by our proposition (4.1), figure 3b shows that the rows of W align with a, thus giving rank structure. The experiments were run on 16-GB RAM Apple MI mac with OS Ventura 13.3.1.
[bo all [br all
122
Cammutde
40
10-3
10
10
202
Ieeetione
(a)
(b)
Figure 3: Three parallel runs of the noise dynamics Eq. (4.1) for 2. (a) The evolution of singular values with do inereasing and d decaying. (b) Measuring the norm of the commutator again as predicted by Proposition 1.
B
Main Proofs
In this section, we present the proofs of the theorem discussed in Section 3. Theorem 3.1. Let (W1, Wz) be the process that follows the GF equations (2.4a)-(2.4b), initialized according condition L,p for some 7 >> 0. Then
(i) The parameters converge global optimum of the loss
lim (Wi(t), Wz(t)) (WPS W Io. 61+00
14

---

(ii) The linear predictor B converges to the minimum E2-norm interpolator
lim (())
argmin
|
B..
XB=YY
(ii)) We have the following variational characterization of the limiting parameters
(WT: W
argmin W,
W2 | +
will | -- 7 7og (det (WiW))). (3.1)
2
Proof. We initialize such that Wi P, W 0. Lemma 3.1 states that the dynamics of gradient flow is restricted to a subspace and can be equivalently described by,
Ż RZ,
22=2,R
(B.1)
where : - 22)
(i) To show the convergence, we track the evolution of tr (R R) and use the following descent inequality to show that it converges to O. With Amin (X"x) denoting the smallest eigenvalue of the X X, the descent inequality (C.13) is as follows,
tr Refer to Lemma C.6 for the detailed proof.
S 27 (X X)tr ( R).
(ii) To show that B + B. in the limit, we show that ,3 E span (X ). -= for some
X. This satisfies the KKT conditions required for the following minimization problem.
1 _ |133|= B.. XB=Y
Bxx E argmin
The complete proof can be found at Lemma C.6.
(iii) For the limit of the projected dynamics (Z Z3a) = lim (Z1(t). Z2(()), Lemma C.4 shows the following.
1+000
(Z 2 , argmin 2 22 1 log (det (Z 2,2z
X2, Z,=y
Using the transformation from Eq. C.7, we have,
Wi P.P.
W₂ = PT Zz.
Thus,
W/ || 7
W2 = =
2, W.W.,
2, W.W,.
Therefore,
|| + - Wi 1| 7 log (det (W, WT),
(WP', W E argmin
Xw, W.W.
This hold on the set ((W,, W2) WT.TT = 0, P. W 0} which is ensured from gradient flow from Lemma 3.1.
Corollary 3.1. [Singular values at the limit] Using the same quantities in Theorem 3.1, denote and (a (B.). (3.),...,0(((3.)) the singular values of W B. are
a, (WT)) (Vo, (B.)² + +2 + 1/2
for 1 Sisk.
(WP (27) for k sis d.
15

---

Proof. Using the transformation W1 P from Lemma 3.1, we obtain.
W, = -
(B.2)
Thus, W1 and Zi share the same singular values, similarly W₂ and Z2 share the same singular values. The expressions for the singular values of Z (similarly Z2) can be found at Lemma C. This along with the fact that 3 - B, proves the Corollary 3.1
Theorem 3.2. [Dynamics of the flow] With the same notations as in Theorem 3.1,
(a) Mirror on singular values: The singular values of [3, denoted by follow the mirror flow
V (Dg) -D dt,
where the potential writes V, (Dg) (2D (Da/1)- D + 77))
(b) Mirror on B. If = 1, the dynamics of B can be characterized mirror flow
1/2
dV thy (B) - VIIBI2 +92
VL dt,
(3.3)
3/2 where the potential writes v/ (B) = [Vi18]] + + +z]"² 27
[VIIS]] + + + +1] 1/2
+))²?
+
Proof. The equivalence with mirror flow for scalar regression is shown in Lemma C.7. The continuous time mirror descent for singular values / is derived in Lemma C.13.
C Supporting Lemmas
This section contains all the technical lemmas and definitions used in the proofs in the section before.
Lemma 3.1. Consider the dynamics of the gradient flow (2.4) initialized at (W, W2
(V27P, 0). Let Zı := W,PT, Zn == -P and the residual R == X T (Y - -x Z2), then the evolution of (Zi, Z2) is governed by the following ODE
RZ
Ż 2
(C.1)
Furthermore, the dynamics of gradient flow (2.4) equivalent (3.4), iee, (W1()), W2()) (Z (t) Zz()) at any time t.
Proof. We choose P1 [R (g(1-d) d)xt such that P1 PT =0. Using the fact that P,P1 orthogonal and span the entire [
=
Denoting (Σ1) WiP. = P1 W2, we have,
W W, P = W.PTPPPP P P wiP Z,P 2 PL.
(C.2)
W₂ = [P' P W. = PT PW, P PL W2 = PT Z2 (Zz)1
(C.3)
Recalling the evolution of gradient flow (2.4) on the loss
W x X -Y W,
W = W (Y XW,W2).
Multiplying the gradient flow updates with PT, P from the right, left (resp.) for the above equations
W xT (Y W,, WT PT. PW, PW (Y x2.
Similarly multiplying with PI, Pl, we have,
- (Y - Xw, WT WT. P.W, = P.W, X -WWW,)
16

---

Using the above, we have,
W = P P1] W2,
-WiPT PW2 + W.
Wi W. = 2,22 + (21)1 (Z2)1.
(C.5)
Rewriting the evolution in terms of Z1, (Z1)_ (Z
2 = = -Y (Z)) (Z2) Z Z2 Z x (22) (zi), = x (Y - 2 (2)), (22) [2)], (2)) - x (Y - x 2 + (21), (Z2)
This is exactly equivalent the gradient flow under the linear transformation which maps W1, W2
to 1 (Z1) Zg, ( Now taking into consideration the evolution of (Z (z) 1. we have that (0,0) equilibrium point for the dynamics. From our initialization Ly.
(Z t= W1 p = 22 = 0,
(22) 1== P1 W2 |=1== 0.
As we initialized at the equilibrium of the dynamics we have ((Zi)_, (0,0) for any time From Eq. (C.5), we have,
W,W, = Z2
(C.6)
The gradient flow (2.4) is equivalent to
= (Y 22, 2
1 x (Y -XZIZ2).
where (1 22TPPT Ia, (0) 0. Furthermore, from Eq. C.2, C.3, we have the following.
W Z.P
W PT=
(C.7)
This finishes the proof the lemma.
L
Lemma C.1. For the projected matrices given in (C.1), we have the following invariant,
z Zi
ZzZ2 2
(C.8)
Proof. Recalling the dynamics (C.1)
Ż RZ,
Ż, = 2 R.
2 Z (Z1) Z1 + 2 (1)) ZRR Z 2 RZ,
Similarly,
ZzZ.j ZRR Z 2 RZ,
Hence, Z %i ZzZ =0. 0. This implies,
2 21 - Z,Z, = 1 Z1 ZzZ,
P = =27Is-
1=
Lemma C.2. Let a Zi Zzz we have the following time evolution of parameters:
à R - aR and == 27 (I aa 7)
17

---

Proof. Taking the time derivative of a,
à Zi' Zz + Zi Z2 -Z Z
Z2 2 Z.
- Z2R'' Zi Z2 R.
-aR a a+b R.
The evolution of Z, 2, aa
(Z)2,) Z aR + Ra 2
(C.9)
ax
(R - aR a +a R - al' a)
Rox T al' aa +aRl T - aa Ro
(I (I-aa' T Ra +aR' (I - aa ).
Computing the evolution of (I - aa )-
(I (I-aa oa )==)a=== - (I-aa) aa ) a0 (I - --a )-
(C.10)
= Ra' T (I- [1-a0] ac ) + +1- (1-ace (1-aa)) aR
(C.11)
Let C. = 2y (I- axx [1-ax)]- and Cz == 2,2,, 80 we have,
C. C.aR + Ra Ca.
Cz = CzaR Ro Cz.
Since, initialization, C. Cz (0), we have Ca Cz. for any time f. Therefore, we have,
Z, = 27 (I-aa) T)-
(C.12)
Using , Z, = Z a and the above invariant, we obtain B = 2y (I-aa)
aa 7)- a.
Lemma C.3. The following property holds for
lim att) E span((*)) t+330
ac
Proof.
From the evolution of a, we have
à= R aR a.
Let U Raxd be the matrix projection each column of the column span of XT, i.e., span(xi..... In), U1 E gdxd be the matrix projection on the orthogonal space, i.e., Ker (X )). a = Ua + U1a. Note that R = (Y - x2 a). So U.R = 0, since U = 0. The evolution of U, that
= (Ua +Ua+U_a)
Again, ULa is the equilibrium point and our initialization a 0 ensures that it stays at this cquilibrium. This proves the lemma.
Lemma C.4. Let (Z )
lim (t), Zz())) the limit of the gradient flow dynamics. Then,
(2 2 argmin
Z2
+
1
Zı 12 7 log (det (Z,2T))
XZ. Z,=Y
18

---

( 2 span so the condition (P2) from Proof. From Lemma C.3, we have that the
Proposition C.5 holds. Note that from (P2) of Proposition C.5, we have,
V (2 = [ 2 (Z2))
(Z))) (1Z))
z 271):
z - 27 (2)0)
which is satisfied by the potential
1 2
2i 1|20 - log (det (Z,Z,)).
v 11
=
When the imbalance is not isotropic, i.e., Z Z Z D, where D is some diagonal matrix (f cI, for any constant c). It this case,
VV, (z = - D (Z
and there exists such function V and the proof breaks.
Proposition C.S. Let (Zi, Z2) satisfy the following minimization problem
(Zi, Zg) argmin V1 (Z1) + 2/2 (Z2),
22
for some non-negative potential functions 1 (Z1), 1/2 (Z)). Then, (Zi. Zz) satisfies
(PI) (Zi) TT V22 E span (X), (P2) VV [ = (zi) 2 (Z))
Proof. The Lagrangian for the minimization problem above is,
L (Z Zz, A) 1 (Zi) V2 (Z2) XA, v).
Taking derivatives w.r.t. to Z1 Zz we get,
22, Zz A) VV, (Z)) (xTN) X) ZT_ Vz,L(Z, Z2 VV, (Zz) + 27 TXTT A).
As Zi: Zž should satisfy VL=
(2j) (V2(2)) = A v (Zi) (Z)))
Lemma C.6. Let (Zi, Z2) be the process that follows the GF equations (3.1), initialized according condition I,, for some Y 0. Then
(i) The parameters converge to a global optimum of the loss
lim Y - X (1) 0.
(ii) The linear predictor 3 converges to the minimum by-norm interpolator
||B|| def
lim B(t) = argmin 100
B..
xØ=Y
Proof. The evolution of R writes,
R = (x X) 2 R RZ, Zz]
tr
R R = 2tr R
(X 22 R)
2tr R (X (XTX)RZ, 2z)
19

---

Note that for any three PSD matrices ABC. Using Lasserre [1995], [see Lemma 7, Min et al., 2021])
tr (AACC) 2 Amin (A) (B)tr C.
From the invariance, we have,
7 1 2 Zr2 Zz2
Z Z > 27In.
Z
Since Z 1 and Zí share the same eigenvalues, Z _ 2 2yIa. Therefore Aman' (Z₁ Z 2 Let Amiin (X X) be the smallest non-zero eigenvalue of X X. Although (X X) have eigenvalues, can always restrict the evolution to the span (XT), which Rt belongs to and without loss of generality assume that all the eigenvalues are non-zero. Using this for the first term, we have,
R (X Z) R) tr ((X) x) 2 R.R.)
2 Min (X X 2 tr (RR.) 2 2Amin X)ytr 33)tr(RR)) (RR").
For the second term, ZØ Zz > we have,
tr (R' (X X RZ Z2) tr ((X X) RZg Z²R
2
X)tr Ammin X)r (RZ,) Z,R
M 0.
Combining both,
( S 27 (X ( (RR.).
(C.13)
Thus, using Gronwall we can show that R r decays exponentially to zero, thus we have X The last step is due overparameterization,, i.e., d n and existence of interpolating solution.
Implicit bias of G. We know that a E spannX ) and 3 (I - -xa T) a. Using, Woodbury matrix identity,
(I aa 1) = a (I a a) aT, (I- ax a= (I a Ta) TT.
Therefore B E span(XT). Hence, this satisfies the KKT conditions for
1 Box argmin _ XB=Y 2 ||31
Lemma C.7 (Mirror flow(k=1).). For t = 1, the dynamics of ß can be characterized as a mirror flow
1/2 VL dt,
dV (B) = - V 11311 + +
(C.14)
where the potential writes
3 V 1181 + +7] + } -27
3/2
V 1181 + + +7] 1/2
by (B)
Proof. Here consider the case & 1, r = XT X13). It is denoted lowercase since here is vector. The gradient flow C.I now in Z1, a can be now can be written as,
à = (I (I-00) aa
1 = T,
B = == 21a) = =(2T 2i a T 2 Zaar : -= X) (2 + 2 Z arrr.
20

---

Recall from Lemma 3.3 that
3= 27(1 aa aa
Using
0:10 T
(I - --a ((ae))=== =
1 1 ||0] 2'
B = 27 (T ac )ー
27
I
a= 2
a
1
|
a
Taking the norms, we get,
|| = 27 [
||0|
2
+
27
a
- = 0.
||||
[|
|||||
7
+
V
+ +1.
|8|
||3||
|||l
?
+
[
2
2 7
+1 +1-2 1
||al|
[|
|||||
2-y
V>2 + ||3|| +7
Now consider the evolution of
á
[ - ax [|,,
aa
a
Q0
I
2
à rr
One cannot write
I 1-
Hessian of any function of a. However, multiplying on both
sides with (1
||a
[3) /
1/2 we obtain,
I
Qx T à
r
(1 ||al²²)
1/2
||0| 3/2
(1 (1- - ||o (3) 1/2
d
a
r
dt
(1 - a [3) 1/2
(1 - || (3)
1/2"
We tackle the left hand side as follows,
27a
1 || 2
1/2
B
(1 - (1-1|a|?))
1/2
2
(1- || ()
1/2
24
24
V/2 + ||/3|
3
227
(V?2 +
||3|| +7)
1/2
21

---

Substituting the above expression and also (1 - ||a|12)
1/2
on the RHS gives us
d
B
= (V + +|1111 +7) 1/2
d
(V + ||01 +7)
1/2
Define the mirror potential,
2 Φ (3)
V 11181 + + 7) 3/2
2y |1133] |// +7/ 71/
(C.15)
Using the fact that,
d 2
V22 + +
3/2
29 p2² ///
+ +
1/2
[V22 x2
p2 +
11/2
2.
dæ
We have,
-1/2
w (B) = (7 + v ||3|| + +
B.
Thus, we can write as a continuous-tin time mirror descent.
||0| + th (B) = (18
1/2
(G
Lemma C.8 (Convergence rate of mirror flow). For the dynamics of the mirror flow given by Eq. (C.14), we have the following rate of convergence,
(33(t) Do. (B)
/2 t
Proof. Note that b. is convex. Let Des, (...) be the Bregman divergence defined with the potential (.). Taking the time derivative of the Bregman divergence, we get,
d D., (3.,
(3), B+)
dt
=- (V/r+ 1 +1|3112 +7) 1/2 (VL - Bz).
We have,
(v + ||3|| +7) 1/2
S -2
Using the convexity of L (.) in G, we have
(B) A+) S - L3))
The last step is using the existence of an interpolating solution. Substituting the above two inequalities, we get,
d D (B.,3) -V2
Integrating, we get,
1
Dv. (Be, Bu)
L. (B(s)) ds < -2
De, (B. - D&s, (Br, <
V2y
22

---

The last step is using the fact that Bregman divergence is positive for convex functions. We will now show that the loss is decreasing along the trajectory.
c(3) (V)(
Viyy (3).
(V>2 + 113112 +7)
1/2
V Vuly (B) 2 (3) B
Vel (3), (v'y), (B) (B) 3,
2 0,
(Using convexity).
Thus L (() < 0. using this,
Dai, (B.,Bo)
tc (B)
C (B(ss)) ds
V2y
This completes the proof.
Definition C.9. [Singular value decomposition and notation.] The singular value decomposition of a given matrix A E Baxk be denoted U UADAV. V where UA Rddx D paxk VA Rkxk
Lef k < d and consider the products AAT and AT A. We will refer to the diagonal matrices A E k and D DAAT E IRddd as D and D3 in order to ease notation. Note that D' is block diakkndaagona matrix containing D3 as the first diagonal block and as the second diagonal block.
Finally, if then D33 D2
Lemma C.10 (Singular values of Zı and (). Using the notations in Desinition C.9, the singular values of Z Zz and (3r are related as the following,
D = =I 2 + and D2 yla + V 2 +D:
Proof. From the invariance (C.8) we further deduce (by appropriately multiplying left and right with Z1, Ze and their transposes):
B Z 22) 2T Ze and
-
-33 2
which implies that [ 3, Z, Zz = 0 and [33T, z = 0. Hence , and Z,j Z2 commute and can be simultancously diagonalizable, same is the case with 333 Z Therefore, the following relation holds (elementwise):
D = - 2 + D and
D2 = >Ia 1 7 + D3.
(C.16)
where the appropriate dimensionality of the diagonal matrices is evident from the indexing of the identity matrices.
Theorem C.11 (Singular vectors are static under orthogonal data). Let a = Z Zz. Then, it holds that
U Uz, Ua. V = Vz, V.-
Furthermore, X X La. it additionally holds that
U. = Un Ui V VR Vg
Proof. From invariance (C.8), by multiplying left and right with VT, and Vz,, respectively, we get that:
D2 VT ZzZ, Vz, = 27Id.
23

---

which implies that Vz, also diagonalizes ZzZ2 (i.e., the left singular vectors of Z₁ are the same as the right singular vectors of Zg for all > 0). As a result, , = Uz, Dg V2, and thus we have shown that U Uz, and V Vz.
Zg
Next, by definition a Z Zz, Zz and Z and z have the same left and right singular vectors (since Z, Vz, D2 U22) This, in conjunction with the above proves Uz, Ua and Vz, Va.
To show that Ua = Ur and Va = VR. we need to do a bit more work. We will show that:
R R ax and
(C.
R Ro
(C.18)
Once we prove (C.17) and (C.18), it directly follows that (aax) (R.. (RR*) (axx T) and
(a a R RR (RJR) (aTa)
(C.19)
which gives the desired result.
To prove (C.17) will show that R a R = 0, Vt 0. First, recall some relations that will be needed. Under orthogonal data, it holds that
R = (3" 3) and
R = -
(C.20)
Furthermore, using the reparametrization in terms of a which induces identity (C.12) and the invariance (C.8), the original dynamics of 3 given by 3 = RZZ Zz + Zi R rewrites as
f 2y [Ra (Ia axa ) (Ia ax ) -R]
(C.21)
Now, using (C.21), it holds that
R = R T R R TR) - 2 (Ia
aa 1 2-a T (Ia
0a Ta
R R R RR a) -22R aTT; a a) a) - -2aa at a T RRa
R R TRRRR'S
a)) a -2,R a(T a a a1-1 - -2 a aa) R a T R
were we used the push-through identity (1+ UV) u U(1 VU)-1
Denoting C = R a. computing the transpose version of the above ODE and using the fact that
C = (C + CT) CC- + (C- - CT) (C+ CT), the following holds:
C - C = - - 2 - a aTaa= - -2764 a aa) aaTaa'CCo CC+
+ a aa)
aa) c +2,0" (4x -aa
= [C2 - (CT)²] + +C- 27 - a aJa) (C- CT) - -2 - -C-CC)))l -a
[1 [I I-- c C- 2n - a -Taa (C- (C-CT)
=
2
- -CC- (1- -7 -
Finally, since - 0, it is a fixed point of the above equation and implies that C = CT, Vt. Thus, identity (C.17) is proven.
We proceed similarly for identity (C.18).
Ra -2,RR T (La
aq T) 1 TaTTaaaa
27(1a aq -)--aaa RR Ra -RR T).
Denoting = Ro T and computing B BJ we get:
B
2 [Ia B B-B 2y(Ta - ar )-
(B (B-B")
+ (B - B
1
(Ia -B B- BT) B-B
- ax (I T)-1
2
24

---

Since (B(0) B - (0)) 0 is a fixed point of the above equation, identity (C.18) is proven.
Finally, show UR Up- and VR Vi3. using yet another invariance, namely
z RZ Z5R Z, VI O.
(C.22)
We proceed by computing the associated time derivatives and showing that their difference is null:
2 2 ZZ z N.. Z2RZ
(C.23)
- RZ, Zz 2, RZ -Z2Z1 Z RZ.
(C.24)
+ +ZzZ, ZRR Z1 +2 Z - ZRR'' Z,
(C.25)
(Z 2, 2 RZ, ZzZZ Z2Z (Z2R 2 RZ,
(C.26)
-2y (2, RZ - Z₂R Zi)
(C.27)
where used (C.20) and the invariance (C.8). Since Zg (0)R (0) ZZ (0) R(0))2(0)) 0.
it is a fixed point of the above equation and we have shown (C.22).
Finally, it holds that
z RZ ZRR. Zn
Z 3Z, Zzz Zi Z2
Z Bz = -2223" Zi
- 3' = =0
(C.28)
where the second equivalence comes from the fact that Z Z₁ and ZzZ are simultaneously di- agonalizable and thus commute (from invariance (C.8)), and the second equivalence comes from multiplying left and right with Z and Z respectively and the definition of a. Furthermore,
= TT = 3Ta a a T
3 a = a 3r,
(C.29)
where the second line comes from the fact that T a From repeated applications of (C.28) and (C.29) obtain that
2 a = a TT.
(a oz) (B = 3r) (aT and (aa T) (B = (B" (axx), (C.30)
and the final identity follows.
Theorem C.12 (Time evolution of singular values orthogonal data). Assume that X = I Then, the ith singular value of B3 evolves over time as
-
- 28 02
ty 21 02 +
27
+ fo t2 texp
tx² (V
(t)
2'
2 Vot, +2 mxp
2t vn +
+2 x
1
1+ Vo + texp 21 /
As a consequence, under vanishing initialization 7 + and with a rescaling of time as t - In(1/7)tt the ith singular value is learned at time T; 1/20
lim (.. = 00 if i <<< 1/200..
o otherwise.
Proof. From Lemma 3.3 have that à R -oRo In light of Theorem C.11, identity (C.18)
and it holds that:
Da (ls D2 Dg. 2yD
25

---

For the 1t singular value of a it holds that:
(t, -r, 2y0a,:tt) + 0z,i
- (Dai(()) + 0x,3
22,)²
(1 -1+
(C.31)
where o, is the jth singular value of B Equation (C.31) is a Ricatti ODE and is separable. For case of notation, let p = o.. = of, and r = - Then, we need to solve the IVP:
(1) Ta,(t)) =r r (1a + 2)² (0, = 0,
(C.32)
(2r (2))t())(((+)))))) 1+oxp (2(/((++())
for which we get (1) = V
p and solving for the initial value gives us Replacing p.q.r and rearranging we finally obtain:
- log
v0-P) Ve+p/
a 2r
Y + v - + 72 exp (-21 Vo + 1²)
(Vo2, + 72 -7)
(t, 1 1+
72
?
02
r + o²,, 72 exp (-2t v /², + +2) (Vo2, + +2
o,,
(C.33)
order to obtain the dynamics for the singular values of 3 Z a we recall the relation given in Lemma C.10:
D2 - + 7 D
where the appropriate dimensionality of the diagonal matrices is evident from the indexing of the identity matrices. Therefore, the singular values (t) are the solutions to the equation
(() = 22. oa,(() [ 122 + +33.00 (1)
and have the following expression:
it)
(t)(()= 1 (t)
27
+
3
+ v o2 t2 exp
- 21 v o0 ,,
7
1+ /o +exp -21 /0,,
2'
2+ / + +2 exp 21 / Vo +
(oo, 7
1+
o,
1+ / 00 4exp 21 /at +
5.,
Note that lom (1)
(Vo2, +2 and therefore we can verify that 3 1 looking at the singular values, since the singular vectors are static (Lemma C.11).
,r by
2y (V 03 +2 - 7)
lim (g.((1= t00
27
Vo,
+22 -7)
0
We can now derive asymptotic transition times at which the singular values are learned. Now, we further process expression (C.33). Let
Vol, +22 7. = Vo + +
then
26

---

w = =
2//0., + ==2==2/,
and 24 We re-write:
(t)
1
(w + v)
w - exp (-1(v + w:)
- (w - v)
20,,i
w + exp( 1/1 + w))u
w exp (-t(w + w))
-1) + -w- -xeexp + w)) v
w
+
20,,
+ exp +(v + w)}u
+ exp( -tee + )
exp (-t(u w))
ew
w + exp( 11v + w))u w + expp-t((x t(v w))
(v w)+ w) exp (-t(v + w))
(C.34)
0x,3
w + vexp( 1(1 + w))
=:h(1)
Since from before we have that (1) 2x)aa
24
and we can write (1)
(1)
blt( (2 h())) the singular values of /3 become:
27)0(;()) (1,((1) 1 (t)
207
[ h(t)
xx.i [1 h(O)
T..i
270
o (2h(t)(2-(0) M(0)
1
(2 b(t)
27
We wish to study the limit of infinitesimal initialization y 0. We introduce constant > 0 and rewrite y = exp(-c) and have c In(1/7). Rescaling time - at and taking the limit c - oo we have
20. [1 h(ct)]
lim (ct((ct
lim
co 2 + vhiel exp(e) (2 h(ee))
σ+,i
1 + lime vh(ct) exp(c)
0, ift < 1/200..3:
d.. if t. = 1/20
1+26,,;
Ox,i if > 1/20
(o + exp - 1(v + w)))
since lim exp(c)h(ct) = lim c1000
w + vexp(-cl(v + w))
2 2 + expp(=222 ex) (cl 21o,,,)) exp
2ct expl 2c)
o., / ; texp(-20)
lim
c000
Voz, exxp(-20 + exp(-c) vexp(-ct(1) +w)
00, ift < 11/20..i
2, if 1/20.,i
0, ift > 1/20,,5-
Lemma C.13 (Mirror on singular values). With the same notations as in Theorem 3.1, The singular values of B, denoted by D&, follow the mirror flow
d (Da) -v dt,
where the potential is V, (Da) = tr ( sinh -1 (Da/n) D3 +
27

---

Proof. We recall the dynamics induces on 2 T
Z2 given in Lemma 3.3:
d= R - R a.
(C.35)
Writing the SVD decomposition of as U U.D.Vo Uz, Da V we have that:
aVz. UU2 Uz, Da V2, Zz
UZ Üz. Da V 2, Vz, + UZ Uz, D. V Vz, + +Z Uz, Da Vz, Vz
U2 Uz Da D. + Da V 22 Vz.
Since Uz, Uz, + 2, Uz = 0 and V Vz. - V2, Vz = 0, the matrices U Üz. and
V2, Vz, are skew-symmetric (have diagonal), and therefore the principal diagonals of the products
U Uz, Da and Da V Vz, are also 0.
We define the linear operator diag: Rd*k - Rdxk which maps any A Rgd to matrix of the same dimensions whose principal diagonal contains the elements on the principal diagonal of and zeros otherwise. With this notation, we have that diag(U)) V2,)) = Da
We similarly apply the orthogonal matrices the left and right of the RHS of C.35, letting R' U, 2,
RVz. 2 E IR dx
diag (U%, R aR T a)) Vz.) = diag R DaR Da)
diag (R') diag (D. R D..)
diag (R') - Dadiag (R') Da
(la D2 diag (R').
Therefore it holds that:
Da = (Ta D& diag (R')
(C.36)
We now wish to arrive at an expression in 3. From the definition of a holds that a = (Z, ) B and, in conjunction with the first part of Theorem C.11 holds that Da = (D2,)* J'Dg. Therefore, the LHS of (C.36) becomes:
= 'D3 (D%,) )
(C.37)
For computing (D%, -1 it is perhaps easiest to proceed did with
UT 2, (2 Z Uzz -UZ =-DD2, ) (D&R' +R'D)) (D%,)
Z 2 Uz,
and therefore
D -1 1 diag
U, [ Z,T -UZ
= =-D22 )11 DDadiag (R') diag (R)) (D2,)* )
=-2(D2,) - D2 diag (R')T.
Putting everything back together in (C.37) we have that
Da -2((D2)) -DD diag (R') D (D DD -22(D22 ) 2D diag (R') + (D%, -
(C.38)
DD..
(C.39)
28

---

Finally, for dealing with the RHS of (C.36), we recall equation (C.12) which is implies that (D 2 ) la Da Putting together (C.39) and (C.36), we get that:
2(D22 2D3 diag (R') + (D2,) LB 27 " (R'),
(C.40)
Therefore, we have the following string of implications:
(C.40) = 2 D3
Dg diag (R')
D22 [T(D2, + +D
Da diag (R')
2
1 (al4+ + v22la D
22 D 7 72 D
D3 diag (R')
2
1
-1
2 2 D
D3 diag (R')
d sinh (2)
diag (R').
dt
As a final step, we
remark that diag (R') = Dz ((3) and that
VD tr 2 sinh (Da/)) VD + sinh Da)
C.1 Extensions and further experiments
In this subsection, we describe how the insights from our analysis can be extend to relaxed assumptions on initialization, discrete gradient descent and non-linear activations.
Perturbations from assumption initialization. Our analysis is contingent upon two assump- tions regarding the initialization shape. The first assumption concerns the orthogonality of the initial feature layer W, (0). In Figure 1a, explored scenario where the feature layer is initialized with a random Gaussian matrix, yet the evolution of singular values closely aligns with our theoretical analysis for orthogonal initialization. In Figure the weight layer W2. We maintain the same experimental setup as in Figure 1a, but employ a random Gaussian initialization for W₂ with variance scales of 10-2 and 10-3 while initializing Wi with variance of 10-3, In this context, the evolution of singular values continues to adhere to the predicted trend from our analysis.
demonstrate the impact of zero initialization for
10
Singular Values
10-1
dg
10
10a
10a
10
10
10
Iterations
WG(0) -NO. .1))
Wa(0) N((0.001)
Figure 4: The time evolution of singular values of the hidden layer weights of a 2-layer linear network when trained with gradient flow initialized with Gaussian random variables and non-zero W2.
Discrete step size. Here we present a simpler problem to show how we can go beyond continuous time analysis. Consider the problem with / = d, k 1, Wi, W2 W, and W is initialized with
29

---

I. The evolution of W, with a learning rate can be written
W+1+- W + n,
a4+1 = a, + +, ac.
R R² (X X) (W.W," + +|aa) -1² (a, W;R:) (a
(xT R.
we further assume that X X = I, then it can be shown that R2, a only grow in norm and do not change direction. So the update of W, is always aligned with the rank-1 matrix Roa. Hence for small initialization, the final Wo is approximately a rank-1 1 matrix. This presents a way forward for the discrete step-size case and orthogonal data. With further analysis, we think this can be generalized to any data matrix satisfying the RIP conditions. This not included the main paper due to restrictive assumptions the data but will include these comments in the appendix.
Non-linear activations. We consider the same teacher-student setup in the Figure 2. To com- pletely characterize the dynamics and the final weight parameters a challenging problem. The Characterization of the dynamics at the small scale of initialization is also absent (for any general data matrix), Boursier et al. [2022] solves this in the case of orthogonal data. To study even this simple case of two neurons, one has to study the dynamics in phases where jumps from zero initialization (a saddle point) to another saddle and then further converge to zero train loss as seen in the Figure 5a. It is difficult to precisely characterize this intermediate saddle. With some careful additional work, we believe that our analysis can capture Phase (where you jump to the first saddle) of the dynamics where you approximately learn rank matrix, see Figure 5b, for general data matrices extending the current understanding. We hope this briefly sketches a way forward for ReI.U networks.
Evolution Singular Values of ReLU network Phase
10a
10-1
-Hist Saddle
Training loss
10
Second Sadile
Singular Values
10-2
10
10-3
2000
30100
10a
16a
10a
10P
Iterations
Iterations
(b)
Figure 5: (a) The training curve of the reacher-studen network which follows saddle-to-sadde dynamics. (b) The time evolution of singular values of the hidden layer weights of a 2-layer ReLU network when trained with gradient flow. The plot represents Phase of the training where you first leamn (approximately) rank-1 hidden layer.
D Noise Dynamies
Noise Model. Here we consider the scalar case = 1), abusing the notation W W,, Wz. The gradients of the loss L (W. a) are
VWL = x (Y XWa)a V.C -T=-wTXTT (Y XWa).
When the labels or outputs are doped with noise of magnitude å > 0, i.c., adding ~ (N, Now the gradients computed after doping with this label noise
- (Y +E- XWa)a v.c = (Y
- V = V.C WTX..
30

---

Thus label noise gradient descent with step size and with added label noise Eg at each iteration writes
W++1 = Wi (Vw (W, a.) - X -x)a,a))
The continuous time version of this SDE writes,
dW V (W, a) dt + no
Now, we consider the large noise regime, where the dominating term in the SDE is the diffusion term. Therefore we consider the SDE,
dW = nóX dB,aa
where B, is n-ddimensional Brownian motion. Note that can get rid of the term nÿ by re-scaling time by a constant factor. Similar steps for the evolution of a gives the SDE,
dw X d da la=Wa x
Now consider the compact SVD decomposition of X, i.e. x UDVT, where U, E Ruxn V Rdxn and UU U U = V = L.
dW VDU dB, da W dB.
d W D (UT (B.)B.) aT, da (v=(VTW)) D (UTdB.).
Using Levy's ' characterization Y, U B, a Brownian motion, since U is an orthogonal matrix. Let W V W, then
dW = D da W DdY,
Here we consider D = I and change the notation to get the SDE 4.1 below. Our results can be extended to any diagonal matrix D.
dw (dB:)a da WT
Proposition 4.1. The dynamics (4.1) has the following convergence properties
(a) Variance explosion. The variance of the norms of W, explode, iee,
lim
[[(((())))
lim t1000
[|a(t)]** + 00.
(b) Scale divergence. For 2 for any > O, we have that,
lim E [+)] + lim E t1000
[|| ||(t)||*]] + c0.
where = e- eaa(s)ds is the exponential moving average of a
(c) Alignment spectral bias. Denote the th row of W as w, Using [ws a] del wna aw
lim E [[ al]] + 8000
Proof. Consider the noise model Eq. (4.1),
dw = (dB:) a
da = TdB..
(D.1)
Variance. Let w, be 7th column of W and a; be the 7th coordinate a. For any E [n]. the diffusion of the quantities can be separately written
du,, = a, da; = (wi; dB:)
Using the Itô chain rule,
d w: || (w,, dw,) + (dw, dw:) = a||a:11 dt + 2a; (w,, dB.).
31

---

Similarly,
da = 2a,da; da,da; = 2a, (B,) + +||2s|| |/tos]/tt. dt.
Note that (w; dB..)
Iw |aB; for some one-dimensional Brownian motion (Bt)
da |12 || dt + wa:||m;||d23 dl| we || = dl al]aaa]]] 2a; wi dB,.
Let t ag
w> using this notation get,
du d BwndB,, du = d' + 2uvdB. Let ug = E [²]] Uo = [ Using the Dynkins formula, taking the expectation, we get, =
(D.2)
dE uu") = dE [00] dt, dE = [u"]
dug = dund/, dug undl.
This system can be transformed into,
a(uo + vdvo) = Vdl (uo + vTvo)) dl.
Solving the above ODE we get,
(uo vdvo) Cbe caeva where co =o((0) + (0).
Similarly,
auuo Vdvo) Vd (uo vdvo) de,
(u0 vāvo) cie Vat where = uo - (dv
2 [coevat + cie cge-vat]
eo 1
2vd [ooe C16 c1e-va vdt
Taking the limit proves the first part of the result.
Scale. From Eq. D.2, we have,
d(u² + vno²)= = Vn(u²² + Vdo") + uvdB.
Again using the Itô chain rule, we get,
d( u² + vdo²)" = a u2 + vdu²) a-1 d(u² + Vdv²²
+ 2
22aa 1)
(2( (2² + vob²) d(u² + vno")d + vno²).
a (u² aauuu + Vdue) a-1 VVllu² Vou +2 (vd++) 1) wodß.
a7aaa 1) (u² + Vio2) a-2 V (4 (vd+1) )² uuz²²² ) dl,
+ 2
The drift term is
avT u² vdoo)) + 2a( 1) (Vd+ 1) 2 (u² 2 v a- u²0²
u20²
=a au² + v Vd+2(a - 1) (Vd+1)²
(u2 vidni)
Again using the Dynkins formula for the evolution of expectation, we have,
2²0²
dIE [u² v E a (u² + vdu2)
Vd+ 2(a 1) (Vd+1)²
dt.
(u2 + /d22)
7
(D.3)
32

---

For any function,
g(v)
y
(u+ Vi) 2'
attains its maximum value at y = Vd. i.., g(vn) = 1/(4vd). Note that a < 1 and we have
u²v²² 2(xx 1) (Vi+ (u² + Vnoo)²
2 2( 1 1) (vī g(vn) (a 1) (VT + 1)²
2/mn
The drift can be lower bounded as the following.
ayn (u² + vdo²)= + 2a(a- - 1) (1) 2 (u² + vno²²) ax2 2²²²
va - 1) (Vi + 1)²
2 2a aal" + v
2vd
For d 2 and 0 < 1, we have co > 0,
Vd (vd+1)² 27
> >-
2/d
Using the above expression in Eq. D.3,
dE [(u2 + vdo²) a 2 acyB [[u² + vdv2) ag di.
Taking the limit,
lim E
[n22 + vito²) +
t
From the SDE (D.2), we obtain the following process with only diffusion,
d(u² - v²) (dv² u2 ddt
d(u² - v²) + (u² - o²) dt (n -1)0², e'd(u² - 0²) c¹ (u² v²) (n 1)e' ue'
de' (u² - 0²) = (n - 1)e'u 'dd
- (12(12 - v² (12) = e (12- (u((1)-)0((1)) + (n--1)
((a-1) 1) "
(u² + (no2(t2)) = =e = (t2-tt) (u²(t))=o'(tt)
- v²(tt)) + / + (n - 1)e t2
e'u? (t)dt,
Denote ((t)===²] + (t))(t), using the fact that ut1 - r²(tt) I C(tz).
C (12) se- (2- C(n) vovo((a) + (d- e'o2(1),
With f = 0, using the notation ([t)==e+ f 0 e' to denote the exponential moving average.
For any time 1> 0, we have,
C (0)====0000))00000+++0)00))00011
Raising to the power of a.
(t) < (0000
33

---

For a. b 0 and 0 < < 1, we have (a+ by a + br. Using the inequality,
c(1) se Se (0)) + d' (
Taking the expectation,
(C0))=" (*
Now, we proceed by taking the limit,
lim E (( 10
lim E[C
d lim 1000
E (t) +B (t) (bb))))))))))
1o
Thus, obtain,
lim 100
EE [t)) + [()) +
where (t) e (() is the exponential moving average.
A similar computation for u will yield,
lim E (u²(t)) + co. 00
Therefore for the limit of a; | we have,
lim
E
w
(t)
20
co,
lim E (a) "] E (a, { (t)))))))) co.
Now we proceed to combine the above results and obtain the result on W. Note that for 0 < < 1, 2 is concave. Further using the Jensen's 's inequality, we have,
[| [1]
Σ ue |1²)
w
(1]l
Similar expression for and taking the limit, we obtain, the result
lim E [[]]] + xo 100
lim
+|000|*+|[(([]]])) + xc.
where a = e (s)ss)ds is the exponential moving average of a.
Alignment. Let æ1, be the rows of the matrix W. Now the evolution of the rows be written as
2; a(dB))
d da Σ zjdBi
j=1
For any two matrices, with same dimensions define [u, v] def
uo vu
d 21,; a] = =d(zaa) d(az,))
34

---

Using the Itō chain rule,
d(z;a² dzza T + z,db' T + dz,da,
aa aaT (dB;) 2i ΣxjdBy dB! (a(dB:))
dB;
aa (dB:)
Σ 212, dBi az; Tt, dt, f=1
d(az,] ) aa T (dB;) + z,a dt + Σ 2 dB/,
d[z; a] - [zz, al dt +
Iz:, ツ,| dB
1=1
From the above evolution, we have that,
dE
[|,, al] = -E [[2 al] dt.
Hence, we have,
E | a = [=)
Let def [zi, al [k, be any (kt) entry of the matrix. Similarly, let Gjj def
[z 25] (k: 11.
de; -c,dt +
EgddB/
de? = -2; edtt + Σ Ci
+ Σ dt,
dež 2ee + Σ c2y 2e; C.
Again using the Ito formula and computing
(e?) for some E (0, 1). get,
d(c))) a = a (e:) de + 2 za 1) (E:) a-2 de?
d(c)) a = a (c:) a-1
-2e + Σ Eddt 2e, Σ cyd
+ = -) ( (E) a-2 te?
Σ
2
-2 (e dt + a (et) a-1 Σ Ejtt - 2a (c)) 1a ΣendB CndBB + jala - 1) (e.) a-2 4e²
Σ
c
dt,
2
= -2a dt + a(2a - 1) (c a-1 Σ Eddt 2a (e7) a-1 es Σ ciddB
Taking a= 0.5,
dlei = -le/ddt - lesl e, Σ cij dB
Taking expectation, we get,
dE [|e]]] -F [|e:|] dt.
Hence,
E || all] = (0). a (0))
35


Reviews:
[{"review_id": "qOBIW2rFEH", "reviewer": "Reviewer_ynT5", "summary": "The paper studies a fairly simple model of deep learning, with a view to understanding the effect of initialisation scale on the spectral structure of the trained network. The model is a two-layer network with out nonlinearity, trained via gradient flow and square loss.\n\nThe authors are able to make explicit statements about the flow followed by the singular values, as well as more implicit statements about potentials that the flow minimises after time $t \\to \\infty$.\n\nThis could be useful for building an understanding of the spectral bias of learning algorithms in more realistic situations.", "questions": "Some questions about extensions:\n- How hard is it to extend this analysis to deeper networks, and why?\n- How hard is it to include a nonlinearity, and why?\n- What's the intuitive reason for using orthogonal init?\n\nOn that vein, I think it might help followup work if you added a section that explains what the main difficulty is in extending these results to a more realistic scenario along various axes.\n\n- have you thought about a rescaled gradient flow, where the flow for one layer is adjusted relative to the other by a constant? This would correspond to the idea that in the discretised setting, both layers would do not need to have the same learning rate\n- does the work suggest any insights about how to do explicit regularisation better?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "My overall impression of this paper is very good. The authors bite off a relevant albeit idealised problem, and derive solid and technically advanced insights about it. The paper is aware of its own limitations, seems to be aware of related work (although I'm not an expert here) and the paper is very well written. I imagine that the paper could inspire future advances towards a more complete picture of the spectral bias of learning.\n\nBy the way, it doesn't really bother me that the authors focus on orthogonal init. In some sense the choice of Gaussian init in practice is somewhat arbitrary, and orthogonal init works well too. See, for example, here: https://github.com/pytorch/pytorch/issues/48144", "weaknesses": "## Minor issues:\n- please label the axes of your plot!\n- please explain more what \"mirror flow\" means\n- typo in various places \"dependant\" --> \"dependent\"\n- typo in line 90:  \"towards minimum l2 linear predictors for as vanishingly small initializations\"\n- remind reader what $P$ is in Lemma 3.4\n- different font for $P$ used in line 284", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and interesting questions which help us in clarifying and improving the paper. \n\n> What's the intuitive reason for using orthogonal init?\n\nThe orthogonal initialization captures the high-dimensional asymptotics of initialization with i.i.d. Gaussian random vectors, as such random vectors are almost orthogonal in high dimensions. Specifically, as the width $l \\ggg d $, $W_1 W_1^{\\top} \\sim I_{d \\times d} $.  Furthermore, technically speaking, this choice eases the theoretical analysis because the invariant (the quantity in Lemma 3.5) becomes isotropic and it gains some commutativity properties (the layers share singular vectors), a fact which is discussed in the comments following Theorem 3.2 and in the Appendix in Theorem C.11.\n\n> Have you thought about a rescaled gradient flow, where the flow for one layer is adjusted relative to the other by a constant? This would correspond to the idea that in the discretised setting, both layers would do not need to have the same learning rate\n\nThank you for this interesting question. Our analysis does support the rescaled gradient flow. Consider the following dynamics:\n\\begin{align*} \\dot{W_1} = \\epsilon R W_2^{\\top} , ~~ \\dot{W_2} = W_1^{\\top} R, \\end{align*}\nwhich induces the projected dynamics (analogue to eq. (3.4) in  the paper)\n\\begin{align*} \\dot{Z_1} = \\epsilon R Z_2^{\\top} , ~~ \\dot{Z_2} = Z_1^{\\top} R,  \\end{align*}\nfor which the rescaled invariance equations become $Z_1^{\\top}Z_1 - \\epsilon Z_2Z_2^{\\top} = 2 \\gamma I$ (analogue to Lemma 3.5 in our work). Building upon this invariance, we can derive expressions for the singular values at convergence in a much similar way as done already in the manuscript. In the limit of $\\epsilon \\to 0$, the limit $W_1^{\\infty}$ is of approximately full-rank, while for large $\\epsilon$,  $W_1^{\\infty}$ has a bias towards low-rank. This behavior makes sense since for $\\epsilon \\to 0$, the feature layer $W_1$ is trained very slowly and feature learning does not happen, while for large $\\epsilon$ the feature layer is trained faster and learns the appropriate features. \n\n> Does the work suggest any insights about how to do explicit regularisation better?\n\nRegularization with $\\ell_2$-norm of the parameters should induce the low rank behaviour (if a global minimum of the regularized objective is reached). However, with explicit regularization the dynamics will not converge towards a solution that interpolates the training loss. \n\n > How hard is it to extend this analysis to deeper networks, and why?  How hard is it to include a nonlinearity, and why?\n\nWe kindly point the reviewer to the general comment where we discuss the technical difficulties of extending the analysis to deeper and non-linear networks. \n"}, {"review_id": "Paw83I65mK", "reviewer": "Reviewer_mL5m", "summary": "The paper provides an analysis of an implicit bias in the hidden layer for two-layer linear networks under a small-scale initialization regime. ", "questions": "Do you have a clue to extend your results to more complex neural networks? \nDo you think the phenomenon you observe within the two-layer linear network can be extended to a wider range of networks/models?", "limitations": "The analysis only applies to two-layer linear networks while there are a lot of large-scaled neural networks in practice. ", "rating": 5, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper shows the hidden layer is biased in a small-scale initialization under over-parameterized two-linear networks. ", "weaknesses": "The novelty and results significance are not enough given there are a lot of classic papers in this area such as https://arxiv.org/pdf/1806.08734.pdf. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback.\n> The novelty and results significance are not enough given there are a lot of classic papers in this area such as https://arxiv.org/pdf/1806.08734.pdf.\n\nWe respectfully disagree with the reviewer's opinion that \"the novelty and significance are not enough given there are a lot of classic papers in this area\". More specifically, we would like to emphasize the following: \n    \n- The study of the implicit bias of training (even linear!) networks has been very intense in the last five years, and it is probably one of the most promising routes to explain the success of deep learning at present. Its relationship with *capacity of predictor space and norms* makes it amenable to provide generalization bounds. Finally, theoretically characterizing it is a hard task, thus every step forward in this direction is important. \n\n- Despite the close similarity of the title, the paper mentioned by the reviewer [1] is **completely orthogonal to ours**. More specifically, [1] studies the spectral bias of the **functions** learned by ReLU networks (and not of the weights!), through the lens of Fourier analysis. In particular, the paper presents the Fourier decomposition of ReLU networks and provides (solely) empirical evidence that lower frequency functions are learned first. We emphasize that the bias in the spectrum of the weight matrices $W_{1, 2}$ studied in our manuscript is of a **completely different nature** than the bias in the Fourier spectrum considered in [1].\n\n- Adding to the previous point, we present a rigorous and tangible argument for the learning process' incremental nature, by looking at the directly-available weight matrices -- as opposed to the more abstract representation provided by the Fourier spectrum in [1]. Our paper thus provides a view which is complementary to that of [1].\n\n- Please also refer to the general answer on possibilities to extend our analysis beyond linear architectures. The novelty of our results is discussed in detail in the comments after Theorem 3.1,3.2 and after Proposition 4.1.  \n    \n[1] Rahaman, Nasim, et al. \"On the spectral bias of neural networks.\" International Conference on Machine Learning. PMLR, 2019.\n"}, {"review_id": "dnRzbYPLq5", "reviewer": "Reviewer_zebK", "summary": "This paper explores the behavior of two-layer fully connected networks with linear activations trained using gradient flow on the square loss. The authors reveal that the optimization process is influenced by the initial parameter values and exhibits an implicit bias. Notably, in the case of small-scale initializations, the hidden layer of the network tends to have a low-rank structure. The paper introduces a variational characterization of the loss minimizers obtained under specific initialization conditions and presents a hidden mirror flow that tracks the evolution of singular values in the weight matrices. Numerical experiments are provided to support the findings and demonstrate the observed phenomena.", "questions": "$\\bf (1)$ Can the insights obtained from this specific setting be applied to deeper networks or networks with non-linear activations? \n\n$\\bf (2)$ How does the analysis technique employed in this work compare to existing techniques in the literature? Does the paper present a novel approach or build upon existing analysis methods? It would be beneficial for the authors to clearly highlight the unique aspects of their analysis and how it differs from or improves upon existing approaches in terms of theoretical insights or methodological advancements.\n ", "limitations": "Not applicable", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "$\\bf Originality:$ First, it introduces a novel approach by studying overparameterized vector regression problems in two-layer fully connected linear neural networks. The focus on linear activations and gradient flow training provides a unique perspective. Additionally, the variational characterization of loss minimizers and the introduction of a hidden mirror flow to track singular values contribute to the originality of the paper.\n\n$\\bf Quality:$ The paper exhibits a high level of quality in terms of its theoretical analysis and methodology. Theorems 3.1 and 3.3 provide rigorous mathematical proofs, establishing the relationship between gradient flow training and the potential of the neural network's parameters and singular values. The explicit expressions derived for the singular values offer concrete insights.  \n\n$\\bf Clarity:$  The paper is well-written and effectively communicates its ideas. The abstract provides a concise overview of the paper's content, and the contributions are clearly stated. Theorems and propositions are presented with clarity, and the mathematical derivations are logically organized and explained. \n\n$\\bf Significance$:  The findings regarding the implicit bias induced by the initialization scale and its impact on the low-rank structure of the hidden layer contribute to our understanding of optimization processes in linear neural networks. The variational characterization and hidden mirror flow provide valuable insights into the geometrical structure of training dynamics. The proposition on stochastic noise-induced low-rank structures expands the understanding of how noise affects network weights. These contributions have implications for network design, initialization strategies, and the optimization landscape of linear neural networks.\n", "weaknesses": "$\\bf (1)$ Restrictive setting: In this work, the theoretical analysis is conducted on a two-layer linear network, which is too restrictive. Many existing works have more general settings, see, e.g., Ref [1].\n\n[1] S. Arora, N. Cohen,W. Hu, and Y. Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019b.\n\n$\\bf (2)$ Limited discussion on practical implications: While the paper provides valuable theoretical insights into the behavior of linear neural networks, it could benefit from a deeper exploration of the practical implications and applications of the findings. How can the observed low-rank structure or the understanding of gradient flow dynamics be utilized in practical scenarios? Discussing potential applications or implications for network design, initialization strategies, or optimization algorithms would enhance the practical relevance of the paper.\n\n$\\bf (3)$ Experimental scope and validation: The paper's experimental section is based on a toy example, which may limit the generalizability and applicability of the findings. The authors should consider conducting more extensive experiments to strengthen the empirical support for their theoretical results.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for the valuable feedback and interesting questions which help us in clarifying and improving the paper. \n\n> Restrictive setting: In this work, the theoretical analysis is conducted on a two-layer linear network, which is too restrictive. Many existing works have more general settings, see, e.g., Ref [1].\n\nWhile [1] indeed studies deeper linear networks, it is still subject to stringent restrictions. In particular, the extension to deep networks is **solely possibly** due to the assumption of fully-balanced weight layers:\n$$\nW^\\top_{j+1}(0)W_{j+1}(0) = W^\\top_{j}(0)W_{j}(0) \\; \\text{for all layers } j,\n$$\nAlso, the implicit bias towards nuclear norm is shown only  with a very small scale of initialization. In comparison, our initialization is more permissive as it is unbalanced and captures the high dimensional behaviour of the standard random initialization. Yet we still provide **explicit formulas of the weights' implicit bias** (instead of the product matrix $W = \\Pi_j  W_j$) as a function of the initialization scale, **for any initialization scale**. \n\n\nIn addition, the implicit bias towards low-rank in general settings studied in [1] is only backed-up empirically (see Section 3.1 of the paper) and the theoretical arguments of the paper do not address this phenomenon head-on, but only touch upon aspects of it (see section \"Theoretical illustration\"). On the other hand, we provide a rigorous characterization of singular values throughout the entire optimization trajectory (for orthogonal data). Also, the derivations in [1] solely showcase the dynamics' dependence on the network's depth and not on the initialization scale, as in the present manuscript.\n \n\n> Limited discussion on practical implications: While the paper provides valuable theoretical insights into the behavior of linear neural networks, it could benefit from a deeper exploration of the practical implications and applications of the findings. How can the observed low-rank structure or the understanding of gradient flow dynamics be utilized in practical scenarios? Discussing potential applications or implications for network design, initialization strategies, or optimization algorithms would enhance the practical relevance of the paper.\n\nThank you for bringing up the practical implications of our work. We will now address these concerns.\n\n- Learning a low-rank network indicates the model may be represented by a much smaller network. This has implications in model compression [1], knowledge distillation [2], the lottery ticket hypothesis [3], and feature quantization [4]. \n- A rigorous characterization of the trajectory of gradient methods can give insight into whether early-stopping is beneficial, and what implications it has for the learned model.\n- In terms of  initialization strategies, we complement existing works showing the optimization benefits of initializing the layers with orthogonal matrices by demonstrating their implicit regularization guarantees. Regarding the training algorithm, we show how stochastic training can still exhibit learning in the rich regime even when the initialization scale is not low. This indicates the benefits of stochastic training. \n\n[1] Denton et. al., Exploiting linear structure within convolutional networks for efficient evaluation. NeurIPS 2014.\n\n[2] Hinton et al. Distilling the knowledge in a neural network, 2015.\n\n[3] Frankle, J. and Carbin, M. The lottery ticket hypothesis:  sparse, trainable neural networks, 2018.\n\n[4] Ramesh et. al.. Hierarchical text-conditional image generation with CLIP latents, 2022.\n\n\n> Experimental scope and validation: The paper's experimental section is based on a toy example, which may limit the generalizability and applicability of the findings. The authors should consider conducting more extensive experiments to strengthen the empirical support for their theoretical results.\n\nPlease refer to the additional experiments and the general response showing that the described phenomenon extends empirically to the case of non-linear ReLU activations.\n\n>    How does the analysis technique employed in this work compare to existing techniques in the literature? Does the paper present a novel approach or build upon existing analysis methods? \n\nThe novelty of the analysis technique can be summarized by the following points (which will be added at the start of the appendix): \n -  Existing techniques target scenarios with either fully-balanced initialization (see above for the mathematical expression) or initializations with a very small scale. Our choice of initialization removes these requirements.\n - To give a variational formulation of the weights at convergence, the existing analyses (for the diagonal linear networks) crucially depend on the existence of a mirror descent equivalent of the gradient flow. Our work takes a different approach, where we first assume an \\textit{ansatz} on the potential. We further derive the required conditions for the potential to be minimized and show that these conditions hold along the gradient flow trajectory by studying the evolution of a key quantity ($\\alpha$). \n- The mirror descent tracking the dynamics for scalar regression ($k=1$) is new. Earlier analyses only hold for the case of $l=1$ (the width) [2]. This result was enabled by studying the evolution of $\\alpha$ and writing $\\beta$ as a function of $\\alpha$ (see Lemma C.2 for details).\n- Regarding stochastic dynamics, we present a **new** analysis aimed at understanding the role of stochastic noise in the algorithm. Leveraging a continuous-time model, we shed light on how the noise further biases towards low-rank structures.\n\n[1] Arora et. al. Implicit regularization in deep matrix factorization. NeurIPS 2019.\n\n[2] Azulay et. al. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. ICML 2021.\n"}, {"review_id": "YmRTW1fDRv", "reviewer": "Reviewer_itua", "summary": "The manuscript studies the evolution of weights in a two layer *linear* NN under gradient flow subject to certain initialization. Results are provided that show how different initialization scales result in qualitatively different behavior of the learned weights.\n", "questions": "- Where are the characterizations of singular values of $W_2$? Appendix B.1 seems to be about the noise model?\n\n- To make a comment in the weakness's section more concrete: how sensitive, if at all, is (iii) of Thm. 3.2 to the fact that the an interpolating solution with minimal norm is found? ", "limitations": "- The restriction to a linear model is clearly quite limiting, though it is also clear that direct analysis of the more general setting is quite challenging. Whether or not analysis of the linear setting helps make progress when the activation function is something else is hard to predict. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main points of the manuscript are relatively easy to follow and the variational characterization of the weights (Thm 3.1 part (iii)) is interesting, as is Corollary 3.2. The manuscript also seems to do a relatively good job positioning itself withing the broader landscape.\n", "weaknesses": "The main weaknesses of the manuscript (in the broader context, not in its self defined scope) are: (1) the assumptions and (2) the restriction to the linear activation function and gradient flow.\n\nRegarding the former, I agree that the orthogonality condition on $W_1$ could likely be removed and is not essential. However, it seems that at least point (ii) in Thm 3.1 may crucially depend on the 0 initialization of $W_2.$ Certainly this is the case for analysis of the random features regime as Landweber iteration. What is not so clear in the manuscript is how sensitive (iii) in Thm 3.2 (which I think is the most interesting result) is to finding a solution of minimal norm. More generally, I think the manuscript would benefit from a bit more analysis/discussion of the sensitivity of the results to the assumptions.  \n\nI understand the restriction to linear activation functions is, at the moment, necessary to have results. Nevertheless, it is undoubtedly a weakness unless a more concrete connection to how it might enable analysis of the more general case is presented (which would also help articulate the potential impact). Such a connection could substantially strengthen the paper since it would help paint a path forward. Similarly, a brief discussion of the difficultly (or lack thereof) of analyzing gradient descent (rather than gradient flow) would be beneficial. (A quick/rough calculation suggests that there may be substantial difficulties based on additional terms that arise in the per step iteration, but maybe this is missing something.)  \n\nA more minor point:\n\n- Characterizing the initialization of $W_2$ at 0 as a limiting case seems odd since there is potentially a substantiate difference between having no component in the null space of $XW_1$ and having something non-trivial. While in the \"only learn $W_2$\" setting the part in the null space would remain small that is less clear here. If $W_2$ is not initialized to 0 would an interpolating solution with near minimal norm be found or could the joint optimization over $W_1$ and $W_2$ inflate such a component.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and interesting questions which help us in clarifying and improving the paper. \n\n > However, it seems that at least point (ii) in Thm 3.1 may crucially depend on the 0 initialization of $W_2$. \nCertainly this is the case for analysis of the random features regime as Landweber iteration. What is not so clear in the manuscript is how sensitive (iii) in Thm 3.2 (which I think is the most interesting result) is to finding a solution of minimal norm. More generally, I think the manuscript would benefit from a bit more analysis/discussion of the sensitivity of the results to the assumptions.\n\nThank you for this comment. We agree that to get precise formulations for Theorem 3.1 (ii) and (iii) the initialization of  $W_2$ to zero is required. Regarding how sensitive these results are to the condition $W_2(0) = 0$, both the bias towards low rank and the convergence to a minimum-norm interpolator are similarly sensitive to perturbations of W_2(0) away from 0. A brief sketch to support this claim is given below:\n\nAssume $l=d$ for simplicity. From our analysis, we have that $\\alpha_{\\infty} - \\alpha_{0} \\in \\mathrm{span}(X)$. Using this property, we get the following expression for $\\beta_{\\infty}$:\n$$ \\beta_{\\infty}  \\in \\left(W_1^{\\infty}(W_1^{\\infty})^{\\top}\\right) \\mathrm{span}(X) +  O(\\|W_2(0)\\|). $$\nLooking at the updates of $W_1$, we get that $W_1^{\\infty} = W_1(0) + \\mathrm{span}(X)$. So, $W_1^{\\infty}(W_1^{\\infty})^{\\top} = W_1(0) W_1(0)^{\\top} + \\mathrm{span}(X)  $. We initialize such that $W_1(0) W_1(0)^{\\top} = 2 \\gamma  I$, hence $\\beta_{\\infty} \\in \\mathrm{span}(X) + O(\\|W_2(0)\\|)$. Using the invariance of the dynamics from Lemma~3.5,\nwe can say that $W_1^{\\top}W_1 = W_1(0)^{\\top}W_1(0) + W_2W_2^{\\top} - O(\\|W_2(0)\\|^2)  $. We will extend  the above pertubation analysis both empirically and analytically and discuss the implications in the revised version of the manuscript.\n\n\n> A brief discussion of the difficultly (or lack thereof) of analyzing gradient descent (rather than gradient flow) would be beneficial. \nWe agree with the reviewer that a general extension to GD is difficult. Note that, with some restrictive assumptions on data, it is however possible to study the dynamics and characterize the limiting behaviour of the weights, in a manner that goes beyond tracking discretization errors. We provide here a sketch.\n\n\n\nConsider a simpler problem with $ l=d, k=1, W_1,W_2 = W,a$ and $W$ is initialized with $I$. The evolution of $W,a$ with a learning rate $\\eta$ can be written as\n\\begin{align*}\n    W_{t+1} &= W_t + \\eta R_t a_t^{\\top}, \\\\\\\\\n    a_{t+1} &= a_t + \\eta W_t^{\\top} a_{t}, \\\\\\\\\n    R_{t+1} &= R_t - \\eta \\left(X^{\\top}X\\right) ( W_t W_t^{\\top} + \\|a_{t}\\|^2 ) R_{t} - \\eta^2 \\left( a_t^{\\top} W_t r_t \\right) \\left(X^{\\top}X\\right) R_t. \n \\end{align*}\n If we further assume that $X^{\\top}X = I$, then it can be shown that $R_{t},a_{t}$ only grow in norm and do not change in direction. So the update of $W_{t}$ is always aligned with the rank-1 matrix $R_{0}a_{0}^{\\top}$. Hence for small initialization, the final $W_{\\infty}$ is\napproximately a rank-1 matrix. This analysis presents a way forward for the discrete step-size case and orthogonal data. With further work, we think this analysis can be generalized to any data matrix satisfying the RIP conditions. This result was not included in the main paper due to the restrictive assumptions on the data but we will include these comments in the appendix.  \n\n  > Characterizing the initialization of $W_2$ at 0 as a limiting case seems odd since there is potentially a substantiate difference between having no component in the null space of $XW_1$ and having something non-trivial. \n If $W_2$ is not initialized to 0 would an interpolating solution with near minimal norm be found or could the joint optimization over $W_1$ and $W_2$ inflate such a component.\n\nThe perturbation analysis in response to the first question of the reviewer may shed some light on why initialization of $W_2$ at 0 can be seen as a limiting case. An alternative explanation involves the trajectories of gradient flow. Let $(W_1(t),W_2(t))\\_{t \\\\geq 0}$ and  $(W_{1}^{\\\\epsilon}(t),W_{2}^{\\\\epsilon}(t))\\_{t \\\\geq 0}$ be the trajectories of gradient flow when initialized at $(\\\\sqrt{2\\\\gamma}P, 0)$ and $(\\\\sqrt{2\\\\gamma}P, \\\\epsilon U )$, respectively, for orthogonal matrices $P,U$ (of appropriate dimensions). Let $\\\\theta(t) = (W_1(t),W_2(t))$ and $ \\\\theta^{\\\\epsilon} = \\\\left( W_{1}^{\\\\epsilon} (t),W_{2}^{\\\\epsilon} (t) \\\\right) $. For time $ t < \\infty $, the difference in trajectories is tracked by:\n\\begin{align*}\n\\frac{d}{dt} \\|\\theta(t) - \\theta_t^{\\epsilon}(t) \\|^2 &= (\\theta(t) - \\theta_t^{\\epsilon}(t))^{\\top}\\left( \\dot{\\theta(t)} - \\dot{\\theta_t^{\\epsilon}(t)} \\right), \\\\\\\\  &= (\\theta(t) - \\theta^{\\epsilon}(t))^{\\top} \\left( \\nabla L{(\\theta^{\\epsilon}(t))} - \\nabla L{(\\theta(t))} \\right), \\\\\\\\\n&\\leq \\left( \\sup_{x \\in \\mathcal{B}}{\\|\\nabla^2 L\\|} \\right) \\|\\theta(t) - \\theta_t^{\\epsilon}(t) \\|^2. \n\\end{align*}\nThe inequality rests on the assumption that the trajectories stay in a bounded region $\\mathcal{B}$. Using the above evolution and by Grönwall's inequality, we have $\\|\\theta(t) - \\theta_t^{\\epsilon}(t) \\| \\lesssim \\mathcal{O}\\left(\\varepsilon\\right), \\; \\forall t < \\infty$. \nThis bound establishes the closeness of the trajectories. Regarding the second part of the question, we do not think the joint optimization would inflate the component in the kernel space of $X$ but it will not diminish such a component either. \n\n > Where are the characterizations of singular values of $W_2$ ? \n\nThe expressions for the singular values of $W_2$ can be found in Lemma C.10 in the appendix. "}, {"review_id": "ztkQeGkT6c", "reviewer": "Reviewer_EWXT", "summary": "The paper explores the convergence of gradient descent for two-layers linear neural networks.", "questions": "No questions", "limitations": "See above", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper substantially advances the knowledge on the topic.", "weaknesses": "Missing reference: E. Boursier et al, Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs. \nPlease mention this reference as they already highlight the low-rank structure of the optimal solution in the nonlinear case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback. \n\n> Missing reference: E. Boursier et. al., Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs. Please mention this reference as they already highlight the low-rank structure of the optimal solution in the nonlinear case.\n\nWe thank the reviewer for pointing out the omission and will update our paper with reference [1] and its relation to our work. More precisely, we will include the following points: \n - Our analysis can be viewed as a theoretical characterization of Phase 1 of [1], in which the neuron alignment happens. The main difference with the ReLU networks considered in [1] is that in the latter, the neurons are required to align along two distinct directions -- one for data with positive outputs and one for negative outputs. In our case (for linear activation), aligning in a single direction is enough (for k=1, i.e., scalar regression).\n  - On a technical note, our analysis broadens the setting of Phase 1 from [1] (in the context of linear networks) in the following way:  [1] works with a very small initialization scale, while our analysis holds for any scale of initialization and, in addition, we do not require the data to be orthogonal. \n\n [1] E. Boursier et al, Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs\n"}]

Author rebuttal:
We thank the reviewers for the detailed feedback and positive evaluations. We address general concerns of the reviewers below. 


## Networks with more depth. [Reviewers ynT5, zebK, mL5m ]
 A first thought to tackle this problem is to simply extend our result by initializing the first layer as an orthogonal one and all other layers as zero. With a similar analysis as the one for two-layer network, we would obtain similar results.  However, such extensions would not contribute to any further understanding or insights as layers other than the first one will be balanced and existing works focusing on balanced initialization already capture this effect. 

For general unbalanced initializations, a rigorous characterization of the implicit bias is still an open problem (even in the two-layer case). In the current literature, the only architecture where a rigorous characterization is proven lies in the case of 2-layer diagonal networks. The characterization here depends on the existence of a hidden mirror flow. For general linear networks, such an equivalence does not hold. In our results, we overcome this difficulty by an ansatz on the potential function and show that gradient flow ensures some properties of the relevant quantities ($\alpha$) (more precisely that $\alpha$ stays in $\mathrm{span}(\{X\})$). 

The hard part to extend these results to deeper networks is to come up with an ansatz for the potential and identify key quantities which can be tracked as we track $\alpha$. 

---------------------------------------------------------------------------------------------

## Networks beyond linear activations. [Reviewers ynT5, zebK, mL5m]
We believe the low-rank phenomenon to be universal and to extend to (homogeneous) non-linear activation. Rigorous experiments with deeper networks demonstrating a low rank phenomenon can be seen in [1,2] but theoretically unravelling this behaviour is still an important open problem. We believe our analysis can serve as a stepping stone to analyze non-linear settings.


To showcase our result's expected extension, we add a toy experiment for ReLU networks (see attached pdf). More precisely, we consider a scalar regression problem modelled through a ReLU network in a teacher-student setup. We generate a training set of size $10$ sampled from a random Gaussian distribution in $\mathbb{R}^{5}$. The training data $(x_i,y_i)_{i=1}^{10} \in \mathbb{R}^{5} \times \mathbb{R}$ is generated by a teacher ReLU network with $2$ neurons $(w_0,w_1)$, i.e.,\begin{align*} y_i =  a_0 \sigma(w_0^{\top}x_i) + a_1 \sigma(w_1^{\top}x_i) , \end{align*} where $\sigma$ is the ReLU non-linearity. We train a student network with 20 hidden neurons. Note that there are two relevant directions $w_0,w_1$ for the student network to learn, therefore we expect the hidden layer to represent these two directions (i.e., a rank 2 hidden layer, and a singular value decomposition with two non-zero singular values). This property is empirically verified in Figure 1 of the attached pdf where we plot the time evolution of singular values. When initialized at low-scale the network converges to an approximately rank-2 matrix. When initialize at large scale, the network is high rank and the neurons do not learn the teacher directions. Thus this new experiment demonstrates the presence of low rank structure in non linear networks and its dependence on the initialization scale. 

However, providing a theoretical characterization of the dynamics or of the final weight parameters is very challenging. To study even this simple case of two neurons with a very low initialization scale, one has to study the dynamics in two phases where the iterate jumps from zero initialization (a saddle point) to another saddle and then further converge to zero train loss as seen in Figure 2a. It is difficult to precisely characterize this intermediate saddle. With some careful additional work, we believe that our analysis can capture Phase 1, in which the dynamics progresses towards the first saddle and a rank-1 matrix is approximately learned. This stage is depicted in Figure 2b. We believe the strategy discussed above can help extend the current understanding of ReLU networks to general (i.e., not orthogonal) data matrices, demonstrating how our analysis of the linear case fits in the broader context of non-linear networks.

[1] M. Huh, H. Mobahi, R. Zhang, B. Cheung, P. Agrawal, and P. Isola. The low-rank simplicity bias in deep networks. TMLR, 2021.

[2] M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion. SGD with large step sizes learns sparse features. In ICML, 2023.

-----

We again thank the reviewers for their valuable comments and remain open to further discussions on any other points they may raise.


Meta-review: This paper investigates the behavior of two-layer fully connected networks with linear activations when subjected to gradient flow training using the square loss. The authors examine how the optimization process introduces an implicit bias on the network parameters, which is contingent upon the scale of their initialization.The central contribution of the study lies in presenting a variational characterization of the loss minimizers achieved by the gradient flow under specific initialization conditions. This characterization unveils that, within the context of small-scale initializations, the hidden layer of the linear neural network exhibits a bias toward adopting a low-rank structure.

All reviewers thought the paper had interesting ideas and recommended acceptance. They raised some technical concerns most of which were addressed. Therefore I recommend acceptance.
Decision: Accept (poster)