Title: SGD vs GD: Rank Deficiency in Linear Networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: SGD dynamics

label noise

implicit bias

low rank
Abstract: In this article, we study the behaviour of continuous-time gradient methods on a two-layer linear network with square loss. A dichotomy between SGD and GD is revealed: GD preserves the rank at initialization while (label noise) SGD diminishes the rank regardless of the initialization. We demonstrate this rank deficiency by studying the time evolution of the *determinant* of a matrix of parameters. To further understand this phenomenon, we derive the stochastic differential equation (SDE) governing the eigenvalues of the parameter matrix. This SDE unveils a *replusive force* between the eigenvalues: a key regularization mechanism which induces rank deficiency. Our results are well supported by experiments illustrating the phenomenon beyond linear networks and regression tasks.

SGD vs GD: Rank Deficiency in Linear Networks

Abstract

In this article, we study the behaviour of continuous-time gradient methods on a two-layer linear network with square loss. A dichotomy between SGD and GD is revealed: GD preserves the rank at initialization while (label noise) SGD diminishes the rank regardless of the initialization. We demonstrate this rank deficiency by studying the time evolution of the determinant of a matrix of parameters. To further understand this phenomenon, we derive the stochastic differential equation (SDE) governing the eigenvalues of the parameter matrix. This SDE unveils a replusive force between the eigenvalues: a key regularization mechanism which induces rank deficiency. Our results are well supported by experiments illustrating the phenomenon beyond linear networks and regression tasks.

1 Introduction

Deep neural networks have significantly advanced machine learning in recent decades. A key attribute of these models is their ability, despite being heavily overparameterized, to learn effective representations which generalizes well across different tasks. This capability has sparked substantial interest in understanding how neural networks learn internal representations for specific tasks [Bengio et al., 2013]. Gaining deeper insights into these mechanisms is crucial for enhancing model interpretability and refining training and application methodologies in real-world scenarios.

The success in learning these representations is often attributed to the gradient methods used in training. These methods navigate complex non-convex landscapes, finding solutions that not only minimize the training objective but also yield effective representations. They achieve this generalization while avoiding the spurious features that could potentially arise from the models' large number of parameters. Empirical studies have shown that the stochastic noise in gradient algorithms enhances generalization [Keskar et al., 2017] by favoring solutions with simpler structures that mitigate spurious features [Andriushchenko et al., 2022]. This paper address the overarching question:

How does stochasticity facilitate the discovery of solutions with simplified structures?

We explore this question using a simplified model: a single hidden-layer linear network. Despite lacking non-linearity, such networks capture some intricate phenomena of real-world deep networks and have been extensively studied to understand convergence [Arora et al., 2019a, Min et al., 2021], learning dynamics [Saxe et al., 2014], and the implicit bias of optimization algorithms [Gunasekar et al., 2017, Soudry et al., 2018]. Our work builds on this foundation by comparing stochastic algorithms with their deterministic counterparts, focusing on how these differences influence the learning of simpler structures.

Specifically, we analyze vector regression on two-layer linear networks trained with both gradient flow and stochastic gradient flow methods. Our contributions include:

In Section 4, we track the evolution of the determinant of the parameter matrix under gradient flow and stochastic gradient flow. We show that stochastic gradient flow drives the determinant towards zero, effectively removing irrelevant direction(s).

In Section 5, we derive a stochastic differential equation that describes the behavior of the eigenvalues of the parameter matrix. This analysis reveals a repulsive force between eigenvalues that pushes them apart and a geometric Brownian motion that pulls them toward zero.

In Section 6, we discuss the generalizability of our approach beyond square loss and various noise models, including discrete step sizes. Finally, we present experimental results in Section 7 that support our theoretical findings.

2 Related Work

Effect of SGD on generalization. The relationship between the stochasticity of SGD and its generalization capabilities has been extensively examined [Mandt et al., 2016, Jastrzebski et al., 2018, He et al., 2019, Hoffer et al., 2017, Kleinberg et al., 2018]. Notably, SGD tends to yield models with superior generalization compared to gradient descent [Keskar et al., 2017, Jastrzebski et al., 2018, He et al., 2019]. Various explorations into this phenomenon have been conducted through various approaches: hypothesizing that SGD favors flatter minima linked to better generalization, as opposed to sharp minima associated with poor generalization [Hochreiter and Schmidhuber, 1997, Keskar et al., 2017, Andriushchenko et al., 2023], using a random walk on a random landscape model to understand the impact of stochasticity [Hoffer et al., 2017], proposing that the inherent noise in SGD smooths the loss landscape [Kleinberg et al., 2018], and exploring the implications of dynamical stability [Wu et al., 2018].

Stochastic dynamics and Label Noise. Recent literature has explored label noise-driven Gradient Descent as an effective method to probe the beneficial impact of stochasticity on generalization, with two distinct perspectives emerging. Firstly, an asymptotic view on general model parametrization is considered, where Blanc et al. [2020], Damian et al. [2021] suggest that stochastic dynamics preferentially optimize a hidden objective linked to the curvature of the loss. In a related vein, Li et al. [2021] demonstrates appropriate limiting dynamics on the manifold of interpolators through time rescaling. Secondly, specifically for diagonal linear networks, HaoChen et al. [2021], Pillaud-Vivien et al. [2022] observe a similar collapsing effect due to label noise but with a finer characterization of the limiting process. Finally, in the absence of label noise, Pesme et al. [2021], Even et al. [2023] have characterized the outcomes of stochastic GF and GD for diagonal linear networks as the solutions to an implicit regularization problem that results in sparser solutions than without stochasticity. Recently, Ghosh et al. [2023] further exhibit a similar sparser features effect for single-neuron autoencoder. Chen et al. [2023] provides a condition under which an invariant set is attractive for SGD - characterizing the local behavior around these sets. The paper also studies linear networks in a teacher-student setup, however due to structured label-noise [Chen et al., 2023, A2 in p.30], the analysis falls short of capturing the repulsive force in the singular values.

Linear Networks. The study of two-layer linear networks has been explored extensively, particularly when optimized using gradient flow on the square loss, across various settings including zero-balance initialization and whitened data Fukumizu [1998], Saxe et al. [2014, 2019], Braun et al. [2022]. Early work by Saxe et al. [2014, 2019] elucidates the temporal changes in the singular values of the predictor, assuming decoupled dynamics and a specific data-dependent weight initialization. This condition is broadened by the analyses of Fukumizu [1998] and Braun et al. [2022], Tarmoun et al. [2021], who apply solutions from a matrix Riccati equation to characterize the weights dynamics under full-rank network initialization. Furthermore, Gidel et al. [2019] extends the existing framework by relaxing the whitened data assumption, conducting a perturbation analysis, and discussing the temporal evolution of the weight matrices' singular values. Additionally, Varre et al. [2024] eliminates the need for zero-balanced and full-rank initializations. Their study provides detailed formulas for weight evolution as a function of the initial scale , also studies a simple version of a stochastic flow without the drift. Wang and Jacot [2023] studied the implicit bias of SGD with ℓ 2 -regularization.

Matrix valued stochastic process and their eigenvalues. Stochastic process on the space of symmetric (or Hermitian) matrices and the evolution of their eigenvalues are well studied since Dyson [1962]. These techniques were further developed by Bru [1989, 1991] to study perturbations of principal component analysis and the eigenvalues of Wishart processes. Norris et al. [1986], Graczyk and Małecki [2013] applied SDE-based techniques to study the eigenvalues and eigenvectors of Brownian motion on ellipsoids.

3 Linear networks and continuous-time gradient method

