SGD vs GD: Rank Deficiency in Linear Networks

We thank the reviewer for their valuable comments, and for pointing out the missing reference.

The result in Theorem 4.2 also highlights how large step sizes can be beneficial for generalization.

Yes, the presence of $\eta$ does control the rate of decrease of the determinant. Hence, the larger the step size, the greater the regularization effect, which can be beneficial for generalization. Note however that due to our choice of continuous-time modeling of gradient methods, as the step size increases, the discretization error makes it difficult to accurately capture the trajectories of the discrete methods with their continuous-time counterparts, leading to potentially different behavior. In the revision, we will add a statement indicating that our model captures the role of step size up to a certain order.

Initialization, balancedness and presentation.

Thank you for the suggestion and the comments on the presentation of the initialization and its balancing effects. We address these points here and will revise the manuscript to incorporate these comments, making the presentation clearer.

We mentioned initialization in the problem setup because it is an important ingredient in studying the implicit bias of homogeneous networks, including linear and ReLU networks. Previous works offer various initialization strategies for gradient flow, where certain structures emerge due to the balancedness properties of gradient flow. In contrast, with stochastic methods, the initialization does not matter. The emergence of low-rank structures is observed despite general initialization. This behavior is exemplified in Theorem 4.1 and 4.2, which hold for any initialization. In addition, we discuss these different initialization strategies to highlight the challenges for stochastic gradients, as the balanced quantities studied for gradient flow are insufficient.

Limited experimentation, more practical networks.

We kindly refer the reviewer to the general comment regarding this limitation.

Detailed comparison with [1]

Thank you for pointing out this work. We were aware of it but unfortunately missed including the reference and detailed comparison in the manuscript. We will revise it accordingly.

[1] investigates a phenomenon similar to ours, referred to as stochastic collapse, where the noise in the gradients causes the iterates to collapse onto certain invariant sets. [1] also uses continuous-time modelling of SGD as an SDE. The differences between the works are outlined as follows:

• In the first part of their work, [1] provides a condition under which an invariant set is attractive for the SDE they consider. Methodologically, their work characterizes the local behavior around the invariant sets. Specifically, once the iterates enter a small $\epsilon$-neighborhood of these sets, they are attracted towards the sets. However, their does not address whether the iterates ever enter such an $\epsilon$-neighborhood for a general initialization.

  In contrast, we offer a global guarantee, albeit within a simplified model: low-rank behavior is observed for any initialization. To establish these global properties, we developed a novel approach by tracking a specific quantity (the determinant). This approach enables us to characterize the global behavior and reveal the dichotomy between the presence and absence of noise.

• The paper also studies linear networks in a teacher-student setup. Under a set of assumptions A1-A4~[1, p.30], they derive the evolution of singular values. However, due to the balanced spectral initialization (A3-A4) and structured label-noise (A2), the analysis falls short of capturing the repulsive force in the singular values. On a technical front, deriving the SDE for the singular values is simpler when the singular vectors remain stationary, as their assumptions ensured. When the singular vectors do not move, all the matrices can be simultaneously diagonalized, which simplifies deriving the singular values' dynamics. In our problem, however, the singular vectors move and follow a stochastic process, which makes tracking them significantly more challenging.

[1] Feng Chen, Daniel Kunin, Atsushi Yamamura, Surya Ganguli. "Stochastic Collapse: How Gradient Noise Attracts SGD Dynamics Towards Simpler Subnetworks". NeurIPS 2023.

We would like to thank the reviewer for their insightful suggestions.

Low-rank solutions and prediction error.

We agree with the reviewer's comment that the implicit regularization might only be useful for generalization if it is aligned with the ground truth model (i.e., in the classical example of sparse regression, $\ell_1$-regularization only improves generalization of ERM if the ground truth regressor is sparse). We also agree with the reviewer's comment regarding convergence to a sub-optimal low-rank solution which can happen due to possible over-regularization. We will discuss these issues in the conclusion of our manuscript.

Intuition for non-linear network.

Empirically, the low-rank phenomenon is prominent for non-linear networks when trained with SGD, see Figure 2,3 in the attached pdf. For piecewise linear non-linearity like ReLU, each neuron's weights exhibit a multiplicative structure. Let $h (a,w, .) = \sum_{j} a_j \sigma(\langle{w_j}{.}\rangle)$ be the parametric model where $\sigma$ is the ReLU non-linearity and $(w_i,a_i)$ are the input and output weights of a neuron. With $\theta_i = \begin{bmatrix} w_i^{\top} & a_i \end{bmatrix}$, the dynamics of the neuron can be rewritten as $\mathrm{d}{\theta_i} = \theta_i \mathrm{d}{J_i}$, for some curated matrix stochastic process $J_i$. For linear networks, these matrix-valued processes are identical across neurons, which makes the analysis tractable. In contrast, non-linear networks do not share this property. However, during the large noise phase, the stochastic matrices $J_i$ are dominated by noise even for the ReLU dynamics, which can result in some similarities between these matrices. This may vaguely explain the observed low-rank behavior.

