Paper: Improving Convergence and Generalization Using Parameter Symmetries (ICLR-2024)

Reviewer comment (Novelty):
My foremost concern is that, compared to existing work on parameter symmetry [1], this work seems more or less incremental in theory, basically extending the analysis from [1] to SGD and Newton's method. Moreover, the extension to SGD is done not based on standard noise assumption. The new results on generalization are justified mostly by experiments, rather than in theory.

Author response:
> compared to existing work on parameter symmetry [1], this work seems more or less incremental in theory, basically extending the analysis from [1] to SGD and Newton's method. Moreover, the extension to SGD is done not based on standard noise assumption.

We respectfully disagree this is an incremental contribution in theory. Theorem 3.1 is the first global convergence theory that allows for teleportation. It was not trivial to identify that such a proof is possible, nor was it trivial to establish. We have also never seen a proof that guarantees that all points on the orbit converge before. There were three main technical challenges/innovations that set this theorem apart from a standard SGD theorem 1) finding the right proof structure and assumptions, 2) introducing a weighted telescoping, and  3) controlling an exponential terms with a cyclic dependency. Below we detail these three challenges. 
1) Before proving Theorems 3.1, we first had to find a proof structure and assumptions that could work with teleportation. Standard proofs based on convexity, or strong convexity, were not amenable to a teleportation step. The reason being that the teleportation step can move far away from the current iterate, which in turn breaks the recurrence relations required by these proofs. We found that only non-convex type proofs could allow for teleportation.

2) After determining a smoothness type proof could work, one difficult point was to handle terms that do not telescope. That is, if you look to eq (28) in the supp. material, the suboptimality terms $(L(w^t)-L(w^*))$ and $(L(w^{t+1})-L(w^*))$ do not telescope because they have different constants weighting these terms. To resolve this, we introduced artificial weights $\alpha_t$ so that the suboptimality terms in eq (29) would telescope. This step is very different from a regular SGD proof.

3) After telescoping eq (29), the weighted telescoping generated terms that grew exponentially, see the term in eq (33). We had to control the growth of this exponential term by making $T$ big enough, but this created a circular dependency between $T$ and the step size $\eta$. Fortunately we were able to resolve this dependency, and still arrive at a. $1/\sqrt{T}$ complexity result.


Regarding the noise assumption, note that our theorem makes no assumption on the noise, which makes it more general than theorems that suppose a certain noise structure. The more classical SGD proofs assume that the stochastic gradients are bounded (equivalently the loss if Lipschitz), that is ||g_t|| <= D which is a fixed constant. We make no explicit assumption on the norm of the gradient, instead, we bound E||g_t|| as a result of assuming the loss is a smooth function, see Lemma A.1 in our paper. This Lemma A.1 and approach for bounding the gradient is standard when assuming smoothness. 

> The new results on generalization are justified mostly by experiments, rather than in theory.

While the main text contains mostly experimental results, we discuss a possible theoretical approach in Appendix F.2. In short, we consider the shift of the minimum under data distribution shifts, and examine how curvature affects the average distance between the old minimum and the new minimum. The correlation depends on the shape of the minimum and may be more complicated beyond a two-dimensional parameter space. We leave a more systematic investigation to future work.

How sharpness affects generalization has been studied extensively but not yet fully understood [3]. We expect the theory on the curvature of minimum to be comparably complex and intriguing.

