Title: On the spectral bias of two-layer linear networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: linear networks

spectral bias

low rank

singular values

mirror flow
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

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.

1 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]. Understanding such phenomena is 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 as the study of the inherent regularization properties (or implicit biases) of training algorithms [Neyshabur et al., 2014, Zhang et al., 2016]. These observations laid the foundation for a new line of work [see, e.g., Vardi, 2022] whose driving question is 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 is capable of learning rich feature representations 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 of representations reflects in the low-rank structure of the hidden layers. Our work aims to precisely explain this phenomenon and quantify the impact of the initialization scale on feature learning.

Unfortunately, studying such phenomena for the types of neural networks used in practice is mathematically challenging at present due to the non-linearity of their activations. Their less expressive linear counterparts, however, are more tractable and represent a good proxy due to their non-convex

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 fullyconnected linear neural networks. We show the following results when the network is trained with gradient flow (GF).

(i) In Theorem 3.1, we 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 a 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 to a mirror flow on the linear predictor. These characterizations give the geometrical structure of the training dynamics of linear neural networks.

(iii) In Proposition 4.1, we design a 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 is the impact of initialization scale: large initial weights induce a learning regime in which the parameters travel a short distance to convergence and feature learning fails to happen ( lazy regime), while small initialization effects a 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 in 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 a preference towards sparse representations is shown [Woodworth et al., 2020], and a restricted setting of the matrix factorization problem, where the implicit bias leads to low-rank representations [Gunasekar et al., 2017, Arora et al., 2019b, Li 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 a 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 a mirror flow is not always guaranteed [Li et al., 2022]. To partly alleviate this issue, Azulay et al. [2021] introduce a nonlinear time-rescaling technique and show that for scalar least-squares regression on a two-layer fully connected network with zero-balance initialization, the implicit bias selects low ℓ 2 -norm predictors. We prove a similar result under imbalanced initialization controlled by a scale parameter, and characterize the weight matrices independently at convergence, thus presenting a higher-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 ℓ 2 linear predictors for vanishingly small initializations. For classification in the case of linear networks, Ji 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-balance 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 a 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 dynamics in the case where the network initialization has full rank. Finally, Gidel et al. [2019] loosen the whitened data assumption through a 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 a 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 t in A ( t ) and simply denote it as A . The time derivative of such variables is denoted d dt A ( t ) as · A .

Vector Regression. The set-up is that of standard vector regression problems with inputs ( x 1 , . . . , x n ) ∈ ( R d ) n and outputs ( y 1 , . . . y n ) ∈ ( R k ) n in the so-called overparametrized regime where d ≥ 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 to learn the input/output rule, we minimize the square loss over a class of parametric models H = { f θ ( · ) : R d → R k | θ ∈ R p } which we specify in the next paragraph. The train loss therefore can be written as

Parameterization with a Linear Network. We consider the parametric model of two-layer linear neural networks of width l ∈ N ∗ : this corresponds to the parametrization θ = ( W 1 , W 2 ) , W 1 ∈ R d × l , W 2 ∈ R l × k and f θ ( x ) = W ⊤ 2 W ⊤ 1 x . The model is linear in the input x , and in terms of expressivity, it is equivalent to the linear class of predictors given by f β ( x ) = β ⊤ x , with β = W 1 W 2 . Wehenceforth use the symbol β to denote the associated linear predictor of the network. An important consequence of this reparametrization is that the prediction function f θ is positively homogeneous of degree 2 in θ : ∀ λ ∈ R , it holds that f λ θ = λ 2 f θ , as it is the case for two-layer ReLU networks. This property has important consequences in the loss landscape through which θ goes.

Train loss. Assume momentarily that k = 1 and denote ϕ ( x ) = W ⊤ 1 x ∈ R l . It is then clear that the predictor rewrites as f θ ( x ) = ⟨ ϕ ( x ) , W 2 ⟩ . For this reason, we call the hidden layer W 1 the feature layer and the last layer W 2 the weight matrix . We study the overparametrized setting where l ≫ d . Letting X ⊤ := [ x 1 , . . . , x n ] and Y ⊤ := [ y 1 , . . . , y n ] , the loss becomes

For brevity, we ignore the N in Eq.(2.2) by implicitly rescaling the data as ( X,Y ) ← ( X / √ N , Y / √ N ) .

Interpolators. Note that when Y ∈ span( X ) and X is non-degenerate (which occurs with probability one if e.g., X,Y are Gaussian and d ≥ n ), there always exists a solution which attains zero loss, i.e., β ∗ ∈ R d × k such that X β ∗ = Y . We emphasize the fact that there are two levels of overparametrization here: on one hand, when d > n , the set of zero loss linear predictors I β := { β ∈ R d × k | X β = Y } is typically an affine set of dimension ( d -n ) k . On the other hand, since we also reparametrize β as a linear network of width l ≫ d , the manifold of interpolators in the reparametrized space of θ , defined by I θ := { θ = ( W 1 , W 2 ) | W 1 W 2 ∈ I β } , is of dimension l ( d + k ) -nk . A natural question, therefore, is to which of these interpolators θ ∗ ∈ I θ does a given optimization algorithm converge. This concept is referred to as the implicit bias of an algorithm. The aim of this work is to study that of gradient flow.

Gradient Flow. The dynamics induced in parameter space by running gradient flow on (2.2) is given by

We wish to describe the implicit regularization properties of this continuous-time process, which is 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 W 1 , W 2 the dynamics translates to

We emphasize a crucial point: even if the function β →∥ X β -Y ∥ 2 is convex, its reparametrization in terms of W 1 , W 2 is not. Non-convexity and non-linearity makes the analysis challenging and a priori it is not even clear whether the time evolution of β can be expressed as a 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 I γ .

(a) Orthogonal feature layer: We initialize the inner layer such that the rows of W 1 are orthogonal and scale with parameter γ > 0 . Mathematically, this translates to W 1 (0) = √ 2 γP for P ∈ R d × l in the Stiefel manifold V d ( R l ) := { P ∈ R d × l , such that PP ⊤ = I d } . Initializing with an orthogonal matrix is studied by Pennington et al. [2018], Hu et al. [2020], however from an optimization perspective. Note that when l is very large, this setting approximates the real-world scenario of initializing the hidden neurons with d i.i.d. Gaussian vectors in R l , which are known to be almost orthogonal.

(b) Zero weight layer: In order to remove any initialization bias from the linear layer, we initialize it at W 2 (0) = 0 . This can be seen as the limiting case of initializing the weight layer with a very small relative scale γ ≪ γ .

As already mentioned in Section 1.1, existing studies on linear networks assume a 'zero-balance initialization', namely that W ⊤ 1 (0) W 1 (0) = W 2 (0) W ⊤ 2 (0) [Saxe et al., 2014, Arora et al., 2019b, Azulay et al., 2021]. This condition introduces the invariant W ⊤ 1 W 1 = W 2 W ⊤ 2 [Du et al., 2018], which holds for all times t ≥ 0 . This balancedness can be seen as a degeneracy assumption on the flow, since it implies that W 1 has at most rank k during the entire process, irrespective of the scale of initialization γ . In contrast, we show that depending on γ the feature layer W 1 is biased (or not) toward a low-rank predictor, thus unveiling a 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 a characterization for GF in the setting of linear networks.

Theorem 3.1. Let ( W 1 , W 2 ) be the process that follows the GF equations (2.4a) -(2.4b) , initialized according to condition I γ , for some γ > 0 . Then

(i) The parameters converge to a global optimum of the loss

(ii) The linear predictor β converges to the minimum ℓ 2 -norm interpolator

(iii) We have the following variational characterization of the limiting parameters

Interpretation of the theorem. The theorem is divided into three parts which state that (i) the matrices converge to a zero loss solution, which is a priori non-trivial since the loss in non-convex; (ii) among all the interpolators in I β , β converges to the minimum ℓ 2 -norm interpolator for all γ > 0 ; and (iii) among all the interpolators in I θ , ( W 1 , W 2 ) converge to the ones that minimize a γ -dependent potential. To fully capture the richness of this result, we observe that in the limit of γ → 0 , problem (3.1) informally translates to [Attouch, 1996]

This is equivalent, in the space of linear predictors β to the minimum nuclear norm solution β ∈ argmin X β = Y ∥ ∥ β ∥ ∥ ∗ (which is also the minimum ℓ 2 -norm interpolator for the problem we study). We informally derived this interpretation by taking first the limit t →∞ and only after γ → 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 a stationary point of the flow.

With increasing γ , we move towards solutions with a large log det ( W 1 W ⊤ 1 ) , 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 et 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, which characterizes the rich and lazy regimes for simpler diagonal linear networks.

Comparison with works on implicit bias of β . Azulay et al. [2021], Min et al. [2021] also study the implicit bias phenomenon in linear networks, however, these results only address the structure of the final predictor β and not that of the factorized problem ( W 1 , W 2 ) . As shown in Theorem 3.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 = 1 ) for which both Theorem 3.1 and [Azulay et al., 2021] show that β is biased towards lowℓ 2 interpolators. This view is not complete, since there exist many pairs ( W 1 , W 2 ) such that W 1 W 2 = β . Theorem 3.1 goes one step further and provides variational characterization of ( W 1 , W 2 ) at convergence. Moreover, it shows that when γ → 0 all columns of W 1 align in the direction of β , thus creating a rank one hidden layer. This is an example of rich representation learning, where W 1 is learning the only feature needed to make a prediction.

