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

Reviewer comment (Experiments):
The implicit bias observed in this paper could be a result of a deliberate choice of parameters, instead of a natural property of SGD and GD, and this point needs clarification from the authors.

Author response:
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.

Reviewer comment (Theory):
It is unclear whether the proof in this paper for Theorem 3.2 also works for GD; specifically, is $B_r$ also absorbing for GD?

Author response:
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.

Reviewer comment (Theory):
Although the main paper is well written, the proofs are not reader-friendly at all. The writing of the proofs is very terse and laconic, omitting many details. While this style can be reminiscent of certain pure mathematics papers, the reviewer believes it should only be acceptable if there is no room for any errors or inaccuracies. The reviewer found at least a few typos, minor errors, and imprecisions in the subset of proofs examined. Since much information is left out for the reader to deduce, even a small number of actual errors greatly expands the difficulty for the reader. The reviewer recommends a substantial revision of the paper with more detailed and careful proofs, and maintaining their score is conditional on that.

Author response:
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.

Reviewer comment (Writing):
On page 12, point "(0)", the definition of the \( U_i, V_i \) is not consistent: the index under the \( U, V, S \) is used to mean both the iteration step in the sequence and the position in the product \( W^L \ldots W^1 \).

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

Reviewer comment (Theory):
On page 13, the authors state: "as \( \lambda \rightarrow 0\), the critical points of the loss move continuously. Consider a continuous path of critical points, as \( \lambda \rightarrow 0\), it converges to...". While the argument makes intuitive sense, filling in the gaps with rigorous proofs is beyond what can be expected from the reader. At least some citations should be provided. The reviewer doubts that simple continuity is enough to guarantee convergence, as even if a subsequence converges, the path could oscillate widely. Probably the only way to rigorously prove the statement is to use a quantitative version relying on calculus of variations.

Author response:
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.

Reviewer comment (Theory):
In the middle of page 14, the equation \( U_{\ell,i}(\lambda)U_{\ell-1,i}(\lambda) \) seems incorrect without at least a missing transpose. The lack of a rigorous and consistent definition of \( U_{\ell,i} \) exacerbates this issue.

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

Reviewer comment (Writing):
At the bottom of page 14, the line starting with "other directions" ends with " \( L-1,)\)", and a few lines down an equation reads \( U_\ell^\top dU_{\ell}+ = -dU_\ell^\top U_\ell \). The reviewer asks what "+=" means here. This same issue appears in other parts of the paper, including the third line of text on page 15.

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

Reviewer comment (Theory):
In the middle of page 13, the authors use the fact that "a matrix cannot be approached with matrices of strictly lower rank", which is true but should probably be cited, especially since the equivalent statement is not true for tensors.

Author response:
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$).

Reviewer comment (Theory):
The proof of Proposition A4 is difficult to follow without additional information. The first sentence is "'let \( A(\lambda) \) be path of global minima restricted to the set of matrices of rank \( r^* \) or less." It is unclear how to construct this path. Even for \( L=2 \), there can be a continuous set of global minima of local intrinsic dimension higher than 2, so how does one use the axiom of choice to construct a path? Sentences such as "going along directions that increase the rank of \( A(\lambda) \), the regularization term increases at a rate of \( d^{2/L} \) for \( d \) the distance" need more mathematical detail.

Author response:
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.

Reviewer comment (Theory):
The statement about \( \phi \) being differentiable in the directions which do not change the rank should be made more precise. A citation to the work by G.A. Watson (1992) is suggested.

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

Reviewer comment (Theory):
For Proposition A.5, the proof begins with "We know that L2 regularized GF \( \theta_\lambda(t) \) converges to unregularized GF \( \theta(t) \) as \( \lambda \rightarrow 0 \)". There are two parameters here, \( \lambda \) and \( t \). The reviewer asks whether the convergence is uniform over all \( t \).

Author response:
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.

Reviewer comment (Theory):
As a minor point, even the statement that the point 0 is a critical point was not immediately obvious to the reviewer as a relative outsider to optimization literature. Perhaps either a calculation of the gradient or a mention of the fact that \( L>1 \) would help.

Author response:
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.

Reviewer comment (Writing):
At the bottom of page 13, the equation before equation (1) is presumably the end of a sentence, so the next line should be rewritten. Additionally, "no such thing happen" should be "no such thing happens".

Author response:
You are right with the notations of proposition A.2, we have changed them to $\tilde{U}_i,\tilde{V}_i$, everything should be fine now. We also fixed the sentence.

Reviewer comment (Theory):
Question: On the third line of page 13, the authors mention that the quantity in the limit is strictly positive but possibly infinite. The reviewer asks for the reasoning behind this, as it is not obvious to them.

Author response:
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.

Reviewer comment (Theory):
Question: At the beginning of the proof of Proposition A5 in the first equation, should the infimum run over \( \text{Rank} \, A > r \) instead of \( \text{Rank} \, A < r \) as written?

Author response:
Regarding your questions:
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^*$.

Reviewer comment (Theory):
I feel the statement of theorems is not very clear. It uses a lot of ''small enough'', "large enough". I think the statement should be more rigorous.

Author response:
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.

Reviewer comment (Theory):
For the function $f_\alpha$, it takes a very specific form. The authors claim that changing $f_\alpha$ with similar properties should not affect the results. I don't see the reason why the condition cannot be extended to more general functions. I believe it could improve the results.

Author response:
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.

Reviewer comment (Writing):
A small suggestion is that the proof in the appendix is not easy to read. I think it can be more organized and add more explanation.

Author response:
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.

Reviewer comment (Theory):
In the remark before section 3.2, it's said that it's possible that GD can recover the minimal rank solution easily. Can you say something about this case?

Author response:
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.

Reviewer comment (Theory):
I have a concern about the constant in Theorem 3.2. It is said $\lambda$ and $C$ are large enough, but an example of acceptable rates in terms of $\lambda$ is $C \sim \lambda^{-1}$. Is it contradictory?

Author response:
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.

Reviewer comment (Theory):
In the proof of Theorem 3.2, r columns with the most observed entries are taken. What if all the columns have the same number of observed entries? Will the $d_{out}-r$ other columns of $W_L$ decay exponentially? I don't see the relation between rank $r$ and the number of observed entries here.

Author response:
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).

Reviewer comment (Theory):
As you point out, the common Langevin-based models predict a transition among each pair of points and thus fall short to capture the phenomenon shown by your result. Do you know if they would nevertheless exhibit a systematic low-rank bias (i.e. the transition toward lower-rank solutions being much more likely than toward higher-rank solutions)?

Author response:
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$).