Reviewer comment (Theory):
The main assumptions are not clearly stated. I would suggest separating the assumptions rather than stating them at the beginning of each result. Moreover, the main result switches between stochastic (SGD) and deterministic settings (Newton's method), which makes it less accessible to the readers.

Author response:
> The main assumptions are not clearly stated. I would suggest separating the assumptions rather than stating them at the beginning of each result. Moreover, the main result switches between stochastic (SGD) and deterministic settings (Newton's method), which makes it less accessible to the readers.

We have added a clarification on the settings in Section 3.2. We plan to add a section dedicated to preliminaries and assumptions in the final version of the paper.

Reviewer comment (Theory):
I'm a bit confused about the claim that teleportation accelerates the convergence rate of SGD. The intuition part makes sense to me since it has the quadratic error term that typically arises from second-order optimization but the convergence rate in Theorem 3.1 is still $O(\epsilon^{-4}$ ), which is the same as SGD. The convergence guarantee is slightly stronger but I don't understand how we can claim teleportation accelerates the convergence rate of SGD.

Author response:
This is a good question, and we agree that “accelerates” is not the correct word here. Instead, we give a stronger convergence guarantee by showing that all parameters on the orbit of the iterates converge to a stationary point at a O(\epsilon^{-4}) rate, as opposed to the standard SGD theory which shows only the iterates converge at this rate. Thus the speed of convergence is the same, but the object that is converging (all points on the orbit) is a stronger notion of convergence.  We will re-word our contribution here accordingly, and thank the reviewer.

Reviewer comment (Evaluation):
Even though the claim is teleportation improves the convergence rate for Adagrad, SGD with momentum, RMSProp, and Adam, the only clear improvements that I could see from Figure 5 is Adagrad. The other graphs seem to have similar performance for algorithm with or without teleportation.

Author response:
We agree that the improvement brought by teleportation is less significant for algorithms other than SGD, perhaps due to how they handle gradients of different magnitudes. However, for all algorithms, the loss curve drops visibly faster right after teleportation, near the beginning of the training. Since teleportation does not have a significant impact on total runtime (last Figure in Appendix), we believe that it is still beneficial to integrate it with various optimizers.

Reviewer comment (Theory):
Does cross-entropy loss satisfy the condition for one teleportation in section 3.3? The experiment on Cifar 10 uses one-time teleportation but the remark in section 3.3 only mentions quadratic loss.

Author response:
It is difficult to analytically check whether cross-entropy loss satisfies the condition in section 3.3, although we observed in the experiment that one teleportation is usually sufficient to improve the convergence speed.

Reviewer comment (Experiments):
Does it matter at which epoch we perform the teleportation? If it does, how do we pick the best epoch?

Author response:
We did not conduct a systematic investigation of when to teleport in this paper, but results in a previous paper [1] suggests that for the purpose of improving convergence, teleporting around the beginning of training leads to the most significant improvement. For improving generalization, we teleport when SGD is close to being converged, because the sharpness and curvature metrics are designed for when loss is near 0. Additionally, these quantities can change during training due to the implicit bias of optimization algorithms, so optimizing them early in the training may not be effective.

Reviewer comment (Evaluation):
In page 7, it says teleporting to points with smaller curvature helps find a minimum with lower validation loss but in the left graph of figure 4, the loss when we move to a place with decreased curvature is actually higher. Am I missing something here?

Author response:
This is a typo in the text and has been fixed. Thanks for pointing it out!

Reviewer comment (Novelty):
One weakness is that some of the theoretical results are not especially novel. For example, the results on teleported SGD (Theorem 3.1) and the Newton steps (Prop. 3.2) seem to be minor modifications of standard proof techniques.

Author response:
Thank you for your comments and positive feedback! 

We respectfully disagree that the theoretical results in the first part of Section 3 are not novel. Theorem 3.1 is the first global convergence theory that allows for teleportation. It was not trivial to identify that such a proof is possible, nor was it trivial to establish. There were three main technical challenges/innovations that set this theorem apart from a standard SGD theorem 1) finding the right proof structure and assumptions, 2) introducing a weighted telescoping, and  3) controlling an exponential terms with a cyclic dependency. Below we detail these three challenges. 

1) Before proving Theorems 3.1, we first had to find a proof structure and assumptions that could work with teleportation. Standard proofs based on convexity, or strong convexity, were not amenable to a teleportation step. The reason being that the teleportation step can move far away from the current iterate, which in turn breaks the recurrence relations required by these proofs. We found that only non-convex type proofs could allow for teleportation.