Implicit bias of the singular values. Proceeding with the description of the spectral bias, we provide a characterization of W 1 's limiting singular values which highlights their dependence on γ .

Corollary 3.1. [Singular values at the limit] Using the same quantities as in Theorem 3.1, denote ( σ 1 ( W ∞ 1 ) , . . . , σ d ( W ∞ 1 )) and ( σ 1 ( β ∗ ) , . . . , σ k ( β ∗ )) the singular values of W ∞ 1 , β ∗ are

Discussion. Similar expressions for the singular values of W 2 can be derived and are deferred to the Appendix B due to lack of space. In the vanishing initialization limit, only the first k singular values are activated and W 1 resembles a rank k matrix. Conversely, for large γ all the singular values are approximately equal in scale and W 1 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 γ , only one singular value grows while the others remain small constants. More precisely, when γ ∼ o ( ∥ ∥ β ∗ ∥ ∥ ) , the training model is approximately rank one, with only one spiked singular value. Conversely, when γ ∼ Ω( ∥ ∥ β ∗ ∥ ∥ ) 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 1 .

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 on 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 a very particular property: it satisfies a mirror flow [Alvarez et al., 2004] with a mirror potential that can be written explicitly.

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 β , denoted by D β , follow the mirror flow

(b) Mirror on β . If k = 1 , the dynamics of β can be characterized as a mirror flow

Mirror on Singular Values. For vector regression, GF on the parameters induces a 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 Ψ γ ∼ ( -ln γ ) ∥ ∥ β ∥ ∥ ∗ . Thus, informally, for small γ the gradient flow on parameters is approximately equivalent to a 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 saddle-to-saddle dynamics [Li et al., 2021, Jacot et al., 2022, Berthier, 2022].

In the case of whitened data, i.e., X ⊤ X = I , we show that the singular vectors of β are stationary (see Appendix C.11 for details). Therefore, the mirror flow on the singular values characterizes the entire system. In this case, we can even provide an exact expression for the evolution of the singular values

(Appendix C.12) by solving a matrix Ricatti equation [Bittanti et al., 1991]. In the limit of γ → 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 γ → 0 and with appropriately rescaled time, the limiting trajectory for the i th singular value σ i, β can be seen as the jump process

where σ i, β ∗ is the i th singular value of β ∗ . Each singular value is activated at time -ln ( γ ) (2 σ i, β ∗ ) -1 . 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 ( W 1 , W 2 ) implies a mirror flow on the predictor β with the potential ψ δ . To be more precise, the evolution is governed by a mirror flow with the time scaled as a function of ∥ ∥ β ∥ ∥ . This technique of time-warping was proposed in Azulay et al. [2021] for the case of a linear network with a single neuron ( l = 1) with balanced initializations. In contrast, with a specific initialization shape, we show the existence of a mirror flow for an arbitrary number of neurons and unbalanced initialization of any scale. The existence of a 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 to show that L ( β ( t )) = O (1 /γt ) (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 et al. [2019b]. Note that with decreasing initialization scale γ , the convergence speed diminishes, while according to the results in Theorems 3.1 better implicit bias is 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 is 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 a 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 1 (0) , W 2 (0)) = ( √ 2 γP, 0 ) . Let Z 1 := W 1 P ⊤ , Z 2 := P W 2 and the residual R := X ⊤ ( Y -X Z 1 Z 2 ) , then the evolution of ( Z 1 , Z 2 ) is governed by the following ODE