Notation We use ⟨ ., . ⟩ to denote the inner product, i.e., ⟨ u, v ⟩ = u ⊤ v for vectors, and ⟨ A,B ⟩ = Tr( AB ⊤ ) for matrices. I d denotes the identity matrix of dimension d and 0 p × k denote the matrix with all zero entries of dimension p × k .

Vector regression. We study the vector regression problems with inputs x 1 , . . . , x n in ( R p ) n and outputs y 1 , . . . , y n in ( R k ) n . We consider the minimization of the square loss over a class of parametric models H = { f θ ( · ) : R p → R k | θ ∈ R d } specified in the next paragraph. The train loss therefore can be written as L ( θ ) = 1 2 n ∑ n i =1 ∥ y i -f θ ( x i ) ∥ 2 .

Parameterization with a linear network. We focus on two-layer linear neural networks of width l ∈ N ∗ . The model is described by the parameterization θ = ( W 1 , W 2 ) , where W 1 ∈ R p × l and W 2 ∈ R l × k , and the function f θ ( x ) = W ⊤ 2 W ⊤ 1 x . This model is linear with respect to the input x . In terms of expressivity, it is comparable to the linear class of predictors, represented as f β ( x ) = β ⊤ x , where β equals W 1 W 2 . Throughout our analysis, we denote the equivalent linear predictor of the network as β . A key aspect of this parametrization is that the prediction function f θ is positive homogeneous of degree 2 with respect to θ : specifically, for any λ ∈ R , f λθ = λ 2 f θ . This property mirrors that of two-layer ReLU networks and significantly influences the loss landscape navigated by the parameters θ . It is important to note that this parameterization introduces some redundancy, a single linear predictor β can have multiple representations W 1 , W 2 such that W 1 W 2 = β . Some representations have a rich structure whereas other resemble random features. For example, consider the case of scalar regression ( k = 1 ), for a vector β there exists rich parameterizations where all the neurons, i.e., columns of W 1 align with β and also some lazy structures where W 1 resembles a random matrix [Chizat et al., 2019, Varre et al., 2023].

Train loss. By defining X ⊤ = [ x 1 , . . . , x n ] and Y ⊤ = [ y 1 , . . . , y n ] , the loss function is given by:

For simplicity, we adjust for the normalization factor n by rescaling the data to ( X,Y ) ← ( X / √ n , Y / √ n ) , thereby implicitly considering it in the loss function without directly mentioning n in the formula. Note that the loss is non-convex in W 1 , W 2 .

Gradient flow. The dynamics induced in parameter space by running GF on Equation (3.1) is given by

Introducing the block matrix, Θ = [ W ⊤ 1 | W 2 ] ∈ R l × ( p + k ) and denoting the residual matrix by R = X ⊤ ( Y -X W 1 W 2 ) , the evolution of Θ can be written as

The gradient flow can therefore be compactly written as

The gradient flow (GF), when expressed in this form, reveals an inherent multiplicative structure with respect to Θ in the gradient of the loss. As we see in subsequent sections, this representation of the gradient flow with block matrices proves to be very convenient.

Label noise gradient descent. Label noise gradient descent (LNGD) is a theoretically studied alternative to SGD that mirrors its practical behavior by sharing the geometric properties of the noise Blanc et al. [2020], Damian et al. [2021]. Let ε t ∈ R n × k , where each entry of ε t is an independent Gaussian random variable. At iteration t , the labels are perturbed with this Gaussian noise at an intensity δ , i.e., ˜ Y = Y + √ δε t . The LNGD algorithm updates the iterates with a step size η in the direction of the gradient computed after the labels have been perturbed, as follows:

where, by an abuse of notation, L ( Y, X , W 1 , W 2 ) = 1 / 2 ∥ X W 1 W 2 -Y ∥ 2 . The iterates can then be restructured into a block matrix:

and J t is defined as in Equation (3.4).

Stochastic gradient flow (SGF). We aim to model the aforementioned LNGD in continuous time using an appropriate SDE. Stochastic continuous-time counterparts of discrete stochastic gradient algorithms are favored for their enhanced amenability to theoretical analysis. We propose the following stochastic differential equation (SDE) to model LNGD in continuous time:

where B t denotes a matrix Brownian motion in R n × k . LNGD as defined in Equation (3.5), can be interpreted as the the Euler-Maryama discretization of the above SGF with a stepsize η . Although the inclusion of step size in the continuous-time modeling of an SDE may seem counter-intuitive, it is a necessary component [Li et al., 2019b]. As all the terms of the SDE in Equation (3.6) are polynomial in Θ , both the drift and diffusion terms are locally Lipschitz continuous. Hence, the solution of the SDE is uniquely defined up to the explosion time τ ∞ [see, e.g., Khasminskii, 2012]. Furthermore, the explosion time can be proven to be infinite ( τ ∞ = ∞ almost surely), by using that the GF does not diverge and applying the techniques outlined by Pillaud-Vivien et al. [2022, Proposition 10].

Initialization. The dynamics of gradient methods on homogeneous models are significantly influenced by initialization, which determines the regime they operate in-specifically, the lazy regime for large initializations and the rich regime for small ones [Chizat et al., 2019, Woodworth et al., 2020]. Thus, the scale of initialization has garnered significant interest, particularly its impact on the training of linear and non-linear networks with GD [Woodworth et al., 2020, Boursier et al., 2022]. It is observed that stochastic methods eliminate the dependence on initialization [Pesme et al., 2021].

Conserved quantities and balanceness. Gradient flows follow specific conservation laws along their trajectory [Marcotte et al., 2023], maintaining characteristics of the initial conditions. For linear networks, this conservation manifests as the balanceness property [Du et al., 2018], described by:

As a result, Saxe et al. [2014], Arora et al. [2018, 2019b] have adopted balanced initialization , where ∆ (0) = 0 , to ensure that weight matrices remain low rank throughout the trajectory. However, unbalanced initialization do not preserve these simple low-rank structures, as aspects of the initial conditions persist.

In contrast, stochastic methods do not adhere to these conservation laws [Ziyin et al., 2023] and the evolution of the imbalance ∆ for SGF is

While there is no diffusion term in the derivative, the matrices remain stochastic and no definitive conclusions can be drawn from this. However, in the case where k = p and X ⊤ X = I p , it can be shown that W ⊤ 1 W 1 -W 2 W ⊤ 2 → 0 , indicating that the stochastic noise eliminates initial imbalance.

Conclusion. Understanding how stochastic methods mitigate dependency on initialization requires exploring beyond the evolution of the imbalance ∆ . To this end, we identify and discuss other conserved quantities, such as the determinant of the block matrix Θ ⊤ Θ in the following sections.

4 Separation between Gradient Flow through determinant

Here, we present our first separation result between GF and SGF. While the determinant of the parameters is preserved in GF, it is driven to zero by the stochasticity of SGF, leading to a simplistic low-rank structure.

4.1 Determinant evolution of the gradient flow

The theorem below demonstrates that the determinant of the parameters is preserved in gradient flow.

Theorem 4.1. For the gradient flow defined in Equation (3.4) , the following property holds,

