Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper aims to explain why stochastic weight averaging (SWA, Izmailov et al., 2018) helps generalization. The paper adopts the framework of Kleinberg et al., 2018 and shows that SGD implicitly minimizes a regularized objective, where the regularizer depends on both Hessian of the loss and the covariance matrix of the stochastic gradient, under the assumption that the regularized loss is one-point convex. The authors show that the regularization strength is proportional to $O(\eta^2)$ where $\eta$ is the learning rate and SGD can reach a solution which is $O(\sqrt{\eta})$-close to the minimizer of the regularizers loss. The authors further explain the benefit of stochastic weight averaging by showing that the averaged iterate better optimizes the regularized loss, that is, SWA can get $O(\eta)$-close to the minimizer. 
strength_and_weaknesses: ##Pros:

1. The paper is well-written and easy to understand. The proofs of the main theorems (Theorem 1 and 2) seem correct to me.

2. The framework (Kleinberg et al., 2018) is originally developed for understanding optimization under assumptions weaker than convexity, that is, one-point convexity. The adoption of their framework towards understanding the implicit bias of SGD and SWA is quite novel, as well as the bias-optimization tradeoff view. 

3. The authors successfully explain the efficacy of SWA using the bias-optimization tradeoff, that is, though SWA implicitly penalizes the same objective as SGD does, SWA optimizes the objective better than SGD.



##Cons:

1. Though the theory proposed in this paper looks self-consistent, the authors don't really justify whether their theory is relevant to practice, like either directly testing their proposed regularizer using other optimizers so there is no optimization issue, or developing algorithms with better generalization using their insight. Technically speaking, it is not clear to me whether there exists a suitable learning rate, such that the optimization-bias tradeoff gives non-trivial improvement compared to the empirical minimizer of the original training loss.

	To be more specific, under Assumption 2, it is fairly easy to see any stationary point $w^*$ of the original loss $f$ satisfies that $\eta^2 \nabla_{w^*} {Tr}(\nabla^2 f(w^*) \Sigma(w^*)) \ge c\|v^*-w^* \|$ , where $\Sigma(w^*)$ is the covariance matrix of noise in Eq. (4). By Assumption (3), $\nabla^3 f(w^*)$ is of constant magnitude, which implies that $\|w^*-v^* \| = O(\eta^2)$.  Immediately we see the optimization guarantee given by both Theorem 1 and 2 are vacuous in terms of the implicit bias. This is because we know even getting $O(\eta^2)$-close (like $w^*$) to the minimizer of the regularized loss $v^*$ is not enough ($w^*$ can be understood by the solution found SGD with tiny LR)
	
	That being said, I don't think the view presented in this paper is wrong. It is still valuable. It could be analysis are not tight or assumptions are too strong. But I do think the authors should address the above issue before publishing this work, which can be viewed as a sanity check to the theory.
	
2. Regarding the relevance to the deep learning practice, another thing I think the authors should reconcile is why in practice typically the learning rate of SGD needs to be annealed and the generalization is still good. 

3. The details of the setting for experiments are missing, e.g. Figure 4. Why does ASGD have higher training loss than SGD in Figure 4? Does the same thing happen for Figure 3 (training loss is not reported there) Intuitively by averaging you should get a smaller training loss, e.g., think about SGD on a quadratic loss with isotropic gaussian noise.

**Typos**: 

1.page 17. "here we prove Theorem 1 which ..". should be Theorem 2.  
2.page 17. All $\nabla R$ should be $R$.