Furthermore, the dynamics of gradient flow (2.4) is equivalent to (3.4) , i.e., ( W 1 ( t ) , W 2 ( t )) = ( Z 1 ( t ) P, P ⊤ Z 2 ( t )) at any time t .

Lemma 3.1 derives an equivalent dynamics to equations (2.4). It shows that weights W ⊤ 1 , W 2 always stay in the column span of the initialization P , thus restricting their evolution to 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,

This invariant ensures that Z ⊤ 1 Z 1 and Z 2 Z ⊤ 2 commute which is a crucial ingredient in the proofs of Theorems 3.1, 3.2. Now, we derive the evolution of α := Z -⊤ 1 Z 2 , which turns out to be the central quantity enabling our result. The lemma below describes certain properties of the evolution of α .

Lemma 3.3. Let α := Z -⊤ 1 Z 2 , we have the following time evolution of parameters:

An outline of the proof of Theorem 3.1. With an ansatz on the potential that it is decomposable in terms of Z 1 , Z 2 , we derive KKT conditions for the constrained optimization problem

Using Lemmas 3.3, we show that α stays in span( X ) . We use the isotropic property of the imbalance from Lemma 3.2 to find appropriate functions ψ 1 , ψ 2 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 a deep-dive into the dynamics of the gradient flow, which we complement here with a 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 et al. [2020].

4.1 The role of noise

It was observed that the noise in stochastic gradient descent has a parameter-dependent shape that induces certain regularization properties [HaoChen et al., 2021, Blanc et al., 2020]. Here, we study the properties of the noise shape induced in the case of parameterization with linear neural networks.

Inspired from the analysis of HaoChen et al. [2021] and the large noise regime described by PillaudVivien 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 l = d and, through an abuse of notation, denote W 1 = W , and W 2 = a . The noise-driven process which we consider is:

where B t is a d -dimensional Brownian motion. Details on how this SDE captures the noise of SGD are deferred to Appendix D. We show that, similarly to the rich regime of the gradient flow ( γ → 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 D.

Proposition 4.1. The dynamics (4.1) has the following convergence properties

(a) Variance explosion . The variance of the norms of W , a

explode, i.e.,

(b) Scale divergence. For d 5 , for any α > 0

where ¯ a := e -t t ∫ 0 e s a ( s ) d s is the exponential moving average of a .

(c) Alignment - spectral bias . Denote the i th row of W as w . Using [ w , a ] def = w a ⊤ a w ⊤ ,

i i i -i .

For any two vectors u, v , [ u, v ] denotes the commutator of the vectors: remark that if [ u, v ] = 0 , then u, v are aligned, i.e, u = cv , for some scalar c ∈ R . 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 i , a ] goes to zero, indicating that all the rows of w i align towards a . Overall, similarly to the gradient flow in the rich regime, this induces a low rank structure in W . This phenomenon can be further seen through the evolution of singular values, where the top singular value of 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-dependent noise which implicitly biases the solutions towards having a low-rank structure of the hidden layer [Andriushchenko et al., 2022].

10

0

10

1

10

2

10

3

10

4

10

5

Iterations

10

-

2

10

-

1

10

0

Singular Values

γ

= 1

γ

= 10

-

4

N

(0

, .

01)

σ

0

σ

1

σ

2

σ

3

(a)

10

0

10

1

10

2

10

3

10

4

10

5

Iterations

10

-

2

10

-

1

10

0

Singular Values

LNGD

σ

0

σ

1

σ

2

σ

3

(b)

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

where δ > 0 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 a regression problem on synthetic data, with n = 5 samples of Gaussian data in R 10 ( d = 10) and the labels in R 3 ( k = 3) generated by a ground truth β ∗ ∈ R d × k . We consider a network with width l = 200 . In Figure 1a, we show the evolution of the top4 singular values of the hidden layer W 1 . We use orthogonal initialization for the network with the two scales of initialization γ = 1 , 10 -4 . Note that, as depicted by Corollary 3.1, for the smaller scale only the first k = 3 singular values are significant in comparison to the remaining d -k singular values. This shows that the matrix is approximately rank k and the neurons align along three directions. In contrast, for the larger scale γ = 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 R l . As described when l ≫ d , the initialization is close to the orthogonal initialization. Hence, in this case, we 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 5 ( n, d = 5 ) and a network with l = 200 . Further details on hyper-parameters can be found in the 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 a toy experiment for ReLU networks (see Figure 2 and further details in Figure 5). Consider a scalar regression problem in a ReLU teacher-student setup. We generate a training set of size 10 sampled from a random Gaussian distribution in R 5 . The training data ( x i , y i ) 10 i =1 ∈ R 5 × R is generated by a teacher ReLU network with 2 neurons ( w 0 , w 1 ) , i.e.,

where σ 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 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 a case for going beyond the characterization of implicit bias in the functional space [Parhi and Nowak, 2022] and further studying the implicit bias in the parameter space. Such characterizations facilitate the identification of crucial ingredients in training algorithms that enable effective feature learning.

Limitations and Future Work. This paper tackles 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 to their activations and rely on various heuristics to achieve state-of-the-art performance, thus being difficult to grasp mathematically.

10

0

10

1

10

2

10

3

10

4

Iterations

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Singular Values

γ

= 10

-

4

γ

= 10

-

2

σ

0

σ

1

σ

2

σ

3

σ

4

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 of 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 a desired objective of future investigations. Further discussion on these aspects is presented in Appendix C.1.

Acknowledgments and Disclosure of Funding

AV is supported by Swiss data science fellowship.

References

Alvarez, J. Bolte, and O. Brahic. Hessian riemannian gradient flows in convex programming. SIAM journal on control and optimization , 43(2):477-501, 2004.

Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion. Sgd with large step sizes learns sparse features. arXiv preprint arXiv:2210.05337 , 2022.

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.

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 , 2019a. URL https://openreview.net/forum?id=SkMQg3C5K7 .

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, Z. Li, and R. Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning , pages 322-332. PMLR, 2019c.

Atanasov, B. Bordelon, and C. Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=1NvflqAdoom .

Attouch. Viscosity solutions of minimization problems. SIAM Journal on Optimization , 6(3): 769-806, 1996.

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.

L. Bartlett, P. M. Long, G. Lugosi, and A. Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences , 2020.

Bengio, A. Courville, and P. Vincent. Representation learning: A review and new perspectives. IEEE Trans. Pattern Anal. Mach. Intell. , 2013.

Berthier. Incremental learning in diagonal linear networks. arXiv preprint arXiv:2208.14673 , 2022.