The proof presented in the App. B.1, is based on straightforward computations of the derivative of the determinant and the fact that the matrix J has zero trace. We note that the simplicity of the proof arises from the strategically chosen block structure of Θ . This result would have been less straightforward with different parametrizations, which likely explains why such a simple finding appears to be novel. The theorem implies that the determinant of M along the trajectory remains equal to the determinant at initialization. If Θ ⊤ 0 Θ 0 is full-rank initially, meaning the determinant is non-zero, the theorem ensures that the determinant of M remains non-zero. Consequently, the rank of Θ does not diminish along the trajectory. When l ≥ p + k , i.e., the hidden layer has a large width and W 1 , W 2 are initialized randomly from a Gaussian distribution, Θ ⊤ 0 Θ 0 has full rank almost surely. The theorem also reveals some implications regarding the impact of initialization scale. Note that λ min ( A ) ≤ n √ det A , indicating that when the scale of initialization is very small, at least one singular value of Θ is small.

4.2 Determinant evolution of the stochastic gradient flow

In contrast, the theorem presented below demonstrates that the determinant of the parameters converges to zero in stochastic gradient flow.

Theorem 4.2. For the SDE, defined in the Equation (3.6) , for t ≤ τ ∞ , the following property holds for the evolution of determinant

Although the evolution of the parameters in SGF is random, the evolution of the determinant is deterministic. The theorem highlights a striking phenomenon: the noise in SGF diminishes the determinant along the trajectory, leading to a simplification of the network over time. The larger the noise and the stepsize, the faster the determinant vanishes. The vanishing of the determinant suggests that the rank of the parameters decreases by at least one, effectively eliminating some components. It holds for any initialization of Θ 0 and indicates how the SGF overrides some aspects of initialization. The proof uses the fact that stochastic Brownian term in the SDE, through Itô's calculus, introduces a negative drift, ultimately driving the determinant to zero (refer to B.3 for the proof).

Limitations. Given the large width of the hidden layer, the determinant converging to zero does not fully reveal the complexity of the situation. It merely indicates that at least one singular value is approaching zero. Furthermore, the theorem provides limited insights when the determinant is already 0 at initialization, det Θ 0 = 0 which happens whenever l < p + k . Next, we explore the mechanisms behind this low-rank phenomenon, suggesting that the repulsive forces induced by stochasticity drive the spurious singular values to zero as seen in the right plot of Figure 1.

5 Mechanism behind the low-rank phenomenon

In this section, we investigate the evolution of singular values under stochastic training to gain deeper insights into the low-rank phenomenon. To simplify the discussion, throughout the section

we consider the case where k = 1 and for notational convenience, we let W 1 = W , W 2 = a . Additionally, we assume that l ≤ p , however the results can be extended to any l .

Warm-up: Comparison with diagonal networks. Let W = U Σ V ⊤ be the singular value decomposition (assuming l ≤ p ). The predictor β can be expressed as

This expression reveals a Hadamard product between σ and c , reminiscent of diagonal networks which are widely studied to understand the nonconvex dynamics of gradient algorithms [Woodworth et al., 2020, Pesme et al., 2021, Pillaud-Vivien et al., 2022]. In the context of diagonal networks, SGD is known to provably induce sparsity in predictions. Similarly, for linear networks, SGF may induce sparsity in terms of the singular value σ . We next derive the SDE governing the evolution of the singular values Σ of the weight matrix to gain a clearer understanding of the low-rank phenomenon.

Scalar Regression. We assume that the data is isotropic, i.e., X = I p . Under these conditions, the loss function for scalar regression can be written as

We train the above objective with SGF, formulated as follows,

where B t is the standard Brownian motion in R p . For analytical convenience, we rescale the time t → t / ηδ and use the process d X = 1 / ηδ ( y -Wa ) d t + d B t . The SGF can then be rewritten as,

Our focus is on understanding the evolution of the singular values of the matrix W . This aim is facilitated by considering the symmetric matrix M = W ⊤ W , whose eigenvalues are the squares of the singular values of W . Taking the derivative of M , we find

Note that d x d y represents d [ x, y ] for any continuous semi-martingales x, y [see, e.g., Ikeda and Watanabe, 1981, chapter 3 for reference].

Eigenvalues of a matrix-valued stochastic process. We leverage tools from the study of eigenvalues of matrix-valued stochastic processes [Bru, 1989, Graczyk and Małecki, 2013] to derive the evolution of the eigenvalues of M in the theorem that follows.

Theorem 5.1. Let s 1 > .. . > s l be the order of the eigenvalues of the matrix M defined by Equation (5.4) . Let the collision time for the eigenvalues be defined as

̸

For t ≤ τ , the eigenvalues are semi-martingales given by the solution of the following SDE

̸

where c = V ⊤ a and Ä d ˜ X ä i = 1 / ηδ Ä ⟨ u i , y ⟩ -√ s i c 2 i ä d t + d ε i with u i being the i th column of U and ( ε 0 , . . . , ε l -1 ) is the standard Brownian motion in R l . The evolution of c i and U are presented in the appendix B.5.

This theorem can be interpreted as the stochastic counterpart to the evolution of eigenvalues previously described for linear networks by Arora et al. [2019c], Varre et al. [2023]. The derivation of the eigenvalues is inspired by the work of Bru [1989].

The evolution of the eigenvalues features a key term highlighted in Equation (5.6) consisting of the sum of skew-symmetric elements s i c 2 j + s j c 2 i / s i -s j . For a pair of indices ( i 0 , j 0 ) with i 0 < j 0 and thus s i 0 > s j 0 , the term s i 0 c 2 j 0 + s j 0 c 2 i 0 / s i 0 -s j 0 positively influences the evolution of the larger eigenvalue d s i 0 and negatively affects the smaller eigenvalue d s j 0 . Therefore, this force is repulsive,

driving the eigenvalues apart and increasing their gap. Another factor influencing the dynamics is the presence of Geometric Brownian motion, where the singular value σ i multiplicatively influences the Brownian motion as √ s i c 2 i Ä d ˜ X ä i , similar to what is observed in diagonal linear networks (refer to the previous discussion for similarities). This effect tends to pull the singular values toward zero. Together with the fact that ( s i , c i ) = (0 , 0) represents a fixed point of the dynamics, these two forces collectively push redundant singular values toward zero.

To further understand the interplay of repulsive forces and geometric Brownian motion, we consider the evolution of the smaller singular value s p for l = p . Using the Ito chain rule, we analyze the evolution of log s p , expressed as,

̸

̸

Using that s p c 2 j + s j c 2 p / s p -s j < -c 2 p , for all indices j , the repulsive force accumulates to -( p -1)( c 2 p / s p ) and the Ito correction term from the logarithm contributes an additional -2( c 2 p / s p ) (the GBMcomponent) thus offsetting the positive drift of p ( c 2 p / s p ) . In the case of l = p , considering a polynomial x α with an appropriate α would demonstrate similar behaviour. This discussion outlines the forces at play, yet a complete characterization of the solution of the SDE Equation (5.6) remains missing. Moreover, we have not established that the eigenvalues avoid a.s. collision, i.e., the explosion time τ ∞ = ∞ which is in itself a significant challenge [Bru, 1989, Graczyk and Małecki, 2014].

Asimplified two-vector problem. To enhance our understanding of the SDE governing the evolution of the eigenvalues detailed in Equation (5.6), we consider the large noise limit. In this scenario, the process described in Equation (5.3) simplifies to a purely noise-driven process without drift:

This SDE exhibits notable symmetry; allowing for an analysis using a matrix with sub-sampled columns. Let S be any subset of 1 , . . . , l , with ( w i ) l i =1 representing the columns of W . We define W S ∈ R p ×| S | as the subsampled matrix obtained by selecting columns w i where i ∈ S , and similarly, we define a subsampled vector a S by selecting the corresponding coordinates. The SDE restricted to the set S is structured as follows:

To demonstrate that the columns of W align, we leverage the symmetry of the SDE by examining the restricted problem on every pair of rows S = { i, j } , and proving alignment within this subset. This approach leads us to consider the two vector problem ( l = 2 ), where W = [ w 1 | w 2 ] and w 1 , w 2 ∈ R p , a ∈ R 2 . We describe the behavior of the eigenvalues for this two-vector problem in the theorem below.

Theorem 5.2. In the large noise limit, let s 0 > s 1 be the eigenvalues of W , the following properties hold, for t ≤ τ defined by τ = { inf t : s 0 ( t ) = s 1 ( t ) } ,

(a) s 0 , s 1 are greater than zero almost surely,

(b) for α = ( p -3) / 2 , s -α 0 is a super-martingale while s -α 1 is a sub-martingale.

This model for l = 2 mirrors the dynamics of the Wishart process studied by Bru [1991], motivating the exploration of the evolution of an appropriately chosen exponent of s 0 , s 1 . The first part of the theorem arises from the fact that s -α 1 s -α 2 is a local continuous martingale that cannot explode to infinity in finite time. The second part highlights a clear separation between the eigenvalues: one is a sub-martingale that consistently increases in expectation, while the other is a super-martingale that diminishes (note that the eigenvalues are raised to a negative power). This dynamic, coupled with the symmetry argument, suggests that for every pair of columns, there is a component that strengthens the alignment through its increases in expectation. Refer to App. B.6 for the proof.

Conclusion. In this section, we derive the SDE of eigenvalues for the matrix of parameters evolving under SGF. This derivation provides deeper insights into the mechanisms contributing to low-rank behavior. Specifically, repulsive forces drive the eigenvalues apart, while the geometric Brownian motion pulls them towards zero. These forces, unique to training with SGF, highlight the regularization effects of stochastic methods compared to gradient flow. However, fully characterizing the solution of this SDE remains a challenging open problem we let as future work.

6 Generalization to other settings

In this section, we generalize our results beyond the square loss and the label noise gradient flow. We consider the general framework of a loss function over the weight product W 1 W 2 defined as

In this framework, the loss function ℓ combines the prediction loss directly with the parametrized model f θ . This approach applies, for example, to classification problems using linear networks where ℓ might represent any classification loss and f θ = W 1 W 2 . It also directly extends to more complex architectures where f θ = σ ( W 1 W 2 ) for an activation function σ , including settings like a self-attention layer with frozen value vectors. We denote the product by β = W 1 W 2 noting it solely controls the loss. We investigate the evolution of the weight matrix determinant for a general loss across various algorithms, from gradient flow to gradient descent, and demonstrate that a similar separation occurs due to stochasticity.

Warm-up: Gradient flow. The gradient flow on the loss L can be written as the following,

Following a similar proof as in Theorem 4.1, we obtain that d ( det ( Θ ⊤ Θ )) = 0 . For separable classification problem, the gradient flow converges to infinity [Soudry et al., 2018, Ji and Telgarsky, 2019], hence, after appropriate rescaling, the layers are aligned, as shown by Ji and Telgarsky [2019]. Next, we contrast this result with the outcomes observed in stochastic and discrete algorithms.

Continuous modelling of SGD. We consider the SGD algorithm with a batch size B . We denote the mini-batch version of the loss functions L and ̂ L as L B and ̂ L B , respectively. The SGD update with stepsize η can be represented with the following block structure,

We denote the SGD noise as g t = Ä ∇ ̂ L ( β ) -∇ ̂ L B ( β ) ä and the noise covariance as Σ t = E î g t ( g t ) ⊤ ó where the expectation is over all the minibatches. Following Li et al. [2019a], the SGD update can be modelled with the following SDE,

The main difference with SGF is that, in overparameterized problems, the noise covariance is timevarying and decreases to zero upon convergence. Using Theorem B.3, the evolution of the determinant of M = Θ ⊤ Θ is given by d (det ( M )) = -η det ( M )Tr (Σ( t )) d t and can be explicitly solved as

Hence, the decay in the determinant is governed by the integral ∫ ∞ 0 Tr(Σ( t )) d t which is a stochastic quantity. Tr(Σ( t )) represents the strength of the stochastic noise, which, in over-parameterized regression, is proportional to the loss, i.e., Tr(Σ( t )) ∝ L ( Θ ) [Pesme et al., 2021]. Therefore, the rate of decay in the determinant depends on ∫ ∞ 0 L (( Θ ( t ))) d t , with slower convergence leading to a simpler model at convergence, as observed in the case of diagonal networks by Pesme et al. [2021]. The result above also holds for non-separable classification tasks where the noise of SGD drives the determinant to 0 , a scenario not covered by the previous analysis of Ji and Telgarsky [2019].

Discrete gradient algorithms. We can extend the previous results to discrete (possibly stochastic) gradient algorithm. Both algorithms can be written as

for stepsize η and J t the possibly stochastic block gradient matrix defined in Equation (6.1). In the context of discrete algorithms, the determinant is controlled by the following lemma (refer to B.4 for the proof).

10 0

10 1

10 2

10

7

10

5

10

3

10

1

10 1

10 3

10 5

Determinant

=0

=2

10 0

10 1

10 2

10 3

10

4

10

3

10

2

10

1

10 0

10 1

Singular values

0

1

2

3

4

=0

=2

Iteration

Iteration

10 0

10 1

10 2

10 3

Iteration

10

9

10

7

10

5

10

3

10

1

10 1

Singular values

0

1

2

3

4

=0.005

=0.1

10 0

10 1

10 2

10 3

Iteration

10

15

10

12

10

9

10

6

10

3

10 0

Singular values

0

1

2

3

4

5

=0.005

=0.5

10 0

10 1

10 2

10 3

10 4

10 5

Iteration

10

2

10

1

10 0

10 1

Singular values

0

1

2

3

4

=0.0001

=0.001

Lemma 6.1. When l = p + k and η 2 ∥ ∥ J t ∥ ∥ 2 F ≤ 1 , the following property holds for the determinant,

If the factor η 2 ∥ ∥ J t ∥ ∥ 2 F ≤ 1 at every iteration t , the determinant is reduced by the discrete step size. However, there is a tradeoff: the sum S := ∑ ∞ t =0 η 2 ∥ ∥ J t ∥ ∥ 2 F can be finite, indicating that it does not completely drive the determinant to zero. Increasing η to increase S might lead to instability and divergence. Furthermore, since ∥ ∥ J t ∥ ∥ 2 F ∝ L ( Θ t ) , there is an additional tradeoff between convergence and the simplicity of the parameters. This illustrates how step sizes that produce non-convergent training loss patterns, such as the catapult effect [Lewkowycz et al., 2020] or the edge of stability mechanisms [Cohen et al., 2020], can simplify the network's parameters.

7 Experimental evidence

We consider a regression problem on synthetic data with n = 1000 samples of Gaussian data in R 5 ( p = 5 ) with labels in R 2 ( k = 2 ) generated by some ground truth β ∈ R 5 × 2 , the width of the network is l = 10 . We use Gaussian initialization of the network parameters with entries from N (0 , 1) . Experiments details can be found in the appendix C. In the left plot of Figure 1, we show the time evolution of the determinant of matrix M . As suggested by theorems 4.1 and 4.2, in the case without label noise, det ( Θ ⊤ Θ ) stays constant, while with the Label Noise of intensity δ = 2 it goes to zero with time. In the right plot of Figure 1, we demonstrate the time evolution of the top-5 singular values of the matrix W 1 . Note that in the case of Gradient Flow all except the first k singular values ( σ 0 and σ 1 ) stay at the same scale, while adding Label Noise forces smallest d + l -k singular values ( σ 2 , σ 3 , and σ 4 ) to tend toward zero. Further experiments illustrate in Figure 2 the evolution of singular values of parameter matrix W 1 when optimized with SGD, for classification tasks and with ReLU network. These results also confirm that the beneficial effects of stochasticity hold in these contexts.

