Improving Convergence and Generalization Using Parameter Symmetries

Response to Questions

“1. What's the actual cost of one teleportation operation (group action) in practice? The authors conduct experiments on small-scale tasks. For huge deep networks, is it feasible to carry out one single teleportation?”

The complexity for one gradient ascent step in teleportation is the same as one step of back-propagation [1]. The time required to teleport a pair of layers scales quadratically with the hidden layer width and linearly with the mini-batch size. For large models, one can teleport a subset of all weights without affecting the others, as individual teleportation can be performed on each pair of consecutive layers. Since the time complexity of teleporting a pair of layers does not depend on the depth, it is feasible to carry out teleportation in deep networks.

“2. Is it possible to show that teleportation achieves better generalization bounds using tools like algorithm stability [2]? The current results are encouraging but a bit lack of rigourous theoretical proof.”

Thanks for pointing us to a relevant approach for developing a generalization bound. The generalization bounds in [2] build on the theorem that uniform stability implies generalization in expectation, and that stochastic gradient methods are $\epsilon$-uniformly stable. One approach to show that teleportation achieves better generalization bounds is to prove that SGD with a teleportation that improves curvature is $\epsilon_1$-uniformly stable with $\epsilon_1 < \epsilon$, which leads to a tighter generalization bound. However, establishing such a proof appears to be challenging, since it is not clear how teleportation affects uniform stability.

Our main contribution in Section 4 is developing a new approximation method for the curvature of the minimum, and effect of teleporting for sharpness/curvature on generalization. Finding a complete theory is beyond the scope of this paper. However, we have included some theoretical investigations on the correlation between curvature and generalization. We believe that Appendix F.2 could serve as a starting point to develop a generalization bound for teleportation. The expected displacement of minimum under distribution shift is directly related to the curvature of the minimum and provides a way to quantify the generalization gap.

References

[1] Zhao, B., Dehmamy, N., Walters, R., & Yu, R. (2022). Symmetry teleportation for accelerated optimization. Advances in Neural Information Processing Systems, 35, 16679-16690.

[2] Hardt, M., Recht, B., & Singer, Y. (2016, June). Train faster, generalize better: Stability of stochastic gradient descent. In International conference on machine learning (pp. 1225-1234). PMLR.

[3] Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, Nicolas Flammarion. A Modern Look at the Relationship between Sharpness and Generalization. International Conference on Machine Learning, 2023.

Thank you for your comments!

Response to weaknesses

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.

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.

Thank you for the suggestion. We will validate the theoretical time complexity by evaluating teleportation's runtime in larger scale experiments in future work.

Thank you for your comments and positive feedback!

Response to weaknesses

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.

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.

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.

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.

Response to Questions

“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.”

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.

“Does it matter at which epoch we perform the teleportation? If it does, how do we pick the best epoch?”

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.

“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?”

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

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.

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.

Figure 4 seems to show the opposite phenomenon to text description.

This is a typo in the text and has been fixed. Figure 4 shows that teleporting to points with larger curvatures helps find a minimum with lower validation loss, while teleporting to points with smaller curvatures has the opposite effect. This relationship is consistent with Table 1. For the dataset and architecture we examined, less curvature is associated with larger validation loss, which is associated with worse generalization.

Thanks for your thoughtful response. I wanted to add that I don't think the re-weighting trick in the proof of Theorem 1 is necessary. In particular, if you take $\eta_t = 1/\sqrt{T}$, then equation (29) becomes

\min_{t = 0, \ldots, T - 1} \mathbb{E}[ | \nabla \mathcal{L}(g^t \cdot w^t) |^2] \leqslant \frac{1}{T} \sum_{t = 1}^T \mathbb{E}[|\nabla \mathcal{L}(g^t \cdot w^t)|^2] \leqslant \frac{1}{\sqrt{T}} \mathbb{E}[ \mathcal{L}(w^1) - \mathcal{L}(w^{T})] + \frac{\beta^2 }{\sqrt{T}} (\mathbb{E}[ \mathcal{L}(w^1) - \mathcal{L}(w_\star)] + \sigma^2). $$ Then since $\mathcal{L}(w^{T}) \geqslant \mathcal{L}\_\star$, we essentially recover Theorem 1, up to dependence on $\beta$ (and I think the exact statement of Theorem 1 is recovered by taking $\eta_t = 1/\beta\sqrt{T - 1}$ as in Theorem 1 itself).

So from what I can tell, the proof is nearly done once equation (29) is obtained, and in particular there is no need for the re-weighting trick.

I am leaving my score unchanged, and stand my by original assessment that this work clearly merits acceptance.

I understood your approach and it is also interesting. But note that the noise constant you introduced requires assuming that the variance of the gradient estimate is globally bounded, and this is an additional assumption. We choose not to have this additional assumption, because it does not lead to an improvement in the rates (despite being more restrictive).

With respect to the two different noise constants. Your noise constant $\tilde{\sigma}^2$ is more common in the convex and non-smooth setting, where we often assume the loss function is Lipschitz. Indeed if $\mathcal{L}(w,\xi)$ is Lipschitz in $w$, then your $\tilde{\sigma}^2$ is bounded. The noise constant $\sigma^2$ we use is sometimes called the interpolation constant, and it is more common in the smooth setting. It's called the interpolation constant because if the model fits perfectly the data than $\sigma^2 =0.$ Thus $\sigma^2$ measure how close we are to interpolation. This constant has recently appeared in other work in the non-convex setting. We will take the reviewers suggestion and add more references to this constant and explain it's meaning. Thanks again.

Thank you for the detailed comments. In the revised version, we have fixed the writing issues according to the list of minor remarks.

Clarity

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.

Reproducibility

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

Response to Questions

“1. In the experiments, have you tried teleporting more often than just in the first epoch and whether it has a positive effect on convergence speed/generalization vs. the increase in runtime?”

For improving convergence, the runtime added by teleportation is negligible, but we observed that teleporting after the first epoch has almost no effect on convergence speed, which is consistent with the results in [2] (section 6.2). For improving generalization, performing a teleportation to change curvature is more time consuming. Under our current implementation, the extra cost does not justify teleporting more than once near a minimum.

“2. How did you decide on 10 and 1 gradient ascent steps, respectively, for the experiments?”

The number of gradient ascent steps is chosen such that the curvature or sharpness shows clear improvement but does not become large enough to cause divergence in the subsequent gradient descent steps.

“3. Do you have any (preliminary) results on how teleportation affects generalization for other optimizers (e.g., AdaGrad, Adam, etc.)?”

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.

“4. 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.”

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.

[1] Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, Nicolas Flammarion. A Modern Look at the Relationship between Sharpness and Generalization. International Conference on Machine Learning, 2023.

[2] Bo Zhao, Nima Dehmamy, Robin Walters, Rose Yu. Symmetry teleportation for accelerated optimization. Advances in Neural Information Processing Systems, 2022.

Thanks for the response and uploading your code!