Bittanti, A. J. Laub, and J. C. Willems. The riccati equation. 1991.

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 Theory, COLT 2020 , Proceedings of Machine Learning Research. PMLR, 2020.

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 https://openreview.net/ forum?id=L74c-iUxQ1I .

Braun, C. C. J. Dominé, J. E. Fitzgerald, and A. M. Saxe. Exact learning dynamics of deep linear networks with prior knowledge. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/ forum?id=lJx2vng-KiC .

Chizat, E. Oyallon, and F. Bach. On Lazy Training in Differentiable Programming . Curran Associates Inc., Red Hook, NY, USA, 2019.

Du, J. Lee, H. Li, L. 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. Du, W. Hu, and J. D. Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In Advances in Neural Information Processing Systems , 2018.

S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2019b.

Fazel, H. Hindi, and S. 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.1109/ACC.2003.1243393.

Fukumizu. Effect of batch learning in multilayer neural networks. Gen , 1(04):1E-03, 1998.

Ghai, E. Hazan, and Y. Singer. Exponentiated gradient meets gradient descent. In Proceedings of the 31st International Conference on Algorithmic Learning Theory , volume 117 of Proceedings of Machine Learning Research , 2020.

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.

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.

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, J. 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.

Z. HaoChen, C. Wei, J. D. 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.

Hu, L. Xiao, and J. Pennington. Provable benefit of orthogonal initialization in optimizing deep linear networks. In International Conference on Learning Representations , 2020.

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 https://openreview.net/forum?id=bCiNWDmlY2 .

Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

Jacot, F. Ged, B. ¸ Sim¸ sek, C. Hongler, and F. Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity, 2022.

Ji and M. Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id= HJflg30qKX .

S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations , 2017. URL https://openreview.net/forum?id=H1oyRlYgg .

Lasserre. A trace inequality for matrix product. IEEE Transactions on Automatic Control , 1995.

LeCun, Y. Bengio, and G. Hinton. Deep learning. nature , 521(7553):436-444, 2015.

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.

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.

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://openreview.net/ forum?id=k4KHXS6_zOV .

Min, S. Tarmoun, R. Vidal, and E. Mallada. On the explicit role of initialization on the convergence and implicit bias of overparametrized linear networks. In Proceedings of the 38th International Conference on Machine Learning , 2021.

Neyshabur, R. Tomioka, and N. Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614 , 2014.

Oksendal. Stochastic differential equations: an introduction with applications . Springer Science &Business Media, 2013.

Parhi and R. D. 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.

Pesme and N. Flammarion. Saddle-to-saddle dynamics in diagonal linear networks, 2023.

Pesme, L. Pillaud-Vivien, and N. Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems , 34: 29218-29230, 2021.

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.

M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In 2nd International Conference on Learning Representations, ICLR 2014 , 2014.

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 , 116(23):11537-11546, 2019.

Soltanolkotabi, A. Javanmard, and J. Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory , 65: 742-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.

Timor, G. Vardi, and O. Shamir. Implicit regularization towards rank minimization in relu networks. In Proceedings of The 34th International Conference on Algorithmic Learning Theory , Proceedings of Machine Learning Research, 2023.

Vardi. On the implicit bias in deep-learning algorithms. arXiv preprint arXiv:2208.12591 , 2022.

Woodworth, S. Gunasekar, J. D. Lee, E. Moroshko, P. Savarese, I. 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-12 Jul 2020.

Yun, S. Krishnan, and H. Mobahi. A unifying view on implicit bias in training linear neural networks. In International Conference on Learning Representations , 2021. URL https:// openreview.net/forum?id=ZsZM-4iMQkH .

Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03530 , 2016.

Organization

The appendix is organized as follows,

In 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 can 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.11, the discussion on the noise model at D.

Notations. For matrices of appropriate dimensions, we use [ A,B ] to denote AB -BA .

A Experiment Details

Experiments. We discretize the SDE (4.1) with a step-size ∼ 1 / √ t . We simulate three parallel runs and track the evolution of singular values and the evolution of alignment using the commutator of the row w i , a , i.e., [ w i , a ] := ( w i a ⊤ -a w ⊤ i ) . We consider the SDE for dimemsion d = 2 . The evolution is initialized at W (0) = I 2 and a (0) = 0 . As seen in Figure 3a, the noise shape inherently induces a low-rank structure where it intensifies a 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 a rank 1 structure. The experiments were run on a 16-GB RAM Apple M1 mac with OS Ventura 13.3.1.

10 0

10 1

10 2

10 3

10 4

Iterations

10

9

10

7

10

5

10

3

10

1

10 1

10 3

10 5

Singular Values

0

1

(a)

10 0

10 1

10 2

10 3

10 4

Iterations

10

7

10

5

10

3

10

1

10 1

Commutator

|[

w

0,

a

]|

|[

w

1,

a

]|

(b)

B Main Proofs

In this section, we present the proofs of the theorem discussed in Section 3.

Theorem 3.1. Let ( W 1 , W 2 ) be the process that follows the GF equations (2.4a) -(2.4b) , initialized according to condition I γ , for some γ > 0 . Then

(i) The parameters converge to a global optimum of the loss

(ii) The linear predictor β converges to the minimum ℓ 2 -norm interpolator

(iii) We have the following variational characterization of the limiting parameters

Proof. We initialize such that W 1 = √ 2 γP, W 2 = 0 . Lemma 3.1 states that the dynamics of gradient flow is restricted to a subspace and can be equivalently described by,

where R := X ⊤ ( Y -X Z 1 Z 2 ) .