8 Conclusion

In this paper, we demonstrate a distinct separation between GF and SGF when trained on linear networks. This separation is obtained by tracking the evolution of the determinant of the parameter matrix. However, while the determinant is a significant factor, it does not fully capture the implicit regularization effects. Notably, the determinant mirrors the imbalance u 2 -v 2 in diagonal networks represented by u ⊙ v , whose dynamics play a crucial role in attuning the implicit regularization across various algorithms [Woodworth et al., 2020, Pesme et al., 2021, Papazov et al., 2024]. Our analysis presents the initial step in deciphering implicit regularization for stochastic methods in linear networks, yet achieving a complete characterization remains a promising direction for future research.

Acknowledgments and Disclosure of Funding

AV is supported by Swiss data science fellowship. This work was supported by the Swiss National Science Foundation (grant number 212111).

References

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

Andriushchenko, F. Croce, M. Müller, M. Hein, and N. Flammarion. A modern look at the relationship between sharpness and generalization. In International Conference on Machine Learning , pages 840-902. PMLR, 2023.

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, N. Cohen, W. Hu, and Y. Luo. Implicit regularization in deep matrix factorization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019c. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/ c0c783b5fc0d7d808f1d14a6e9c8280d-Paper.pdf .

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

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 .

M.-F. Bru. Diffusions of perturbed principal component analysis. Journal of Multivariate Analysis , 29(1):127-136, 1989. ISSN 0047-259X. doi: https://doi.org/10.1016/0047-259X(89)90080-8. URL https://www.sciencedirect.com/science/article/pii/0047259X89900808 .

M.-F. Bru. Wishart processes. Journal of Theoretical Probability , 4:725-751, 1991.

Chen, D. Kunin, A. Yamamura, and S. Ganguli. Stochastic collapse: How gradient noise attracts SGD dynamics towards simpler subnetworks. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/forum?id=iFxWrxDekd .

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

Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations , 2020.

Damian, T. Ma, and J. D. Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems , 34:27449-27461, 2021.

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.

J. Dyson. A Brownian-Motion Model for the Eigenvalues of a Random Matrix. Journal of Mathematical Physics , 3(6):1191-1198, 11 1962. ISSN 0022-2488. doi: 10.1063/1.1703862. URL https://doi.org/10.1063/1.1703862 .

Even, S. Pesme, S. Gunasekar, and N. Flammarion. (s) gd over diagonal linear networks: Implicit regularisation, large stepsizes and edge of stability. Advances in Neural Information Processing Systems , 2023.

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

Ghosh, S. Frei, W. Ha, and B. Yu. The effect of sgd batch size on autoencoder learning: Sparsity, sharpness, and feature learning. arXiv preprint arXiv:2308.03215 , 2023.

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.

Graczyk and J. Małecki. Multidimensional yamada-watanabe theorem and its applications to particle systems. Journal of Mathematical Physics , 54(2), 2013.

Graczyk and J. Małecki. Strong solutions of non-colliding particle systems. 2014.

Gunasekar, B. E. Woodworth, S. Bhojanapalli, B. Neyshabur, and N. Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems , 30, 2017.

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.

R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy. Nature , 585(7825):357-362, Sept. 2020. doi: 10.1038/s41586-020-2649-2. URL https://doi.org/10. 1038/s41586-020-2649-2 .

He, T. Liu, and D. Tao. Control batch size and learning rate to generalize well: Theoretical and empirical evidence. In Advances in Neural Information Processing Systems , volume 32, 2019.

Hochreiter and J. Schmidhuber. Flat minima. Neural Comput. , 9(1):1-42, jan 1997.

Hoffer, I. Hubara, and D. Soudry. Train longer, generalize better: Closing the generalization gap in large batch training of neural networks. In Proceedings of the 31st International Conference on Neural Information Processing Systems , NIPS'17, page 1729-1739, 2017.

Ikeda and S. Watanabe. Stochastic differential equations and diffusion processes , volume 24 of North-Holland Mathematical Library . North-Holland Publishing Co., Amsterdam, 1981. ISBN 0-444-86172-6.

Jastrzebski, Z. Kenton, D. Arpit, N. Ballas, A. Fischer, A. Storkey, and Y. Bengio. Three factors influencing minima in SGD. In International Conference on Learning Representations , 2018.

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.

Khasminskii. Stochastic Stability of Differential Equations , volume 66 of Stochastic Modelling and Applied Probability . Springer, Heidelberg, second edition, 2012.

Kleinberg, Y. Li, and Y. Yuan. An alternative view: When does SGD escape local minima? In Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 2698-2707. PMLR, 10-15 Jul 2018.

Lewkowycz, Y. Bahri, E. Dyer, J. Sohl-Dickstein, and G. Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218 , 2020.

Li, C. Tai, and W. E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research , 20(40):1-47, 2019a.

Li, C. Tai, and W. E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research , 20(40):1-47, 2019b. URL http://jmlr.org/papers/v20/17-526.html .

Li, T. Wang, and S. Arora. What happens after sgd reaches zero loss?-a mathematical framework. In International Conference on Learning Representations , 2021.

Mandt, M. D. Hoffman, and D. M. Blei. A variational analysis of stochastic gradient algorithms. In Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48 , ICML'16, page 354-363, 2016.

Marcotte, R. Gribonval, and G. Peyré. Abide by the law and follow the flow: conservation laws for gradient flows. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/forum?id=kMueEV8Eyy .

Mayerhofer, O. Pfaffel, and R. Stelzer. On strong solutions for positive definite jump diffusions. Stochastic processes and their applications , 121(9):2072-2086, 2011.

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.

R. Norris, L. C. G. Rogers, and D. Williams. Brownian motions of ellipsoids. Transactions of the American Mathematical Society , 294(2):757-765, 1986. ISSN 00029947. URL http: //www.jstor.org/stable/2000214 .

Papazov, S. Pesme, and N. Flammarion. Leveraging continuous time to understand momentum when training diagonal linear networks. In International Conference on Artificial Intelligence and Statistics , pages 3556-3564. PMLR, 2024.

Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32 , pages 8024-8035. Curran Associates, Inc., 2019. URL http://papers.neurips.cc/paper/ 9015-pytorch-an-imperative-style-high-performance-deep-learning-library. pdf .

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.

Soudry, E. Hoffer, M. S. Nacson, S. Gunasekar, and N. Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research , 19(70):1-57, 2018.

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.

Townsend. Differentiating the singular value decomposition. Technical report, Technical Report 2016, https://j-towns. github. io/papers/svd-derivative . . . , 2016.

Van Rossum and F. L. Drake. Python 3 Reference Manual . CreateSpace, Scotts Valley, CA, 2009. ISBN 1441412697.

V. Varre, M.-L. Vladarean, L. Pillaud-Vivien, and N. Flammarion. On the spectral bias of two-layer linear networks. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/forum?id=FFdrXkm3Cz .

V. Varre, M.-L. Vladarean, L. Pillaud-Vivien, and N. Flammarion. On the spectral bias of two-layer linear networks. Advances in Neural Information Processing Systems , 36, 2024.

Wang and A. Jacot. Implicit bias of sgd in l 2 -regularized linear dnns: One-way jumps from high to low rank, 2023.

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.