Is $M = \theta^{\top}\theta$ in line 174 ? The clarification of $W^{\top}a$ in line 207 ?

Yes, thank you for pointing out this missing definition. We will correct it. In l.207, it should be $Wa$ instead of $W^{\top}a$.

The jump from eq. (5.2) to (5.3).

It is possible due to changing the time, let $t' = (\eta \delta) t$. Now we have that $\mathrm{d}{t} = \frac{1}{\eta \delta} \mathrm{d}{t'}.$ The Brownian motion has the following property $\mathrm{d}{\mathbf{B}_{t'}} = \sqrt{\eta \delta} ~ \mathrm{d} {\mathbf{B}}_t$ (the square root is due to different scaling of Brownian motion while changing time). Substituting these gives us the required jump.

We thank the reviewer for their valuable feedback, encouraging remarks, and meticulous proofreading of our work. They will help to make our work clearer.

It 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 such as the effective rank.

We refer the reviewer to the general comment for the experiments with higher dimensions.

In Section 7, I think I am missing something obvious: why are the smallest singular values equal $\sigma_2, \sigma_3, \sigma_4$. since to my understanding $d + l - k = 18$ ?

Apologies for the typo. As we are measuring the singular values of $W_1$, it should be the $p - k$ smallest singular values instead of the $p + l -k$ ones . Here $p=5$, $l = 10 (>p)$ and $k=2$, hence the $p - k = 3$ smallest singular values are $\sigma_2,\sigma_3,\sigma_{4}$.

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)

Thank you for this question. Indeed, deeper layers cannot be directly written in this multiplicative format. However, if the network has $l$ layers $W_1,\ldots,W_l$, for any adjacent layers $W_{i}$ and $W_{i+1}$, we can form a block matrix, $\theta_{i} = \begin{bmatrix} W_{i}^{\top} & W_{i+1} \end{bmatrix}$ and the multiplicative structure can be seen in the evolution of $\theta_i$. The analysis can then be performed with these block structures. For the stochastic case, however, the noise covariance is not straightforward and needs to be handled carefully. This approach might present a preliminary way forward for deeper networks. Another approach could be to carefully formulate a tensor structure that extends the spirit of our work beyond two layers.

We thank the reviewer for their interesting questions and encouraging comments.

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.

Please refer to the general comment.

Is 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?

The phenomenon that the parameters of solutions found by SGD live in a low-rank manifold is already empirically observed in [1] (note that the noise is more prominent when the step size is large, hence the effects of SGD are prominent with large step size). Given that measuring the rank is too computationally expensive, metrics that measure the similarity between the activations or weights of the neurons would serve as a good alternative. Such metrics are used by [1] and [2].

[1] M. Andruischenko, AV Varre, L. Pillaud-Vivien, N. Flammarion. SGD with large step size learns sparse features, ICML 2023.

[2] Feng Chen, Daniel Kunin, Atsushi Yamamura, Surya Ganguli. "Stochastic Collapse: How Gradient Noise Attracts SGD Dynamics Towards Simpler Subnetworks". NeurIPS 2023.

Does this theory imply that for the simpler vector problem the solution of SGD is a sparse vector where FGD finds dense solutions?

No, the theory does not directly imply it. Consider the vector problem of diagonal networks where $u \odot v$ is the re-parameterization considered. Let $U = \textrm{diag}(u)$, and $V = \textrm{diag}(v)$ be diagonal matrices with diagonal entries $u$ and $v$, respectively. The block structure can be recovered with $\theta = [ U ~ ~ V ]$. However, the dimensions of this block matrix are $p \times 2p$, hence our theory does not have any implications as we need $l > d+k$. However, if a complete characterization can be obtained, then indeed a low rank of $\theta$ would imply a sparse solution. The sparse vector problem with label noise has received much attention [3,4] and the complete characterization in the symmetric case (u=v) is given by [5].

[3]J. Z. HaoChen, C. Wei, J. Lee, and T. Ma. Shape matters: Understanding the implicit bias of the noise covariance. COLT, 2021.

[4] Z. Li, T. Wang, and S. Arora. What happens after sgd reaches zero loss? –a mathematical framework. ICLR 2022.

[5] Pillaud-Vivien, L., Reygner, J., and Flammarion, N. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. COLT, 2022.

If 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?

Please refer to the general comment, particularly the paragraph on ReLU networks, as the experiments indicate the simplicity bias that SGD tends to induce.