2) After determining a smoothness type proof could work, one difficult point was to handle terms that do not telescope. That is, if you look to eq (28) in the supp. material, the suboptimality terms $(L(w^t)-L(w^*))$ and $(L(w^{t+1})-L(w^*))$ do not telescope because they have different constants weighting these terms. To resolve this, we introduced artificial weights $\alpha_t$ so that the suboptimality terms in eq (29) would telescope. This step is very different from a regular SGD proof.

3) After telescoping eq (29), the weighted telescoping generated terms that grew exponentially, see the term in eq (33). We had to control the growth of this exponential term by making $T$ big enough, but this created a circular dependency between $T$ and the step size $\eta$. Fortunately we were able to resolve this dependency, and still arrive at a. $1/\sqrt{T}$ complexity result.

Reviewer comment (Theory):
It seems that Definition 3.3 doesn't exclude *minimizers* of \( w \mapsto \|\partial \mathcal{L}/\partial w\|^2 \). It might help the exposition to mention this.

Author response:
**Response to Questions**
> “It seems that Definition 3.3 doesn't exclude minimizers of $w \mapsto |\partial \mathcal{L}/\partial w|^2$. It might help the exposition to mention this.”

Thanks for the suggestion! We agree that this point is not clear in the original draft and have made this explicit in the revised version.

Reviewer comment (Writing):
**Clarity**: Although the paper is generally well-written, I did find it difficult to follow at times, especially with respect to the overall structure. One suggestion would be to switch the order of sections 4 and 5, as section 5 investigates how teleportation improves optimization, while section 4 is more or less self-contained with respect to generalization. I would also suggest having a (sub)section which is dedicated to introducing the necessary preliminaries and assumptions, with additional pointers to literature (e.g., some of the notation in section 3.3 could be introduced in the preliminaries already). In the appendix, it would be helpful to restate all the needed notation and equation, so the reader does not have to switch between the main paper and the appendix to follow the proofs.

Author response:
We appreciate the suggestions on the overall structure. Section 5 is placed at the end since while experiments in that section focus on optimization, teleportation can be integrated with other optimizers both to improve optimization and to improve generalization. We agree that a section dedicated to preliminaries and assumptions would be helpful, but currently it is difficult to add to the main text due to the page limit. We will add more background in the final version of the paper. We have added the definition of needed notations in the appendix in the revised version.

Reviewer comment (Reproducibility):
**Reproducibility**: If I am not mistaken, there is no reference to or mention of any source code; it would be great to make your code publicly available.

Author response:
We have uploaded our source code together with the revised paper. We will add a github link in the final version of the paper.

Reviewer comment (Experiments):
Do you have any (preliminary) results on how teleportation affects generalization for other optimizers (e.g., AdaGrad, Adam, etc.)?

Author response:
Thank you for the suggestion! We added a new experiment that investigates how teleportation affects generalization for AdaGrad (Figure 12 in the revised paper). We found that similar to SGD, changing curvature via teleportation affects the validation loss, while changing sharpness has negligible effects. Teleporting to points with larger curvatures helps find minima with slightly lower validation loss. Teleporting to points with smaller curvatures increases the gap between training and validation loss.

Reviewer comment (Novelty):
This question goes beyond the scope of the paper, but I would be interested in your opinion on [1] in light of your contribution, a recent paper which challenges the current view on the correlation between sharpness and generalization.

[1] https://openreview.net/pdf?id=VZp9X410D3

Author response:
The paper by Andriushchenko et. al provides an extensive empirical study on the correlation between sharpness and generalization for transformers and ConvNets on large datasets. Their observations that sharpness does not correlate well with generalization and that the right sharpness measure is data-dependent reveals the complexity of the role of sharpness in generalization. Their results may also explain our observation that teleporting in symmetry directions to change sharpness has negligible effect on generalization. We expect curvature to have a comparably complex and intriguing effect on generalization. We believe that the link between sharpness/curvature and generalization is interesting and deserves further investigation, and we hope that teleportation could become useful as a tool to explore the minimum manifold.