clarity,_quality,_novelty_and_reproducibility: The paper is well-written and easy to understand. The proofs of the main theorems (Theorem 1 and 2) seem correct. The adoption of the framework (Kleinberg et al., 2018) framework toward understanding the implicit bias of SGD and SWA is quite novel. More details about experiments are needed.
summary_of_the_review: This paper gives a novel understanding of the implicit bias of SGD, which explains the efficacy of stochastic weight averaging. However, I am not sure how this understanding is related to practice as the optimization guarantee presented in the main theorems is quite weak, not enough to separate SGD and SWA from ERM even. Thus now I tend to reject this paper, but I am willing to increase the score if the authors can address my concerns satisfactorily. 
(A good example is Damian et al.,2021, where they showed that the minimizer of the original loss doesn't meet their guarantee for label noise SGD on the regularized loss)

- Damian, Alex, Tengyu Ma, and Jason D. Lee. "Label noise sgd provably prefers flat global minimizers." Advances in Neural Information Processing Systems 34 (2021): 27449-27461.

================================================

Update after rebuttal: I've read the revision and rebuttal by the authors as well as reviews by other reviewers. My concern that "whether there exists a suitable learning rate, such that the optimization-bias tradeoff gives non-trivial improvement compared to the empirical minimizer of the original training loss" is not well-addressed by the authors and I decide to keep my score. 

The new example in appendix C given by the authors in the revision shows when $\delta$ is sufficiently small, SGD and average SGD will converge to the flat local minimizer (which is also the global minimizer of original loss $f$). This is just an optimization result because any algorithm that finds global minimizer will also find the flat minimizer in this case. To make it a result of the implicit bias of SGD, the author should demonstrate there is a case where the sharp local minimizer is indeed a global minimizer but still gradient noise biases SGD to the flatter area, even if the loss there could be slightly higher.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper aims to understand how averaging the parameters of SGD (averaged SGD) leads to a more stable convergence to flat regions. 
The authors take the alternative view of SGD from Kleinberg et al. (2018), which tracks $v = w - \eta \nabla f(v)$ instead of $w$ itself. Then SGD wrt $w$ can be viewed as a stochastic semi-gradient method wrt $v$ on a smoothed loss, where the smoothening factor depends on eta. Under one-point convexity at the minimizer of the smoothed loss (and also various regularity conditions), they show that both SGD and averaged SGD oscillate around the minimizer, but the distance to the minimizer from the final parameter of averaged SGD is smaller than that from SGD ($O(\eta)$ v.s. $O(\sqrt{\eta})$).
strength_and_weaknesses: ### Strength

1. The phenomenon that averaged SGD leads to a more stable convergence to flat regions is interesting to study, and may deepen the current understanding of the training dynamics in the late training phase.
2. Clear convergence results for both SGD and averaged SGD.
3. The theoretical result indeed provides a clear separation between SGD and averaged SGD in terms of the distance upper bounds to the minimizer of the smoothed loss.


### Weakness

**Major Concern.** The major concern I have is whether the current result indeed shows that larger LR leads to a flatter minimizer. If not, then people may want to use SGD with a smaller LR so that the distance to the minimizer is smaller, and then SGD may beat averaged SGD. Basically, the authors should check whether the minimizer $v_*(\eta)$ as a function of $\eta$ can move a lot when $\eta$ goes from a big value $\eta$ to a small value $\eta'$.

However, the distance between $v_*(\eta)$ and $v_*(\eta')$ should be bounded by $O(\eta^2)$. The reason is that we can apply the implicit function theorem for solving the equation $\nabla F(v) = 0$, which gives us a curve $v_*(\eta)$ wrt $\eta$. Note that the derivative wrt $\eta$ is $-(\nabla^2 F(v_*))^{-1} \left(\frac{\eta^2}{2} \nabla \mathrm{Tr}(\nabla^2 f(v_*) \mathbb{E}[\epsilon' \epsilon'^\top]) + O(\eta^3)\right)$. According to A4 and A5, this should be of order $O(\eta^2)$.

Note that the theorem for SGD shows a distance bound $O(\eta)$. So for SGD with a much smaller LR $\eta'$, the distance from the final parameter to the minimizer of the smoothed loss defined with $\eta$ would be $O(\eta^2) + O(\sqrt{\eta'}) = O(\eta^2)$. But for averaged SGD with LR $\eta$, it is $O(\eta)$, which is actually much higher.

The above argument would suggest that averaged SGD with large LR does not encourage flatness at all.

**Minor concerns.**
1. I understand that one-point convexity makes the analysis possible, but the authors should give some concrete examples of one-point convex functions, and discuss how general this notion is. It should also be stated clearly whether one-point convexity holds in the experiments or not.
2. Comparing the distance upper bounds for SGD and averaged SGD is reasonable, but this does not mean SGD cannot get much closer to the minimizer because the bound may not be tight. I would recommend the authors prove a lower bound in the SGD case to make the separation more concrete.
clarity,_quality,_novelty_and_reproducibility: This paper is well-written and easy to follow. The assumptions and theorems are stated very clearly.

The convergence analysis for SGD in the alternative view in this paper is new and the result is stronger than Kleinberg et al. (2018).
summary_of_the_review: I would like to hear the authors' comments on my major concern because I could be wrong. If $v_*$ can indeed change a lot for different learning rates, I would be rather happy to raise my score. Otherwise, I have to vote for rejection, but I encourage the authors to come up with another argument to motivate their convergence analysis and resubmit.

=================

After reading the rebuttal, I found the 1-dim example given by the author helpful. This simple example shows that the theory indeed predicts the averaged SGD converges to a flatter solution than SGD, so I have raised my score from 3 to 5.

