Simplicity Bias of SGD via Sharpness Minimization

The proof of section 4.3 is not well-written.

Thank you for the feedback, indeed we shifted the proof overview of the second part of Theorem 3 after the overview of the first part.

$g$ and $\gamma$ are not defined.

Thank you for spotting this, the single line that includes the notation $g,\gamma$ is indeed a typo based on our previous notation. We have corrected this in the new version.

Typos/No code provided.

We apologize for the typos, we have fixed them in the new version. We have also submitted our code.

Q1: Where are the assumptions on the activations used?

The first assumption is crucial in obtaining a closed-form for the stationary point of the trace of Hessian, hence it is used in the proof of Theorem 2. The second assumption (Assumption 3) arises naturally in showing the local g-convexity property of trace of Hessian at approximate stationary points, which we crucially use to obtain our convergence rate.

Q2: Why not consider $d=n$?

The point mentioned by the reviewer is indeed correct, the subspace orthogonal to the data points plays no role and does not change during the optimization, hence the proof is similar to the case $d=n$. We preferred to present our result in this slightly more general setting because in applications of neural networks with high-dimensional data, it is often the case that $d > n$. Note that the subspace of the data span, even though it has a lower dimension than the ambient space, is still informative about the ultimate function that we are trying to learn.

Q3: No future directions provided.

Thank you for your feedback, we added two ideas in the conclusion section about future directions that we think are interesting regarding our work.

Q4: How is this work related to "Sharpness-Aware Minimization Leads to Low-Rank Features"

In the paper by writing down the update rule of SAM for a relu network, the authors unravel the local tendency of SAM to make the features sparser at each step. They further confirm this hypothesis empirically for deep relu networks and see that this sparsity also leads to low-rank features. In comparison, our work rigorously proves the low-rank property of the learned features by label noise SGD for a class of two-layer networks. Notably, our result can be generalized to 1-SAM as well by applying the results in Wen et al., 2023.

Given our changes and replies, we kindly ask the reviewer to increase her score.

Q1: The result of Li et al is asymptotic and not general.

Response: We want to clarify that we have already clearly stated that our main theorem (Theorem 1) only holds for a sufficiently small learning rate, label noise SGD, and that with step size $\eta$, the simplicity bias happens in $\tilde O(1/\eta^2)$ iterations. Following the reviewer's suggestion, we also stressed in the main text before Theorem 1, that our theorem " holds for sufficiently small learning rate."

Q2: The proof of the generalization to unequal second-layer weights is not clear.

Our analysis indeed extends to the case when the second layer weights are not all equal, as discussed in the proof of Theorem 2 (page 7 the paragraph "This characterizes...".) In such case, our simplicity bias will be slightly different, in the sense that the $\phi''$ image of the feature matrix becomes rank one instead.

We apologize for the confusion. In the revision, we have removed the phrases "for simplicity of presentation" and have moved the discussion about the case with general weights right after the proof of Theorem 2 at the bottom of page 7.

Q4: More details for Theorem 1.

We explicitly noted in Theorem 1 that the result holds for any level of noise in label noise SGD. We further reminded the definition of label noise SGD to the reader right before stating Theorem 1.

We thank the reviewer for the valuable references provided; indeed obtaining non-asymptotic results for label noise SGD is more demanding and we see it as an important future direction, even for more general networks. We have added the provided references to the paper.

Q5: Typos.

Thank you for spotting the typos. We just want to point out that regarding your comment about page 3 typos there is probably a misunderstanding, the line “the first phase where the algorithm is far from stationary points trace of Hessian…” should be read as “the first phase where the algorithm is far from stationary points, trace of Hessian decreases slowly and...”. We have added the comma after “points” to make that clear.

Given the changes we have made, we kindly ask the reviewer to increase her score.

Q2 Equality of the neurons in the span of data.

Thank you for this comment, indeed we obtain the equality of the neurons in the span of data by the fact that their dot product with the data points are all equal.

Q3 What if we add Weight decay?

Adding tiny weight decay might work ($\lambda\ll \eta$) if we also train for a very long time ( $1/\eta\lambda$ steps), but we don't have a formal proof. Adding constant or moderate Weight decay might prevent training loss from going to 0 or hurt sharpness minimization.

Q1: (typos)

We thank the reviewer for spotting the typos, we have fixed them in the new version.

Q2: Case of unequal weights and connection between $\phi$ and $\phi''$

We only consider the case of equal weights, but in the general case with unequal weights, it turns out that with a similar analysis as in the proof of Theorem 2, the $\phi''$ image of the feature matrix becomes rank-one. We have moved the discussion about the general case which was within the proof of Theorem 2 to the end of Section 4.2. We apologize if this created any confusion.

Q3: The flow might not stay in a bounded region without Assumption 1

We thank the reviewer for pointing out to this typo in the statement of Lemma 15; indeed, Lemma 15 only requires the positivity of the third derivative $\phi''' > 0$, and does not depend on Assuumption 1. Therefore, our argument for obtaining Assumption 1 from a weaker positivity assumption on $\phi',\phi'''$ is not circular. We corrected this typo in the new version that we uploaded.

Given our changes and feedback, we kindly ask the reviewer to increase her score.