Paper: Fast Equilibrium of SGD in Generic Situations (ICLR-2024)

Reviewer comment (Experiments):
They only conduct experiments on MNIST which is almost a linearly-separable dataset, which is not good for deep learning analysis. I suggest the authors conduct the experiments on other more difficult datasets, such as CIFAR-10.

Author response:
We are running an experiment with CIFAR 10 and will add it to the revised version.

Reviewer comment (Writing):
There are some tiny errors in the paper, so I'd suggest the authors to proofread their paper more carefully.

Author response:
We will correct the typos and inaccuracies you pointed out.

Reviewer comment (Writing):
$\tilde{O}(1/\eta\lambda)$ in the abstract may need to be changed as $\tilde{O}(\frac{1}{\eta \lambda})$.

Author response:
You are right about the exponents in the abstract. They should be $e^{O( (\eta\lambda )^{-1}) } $ and $\tilde{O}(\frac{1}{\eta\lambda})$ respectively. We will fix them.

Reviewer comment (Experiments):
One obvious question is "can you have more experiments". This might be less relevant in this particular paper, but perhaps there is something you can do to reach people outside the very technical domain.

Author response:
Thanks for this suggestion. Experiments where the empirically observed rates of convergence are polynomial were contained in the original paper (Li et. al 2020) where the Fast Equilibrium Conjecture was first asked. We will add references to figures from that paper.

Reviewer comment (Experiments):
It would be useful to illustrate the rates known in the literature and the conjecture, perhaps with evidence from some datasets and some networks. A headline figure showing the speed of convergence to the stationary distribution and the exact rate you prove could attract the interest of many people.

Author response:
Thanks for this suggestion. Experiments where the empirically observed rates of convergence are polynomial were contained in the original paper (Li et. al 2020) where the Fast Equilibrium Conjecture was first asked. We will add references to figures from that paper.

Reviewer comment (Theory):
It is a bit confusing going back-and-forward between $O(1/\eta)$ and $e^{O(\eta^{-1})}$ results, which raises a question.

Author response:
Thanks for pointing out the unclarity of the writing here. The difference between the polynomial and exponential times is explained in the response 2 below. The difference between $(\eta\lambda)^{-1}$ and $(\eta\lambda)^{-\frac12}$, is due to the fact that (Bovier et al., 2004; Shi et al., 2020) studied neural networks without normalization, which involved an analysis on $R^d$. With normalization and weight decay together, (Li et al. 2020) proved that the problem, via the polar coordinate decomposition written as (4) and (5) in our paper, can be considered as a problem on $S^{d-1}$. After this polar decomposition step, the effective learning rate is $O((\eta\lambda)^{\frac12})$, which results in replacements of exponents appearing in convergence times.

Reviewer comment (Theory):
From SDE knowledge, convergence to the stationary distribution is generally dominated by the drift. What is the convergence rate to a local minimum if the noise term is dropped? Providing an analysis of this setting could help the reader understand the proof in a simplified context.

Author response:
In this case, the observed convergence to equilibrium is dominated by the diffusion. The drift regulates how fast the trajectories fall near the bottom of a basin (the “descent stage”), i.e. how fast the loss function value converges. However, the trajectories then start a random walk all around the bottom of the basin. This is the “diffusion stage”, whose time scale is regulated by the diffusion term. (The third stage takes an exponentially long time when trajectories leave the basin to nearby basins but this cannot be observed in practice because of the exponential time scale.)

Reviewer comment (Theory):
The SDE studied is a model for SGD in the given setting, but there is concern about the noise structure. Is there some guarantee in previous literature that constant gradient noise in the non-normalized setting translates to the noise projection structure studied in the normalized case? This refers specifically to formula 4-5 in comparison to the discrete update.

Author response:
Even though some earlier papers assumed a constant gradient noise structure, our current paper does not need to assume constant gradient noise. In fact, equations (4) and (5) were proved in (Li et al. 2022) and didn’t require constant noise structure.

Reviewer comment (Theory):
Are all the remaining three assumptions essential to prove the Fast Equilibrium conjecture? Would it be possible to even weaken the assumptions?

Author response:
You are absolutely right that the conditions in Assumption in 1.2 are essential for our analysis and we will add explanations in a revision. To be precise:
 * Assumption 1.2 (i) is essential because without this assumption, the SDE would not be equivalent to a SDE on the sphere $S^{d-1}$ (written as (6) in the original version, the one before current revision), which is crucial to our analysis. Without this assumption, similar analysis can probably be formulated on $R^d$ instead of the $S^{d-1}$ coordinate but there will be new technical obstacles to overcome. Since the original fast equilibrium conjecture was asked for normalized neural nets, we restrict our study to the current setting. 
 * Assumption 1.2 (ii) is important because if not, a trajectory may stay near a critical point (for example a saddle point) for a very long period of time, and thus it would not be able to converge within a polynomial time.  
* The reason why we need Assumption 1.2 (iii) is that if the span is not the whole tangent space, but instead a subspace of the tangent space, then the diffusion will be restrained to this subspace, which a priori may be very fractal and existing mathematical theory is not sufficient to guarantee a unique equilibrium in limit.

Reviewer comment (Writing):
In the first line in Section 4, "Assumption 1.3. (ii)" -> "Assumption 1.3. (i)"?

Author response:
We will fix this typo.

Reviewer comment (Theory):
Referring to the second experiment as "stochastic weight averaging" (SWA) is inappropriate, as SWA averages the model parameters at different iterations along the same trajectory. Conversely, the approach in this paper averages the parameters at the same iterations from different trajectories.

Author response:
We would like to thank you for your valuable comments. Some explanations in response to your remarks are below and we will address these concerns in the revised version accordingly. 
1. We would change the name of SWA to “weight average over trajectories”.

Reviewer comment (Novelty):
Missing discussion of related works, which argue that the iterates stay in the same basin for a significant amount of time when starting from the same initialization. 
- Frankle, J., Dziugaite, G. K., Roy, D., & Carbin, M. (2020, November). Linear mode connectivity and the lottery ticket hypothesis. In *International Conference on Machine Learning* (pp. 3259-3269). PMLR.
- Gupta, V., Serrano, S. A., & DeCoste, D. (2019, September). Stochastic Weight Averaging in Parallel: Large-Batch Training That Generalizes Well. In *International Conference on Learning Representations*.

Author response:
We would like to thank you for your valuable comments. Some explanations in response to your remarks are below and we will address these concerns in the revised version accordingly. 
2. We will add the references you pointed out.

Reviewer comment (Writing):
The paper is abstract in its current form, and it would be beneficial if the authors could provide specific examples where the removed assumptions may be too restrictive.

Author response:
We would like to thank you for your valuable comments. Some explanations in response to your remarks are below and we will address these concerns in the revised version accordingly. 
3. The previous assumption of analyticity is restrictive because the regularity of the loss function is decided by that of the activation function. Even though popular activation functions such as Sigmoid are analytic, a priori one could use smooth but not analytic functions. The one basin assumption is restrictive as we do not see empirical evidence of proof that $L$ only has one basin. In fact, the experiments at the end of the paper suggests that there are multiple basins.