Paper: Simplicity Bias of SGD via Sharpness Minimization (ICLR-2024)

Reviewer comment (Theory):
The conclusion says: "we show that by initializing the neurons in the subspace of high-dimensional input data, all of the neurons converge into a single vector." But, I can't find any part of the text that discusses either this special initialization, or the result that the neurons converge to a single vector. By contrast, the main text proves that for _any_ initialization, all the neurons _agree entirely on the training dataset_, which is weaker than saying that all neurons are identical. I think I do understand why the sentence in the conclusion follows from this: if the neurons are initialized within the span of the input data, then since they always move within the span of the input data, they must end up within the span of the input data; and there is only one possible set of weights that both lies within the span of the input data and matches the necessary targets, and this is a pseudoinverse. But the paper did not discuss this explicitly.

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

Reviewer comment (Theory):
If you added weight decay, would it be true that for any initialization, all the neurons converge to a single vector?

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

Reviewer comment (Theory):
The authors claim that for the ease of exposition, they set the second layer to one. If there is no clear proof that the result can be adapted from this, it is not an easy step to adapt their result to such a setup, even though the qualitative behavior might not change.

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

Reviewer comment (Theory):
Even though Theorem 1 is informative, there is a need for more precision at this stage. There is no clear definition of label noise SGD, no quantitative dependency on the step size and the noise of SGD. Li et al.'s result is asymptotic and non-quantitative in essence. Any quantitative and non-asymptotic result on this program is significantly difficult to obtain (see e.g., **Label noise (stochastic) gradient descent implicitly solves the Lasso for quadratic parametrisation**, L. Pillaud-Vivien et al., COLT 2022 on a specific prediction function).

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

Reviewer comment (Writing):
There are a number of typos. For example, on page 4: "Equation equation 1", "non=degeneracy", "equipp".

Author response:
**Q1: (typos)** 

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

Reviewer comment (Theory):
Page 5: "In Lemma 15, we show that θ(t) remains in a bounded domain along the gradient flow, which means having the weaker assumption that φ′, φ′′′ > 0 automatically implies Assumption 1 for some positive constants ρ1 and ρ2." I am not fully convinced by this claim. Lemma 15 is also established based on Assumption 1. Thus, if we do not have positive lower bounds on φ′, φ′′′, Lemma 15 may not hold and θ(t) may not stay in a bounded region.

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