Wu, C. Ma, and W. E. How SGD selects the global minima in over-parameterized learning: A dynamical stability perspective. In Advances in Neural Information Processing Systems , volume 31, 2018.

Ziyin, H. Li, and M. Ueda. Law of balance and stationary distribution of stochastic gradient descent. arXiv preprint arXiv:2308.06671 , 2023.

A Notations

Notation S d , S + d , S ++ d denote the set of symmetric, positive semi-definite and positive definite matrices in R d × d . We use ⊙ to denote the Hadamard product.

B Proofs

Theorem B.1. For the gradient flow defined in Equation (3.4) , the following property holds,

First, we present a proof of this theorem, based on straightforward computations of the derivative of the determinant and the fact that the matrix J has zero trace.

Proof. Let M = Θ ⊤ Θ . The dynamics of M are governed by the ODE,

Using the gradient of the determinant given in Proposition B.2, the determinant of M evolves as follows,

Given that Tr( J ) = 0 , it follows that d (det ( M )) = 0 .

Proposition B.2. For any matrix M in S ++ d , the first two derivatives of the determinant of M, denoted by det ( M ) are the following

(i) ∇ det ( M ) = det ( M ) M -1

(ii) For 1 ≤ a, b, k, l ≤ d , the second order partial derivative is given by

Theorem B.3. For a stochastic process given by the SDE,

with Tr J = Tr ξ = 0 , the determinant of the M = Θ ⊤ Θ evolves as

Proof. First, we compute the evolution of M = Θ ⊤ Θ using the Ito's product rule,

The last term is interpreted as a derivative of the finite variation and it should be computed using d t. ( d B t ) ij = 0 and ( d B t ) ij . ( d B t ) kl = δ i = k ∧ j = l d t . Using Eq. (3.6),

Using the Ito chain rule, we can compute the evolution of determinant as following,

The first term is

Using the property that Tr( J ) = Tr ( d ξ ) = 0 . We get that ⟨∇ det ( M ) , d M ⟩ = 〈 M -1 , d ξ M ξ 〉 . For the second term

Rearranging the terms in the summation, we get,

Similarly for the other term, we get,

Note that the diffusion part of M -1 d M is d ξ + M -1 d ξ M . Using this

as Tr d ξ = 0 . For the other term,

Putting everything together, we get,

which gives us

Lemma B.4. When l = p + k and η 2 ∥ ∥ J t ∥ ∥ 2 F ≤ 1 , the following property holds for the determinant,

Proof. Note that because of the block structure of the matrix J t , its nonzero eigenvalues come in ± -pairs: ± σ 1 , . . . , ± σ m , moreover, since J t is symmetric, singular values of J t are the absolute values of eigenvalues, i.e. σ 1 , . . . , σ m . Then, the determinant of Θ t +1 can be written as the following,

Using that 1 -x 2 ≤ e -x 2 for all x , we can estimate

We obtain the required inequality by observing that m ∏ i =1 (1 -η 2 σ 2 i ) = ∣ ∣ ∣ ∣ m ∏ i =1 (1 -η 2 σ 2 i ) ∣ ∣ ∣ ∣ since each term 1 -η 2 σ 2 i ≥ 0 when η 2 ∥ ∥ J t ∥ ∥ 2 F < 1 .

Theorem B.5. Let s 1 > . . . s l be the order of the eigenvalues of the matrix M defined by Equation (5.4) . Let the collision time for the eigenvalues be defined as

̸

For t ≤ τ , the eigenvalues are semi-martingales given by the solution of the following SDE

̸

where Ä d ˜ X ä i = 1 / ηδ Ä ⟨ u i , y ⟩ -√ s i c 2 i ä d t + d ε i with u i being the i th column of U and ( ε 0 , . . . , ε l -1 ) is the standard Brownian motion in R l . The evolution of c i and U are presented in the appendix.

Proof. The proof follows the approach of Bru [1989]. Let W = UΣV ⊤ be the singularvalue decomposition (see Def.D.1 involved with r = l and l < p and it will be the rank). Our focus is on understanding the evolution of the singular values and singular vectors of the matrix W . To derive the evolution of Σ , V we can consider the eigenvalues and eigenvectors of the PSD matrix process M . Note that M = VΣ 2 V ⊤ , let D = Σ 2 .

Evolution of D and V Taking the derivative of M , we find

We invoke the theorem D.2 we derived to give the eigenvalues of any matrix valued stochastic process. Note that VV ⊤ = I l , so some terms of the computation are not required.

and the evolution of the eigenvectors,

where you define,

where the matrix S is given by

Note that Σ = diag (( σ 0 , . . . , σ l -1 )) where σ 0 > σ 1 . . . > σ l -1 . Let D = Σ 2 and denote the entires of D as following, D = diag (( s 0 , . . . , s p -1 )) . Note that

Using Levy's characterization U ⊤ d B t is a Brownian motion in R l , lets call that d ˜ B t . The diffusion part of d ˜ M (say d F )

where d m t def = ( σ ⊙ d ˜ B t ) . We are required to compute d F ( S ⊙ d F ) to compute the evolution of eigenvalues. Using the lemma D.4, we get

The element wise computation of this term gives the required result for evolution of eigenvalues. Evolution of c . Note that c = V ⊤ a . Computing the derivative using the Ito's product rule, we get,

Using the lemma D.6, D.5, D.4 and computing the element wise summation, we get the following evolution for d c

̸

Evolution of U . To compute the evolution of U , we invoke the theorem D.2 on the evolution of WW ⊤ = UDU ⊤ . We ignore it here as it does not have much consequence on our results.

Theorem B.6. In the large noise limit, when l = 2 , the following properties hold, for t ≤ τ ,

(a) s 0 , s 1 are greater than zero almost surely.

(b) for α = ( p -3) / 2 , s -α 0 is a super-martingale while s -α 1 is a sub-martingale.

Proof. First, note that in the large noise limit with l = 2 , the evolution of the eigenvalues is expressed as

Using the Ito chain rule, for the evolution of s -α 0 we can write

analogously

and finally for s -α 0 s -α 1

Now, we can show that the drift term in the SDE that describes the dynamics of s -α 0 s -α 1 is zero, which gives us the first part of the result by Mckean's argument [Mayerhofer et al., 2011],

The second part is obtained by noticing that

and hence the drift term of d ( s -α 0 ) is not positive, while the drift term of d ( s -α 1 ) is not negative.

C Experiment details

In all the graphs we plot the values averaged on 20 runs with different random seeds as well as the 95% confidence interval (lightly colored). To numerically emulate GF (Figure 1), we set a stepsize of 1 e -6 in numerical simulation.

In the further experiments, we study the behaviour of the linear network for regression with the same synthetic data and same network initialization as in previous experiment. As seen in the left plot of the Figure 2, when the stepsize is large ( η = 0 . 1 ), singular values exhibit behavior similar to the case of LNGF, while with the small stepsize ( η = 0 . 005 ) the evolution of singular values is closer to GF case. Next, we examine the effect of SGD in the case of classification task with logistic loss, as illustrated in the middle plot of the Figure 2. We consider synthetic data with n = 1000 samples of Gaussian data in R 5 ( d = 5 ) constituting two clusters corresponding to two classes ( k = 1 ). Note that larger stepsize ( η = 0 . 5 ) in this case also forces the smallest singular value to tend to zero, however the effect is not so dramatic for the rest of singular values. Additionally, we study the 2-layer ReLu network optimized with SGD on the same regression task as before. As seen in the right plot of the Figure 2, the decrease of the last singular value σ 4 is much slower than in the case of the linear network, however, the larger stepsize still facilitates divergence of k largest ( σ 0 and σ 1 ) and p -k smallest ( σ 2 , σ 3 and σ 4 ) singular values.