(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 0. With λ min ( X T X ) denoting the smallest eigenvalue of the X ⊤ X , the descent inequality (C.13) is as follows,

Refer to Lemma C.6 for the detailed proof.

(ii) To show that β → β ∗ in the limit, we show that β ∈ span ( X ⊤ ) , i.e., β = X ⊤ λ , for some λ . This satisfies the KKT conditions required for the following minimization problem.

The complete proof can be found at Lemma C.6.

(iii) For the limit of the projected dynamics ( Z ∞ 1 , Z ∞ 2 ) := lim t →∞ ( Z 1 ( t ) , Z 2 ( t )) , Lemma C.4 shows the following,

Using the transformation from Eq. C.7, we have,

Thus,

Therefore,

This hold on the set { ( W 1 , W 2 ) : W 1 P ⊤ ⊥ = 0 , P ⊥ W 2 = 0 } which is ensured from gradient flow from Lemma 3.1.

Corollary 3.1. [Singular values at the limit] Using the same quantities as in Theorem 3.1, denote ( σ 1 ( W ∞ 1 ) , . . . , σ d ( W ∞ 1 )) and ( σ 1 ( β ∗ ) , . . . , σ k ( β ∗ )) the singular values of W ∞ 1 , β ∗ are

Proof. Using the transformation W 1 = Z 1 P from Lemma 3.1, we obtain,

Thus, W 1 and Z 1 share the same singular values, similarly W 2 and Z 2 share the same singular values. The expressions for the singular values of Z 1 (similarly Z 2 ) can be found at Lemma C.10. This along with the fact that β → β ∗ 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 β , denoted by D , follow the mirror flow

where the potential writes Ψ γ ( D β ) := tr Ä 1 2 D β sinh -1 ( D β /γ ) -» D 2 β + γ 2 ä .

(b) Mirror on β . If k = 1 , the dynamics of β can be characterized as a mirror flow

where the potential writes ψ γ ( β ) := 2 3 [ » ∥ β ∥ 2 + γ 2 + γ ] 3 / 2 -2 γ [ » ∥ β ∥ 2 + γ 2 + γ ] 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 of β 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 1 (0) , W 2 (0)) = ( √ 2 γP, 0 ) . Let Z 1 := W 1 P ⊤ , Z 2 := P W 2 and the residual R := X ⊤ ( Y -X Z 1 Z 2 ) , then the evolution of ( Z 1 , Z 2 ) is governed by the following ODE

Furthermore, the dynamics of gradient flow (2.4) is equivalent to (3.4) , i.e., ( W 1 ( t ) , W 2 ( t )) = ( Z 1 ( t ) P, P ⊤ Z 2 ( t )) at any time t .

Proof. We choose P ⊥ ∈ R ( l -d ) × l such that P ⊥ P ⊤ = 0 . Using the fact that P, P ⊥ orthogonal and span the entire R l × l ,

Denoting ( Z 1 ) ⊥ := W 1 P ⊤ ⊥ , ( Z 2 ) ⊥ := P ⊥ W 2 , we have,

Recalling the evolution of gradient flow (2.4) on the loss

Multiplying the gradient flow updates with P ⊤ , P from the right, left (resp.) for the above equations

Similarly multiplying with P ⊤ ⊥ , P ⊥ , we have,

Using the above, we have,

Rewriting the evolution in terms of Z 1 , Z 2 , ( Z 1 ) ⊥ , ( Z 2 ) ⊥ ,

This is exactly equivalent to the gradient flow under the linear transformation which maps W 1 , W 2 to ( Z 1 , ( Z 1 ) ⊥ ) , ( Z 2 , ( Z 2 ) ⊥ ) . Now taking into consideration the evolution of ( Z 1 ) ⊥ , ( Z 2 ) ⊥ , we have that (0 , 0) is a equilibrium point for the dynamics. From our initialization I γ ,

As we initialized at the equilibrium of the dynamics we have (( Z 1 ) ⊥ , ( Z 2 ) ⊥ ) = (0 , 0) for any time t . From Eq. (C.5), we have,

The gradient flow (2.4) is equivalent to

where Z 1 (0) = √ 2 γPP ⊤ = I d , Z 2 (0) = 0 . Furthermore, from Eq. C.2, C.3, we have the following,

This finishes the proof of the lemma.

Lemma C.1. For the projected matrices given in (C.1) , we have the following invariant,

Proof. Recalling the dynamics (C.1)

Similarly,

Hence, · ˇGLYPH<0> Z ⊤ 1 Z 1 -Z 2 Z ⊤ 2 = 0 . This implies,

Lemma C.2. Let α := Z -⊤ 1 Z 2 , we have the following time evolution of parameters:

Proof. Taking the time derivative of α ,

The evolution of Z 1 Z ⊤ 1 , αα ⊤ ,

Computing the evolution of ( I -αα ⊤ ) -1 ,

Let C α := 2 γ ( I -αα ⊤ ) -1 and C Z := Z 1 Z ⊤ 1 , so we have,

Since, at initialization, C α (0) = C Z (0) , we have C α = C Z , for any time t . Therefore, we have,

Using β = Z 1 Z 2 = Z 1 Z ⊤ 1 α and the above invariant, we obtain β = 2 γ ( I -αα ⊤ ) -1 α .

Lemma C.3. The following property holds for α :

Proof. From the evolution of α , we have

Let U ∈ R d × d be the matrix projection each column of α on the column span of X ⊤ , i.e., span { x 1 , . . . , x n } , U ⊥ ∈ R d × d be the matrix projection on the orthogonal space, i.e., Ker ( X ⊤ ) . So α = U α + U ⊥ α . Note that R = X ⊤ ( Y -X Z 1 Z ⊤ 1 α ) . So U ⊥ R = 0 , since U ⊥ X ⊤ = 0 . The evolution of U ⊥ α is that

Again, U ⊥ α = 0 is the equilibrium point and our initialization α = 0 ensures that it stays at this equilibrium. This proves the lemma.

Lemma C.4. Let ( Z ∞ 1 , Z ∞ 2 ) def = lim t →∞ ( Z 1 ( t ) , Z 2 ( t )) the limit of the gradient flow dynamics. Then,

Proof. From Lemma C.3, we have that the ( Z ∞ 1 ) -⊤ Z ∞ 2 ∈ span ( X ) , so the condition ( P2 ) from Proposition C.5 holds. Note that from ( P2 ) of Proposition C.5, we have,

which is satisfied by the potential

When the imbalance is not isotropic, i.e., Z ⊤ 1 Z 1 -Z 2 Z ⊤ 2 = D , where D is some diagonal matrix ( = c I , for any constant c ). It this case,

̸

and there exists no such function Ψ 1 and the proof breaks.

Proposition C.5. Let ( Z ∗ 1 , Z ∗ 2 ) satisfy the following minimization problem

for some non-negative potential functions Ψ 1 ( Z 1 ) , Ψ 2 ( Z 2 ) . Then, ( Z ∗ 1 , Z ∗ 2 ) satisfies

Proof. The Lagrangian for the minimization problem above is,

Taking derivatives w.r.t. to Z 1 , Z 2 , we get,

As Z ∗ 1 , Z ∗ 2 should satisfy ∇L = 0 .

Lemma C.6. Let ( Z 1 , Z 2 ) be the process that follows the GF equations (3.1) , initialized according to condition I γ , for some γ > 0 . Then

(i) The parameters converge to a global optimum of the loss

