Implicit bias of SGD in $L_2$-regularized linear DNNs: One-way jumps from high to low rank

Thank you for the thoughtful review.

First note that we actually do not need $\lambda$ to be large for our theorem, only $C$ needs to be large. We apologize for the confusion, we did not realize that our statement was confusing, and we have replaced it with "for $r\geq0$, any $\lambda\geq 0$, any large enough $C$ and small enough $\alpha,\epsilon_1,\epsilon_2,\eta$" which can only be understood in the correct way.

This should resolve your concern that the effect we observe is only a result of a very large regularization, since our result applies to arbitrarily small ridge $\lambda$.

Note also that an upper bound on $\eta$ is unavoidable, since large $\eta$ could lead to divergence. Now it seems likely that the bound we require on $\eta$ is not tight, especially w.r.t. its dependence on $\lambda$, and this could be improved in future work.

From our theorem, one can identify settings (for $\lambda,\eta$ small enough) where GD can get stuck at a rank-overestimating minima, while SGD will be able to escape it given enough time, showing a fundamental difference between the two regimes.

Regarding your question. Yes, the sets $B_r$ are absorbing for GD too, only the jumps are specific to SGD. But both are needed to guarantee a gradual decrease of the rank.

Thank you for the thorough and detailed review.

Your remarks were really valuable to improve the readability of our proofs. We have made many changes that should address your issues and we are still working on improving readability.

We have replaced $U_i,V_i$ by $V_i,W_i$ to avoid the conflict in notation with the $U_\ell$s.

We have removed the reliance of the proofs on a continuous path of minima as $\lambda \searrow 0$. As you said, some work needs to be done to make such statement rigorous and it turns out to not be necessary.

The "eigenvalue" was a typo, it should have been "singular value", and we have changed this part of the proof in the newer version anyway, since we do not rely on a path of minimizers.

You were right that a transpose was missing and we have added a definition of $U_{\ell,i}$.

The "+=" was a typo, we replaced it with "=".

The saddle-to-saddle regime refers to the two papers cited at the end of the sentence, we have moved the citations to earlier in the sentence to clarify this.

We have added a short explanation for why a matrix cannot be approached by matrices of lower rank (the error will always be at least the $r+1$ singular value of $A$).

We have removed the reliance on a path of minima, instead we just take $A(\lambda)$ to be a global minimum of $C_\lambda$ amongst matrices of rank $r^*$ or less.

Thanks for the reference you provided for the differentiability along directions that do not change the rank, we will add it as a justification.

The convergence is not uniform in $t$, but we only need sequential convergence: for all fixed $t$, we have convergence as $\lambda \searrow 0$. We have made this clearer.

We will add a short explanation: the first $L-1$ derivatives of the unregularized loss $\mathcal{L}$ vanish at $\theta=0$, thus adding a regularization terms $\lambda\Vert\theta\Vert^2$ leads to a local minimum.

We fixed the typos and error you mentioned and added a citation for Fact C.4.

Regarding your questions:

1. In general you get an explosion if $A_{\theta_i}$ approaches $A_{\hat{\theta}}$ from a higher rank direction. Assume that $A_{\theta_i}$ equals to $A_{\hat{\theta}}$ with a small additional singular value: $A_{\theta_i} = A_{\hat{\theta}} + uv^T / i$, then for balanced parameters $\Vert \theta_i - \hat{\theta} \Vert^2 = i^{-\frac{2}{L}}$, leading to an infinite limit.

2. No, we do want the infimum over $\mathrm{Rank}A<r^*$. Since we know it will be strictly larger than $0$, we can always choose $t_0$ large enough and $\lambda$ small enough to be below this threshold. Once we are below this threshold, we know that the rank of $A_{\theta(t)}$ cannot be strictly smaller than $r^*$.

Thank you for the thoughtful review. Regarding your points:

1. The statement "for small enough $x$" is simply a (commonly used) shortcut for "there is a $x_0$ such that for all $x\leq x_0$". They are both rigorous formal statements. Note that exact bounds are given in the appendix, we did not put them in the main to improve readability.

2. Since the loss $\mathcal{L}_\lambda$ is smooth, GD (and SGD) converges to GF as $\eta \searrow 0$. We stated the Proposition for GF to simplify the statement. We preferred to assume an initialization that converges to a global minimum directly for the following reason. There exist many papers that can guarantee convergence to a global minimum (e.g. in the NTK regime, or with specific initialization, we will add references to such results in the main after the statement of this Proposition), and since it seems that we have not yet fully answered the question of convergence of DLNs, we expect more result to appear in the next years. Our Proposition can be applied to any such settings.

3. We feel that our proofs are already quite complex and technical, and considering a general $f_\alpha$ would lead to even more complexity for only a small increase in generality.

4. We have clarified the proofs thanks to the remarks of the reviewers. Please tell us if there is one section in particular that you find difficult to understand, so that we may focus our efforts there.

Regarding your questions:

1. We say that there might not be any rank-overestimating minima (in contrast to the fact that there always exist rank-underestimating minima, e.g. the origin $\theta=0$), in which case GD with a small enough ridge $\lambda$ and learning rate $\eta$ will converge to a minimum with the right rank by Proposition 3.2. Of course we do not know under what condition this happens, and it might be difficult to characterize, a trivial example would be if we observe all entries of the matrix.

2. Sorry, we did not realize that our statement was confusing, we replaced it with "for $r\geq0$, any $\lambda\leq 0$, any large enough $C$ and small enough $\alpha,\epsilon_1,\epsilon_2,\eta$" which can only be understood in the correct way. Namely, we do not need $\lambda$ to be large for our theorem, only $C$ needs to be large.

3. We can guarantee a jump under the event that SGD randomly selects observations only from the same $r$ rows/columns over a large period of time. This could happen for any set of $r$ rows/columns, but we only consider the $r$ rows/columns with the most observed entries since it is most likely to happen to these rows entries. When all rows have the same amount of observed entries, this event is as likely for all subsets of rows/columns (with probability equal to the lower bound given in the statement of the Theorem).

Thanks for the review!

Regarding your question, we would also observe a low-rank bias with Langevin dynamics since minima with a lower rank generally have a smaller parameter norm and thus a smaller regularized loss. It should therefore be more likely to jump from a high loss to a low loss minima than the other way round. Now the relation between parameter norm and rank is not exact and there can be example where the minimal parameter norm solution has a higher rank than another solution with a larger parameter norm (this happens for $\epsilon$ small enough in the setup of our numerical experiments). So the low-rank bias would be in some sense weaker.

But Langevin dynamics do not capture the covariance of the noise of SGD. One can instead study an SDE with matching covariance. But this covariance is not full rank especially inside and around the $B_r$ sets, our intuition is that there is a lot noise along direction that keep or lower the rank, but (almost) no noise along direction that increase the rank. It is still unclear whether a similarcould be proven in this setting (e. .whether the limiting distribution is also supported inside $B_1$).