All experiments are implemented with Python 3 [Van Rossum and Drake, 2009] under PSF license, NumPy [Harris et al., 2020] under BSD license, and PyTorch [Paszke et al., 2019] under BSD-3Clause license.

The experiments were run on a Intel i5-8250U, 8-GB RAM, with OS Ubuntu 20.04.6.

D Supplementary material

D.1 Notations and preliminary definitions

Definition D.1 (Eigen decomposition and Singular Value decomposition) . We discuss the eigen value decomposition for a symmetric square matrix, and the singular value decompostion for any matrix is defined as the following

(a) Eigen decomposition . For any rank r matrix R ∈ S p , R = VDV ⊤ is the eigen decomposition, where V ∈ R p × r , D ∈ R r × r , D is a diagonal matrix and V ⊤ V = I r , however, VV ⊤ is not necessarily an identity matrix unless r = p .

(b) Singular Value Decomposition . For any rank r matrix W ∈ R p × l , W = UΣV ⊤ , where U ∈ R p × r , V ∈ R l × r , Σ ∈ R r × r , Σ is a diagonal matrix and U ⊤ U = V ⊤ V = I r , however the UU ⊤ and VV ⊤ are not necessarily identity unless r = p or r = l respectively.

D.2 Eigenvalues of matrix valued stochastic process

Theorem D.2. For a matrix-valued stochastic process on S ++ p + k ,

where d M is a local martingale process. Let R = VDV ⊤ is the eigenvalue decomposition of the process, the evolution of eigenvalues satisfy the SDE for time t less than the collision time,

where S is defined as per Eq. D.1 and d ˜ M = V ⊤ d MV , ˜ N = V ⊤ NV . The evolution of the eigenvectors,

where you define,

Evolution of eigenvalues for general matrix SDE

Proof. Using the eigen decomposition, we have R = VDV ⊤ ,

The approach we follow is use the jacobian of the evolution of V (see [Townsend, 2016] ) and solve the constrains equations to obtain the Ito correction term as done in Bru [1989]. Let ( s 1 , s 2 , . . . , s r ) denote the diagonal entries of D . Furthermore, we define the matrix S , which plays a notable role in Jacobian w.r.t V , as the following,

For the sake of brevity, we denote the evolution

The evolution of the eigenvectors,

Using the Jacobian of the eigen vectors, we write,

Note that V ⊤ V = I , using this we have,

Coming back to the evolution of singular values,

Note that d D is diagonal, hence, I ⊙ d D = d D .

Note that I ⊙ ( DM ) = I ⊙ ( MD ) = D ⊙ M for any matrix M and diagonal matrix D , using this property, we can simplify the above expression as,

Using Eq. D.2, we have,

Using this,

Evolution of eigenvectors for general matrix SDE. Here, we derive the evolution of eigenvectors, Using Eq. D.2, we have,

Combing these, we get the diagonal and off diagonal terms of Q ∥

Computing of Q ⊥ . Recalling the evolution of the eigenvectors,

Using the Jacobian of the eigen vectors, we write,

Using the fact that ( I -VV ⊤ ) R = 0

Adding the terms up we get,

This gives the expression for Q ⊥ and this ends our computation.

Lemma D.3. For any matrix A ∈ R n × m , B ∈ R n × n , m × n -dimensional Brownian motion d B t , the following results hold on the covariance

Lemma D.4. With S defined in Equation (D.1) , d F = d F = ΣV ⊤ d B t c ⊤ + c d B t ⊤ VΣ and d m t def = ( σ ⊙ d ˜ B t ) .

Proof.

Lemma D.5. With S defined in Equation (D.1) , d F = d F = ΣV ⊤ d B t c ⊤ + c d B t ⊤ VΣ and d m t def = ( σ ⊙ d ˜ B t ) .

Proof.

Now, computing the product,

Lemma D.6.

Proof.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The dichotomy between SGD and GD is revealed by theorems 4.1 and 4.2, the repulsive force between the eigenvalues of parameter matrix is discussed in theorems 5.1 and 5.2. Supporting experiments are discussed in section 7.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations of the work are discussed in section 4 in a designated paragraph as well as in section 5 in the discussion of theorem 5.1.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Proofs of the theorems that are not presented in the main body are presented in appendix B.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The data as well as the models and optimization algorithms used are discussed in the section 7

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code in the form of Jupyter notebook is provided and all the random seeds are fixed for the reproducibility purposes.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The setup is discussed in the section 7.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: All the graphs of the parameters evolution are accompanied with the 95% confidence interval calculated on the 20 runs with different random seeds.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All relevant information is stated in the appendix section C.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research is theoretical and does not suggest any new model that can cause harm. All the data used is synthetic.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: The work investigates the effects of well known algorithms on the simple models and doesn't suggest any new applications.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The work doesn't entail models or datasets releases.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All relevant information is stated in the appendix section C.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an interesting dichotomy between SGD vs. GD, and how stochasticity in the gradients have the ability to implicitly regularize towards low-rank solutions. I believe this point is best highlighted by Theorem 5.1, where the show that the highlighted term in Equation (5.6) has a repulsive force, increasing the gap between larger and smaller eigenvalues. The \"punchline\" of the paper is technically given in Theorems 4.1 and 4.2, where they show that the solutions given by GD cannot have its rank diminished (rather stays at initialization), whereas we can see a decay in SGD (or SGF rather).\n\nSTRENGTHS:\n- I believe this paper fits well within the study of implicit regularization properties of certain optimizers and how they drive generalization. To this end, it also includes a related works section.\n- The results are interesting, albeit a little limiting.\n- I think the result in Theorem 4.2 also highlights how large step sizes can be beneficial for generalization. This is quite a popular topic in the literature. It is easy to see that when $\\eta$ increases, the determinant vanishes faster.\n\nWEAKNESSES:\n- I think the presentation of the paper can be slightly improved. For example, the discussion on initialization on Page 4 was a bit confusing. The authors discuss different initialization tactics and the balancing effect, but they do not clarify which initialization they ultimately chose. I believe the intention was to show how initialization plays a role in Theorems 4.1 and 4.2, but this could be made clearer if that was indeed the intention. Furthermore, the authors could also change the notation to $\\Delta(t)$ between lines 149-150. The authors are also missing a > in line 228.\n- The experiments section is a bit lacking. Are there ways we can leverage the results from this paper to more practical scenarios / networks?\n- My main concern (which is also the weakness) is the similarity between this paper and [1], which I don't believe the authors cited. From what I read, the main takeaway in [1] is that the noise in SGD has the ability to \"push\" or \"attract\" networks towards simpler networks, which they describe to be invariant sets. If you look at Sections 5 and 6 in [1], I believe what they are trying to say is that the noise has the ability to collapse the network to an invariant set that is low-rank, and one in which some singular modes are actually never learned. At a high-level, it seems to me that this might be a stronger result than to say that there exists a repulsive force between the eigenvalues, which is the current result. If I am missing something, please feel free to correct me. Though, at the very least, I think the authors should cite this paper and delineate the similarities / differences.\n\n[1] Feng Chen, Daniel Kunin, Atsushi Yamamura, Surya Ganguli. \"Stochastic Collapse: How Gradient Noise Attracts SGD Dynamics Towards Simpler Subnetworks\". NeurIPS 2023.\n\nQUESTIONS:\nSee Weaknesses. I would be more than happy to raise my score if my concerns can be addressed.\n\nLIMITATIONS:\nI have stated the limitations throughout my review.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