(ii) The linear predictor β converges to the minimum ℓ 2 -norm interpolator

Proof. The evolution of R writes,

Note that for any three PSD matrices ABC . Using Lasserre [1995],( [see Lemma 7, Min et al., 2021])

Since Z ⊤ 1 Z 1 and Z 1 Z ⊤ 1 , share the same eigenvalues, Z 1 Z ⊤ 1 ≥ 2 γ I d . Therefore λ min ( Z 1 Z ⊤ 1 ) ≥ 2 γ . Let λ min ( X ⊤ X ) be the smallest non-zero eigenvalue of X ⊤ X . Although ( X ⊤ X ) can have zero eigenvalues, we can always restrict the evolution to the span ( X ⊤ ) , which R belongs to and without loss of generality assume that all the eigenvalues are non-zero. Using this for the first term, we have,

For the second term, Z ⊤ 2 Z 2 ≽ 0 we have,

Combining both,

Thus, using Gronwall we can show that ∥ ∥ R ∥ ∥ 2 F decays exponentially to zero, thus we have X β -Y → 0 . The last step is due to overparameterization, i.e., d > n and existence of a interpolating solution.

Implicit bias of β . We know that α ∈ span( X ⊤ ) and β = ( I -αα ⊤ ) -1 α . Using, Woodbury matrix identity,

Therefore β ∈ span ( X ⊤ ) . Hence, this satisfies the KKT conditions for

Lemma C.7 ( Mirror flow(k=1) .) . For k = 1 , the dynamics of β can be characterized as a mirror flow

where the potential writes

Proof. Here we consider the case k = 1 , let r := X ⊤ ( Y -X β ) . It is denoted lowercase since r here is a vector. The gradient flow C.1 now in Z 1 , α can be now can be written as,

From the invariance, we have,

Recall from Lemma 3.3 that

Using Sherman-Morrison,

Taking the norms, we get,

Now consider the evolution of α

ñ

∥

One cannot write I + αα ⊤ 1 -∥ ∥ α ∥ ∥ 2 as a Hessian of any function of α . However, multiplying on both sides with (1 -∥ ∥ α ∥ 2 ) -1 / 2 , we obtain,

We tackle the left hand side as follows,

Substituting the above expression and also Ä 1 -∥ ∥ α ∥ ∥ 2 ä 1 / 2 on the RHS gives us,

Define the mirror potential,

Using the fact that,

We have,

Integrating, we get,

Thus, we can write it as a continuous-time mirror descent.

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,

Proof. Note that ψ γ is convex. Let D ψ γ ( ., . ) be the Bregman divergence defined with the potential ψ γ ( . ) . Taking the time derivative of the Bregman divergence, we get,

We have,

Using the convexity of L ( . ) in β , we have

The last step is using the existence of an interpolating solution. Substituting the above two inequalities, we get,

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,

Thus · L ( β ) ≤ 0 , using this,

This completes the proof.

Definition C.9. [ Singular value decomposition and notation. ] The singular value decomposition of a given matrix A ∈ R d × k be denoted as U A D A V ⊤ A where U A ∈ R d × d , D A ∈ R d × k , V A ∈ R k × k .

Let k < d and consider the products AA ⊤ and A ⊤ A . We will refer to the diagonal matrices D A ⊤ A ∈ R k × k and D AA ⊤ ∈ R d × d as D 2 A and ˜ D 2 A , in order to ease notation. Note that ˜ D 2 A is a block-diagonal matrix containing D 2 A as the first diagonal block and 0 as the second diagonal block.

Lemma C.10 ( Singular values of Z 1 and β ) . Using the notations in Definition C.9, the singular values of Z 1 , Z 2 and β are related as the following,

Proof. From the invariance (C.8) we further deduce (by appropriately multiplying left and right with Z 1 , Z 2 and their transposes):

which implies that [ β ⊤ β , Z ⊤ 2 Z 2 ] = 0 and [ ββ ⊤ , Z 1 Z ⊤ 1 ] = 0 . Hence β ⊤ β and Z ⊤ 2 Z 2 commute and can be simultaneously diagonalizable, same is the case with ββ ⊤ , Z 1 Z ⊤ 1 . Therefore, the following relation holds (elementwise):

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 α := Z -⊤ 1 Z 2 . Then, it holds that

Furthermore, if X ⊤ X = I d , it additionally holds that

Proof. From invariance (C.8), by multiplying left and right with V ⊤ Z 1 and V Z 1 , respectively, we get that:

which implies that V Z 1 also diagonalizes Z 2 Z ⊤ 2 (i.e., the left singular vectors of Z 1 are the same as the right singular vectors of Z 2 for all t > 0 ). As a result, β = U Z 1 D β V ⊤ Z 2 and thus we have shown that U β = U Z 1 and V β = V Z 2 .

Next, by definition α = Z -⊤ 1 Z 2 , and Z 1 and Z -T 1 have the same left and right singular vectors (since Z -1 1 := V Z 1 D -1 Z 1 U ⊤ Z 1 ). This, in conjunction with the above proves U Z 1 = U α and V Z 2 = V α .

To show that U α = U R and V α = V R , we need to do a bit more work. We will show that:

Once we prove (C.17) and (C.18), it directly follows that

which gives the desired result.

To prove (C.17) we will show that R ⊤ α -α ⊤ R = 0 , ∀ t > 0 . First, we recall some relations that will be needed. Under orthogonal data, it holds that

Furthermore, using the reparametrization in terms of α which induces identity (C.12) and the invariance (C.8), the original dynamics of β given by · β = R Z T 2 Z 2 + Z 1 Z ⊤ 1 R rewrites as

Now, using (C.21), it holds that

were we used the push-through identity ( I + UV ) U = U ( I + VU ) .

Denoting C := R ⊤ α , computing the transpose version of the above ODE and using the fact that C 2 -( C ⊤ ) 2 = 1 2 ( C + C ⊤ ) ( C -C ⊤ ) + 1 2 ( C -C ⊤ ) ( C + C ⊤ ) , the following holds:

Finally, since C (0) -C ⊤ (0) = 0 , it is a fixed point of the above equation and it implies that C = C ⊤ , ∀ t . Thus, identity (C.17) is proven.

We proceed similarly for identity (C.18).

Denoting B := R α ⊤ and computing · B -· B ⊤ we get:

Since ( B (0) B (0)) = 0

Finally, we show U R = U β ∗ and V R = V β ∗ using yet another invariance, namely

-⊤ is a fixed point of the above equation, identity (C.18) is proven.