Overall, I'm not against accepting this paper because the convergence analysis still looks good to me. But I feel that the paper is at the borderline due to the following concerns:
1. As I pointed out in the initial review, this paper compares the upper bounds for SGD and averaged SGD (even in the 1-dim example, only the upper bounds are compared). A more solid theoretical result should give a lower bound for SGD to ensure that it doesn't converge to flat minimizes.
2. Though I appreciate the 1-dim example for its simplicity, it is unclear how general it is. It would be better if the authors could have generalized the example to a broad function class. Alternatively, the authors could also make up the generality issue in theory by measuring the empirical loss landscape to see if it is similar to the 1-dim example in some sense. For this reason, I believe that this paper still has much room for improvement.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The manuscript studies the optimization and generalization trade-off caused by stochastic noise, step size, implicit bias, and sharpness of the minimum; and more importantly, why an averaged SGD with a large step size can improve the trade-offs and converge to a flat region more stably than SGD.
strength_and_weaknesses: ## Strength
* The paper is well-structured and well-written (in high quality), with a clear related work section. The reviewer was enjoying reading most part of the manuscript.
* The studied question is of great importance to the optimization community.
* Extensive numerical results are included in the manuscript in terms of various neural architectures and learning rate schedules; besides, some visualization results are very insightful. 

## Weaknesses
1. The manuscript argues that SGD using the step size $\eta$ converges to a distance $O(\sqrt{\eta})$ from the solution, whereas the averaged SGD using the same step size converges to a distance to $O(\eta)$. First of all, the reviewer is unclear on how to conclude the $O(\sqrt{\eta})$ and $O(\eta)$ from Theorem 1 and Theorem 2; and authors are required to elaborate on them. More importantly, compared to Theorem 1, Theorem 2 relies on the additional Assumption 3, and thus could derive a tighter rate; however, it does not make much sense to compare two results under different assumptions, and even highlight it as a key contribution.
2. If the reviewer understands the manuscript correctly, the current convergence theory in the submission cannot capture the training dynamic caused by using a larger step size as the averaged SGD only runs normal SGD and takes the average in the end. Understanding this would be a key step in interpreting the implicit bias in SGD, however, the current draft cannot fill this gap.
3. The choice of hyper-parameters used in the manuscript seems a bit random and no justification can be found. For example, weight decay with a coefficient of 0.05 was used but the reviewer believes many SOTA training recipes use 1e-4 or 5e-4; Similar comment can be applied to the choice of decay factor as well as the milestones for three cases.
clarity,_quality,_novelty_and_reproducibility: check comments above
summary_of_the_review: In general, the manuscript is well-written, and the main concerns of the reviewer are
* unfair comparison of theoretical/numerical results
* unsupported experimental setups.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper presents the convergence analysis for SGD and averaged SGD with fixed step size, showing that iterates converge to the vicinity of the optimizer of a smoothed loss function. Compared to prior work, the author improved existing results with a less strict assumption on the step size, and provide a convergence result for the averaged SGD, showing a better rate over vanilla SGD.
strength_and_weaknesses: Strength: This paper is well written, and it is easy to follow. The theoretical results are novel and the discussions/remarks are very helpful. While the results seem incremental over some existing work on the implicit bias of SGD, there is still a significant contribution, showing an improved rate using averaged SGD.

Weakness: The author mentioned before Theorem 1, "different proof techniques" than (Kleinber et al.), maybe you can elaborate?
clarity,_quality,_novelty_and_reproducibility: This paper is well-written and novel.
summary_of_the_review: Good paper studies a very relevant topic in the theory of machine learning.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: Summary: The manuscript studies the optimization and generalization trade-off caused by stochastic noise, step size, implicit bias, and sharpness of the minimum; and more importantly, why an averaged SGD with a large step size can improve the trade-offs and converge to a flat region more stably than SGD.
	
This paper was considered borderline by the reviewers (even after considering the rebuttal and updating their reviews). Reviewer vTYE recommended acceptance, but did not argue for it in discussion (and provided a fairly short review). All the other three reviewers gave the paper a 5. They acknowledged that it makes an interesting contribution to ICLR, providing a novel understanding of the implicit bias of averaged SGD. However, they felt it had some unclear points that needed to be addressed, in particular, as mentioned by reviewer EP51 and fQ6K, "whether there exists a suitable learning rate, such that the optimization-bias tradeoff gives non-trivial improvement compared to the empirical minimizer of the original training loss".
justification_for_why_not_higher_score: Lack of support among reviewers.
Decision: Reject