comments: SUMMARY:
This paper studies and analyzes how the rank of the parameter matrices evolves for two-layer linear networks when using GD and label noise SGD. The paper basically shows that while GD preserves the rank at initialization throughout its trajectory, SGD reduces the rank, thereby removing spurious directions. A stochastic differential equation (SDE) is derived to characterize the evolution of the eigenvalues of the parameter matrix to formalize the rank-reducing behavior of SGD. Small-scale experiments validate the derived theory.

STRENGTHS:
The paper has a nice insight backed up by solid theory. I'm not an expert in this area but it seems that this is a novel insight.

WEAKNESSES:
**1.** While the rank-reducing property of SGD is interesting, it does not imply that the prediction error (which is the generalization metric most people would care about) of SGD is lower than that of GD. What if SGD converges to a completely erroneous low-rank solution? I understand that the paper is not directly about the generalization abilities of SGD and GD, but I feel this point should be discussed somewhere. 

**2.** Since the results here are only for linear networks (which is completely fine), is there any intuition for why the results and claims of this paper should extend to non-linear networks?

QUESTIONS:
Please see the Weaknesses. Also, I have the following questions:

**1.** In line 174, is $M = \Theta^\top \Theta$?

**2.** In line 207 and the equation thereafter, if $W = U \Sigma V^\top$, then $W^\top a$ should be $V \Sigma U^\top a$.

**3.** I don't understand the jump from eq. (5.2) to (5.3). There seems to be an issue with the $\sqrt{\eta \delta}$ term (with the definition of $d X$).

LIMITATIONS:
Discussed somewhat throughout the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the implicit bias of SGD for two-layer linear networks. The authors study this primarily using (stochastic) gradient flow (S)GF, the continuous time version of (S)GD. To model the stochasticity, they approximate the SGD noise with independent label noise. Under these assumptions, they prove that the determinant of the Gram matrix of the parameters is invariant under GF, whereas it exponentially decays under SGF. Hence, if such a network is initialized with a full-rank parameter matrix, the rank asymptotically drops by at least one under label noise SGF, which lends some evidence to a low-rank bias for SGD not present in GD. The stochastic eigenvalue dynamics are investigated further, and while the exact dynamics are not rigorously pinned down, the paper calls out the repulsive forces that could explain why SGD converges to lower rank solutions. These results are empirically verified using some synthetic data experiments for regression and classification.\n\nSTRENGTHS:\n1. Laying theoretical foundations for the benefit of SGD over GD for optimizing neural networks is an important direction for the theory of deep learning. This paper takes a step towards identifying the implicit bias of SGD.\n2. Previous work has that balancedness of adjacent layers is a conserved quantity for deep linear neural networks optimized using GF. Invariants of the dynamics are an important analytic tool; balancedness is crucial for proving convergence for deep linear networks. This paper identifies a very clean invariant for GF which separates GF from (label noise) SGF: the determinant of the Gram matrix of a certain concatenation of the linear layers ($\\Theta$ in the paper). In particular, although the determinant is conserved for GF, for SGF the determinant decays exponentially (exactly!). This is a very nice contribution to the theory of SGD. \n3. By itself, the determinant decay only establishes that the rank drops by one in SGF. However, the authors also investigate the behavior of the singular values of the first layer $W$ in a simpler scalar regression setting by deriving the SDE for the eigenvalues of $M = W^\\top W$. Although they are not able to characterize the solution to this SDE, they clearly interpret the forces at play and highlight a repulsive effect between the eigenvalues of $M$, which encourage the singular values to decay towards zero.\n4. The authors also consider the large noise limit and cleverly reduce to the 2-dimensional setting, where having a zero singular value is equivalent to the two rows being multiples of each other. They prove that, in expectation, (a power of) the larger singular value does not decrease, whereas the (same power of the) smaller singular value does not increase.\n5. Section 6 contains many interesting extensions to other settings, such as classification, directly modeling the SGD noise as Gaussian with matching covariance, and discrete time methods. \n6. The paper is quite transparent with the limitations of the theory. It also clearly highlights where the future directions are (e.g., the full characterization of the solution for the SDE governing the evolution of eigenvalues).\n\nWEAKNESSES:\nIt was nice to see some experiments validating the theory, but I felt that the dimensionality of the problems should be increased to be more convincing. Right now the experiments are done with $(p, l, k) = (5, 10, 2)$. Even increasing these to being in the hundreds would be interesting, and correspondingly it might be appropriate to use other metrics to measure the rank of $W_1$, such as the effective rank  $\\tr(M)/||M||$.\n\nQUESTIONS:\n1. In line 95, should $d$ and $p$ be swapped? Since the input is $p$-dimensional, whereas the parameter $\\theta$ is $d$-dimensional (for example, the two-layer setting $d = p \\cdot k$).\n2. In line 96, there should be an $\\ell_2$ norm outside of $y_i - f_\\theta(x_i)$. \n3. In line 108, $k=1$ should be in math mode. \n4. There is a minor typo on the equation after line 208: it should read $Wa$ rather than $W^\\top a$.\n5. In the statement of Theorem 5.1 on Line 228, there is a missing comma in the definition of the ordered eigenvalues. \n6. Could the authors add a reference to the proofs in the appendix when the result is stated in the main text (e.g. Lemma 6.1)?\n7. In Section 7, I think I am missing something obvious: why are the $d + \\ell - k$ smallest singular values equal $\\sigma_2, \\sigma_3, \\sigma_4$, since to my understanding $d + \\ell - k = 10 + 10 - 2 = 18$?\n8. What can one hope to say about deeper linear networks? If I understand correctly, if one naively generalizes the construction of $\\Theta$ then the special block structure of the label noise portion will no longer hold (in particular, it seems that $Theta$ would need to contain all the prefix and suffix products of the linear layers)\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an analysis of the gradient flows of two layer linear neural networks with a squared loss borrowing tools from differential equations. The paper's result establishes that the stochastic gradient method generates solutions (limit of the flow as time goes to infinity) with determinant of the parameters x parameters transpose tending to zero. This implies that the parameters of a model optimized with stochastic gradient converge to a low-rank (simple) solution. While the full gradient method's solutions for the same problem do not exhibit the same simple structure. Authors conclude that this finding gives an explanation to the implicit bias of SGD.\n\nSTRENGTHS:\nThe paper is tackling a very interesting problem and in a relatively simple yet representative setup where learnings can be used for deriving broader intuitions on the reasons behind SGD's implicit bias. The paper's methodology looks solid and the conclusions in the proposed context are well presented.\n\nWEAKNESSES:\nThe paper studies a very simple model. Numerical experiments are also on the same and in very small data. It would benefit the paper to run experiments on larger / more complex models to discuss the limits and validity of this theory. The generalization to other settings section 6 is only changing the loss function. The section's title is perhaps overpromising.\n\nQUESTIONS:\nIs it possible to verify the breadth of validity of the result (parameters of SGD based solution live in a low-rank manifold, not the FGD) on different neural network architectures and see whether similar conclusions hold? Or which other property (if low rank is too restrictive) would you test for?\n\nDoes this theory imply that for the simpler vector problem the solution of SGD is a sparse vector where FGD finds dense solutions?\n\n\nIf we add a ReLU non-linearity between W1 and W2, then in what form the simplicity of the solution by SGD will get modified? locally low-rank?\n\nLIMITATIONS:\nThe papers does not discuss how broadly their findings are valid. Numerical experiments could be used for testing whether findings are still valid in slightly more complex models", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}