We proceed by computing the associated time derivatives and showing that their difference is null:

where we used (C.20) and the invariance (C.8). Since Z 2 (0) R ⊤ (0) Z 1 (0) -Z 1 (0) ⊤ R (0) Z 2 (0) ⊤ = 0 , it is a fixed point of the above equation and we have shown (C.22).

Finally, it holds that

where the second equivalence comes from the fact that Z ⊤ 1 Z 1 and Z 2 Z ⊤ 2 are simultaneously diagonalizable and thus commute (from invariance (C.8)), and the second equivalence comes from multiplying left and right with Z -⊤ 1 and Z -1 1 , respectively and the definition of α . Furthermore,

where the second line comes from the fact that β ⊤ α = α ⊤ Z 1 Z ⊤ 1 α = α ⊤ β . From repeated applications of (C.28) and (C.29) we obtain that

and the final identity follows.

Theorem C.12 ( Time evolution of singular values - orthogonal data ) . Assume that X ⊤ X = I d . Then, the i th singular value of β evolves over time as

As a consequence, under vanishing initialization γ → 0 and with a rescaling of time as t → ln(1 /γ ) t , the i th singular value is learned at time T i = 1 / 2 σ ⋆,i :

Proof. From Lemma 3.3 we have that · α = R -α R α ⊤ . In light of Theorem C.11, identity (C.18) and it holds that:

For the i th singular value of α it holds that:

where σ ⋆,i is the i th singular value of β ∗ . Equation (C.31) is a Ricatti ODE and is separable. For ease of notation, let p := γ σ ⋆,i , q := 1 + γ 2 σ 2 ⋆,i and r := -σ ⋆,i . Then, we need to solve the IVP:

for which we get σ α ,i ( t ) = √ q ï 1 -exp ( 2 r √ q ( t + c 1 ) ) 1+exp ( 2 r √ q ( t + c 1 ) ) ò -p and solving for the initial value gives us c 1 = 1 2 r √ q log Ä √ q -p √ q + p ä . Replacing p, q, r and rearranging we finally obtain:

In order to obtain the dynamics for the singular values of β = Z 1 Z ⊤ 1 α we recall the relation given in Lemma C.10:

where the appropriate dimensionality of the diagonal matrices is evident from the indexing of the identity matrices. Therefore, the singular values σ β ,i ( t ) are the solutions to the equation

and have the following expression:

Note that lim t →∞ σ α ,i ( t ) = 1 σ ⋆,i Ä» σ 2 ⋆,i + γ 2 -γ ä and therefore we can verify that β t →∞ -→ β ∗ by looking at the singular values, since the singular vectors are static (Lemma C.11).

We can now derive asymptotic transition times at which the singular values are learned. Now, we further process expression (C.33). Let v := » σ 2 ⋆,i + γ 2 -γ , w := » σ 2 ⋆,i + γ 2 + γ , then

Since from before we have that σ β ,i ( t ) = 2 γσ α ,i ( t ) 1 -σ 2 α ,i ( t ) and we can write 1 -σ 2 α ,i ( t ) = 2 γv σ 2 ⋆,i +

v 2 h ( t ) (2 -h ( t )) σ 2 ⋆,i , the singular values of β become:

We wish to study the limit of infinitesimal initialization γ → 0 . We introduce a constant c > 0 and rewrite γ = exp( -c ) and have c = ln(1 /γ ) . Rescaling time t → ct and taking the limit c →∞ we have

Lemma C.13 ( Mirror on singular values ) . With the same notations as in Theorem 3.1, The singular values of β , denoted by D β , follow the mirror flow

where the potential is Ψ γ ( D β ) := tr Ä 1 2 D β sinh -1 ( D β /γ ) -» D 2 β + γ 2 ä .

Proof. We recall the dynamics induces on α = Z -⊤ 1 Z 2 given in Lemma 3.3:

Writing the SVD decomposition of α as U α D α V ⊤ α = U Z 1 D α V ⊤ Z 2 we have that:

Since U ⊤ Z 1 · U Z 1 + · U ⊤ Z 1 U Z 1 = 0 and V ⊤ Z 2 · V Z 2 + · V ⊤ Z 2 V Z 2 = 0 , the matrices U ⊤ Z 1 · U Z 1 and V ⊤ Z 2 · V Z 2 are skew-symmetric (have 0 diagonal), and therefore the principal diagonals of the products U ⊤ Z 1 · U Z 1 D α and D α · V ⊤ Z 2 V Z 2 are also 0 .

We define the linear operator diag : R d × k → R d × k which maps any A ∈ R d × k to a matrix of the same dimensions whose principal diagonal contains the elements on the principal diagonal of A , and zeros otherwise. With this notation, we have that diag( U ⊤ Z 1 · α V Z 2 ) = · D α .

We similarly apply the orthogonal matrices to the left and right of the RHS of C.35, letting R ′ := U ⊤ Z 1 RV Z 2 ∈ R d × k :

Therefore it holds that:

We now wish to arrive at an expression in β . From the definition of α it holds that α = ( Z 1 Z ⊤ 1 ) -1 β and, in conjunction with the first part of Theorem C.11 it holds that D α = ( ˜ D 2 Z 1 ) -1 D β . Therefore, the LHS of (C.36) becomes:

For computing · ( ˜ D 2 Z 1 ) -1 it is perhaps easiest to proceed as we did with α .

and therefore

Putting everything back together in (C.37) we have that

Finally, for dealing with the RHS of (C.36), we recall equation (C.12) which is implies that 2 γ ( ˜ D 2 Z 1 ) -1 = I d -˜ D 2 α . Putting together (C.39) and (C.36), we get that:

Therefore, we have the following string of implications:

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 on initialization. Our analysis is contingent upon two assumptions regarding the initialization shape. The first assumption concerns the orthogonality of the initial feature layer W 1 (0) . In Figure 1a, we explored a 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 4, we demonstrate the impact of zero initialization for the weight layer W 2 . We maintain the same experimental setup as in Figure 1a, but employ a random Gaussian initialization for W 2 with variance scales of 10 -2 and 10 -3 , while initializing W 1 with a variance of 10 -3 . In this context, the evolution of singular values continues to adhere to the predicted trend from our analysis.

10

0

10

1

10

2

10

3

10

4

10

5

Iterations

10

-

1

10

0

Singular Values

W

2

(0)

∼

N

(0

, .

01)

W

2

(0)

∼

N

(0

, .

001)

σ

0

σ

1

σ

2

σ

3

Discrete step size. Here we present a simpler problem to show how we can go beyond continuous time analysis. Consider the 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 η can be written as

If we further assume that X ⊤ 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 . Hence for small initialization, the final W ∞ is approximately a rank-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 is not included in the main paper due to restrictive assumptions on the data but we will include these comments in the appendix.

Non-linear activations. We consider the same teacher-student setup as in the Figure 2. To completely characterize the dynamics and the final weight parameters is 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 two phases where one 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 1 (where you jump to the first saddle) of the dynamics where you approximately learn rank 1 matrix, see Figure 5b, for general data matrices extending the current understanding. We hope this briefly sketches a way forward for ReLU networks.

0

2000

4000

6000

8000

10000

Iterations

10

-

3

10

-

2

10

-

1

10

0

Training loss

-First Saddle-

-Second Saddle-

(a)

10

0

10

1

10

2

10

3

Iterations

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

Singular Values

Evolution of Singular Values of ReLU network - Phase 1

σ

0

σ

1

σ

2

σ

3

σ

4

(b)

D Noise Dynamics

Noise Model. Here we consider the scalar case ( k = 1) , abusing the notation W = W 1 , a = W 2 . The gradients of the loss L ( W , a ) are

When the labels or outputs are doped with a noise of magnitude δ > 0 , i.e., adding ε ∼ δ N(0 , I n ) . Now the gradients computed after doping with this label noise are

Thus label noise gradient descent with step size η and with added label noise ε t at each iteration writes

The continuous time version of this SDE writes,

Now, we consider the large noise regime, where the dominating term in the SDE is the diffusion term. Therefore we consider the SDE,

where B t is a n -dimensional Brownian motion. Note that we can get rid of the term ηδ by re-scaling time by a constant factor. Similar steps for the evolution of a gives the SDE,

Now consider the compact SVD decomposition of X , i.e., X = UDV ⊤ , where U, D ∈ R n × n , V ∈ R d × n and UU ⊤ = U ⊤ U = V ⊤ V = I n .

Using Levy's characterization Y t = U ⊤ B t is a Brownian motion, since U is an orthogonal matrix. Let ˜ W = V ⊤ W , then

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 .

Proposition 4.1. The dynamics (4.1) has the following convergence properties

(a) Variance explosion . The variance of the norms of W , a explode, i.e.,

(b) Scale divergence. For d ≥ 5 , for any α > 0 , we have that,

where ¯ a := e -t t ∫ 0 e s a ( s ) d s is the exponential moving average of a .

(c) Alignment - spectral bias . Denote the i th row of W as w i . Using [ w i , a ] def = w i a ⊤ -a w ⊤ i ,

Proof. Consider the noise model Eq. (4.1),

Variance. Let w i be the i th column of W and a i be the i th coordinate of a . For any i ∈ [ n ] , the diffusion of the quantities can be separately written as

Using the Itô chain rule,

Similarly,

Note that ⟨ w i , d B t ⟩ ∼ ∥ ∥ w i ∥ ∥ d ˜ B t for some one-dimensional Brownian motion ( ˜ B t ) t ≥ 0 .

Let u 0 := E [ u 2 ] , v 0 := E [ v 2 ] . Using the Dynkins formula, taking the expectation, we get,

This system can be transformed into,

Solving the above ODE we get,

Similarly,

Taking the limit proves the first part of the result.

Scale. From Eq. D.2, we have,

Again using the Itô chain rule, we get,

The drift term is

Again using the Dynkins formula for the evolution of expectation, we have,

For any function,

Taking the limit,

Raising to the power of α ,

attains its maximum value at y = √ d , i.e., g ( √ n ) = 1 / (4 √ d ) . Note that α < 1 and we have

The drift can be lower bounded as the following,

For d ≥ 5 and 0 < α < 1 , we have c 0 > 0 ,

Using the above expression in Eq. D.3,

From the SDE (D.2), we obtain the following process with only diffusion,

Denote C ( t ) := u 2 ( t ) + d v 2 ( t ) , using the fact that u 2 ( t 1 ) v 2 ( t 1 ) C ( t 1 ) ,

With t 1 = 0 , using the notation ˜ v 2 ( t ) := e -t t ∫ 0 e s v 2 ( s ) d s , to denote the exponential moving average. For any time t > 0 , we have,

For a, b > 0 and 0 < p < 1 , we have ( a + b ) p ≤ a p + b p . Using the inequality,

Taking the expectation,

Now, we proceed by taking the limit,

Thus, we obtain,

where ˜ v 2 ( t ) := e -t t ∫ 0 e s v 2 ( s ) d s is the exponential moving average.

A similar computation for u will yield,

Therefore for the limit of a i , ∥ ∥ w i ∥ ∥ , we have,

Now we proceed to combine the above results and obtain the result on a , W . Note that for 0 < α < 1 , x α is concave. Further using the Jensen's inequality, we have,

Similar expression for a and taking the limit, we obtain, the result

where ¯ a := e -t t ∫ 0 e s a ( s ) d s is the exponential moving average of a .

Alignment. Let z 1 , z 2 . . . z d be the rows of the matrix W . Now the evolution of the rows can be written as

For any two matrices, with same dimensions define [ u, v ] def = uv ⊤ -vu ⊤ .

Using the Itô chain rule,

From the above evolution, we have that,

Hence, we have,

E

[[

z

i

,

a

]] = [

z

i

(0)

,

a

(0)]

e

-

t

.

Let e i def = [ z i , a ] [ k, l ] , be any ( kl ) th entry of the matrix. Similarly, let c ij def = [ z i , z j ] [ k, l ] .

Again using the Ito formula and computing ( e 2 i ) α for some α ∈ (0 , 1) , we get,

Taking α = 0 . 5 ,

Taking expectation, we get,

Hence,

d

E

[

|

e

i

|

] =

-

E

[

|

e

i

|

]

d

t.

Peer reviews:
comments: SUMMARY:
The paper explores the convergence of gradient descent for two-layers linear neural networks.

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. 
Please mention this reference as they already highlight the low-rank structure of the optimal solution in the nonlinear case.

QUESTIONS:
No questions

LIMITATIONS:
See above
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n$\\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\nWEAKNESSES:\n$\\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.\n\nQUESTIONS:\n$\\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\nLIMITATIONS:\nNot applicable", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

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

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.

QUESTIONS:
Do you have a clue to extend your results to more complex neural networks? 
Do 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nMy 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\n\nWEAKNESSES:\n## 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\n\nQUESTIONS:\nSome 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-28"}