Improving Convergence and Generalization Using Parameter Symmetries
Authors: Bo Zhao
Robert M. Gower
Robin Walters
Rose Yu
Venue: ICLR 2024 oral (2024)
Keywords: Symmetry
optimization
generalization
Primary area: optimization
TL;DR: We provide theoretical guarantees that teleportation accelerates the convergence rate, show that teleportation can be used to improve generalization, and integrate teleportation into various optimization algorithms such as meta-learning.

Abstract: In many neural networks, different values of the parameters may result in the same loss value. Parameter space symmetries are loss-invariant transformations that change the model parameters. Teleportation applies such transformations to accelerate optimization. However, the exact mechanism behind this algorithm's success is not well understood. In this paper, we show that teleportation not only speeds up optimization in the short-term, but gives overall faster time to convergence. Additionally, teleporting to minima with different curvatures improves generalization, which suggests a connection between the curvature of the minimum and generalization ability. Finally, we show that integrating teleportation into a wide range of optimization algorithms and optimization-based meta-learning improves convergence. Our results showcase the versatility of teleportation and demonstrate the potential of incorporating symmetry in optimization.

Published as a conference paper at ICLR 2024
IMPROVING CONVERGENCE AND GENERALIZATION USING PARAMETER SYMMETRIES
Bo Zhao
Robert M. Gower
University of California San Diego
Flatiron Institute rgower
bozhaolucsd. edu
Robin Walters
Rose Yu
Northeastern University
University of California San Diego roseyu@ucsd.em
wwwters@norrheateeeeeeee edu
ABSTRACT
In many neural networks, different values of the parameters may result the same loss value. Parameter space symmetries are loss-invariant transformations that change the model parameters. Teleportation applies such transformations to accel- erate optimization. However, the exact mechanism behind this algorithm's 's success is not well understood. this paper, we show that teleportation not only speeds up optimization in the short-term, but gives overall faster time convergence. Additionally, teleporting to minima with different curvatures improves general- ization, which suggests a connection between the curvature of the minimum and generalization ability. Finally, we show that integrating teleportation into a wide range of optimization algorithms and optimization-ban based meta- leeanlea improves convergence. Our results showcase the versatility of teleportation and demonstrate the potential of incorporating symmetry in optimization.
INTRODUCTION
Given deep neural network architecture and a dataset, there may be multiple points in the parameter space that correspond the same loss value. Despite having the same loss, the gradients and learn- ing dynamics originating from these points can be very different (Kunin et al. [202]] Van Laarhoven 2017 Grigsby et al., 2022). Parameter space symmetries, which are transformations of the parame- ters that leave the loss function invariant, allow us teleport between points in the parameter space on the same level set of the loss function (Armenta et al.) 2023). In particular, teleporting to a steeper point in the loss landscape leads to faster optimization.
Despite the empirical evidence, the exact mechanism of how teleportation improves convergence in optimizing non-convex objectives remains elusive. Previous work shows that gradient increases momentarily after teleportation, but could not show that this results in overall faster convergence (Zhao al.) 2022). In this paper, we provide theoretical guarantces on the convergence rate. In particular, we show that stochastic gradient descent (SGD) with teleportation converges to a basin of stationary points, where every point reachable by teleportation is also stationary. We also provide conditions under which one teleportation guarantees optimality of the entire gradient flow trajectory.
Previous applications of teleportation are limited accelerating optimization. The second part of this paper explores a different objective improving generalization. We relate properties of minima to their generalization ability and optimize them using teleportation. We empirically verify that cer- tain sharpness metrics correlated with generalization (Keskar et al. (017), although teleporting towards flatter regions has negligible effects on the validation loss. Additionally, we hypothesize that generalization also depends on the curvature of minima. For fully connected networks, derive an explicit expression for estimating curvatures and show that teleporting towards larger curvatures improves the model's generalizability.
To demonstrate the wide applicability of parameter space symmetry, we expand teleportation to standard optimization algorithms beyond SGD, including momentum, AdaGrad, RMSProp, and

---

Published as a conference paper at ICLR 2024
Adam. Experimentally, teleportation improves the convergence speed for these algorithms. Inspired by conditional programming and optimization-bnn based meta-learning (Andrychowicz et al., 2016), we also propose meta-optimizer to learn where to move parameters in a loss level set. This approach avoids the computation cost of optimization on group manifolds and improves upon existing meta- learning methods that are restricted to local updates.
The convergence speedup, applications in improving generalization, and the ability to integrate with different optimizers demonstrate the potential of improving optimization using symmetry. In sum- mary, our main contributions are:
theoretical guarantees that teleportation accelerates the convergence rate of SGD;
quantifying the curvature of a minimum and evidence its correlation with generalization; a teleportation based algorithm to improve generalization;
various optimization algorithms with integrated teleportation including momentum, Ada- Grad, and optimization-base meta-learning.
RELATED WORK
Parameter space symmetry. Continuous symmetries have been identified in the parameter space of various architectures, including homogeneous activations (Badrinarayanan et al. 2015 Du et al.] radial rescaling activations (Ganev et al. 2022), and softmax and batchnorm functions (Kumin
2018).
et al. 202T). Permutation symmetry has been linked to the structure of minima (Şimşek et al.j
2021
Entezari et al. 2022). Quiver representation theory provides more general framework for
symmetries in neural networks with pointwise (Armenta & Jodoin [021) and rescaling activations
(Ganev & Walters 2022). A new class of nonlinear and data-dependent symmetries are identified in
(Zhao et al.) 2023) Since symmetry defines transformations of parameters within a level set of the loss function, these works are the basis of the teleportation method discussed our paper.
Knowledge of parameter space symmetry motivates new optimization methods. One line of work seeks algorithms that are invariant to symmetry transformations (Neyshabur et al. 2015) Meng al. 2019). Others search in the orbit for parameters that can be optimized faster (Armenta et al. 2023 Zhao et al. 2022). We build on the latter by providing theoretical analysis on the improvement of the convergence rate and by augmenting the teleportation objective to improve generalization.
Initializations and restarts. Teleportation before training changes the initialization of parame- ters, which is known to affect the training dynamics. For example, imbalance between layers at initialization affects the convergence of gradient flows in two-layer models TTarmoun et al. 202I]. Different initializations, among other sources of variance, also lead to different model performance after convergence (Dodge et al. 2020 Bouthillier et al. 2021 Ramasinghe et al. 2022). In addition to initialization, teleportation allows changes in landscape multiple times throughout the training.
Teleportation during training re- re-initializes the parameters to a point with the same loss. Its effect can resemble warm restart (Loshchilov & Hutter 2017). which encourages parameters move more stable regions by periodically increasing the learning rate. Compared to restarts, teleportation leads to smaller temporary increase in loss and provides more control of where move the parameters.
Sharpness of minima and generalization. The sharpness of minima has been linked to the gen- eralization ability of models both empirically and theoretically (Hochreiter & Schmidhuber) 1997 Keskar et al., 2017 Petzka et al. 202T Ding et al. 2022 Zhou et al. (2020), which motivates op- timization methods that find flatter minima (Chaudhari al. [2017] Foret al., 202T Kwon al. 2021 Kim et al.) (022). We employ teleportation to search for flatter points along the loss level sets. The sharpness of a minimum is often defined using properties of the Hessian of the loss fune- tion, such as the number of small eigenvalues (Keskar et al., 2017: CCaudharr et al., 2017; Sagun the top eigenvalues (Wu et al., 2017). Alternatively, sharpness can be characterized by the maximum loss within neighborhood of minimum (Keskar et al. 2017 Foret et al. 2021 I Kim et al. 2022) or approximated by the growth in the loss curve averaged over random directions (Izmailov et al. 2018). ). The sharpness of minima does not always capture gener- alization (Dinh et al.) (017) (Andrushchenko et al. (023). Some reparametrizations do not affect generalization but can lead to minima with different sharpness.
cet al., 2017) or the product
2

---

Published as a conference paper at ICLR 2024
3 THEORETICAL GUARANTEES FOR IMPROVING OPTIMIZATION
In this section, we provide a theoretical analysis teleportation. We show that with teleportation, SGD converges basin of stationary points. Building its relation to Newton's o''s method, telepor- tation leads to a mixture of linear and quadratic convergence. Lastly, in certain loss functions, one teleportation guarantees optimality of the entire gradient flow trajectory.
Symmetry Teleportation. We briefty review the symmetry teleportation algorithm (Zhao et al., 2022), which searches for steeper points in loss level set to accelerate gradient descent. Consider the optimization problem
w arg min L(w), C(w) def BEWD [L( E)]
wERd
where 'D is the data distribution, data sampled from D, L the loss, the parameters of the model, and I the parameter space. Let G be group acting on the parameter space, such that
((w) L(g w), Vg E G, Ve E Rd.
Symmetry teleportation uses gradient ascent to find the group element that maximizes the magni- tude of the gradient, and applies g to the parameters while leaving the loss value unchanged:
w - w. g = argmax [VL(g w)ll
gEG
3.1 TELEPORTATION AND SGD
At each iteration N+ in SGD, we choose group element g² E G and use teleportation before each gradient step follows
V w/,
(1)
Here is learning rate, VL(w", (*) is the gradient of L(w', {*) with respect to the parameters w, and 5 DD a mini-batch of data sampled i.i.d from the data distribution at each iteration.
By choosing the group element that maximizes the gradient norm, we show in the following theorem that the iterates equation converge to a basin of stationary points, where all points that can be reached via teleportation are also stationary points (visualized in Figure I.
Theorem 3.1. (Smooth non-convex) Let L(w, {) be 8- smooth and let
0² def L(w") E inf C(w. {)
Consider the iterates wt given by equation where
W
(cM LM)
g' E arg max V VL(g wI,
gEG
which we assume exists.
If = DVT-1 then
W1
VL(g r0²)112
W2
min
E
max .9G
V
T
2,3
Ba²
Figure I: With teleportation, SGD converges to basin where all points on the level set are stationary points.
E [(ww) L(ww))]
(2)
VT-1
VT-I'
where the expectation is the total expectation with respect to the data EL foo 0,...,
This theorem is an improvement over vanilla SGD, for which we would have instead that
233
300
min E [[VV((')) <
E [Lww") L(w")] +
VT
VT-
For instance when G is compact and | C(g un is continuous over G, when the gradient is coercive function and G is bounded.

---

Published as a conference paper at ICLR 2024
The above only guarantees that there exists single point w' for which the gradient norm will eventually be small. In contrast, our result in equation _ guarantees that for all points over the orbit {g- wt Vg G}, the gradient norm will be small. For strictly convex loss functions, maxgeG|l V Lig w)]l 2 is non-decreasing with ((w) In this case, the value of L is smaller after T steps of SGD with teleportation, compared vanilla SGD (Proposition A.2)
3.2 TELEPORTATION AND NEWTON'S S METHOD
Intuitively, teleportation can speed up optimization as it behaves similarly to Newton's 's method. After a teleportation that takes parameters to a critical point on level set, the gradient descent direction the same as the Newton direction (Zhao et al.) 2022). As result, we can leverage the convergence of Newton's 's method to derive the convergence rate of teleportation for the deterministic setting.
Proposition 3.2 (Quadratic term in convergence rate). Let L be strictly convex and let wo Rd.
Let
wn' E arg max [IV((1)) s.t. L(un) ((n))=((wo)- WElld
Let L be the Hessian of L, and Amax (V2 be the largest eigenvalue of V2 If V((w') & O, then there exists λn such that 0 S Xo S Amax (V2((Wo))) and one step of gra- dient descent after teleportation with learning rate 7 > gives
wi w' [((w)) w do 2
Let w = go wo- If S L. C is ju-strongly convex L L-smooth function, and the Hessian is G-Lipschitz, then we have that
G 1000 wo w0 112 +11 >oo 24
L we w
||0 w0ll S
l1
2/.
More details about the assumptions and the proof are in Appendix Note that due to unknown step size λg extra care is needed in establishing this convergence rate.
The above proposition shows that taking one step of teleportation and one gradient step, the result is equal to taking dampened Newton step (equation 3 Hence, the convergence rate has quadrat- ically contracting term [[go ww wI, which is typical of second order methods. In particular, setting 7 1/Ao would have local quadratic convergence. In contrast, without the teleportation step and under the same assumptions, we would have the following linear convergence
|w/ w" s (1 (4) |mw w
for y < 1 using gradient descent. Thus there would be no quadratically contracting term.
3.3 WHEN IS ONE TELEPORTATION ENOUGII
Despite the guaranteed improvement in convergence, teleporting before every gradient descent step is computationally expensive. Hence we teleport only occasionally. In fact, for certain optimization objectives, every point the gradient flow has the largest gradient norm in its loss level set after one teleportation (Zhao et al. 2022). In past work. this result is limited to convex quadratic functions. In this section, we give a sufficient condition for when one teleportation results in an optimal trajectory for general loss functions. Full proofs can be found in Appendix C
Let V M + TM be a vector field on the manifold M, where T.M denotes the associated tangent bundle. Here we consider the parameter space . = R", although results in this section can be extended to optimization on other manifolds. this case, we may write v² Dar using the
component functions v² R + R and coordinates w
Consider smooth loss function L M R. Let G' be symmetry group of C, i.c. Lig w) L(w) for all M and
G. Let X be the set of all vector fields M. Let R = r² à where Td dw, be the reverse gradient vector field. Let X1 {A a' 8
Jus X a'
C (M) and Σ; ((w))r((( = 0, Vw E M3 be the set vector fields orthogonal to R. If G is Lie group, the infinitesimal action of its Lie algebra defines set of vector fields X, X
4

---

Published as a conference paper at ICLR 2024
A gradient flow is curve y R - .M where the velocity is given by the value of R. i... (t) Re(()) for all & R. The Lie bracket [A, R. defines the derivative of with respect to A. Flows of A and commute and only if [A, R affect the convergence rate only if [A, RJ C * O for some A Xg To simplify notation, we write (WW,R[))())) = 0 for a set of vector fields W S X when ([A, (w) = 0 for all W.
(Theorem 9.44, Leee 2013 That is, teleportation can
We consider gradient flow optimal if every point on the flow is critical point of the magnitude of gradient in its loss level set. Note that this definition does not exclude the case where points on the flow are minimizers of the magnitude of gradient.
Definition 3.3. Let f : - R,
te l|z point w E M optimal with respect to a set of vector fields W S 11 if Af(w) = 0 for all W. A gradient flow y M is optimal with respect to W if (() is optimal with respect to W for all i E IR.
Proposition 3.4. point w E . is optimal with respect to a set of vector fields W if and only if RW w) O.
A sufficient condition for one teleportation result in an optimal trajectory is that whenever the function [A, R vanishes at w E .M, it vanishes along the entire gradient flow starting at w.
Proposition 3.5. Let W C R, be set of vector fields that are orthogonal to ac Tw. Assume that for
all w E M such that ([, (w)) O, we have that (R(W, (w)) O. Then the gradient flow starting at any optimal point with respect to W is optimal with respect to W.
To help check when the assumption in Proposition 3.5 satisfied, provide alternative form of R(W, ())(w) when [W, R]((w) = O.
Proposition 3.6. If at all optimal points in {(M y Bus
E Rux" M =MM},
Oc dc L
dc
M Own Owa Ow Oww
Ow
for all anti-symmetric matrices M E R then the gradient flow starting at an optimal point in is optimal in S.
From Proposition 3.6 see that RIW, R))(w) not automatically when [W, [[w( O. Therefore, even if the group is big enough to have its infinitesimal actions cover the tangent space of the level set (3g 11), one teleportation does not guarantee that the gradient flow intersects all future level sets at optimal points. However, for loss functions that satisfy the condition in Proposition 3.5 teleporting once optimizes the entire trajectory. This is the case, for example, when C
BC
C &C for all i, ki j, (Proposition C.3 In particular, all quadratic Daan th Buo
functions meet this condition.
4 TELEPORTATION FOR IMPROVING GENERALIZATION
Teleportation was originally proposed to speedup optimization. In this section, we explore the suit- ability of teleportation for improving generalization, which is another important aspect of deep learn- ing. We first review definitions of the sharpness of minima. Then, we introduce a novel notion of the curvature of minima and discuss its implications on generalization. By observing how sharpness and curvature of minima correlated with generalization, we improve generalization by incorporating sharpness and curvature into the objective for teleportation.
4.1 SHARPNESS OF MINIMA
Flat minima tend to generalize well (Hochreiter & Schmidhuber 1997). typically characterized by numerous small Hessian eigenvalues. Although Hessian- based sharpness metrics are known to correlate well with generalization, they are expensive compute and differentiate through. To use sharpness as an objective in teleportation, we consider changes in the loss averaged over random directions. Let D be set of vectors drawn randomly from the unit sphere d; {d E 12 ||]| 1}. and 7" a list of displacements tj E R. Then, we have the following metric (Izmailov et al. 2018
Sharpness:
$(w, T,
1
Σ TET
Σ
L(w + td).
D
deD
5

---

Published as a conference paper at ICLR 2024
4.2 CURVATURE MINIMA
At a minimum, the loss-invariant or flat directions are zero eigenvectors of the Hessian. The curvature along these directions does not directly affect Hessian-based sharpness metrics. However, these curvatures may affect generalization, by themselves or by correlating to the cur- vature along non-flat directions. Unlike the curvature of the loss (curve ((w) in Figure 2) the curvature of the minima (curve 7) less well studied. We provide novel method to quantify the curvature of the minima below.
c(w)
Assume that the loss function sider the curve YM IR x R > R" where M Lie(G) and (t, w) = exp (fM) w. Then 7(0, w) = w. and every point Yar is in the minimum if w minimum.
a G symmetry. Con-
W2
W1
Let > dr be the derivative of curve y.. The curva- 7') TA'T where T(t) |ل is
Figure 2: Gradient flow (C(w)) and
ture of y is x(7,() f)
the minimum (γ). The curva- ture of both curves may affect general- ization.
the unit tangent vector. We assume that the action map
smooth, since calculating the curvature requires second derivatives and optimizing the curvature via gradient descent requires third derivatives. For multi- layer network with element-wise activations, we derive the group action, 7, and in Appendix D
Since the minimum can have more than dimension, we measure the curvature of point w the minimum by averaging the curvature of & curves with randomly selected Lie algebra elements M, Lie(G). The resulting new metric is
Curvature:
(ww, k))
k Σ <(7MM, w),
(5)
There are different ways to measure the curvature of higher-dimensione manifold, such as using the Gaussian curvature 2D subspaces of the tangent space. However, our method of approximating the mean curvature is easier to compute and suitable as differentiable objective.
(a)
(b)
(c)
(d)
I(w)
[(w)
f
L
W2
W2
('(w)
dist(w"
di xr
AC
1' (w)
AC
W
W
W
w
W1
W1
Figure 3: Illustration of the effect of sharpness (a.b) and curvature (c,d) of minima on generalization. See Figure 2 for 3D visualization of the curves L(w) and 7. When the loss landscape shifts due to change in data distribution, sharper minima have larger increase loss. In the example shown, minima with larger curvature moves further away from the shifted minima.
4.3 CORRELATION WITH GENERALIZATION
Generalization reflects how loss changes with shifts in data distribution. The sharpness of minima is well known to be correlated with generalization. FigureBja)(b visualizes example of the shift loss landscape (C(w)) )), and the change loss AL at a minimizer w" large when the minimum is sharp. The relation between the curvature of minimum and generalization is less well studied. Figure (d) shows one possible shift the minimum (y). Under this shifting, the minimizer with larger curvature becomes farther away from the shifted minimum. The curve on the minimum can shift other directions. Appendix [2.2 provides analytical examples of the correlation between curvature and expected distance between the old and shifted minimum.

---

Published as a conference paper at ICLR 2024
Table 1: Correlation with validation loss
sharpness (+)
curvature (e)
MNIST Fashion-MNNIST CIFAR-10
MNIST
Fashion-MNN
CIFAR-10
0.704
0.790
0.899
-0.050
-0.232
-0.167
We verify the correlation between sharpness, curvatures, and validation loss on MNIST (Deng 2012). Fashion-MNIST (Xiao al., 2017), and CIFAR-10 (Krizhevsky et al. 2009). On each dataset, we train 100 three-layer neural networks with LeakyReLU using different initializations. Details of the setup be found in Appendix E.3
Table shows the Pearson correlation between validation loss and sharpness or curvature (scatter plots in Figure 9 and [0] in the appendix). In all three datasets, sharpness has strong positive correlation with validation loss, meaning that the average change in loss under perturbations is a good indicator of test performance. For the architecture negatively correlated with the validation loss. We observe that the magnitudes of the curvatures are small, which suggests that the minima are relatively flat.
consider, the curvature of minima is
4.4 TELEPORTATION FOR IMPROVING GENERALIZATION
To improve the generalization ability of the minimizer and to gain understanding of the curvature of minima, teleport parameters regions with different sharpness and curvature. Multi-layer neural networks have GL(R) symmetry between layers (Appendix D.I. We parametrize the group by its Lie algebra T. and perform gradient ascent T maximize the gradient norm at the transformed parameters V Lexxp (T) Algorithm 2 in Appendix E.4 demonstrates how to increase curvature ei by teleporting two layers, with hidden dimension h, in MLP. In experiments, we use an extended version of the algorithm, which teleports all layers by optimizing on list of T's concurrently. During teleportation, we perform gradient descent on the group elements to change @ or sb. Results are averaged over 5 runs.
2.0
2.0
SGD teleport(decrease $) teleporttincrease $)
SGD teleport|decrease teleportiincrease y)
Loss
Loss
1.7
1.7
1.4
1.4
20
40
0
20 Epoch
40
Epoch
Figure 4: Changing sharpness (left) or curvature (right) using teleportation and its effect gener- alization on CIFAR-10. Solid line represents average test loss, and dashed line represent average training loss. Teleporting to decrease sharpness improves validation loss slightly. Teleportation changing curvatures has a more significant impact on generalization ability.
Figure 4 shows the training curve of SGD on CIFAR-10, with one teleportation epoch 20. Similar results for AdaGrad can be found Appendix F.4 Teleporting to flatter points slightly improves the validation loss, while teleporting to sharper points has effect. Since the group action keeps the loss invariant only on the batch of data used in teleportation, the errors incurred in teleportation have similar effect
restart, which makes the effect of changing sharpness less clear.
Interestingly, by changing the curvature, teleportation is able to affect generalization. Teleporting to points with larger curvatures helps find minimum with lower validation loss, while teleporting to points with smaller curvatures has the opposite effect. This suggests that at least locally, curvature correlated with generalization. Details of the experiment setup can be found in Appendix E.4

---

Published as a conference paper at ICLR 2024
5 APPLICATIONS TO OTHER OPTIMIZATION ALGORITHMS
Having shown teleportation's potential improve optimization and generalization, demonstrate its wide applicability by integrating teleportation into different optimizers and meta-learning.
5.1 STANDARD OPTIMIZERS
Teleportation improves optimization not only for SGD. show that teleportation works well with other standard optimizers, train 3-layer neural network MNIST using different optimizers with and without teleportation. During training, we teleport once at the first epoch, using mini- batches of size 200. Details can be found in Appendix F2
Figure shows that teleportation improves the convergence rate when using AdaGrad, SGD with momentum, RMSProp, and Adam. The runtime for teleportation is smaller than the time required to train one epoch, hence teleportation improves convergence rate per epoch at almost no additional cost of time (Figure 13 in the appendix).
Adagraad trait Adagpiad Adkggnd
momrnntu tran
RNGproppiig trann
train
INGgrap test
Los
Loss
0.6
Loss
0.3
Epoch
Epoch
Epoch
Epoch
Figure 5: Integrating teleportation with AdaGrad, momentum, RMSProp, and Adam improves the convergence rate on MNIST. Solid line represents the average test loss, and dashed line represents the average training loss. Shaded areas are standard deviation of the test loss across 5 runs.
5.2 LEARNING TO TELEPORT
In optimization based meta-learning, the parameter update rule or the hyperparameters are learned using meta-optimizer (Andrychowicz et al. 2016 Finn al. 2017). Teleportation introduces an additional degree of freedom in parameter updates. We augment existing meta- meta-learning algorithms by learning both the local update and teleportation. This allows to teleport without implementing the additional optimization step on groups, which reduces computation time.
Let w, Rd be the parameters time and V. descent, the update rule with learning rate 7
ac
be the gradient of the loss L. In gradient
lw
Wh+1 = w. V.
In meta-leamiigg (Andrychowicz al. 2016). the update w; is learned using meta-learning optimizer m, which takes V, as input. Here m is an LSTM model. Denote h, as the hidden state in the LSTM and á as the parameters in m. The update rule is
W wi + ft ft he+1
Extending this approach beyond additive update rule, we learn teleport. Let G be a group whose action on the parameter space leaves L invariant. We use two meta-learning optimizers m., mg to learn the update direction f. Ed and the group element ge G:
W++1 = 9t (w, + (s))
fi = m, (V.,,,, h1
9t
m2(V (2).
h2
8

---

Published as a conference paper at ICLR 2024
Experiment setup. We train and test two-layer neural networks C(W1. W2) |Y W2X(WiX) ||2 where W2. Wi, Y E 200000 and a is the LeakyReLU function with slope co- efficient 0.1. Both meta-optimizers are two-layer LSTMs with hidden dimension 300. We train the meta-optimizers on multiple trajectories created with different initializations, each consisting of 100 steps of gradient descent on & with random X, Y and randomly initialized W's. We update the parameters in m and m2 by unrolling every 10 steps. The learning rate for meta-optimizers are 10-4 for my and 10-3 for m2. We test the meta-optimizers using 5 trajectories not seen in training.
Algorithm T summarizes the training procedure. The vanilla gradient descent baseline ("GD") uses the largest learning rate that does not lead to divergence (3 10-4), The second baseline ("LSTM(update)") learns the update ft only and does not perform teleportation (g. I, Vt). The third baseline ("LSTM(Irtelee)) learns the group element ft and the learning rate used to perform gradient descent instead of the update ft- We keep training until adding more training trajectories does not improve convergence rate. We use 700 training trajectories for our approach, 600 for the second baseline, and 30 for the third baseline.
Results. By learning both the local update fa and non-local transformation 9t. our meta-optimizer successfully learns to learn faster. Tigure/J]sowws the improvement of our approach from the previ- ous meta-learning method, which only learns ft. Compared to the baselines, learning the two types of updates together ("LSTM(update.tele)) achieves better convergence rate than learning them sep- arately. Additionally, learning the group element It eliminates the need for performing gradient ascent on the group manifold and reduces hyperparameter tuning for teleportation. As an example of successful integration of teleportation into existing optimization algorithms, this toy experiment demonstrates the flexibility and promising applications of teleportation.
Algorithm Learning to teleport
GD
10
LSTMtir.tele
Input: Loss function L. learning rate 7, number of epochs T. LSTM models m1, m2 with initial parameters Φ1, 622 unroll step tunrolt-
LSTMiupdatel
104
LSTMiupdatettele)
10a
Output: Trained parameters Φ1 and Φ2-
for each training initialization do for t= = 1 to I do
10a
10
fr, h1.1 mi (v hic Φ1)
10
30
Epoch
Jt, h2.4 m2 (V h2. 22)
+ 9t (w +f+)
Figure 6: Performance of the trained meta-optimizer on the test set. Learning both local update ft and nonlocal transformation ge results in better convergence rate than learning only local updates or learning only teleportation.
if mod lunrall = 0 then update Φ1, 22 by mack-propogation from the accu- mulated loss Dimt-taroa ((w.)
end if end for end for
6 DISCUSSION
Teleportation is a powerful tool to search in the loss level sets for parameters with desired properties. We provide theoretical guarantees that teleportation accelerates the convergence rate of SGD. Using concepts in symmetry, we propose a novel notion of curvature and show that incorporating addi- tional teleportation objectives such as changing the curvatures can be beneficial to generalization. The close relationship between symmetry and optimization opens up number of exciting opportu- nities. Exploring other objectives in teleportation appears to be an interesting future direction. Other possible applications include extending teleportation to different architectures, such as convolutional or graph neural networks, and to different algorithms, such as sampling-base based optimization.
The empirical results linking sharpness and curvatures to generalization are intriguing. However, the theoretical origin of their relation remains unclear. In particular, precise description of how the loss landscape changes under distribution shifts is not known. More investigation of the correlation between curvatures and generalization will help teleportation to further improve generalization and take us a step closer to understanding the loss landscape.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work was supported in part by Army-BCASE award W91 INF-23-1-0231, the U.S. Department Of Energy. Office of Science under #DE-SC0022255, IARPA HAYSTAC Program, CDC-RFA-FT- 23-0069, NSF Grants #2205093, #2146343, #2134274, #2107256, and #2134178.
REFERENCES
Osmar Aléssio. Formulas for second curvature, third curvature, normal curvature, first geodesic curvature and first geodesic torsion of implicit curve in n-dimensions. Computer Aided Geometric Design, 29(3-4):189-22 201, 2012.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. modern look at the relationship between sharpness and generalization. International Conference on Machine Learning, 2023.
Marein Andrychowicz, Misha Denil, Sergio Gomez, Matthew Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in Neural Information Processing Systems, 29, 2016.
Marco Armenta and Pierre-Mare Jodoin. The representation theory of neural networks. Mathemat- ics, 9(24), 2021. ISSN 2227-7390.
Marco Armenta, Thierry Judge, Nathan Painchaud, Youssef Skandarani, Carl Lemaire, Gabriel Gibeau Sanchez, Philippe Spino, and Pierre-Marc Jodoin. Neural teleportation. Mathematics, 11(2):480, 2023.
Vijay Badrinarayanan, Bamdev Mishra, and Roberto Cipolla. Symmetry-invarian optimization in deep networks. arXiv preprint arXiv:1511.01754. 2015.
Xavier Bouthillier, Pierre Delaunay, Mirko Bronzi, Assya Trofimov, Brennan Nichyporuk, Justin Szeto, Nazanin Mohammadi Sepahvand, Edward Raff, Kanika Madan, Vikram Voleti, et al. Ac- counting for variance in machine learning benchmarks. Proceedings of Machine Learning and Systems, 3:747-769, 2021.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. International Conference on Learning Representations, 2017.
Li Deng. The mnist database of handwritten digit images for machine learning research [best of the web]. IEEE signal processing magazine, 29(6): 141-142, 2012.
Lijun Ding, Dmitriy Drusvyatskiy, Maryam Fazel, and Zaid Harchaoui. Flat minima generalize for low-rank matrix recovery. arXiv preprint mxXiv:2203.0376 2022.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference Machine Leurning, pp. 1019-1028. PMLR, 2017.
Jesse Dodge, Gabriel Ilharco, Roy Schwartz, Ali Farhadi, Hannaneh Hajishirzi, and Noah Smith. Fine-tuning pretrained language models: Weight initializations, data orders, and early stopping. arXiv preprint arXiv:2002.06305 2020.
Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. Neural Information Processing Systems, 2018.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. International Conference on Learning Representations, 2022.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference machine learning, pp. 1126-1 135. PMLR, 2017.
10

---

Published as a conference paper at ICLR 2024
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations, 2021.
Iordan Ganev and Robin Walters. Quiver neural networks. arXiv preprint arXiv:2207.12777 2022.
lordan Ganev, Twan van Laarhoven, and Robin Walters. Universal approximation and model com- pression for radial neural networks. arXiv preprint arXiv:2107.0255022 2022.
J Elisenda Grigsby, Kathryn Lindsey, Robert Meyerhoff, and Chenxi Wu. Funetional dimension of feedforward relu neural networks. arXiv preprint arXiv:2209.04036, 2022.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. Conference on Uncertainty in Artificial Intelligence, 2018.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. Inter- national Conference on Learning Representations, 2017.
Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher sam: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning, pp. 11 11148- 11161. PMLR, 2022.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Daniel Kunin, Javier Sagastuy-Brent Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. In Inter- national Conference on Learning Representations, 2021.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International Con- ference on Machine Learning, pp. 5905-5914. PMLR, 2021.
John M Lee. Introduction to Smooth Manifolds. Graduate Texts in Mathematics, vol 218. Springer,
New York, NY, 2013.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. Interna- tional Conference on Learning Representations, 2017.
Qi Meng, Shuxin Zheng, Huishuai Zhang, Wei Chen, Zhi-Ming Ma, and Tie-Yan Liu. G-SGD: Optimizing relu neural networks in its positively scale-invariant space. International Conference on Learning Representations, 2019.
Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro. Path-SGD: Path-normalized optimiza- tion in deep neural networks. Advances in Neural Information Processing Systems, 2015.
Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchiseseu, and Mario Boley. Relative flatness and generalization. 35th Conference on Neural Information Processing Systems, 2021.
Sameera Ramasinghe, Lachlan MacDonald, Moshiur Farazi, Hemanth Sartachandran, and Simon Lucey. How you start matters for generalization. arXiv preprint arXiv:2206.0885 220.:0888888888 2022.
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over- parrperreetriieed neural networks. arXiv preprint 17Xiv:1706.04454 2017.
Aleksandr Mikhailovich Shelekhov. On the curvatures of curve in n-dimensional euclidean space. Russian Mathematics, 65(1 1):46-58, 2021.
Berfin Şimşek, François Ged, Arthur Jacot, Francesco Spadaro, Clément Hongler, Wulfram Gerst- ner, and Johanni Brea. Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances. In International Conference on Machine Learning, pp. 9722-9732 PMLR, 2021.
11

---

Published as a conference paper at ICLR 2024
Sebastian U. Stich. Unified optimal analysis of the (stochastic) gradient method. CoRR, 2019.
Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal. Understanding the dynamics of gradient flow in overparameterized linear models. In International Conference on Machine Learning, pp. 10153-10161. PMLR, 2021.
Twan Van Laarhoven. L.2 regularization versus batch and weight normalization. Advances Neural Information Processing Systems, 2017.
Lei Wu, Zhanxing Zhu, al. Towards understanding generalization of deep learning: Perspective of loss landscapes. arXiv preprint arXiv: 10239, 2017.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprini arXiv: (eti:::700.007444 2017.
Bo Zhao, Nima Dehmamy, Robin Walters, and Rose Yu. Symmetry teleportation for accelerated optimization. Advances in Neural Information Processing Systems, 2022.
Bo Zhao, Iordan Ganev, Robin Walters, Rose Yu, and Nima Dehmamy. Symmetries, flat minima, and the conserved quantities of gradient flow. International Conference on Learning Representa- tions, 2023.
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoreti- cally understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems, 33:21285-21296, 2020.
12

---

Published as a conference paper at ICLR 2024
APPENDIX
This appendix contains proofs, experiment setups, as well as additional results and discussions. Appendix A through C contain proofs for theoretical results Section [ Appendix D provides details about curves induced by symmetry and the curvature of the minimum. Appendix discusses possible theoretical approaches experiment details on computing correlations and the algorithm that uses teleportation to change cur- vature. Appendix describes experiment setups and different strategies of integrating teleportation into various optimization algorithms.
relate curvatures and generalization. This section also contains
The code used for our experiments available at: https: bittps//github... com Rose STL--Lb/ Te deeleportat.on-pptiime
A TELEPORTATION AND SGD
This section includes proof for Theorem Additionally, we discuss the theorem's 's implication when the loss function strictly convex.
Lemma A.1 (Descent Lemma). Let L(w, E) be 8-smooth function. It follows that
E (IVL(w,<))F)) < 2 [(w")) + 2B(C(wx") - E inf L(w,
(6)
Proof. Since C(w, E) is smooth we have that
3 wil?
((z,{) [(w,E) {) < (VL(w, E), -w)
Vz, w E Rd.
2
By inserting
v {)
B
into equation
7 we have that
1
[(w (..))) C(w.
23 [C(w,))|1
(8)
Re-arranging we have that
L(w", .)) E) L(w", ,)) inf L(w, {) + inf L(w, 5) L(un, g)
< L(w", ,)) inf L( (w 5) + L(w (1/3) L ((w,)www,)) .), L(w,,
w
equation I
1
L(w", {) inf L( ((w.))
V vL(w.
w
23
where the first inequality follows because inf L(w,E) L(w, E), Vw. Re-arranging the above and taking expectation gives
E [[] E) < 2E ([[(w"EE) inf L(w, + L(w, E) - L(w", {))
< 2,3 [(w", - inf L(w, {) + L(w, L(w", [5] ((3(((w) - -(ww)) + +2,3(((2))) E inf L(w, E)
At each iteration E N+ in SGD, we choose group element g" E G and use teleportation before each gradient step as follows
w ++1 9" w' L/Vg w²,st)
Here 17 is a learning rate, VL(w',') is a gradient of L(w', with respect to the parameters w, and gt D is mini-batch of data sampled i.i.d each iteration.
13

---

Published as a conference paper at ICLR 2024
Theorem 3.1 Let [(w,E) be B 3-smooth and let def [(w")
-E [if ((w, 5)
w
Consider the iterates w given by equation
where g E arg max DEG
VL(g w')l
(10)
If
BVT-I then
23
3o²
min
E
max ggEG
VL(g w')ll
<
VT / [[((w)) ((((w))-L((0 [(w")]
(11)
t=0,
T
VT
Proof. First note that if L(w,E) {) is 3-smooth, then L(w) is also a B-smooth function, that is
[(z) - [(w) (VL(w), = w) S 2 w||?
(12)
Using equation I with x = ++ and w = g" w', together with equation 12 and the fact that the group action preserves loss, we have that
B
[(w S L(gg' ( + (VL(g' w)), w +1 g¹ .w0'
w g¹ u.''11
(13)
2
= =((w)) - t/e (VLig' w²), V w/, ,')) Bi
VLlg' ,t)?
(14)
Taking expectation conditioned w', we have that
Eg [(( < [(ww)) - nel V w)) 221l + 3n E, [[V w -
(15)
2
Now since [(w, {) is smooth, from Lemma A.1 above have that
E [ 1 & 28(((w) [(w")) + +2((((w)) inf L(w, (())
(16)
Using equation 16 with w g" w" we have that
Er [(ww++)) +1)] S [(w') - 1: 1:llV[(g' w²)ll
+ n2 ((('' w') L(w") + L(ww) - inf ((u,, .))) (17)
w
Using that Lig' w²) L(wr), taking full expectation and re-arranging terms gives
1 VL w || < (1 + 3² [(ww) L' - E [(w" -((] -- B330000 (18)
Now re-weighting trick introduced in Stich (2019). Let 04 be sequence such that ar (1 + 32nr² ay-1- Consequently if a_1 then az (1 + 32 .
MMltiplyyng by both
sides of equation 18 by au thus gives
a:T1EE [|V - ? & at-1 E [[(w') - c*] - agE [[(w(+1) - L*] + n202 (19)
Summing up from f = 0. T 1, and using telescopic cancellation, gives
T
anE VL w w < E [C(w)) - -* +
(20)
Let A = Et- Og Dividing both sides by A gives
E [[VV w w²)]13] <
Σ a:T VL(g' w²)12
Er= aml t=0
E [(ww) L + f2 _TT T-
am
(21)
14

---

Published as a conference paper at ICLR 2024
Finally, if 1) then
1 - (1 (1+ 11
Σ 1=
047
11
Σ (1
((+1)
(t+1)
1+ 12 (p-
7
(22)
1=
(1 12
(23)
B²m
To bound the term with the -T power, we use that
1og(2)
(1 + (1+ 2
S T.
log 32122)
To simplify the above expression we can use
S log(1 + x) S x,
for 2 -1
1 ++a x
thus
log(2)
1
5
< T.
log(1 + 11)2)
1211
Using the above we have that
T-1
Σ
1
1 1 12
Oiel 2
for T' 2
2,82n'
82
Using this lower bound equation 21 gives
1 881p
min t= T-
[[V[''' w w²)|l?
< 2,3 [(ww)) - L" + nns002 for T
Now note that
1 12
T
BT
3212
8 3TT-1
Thus finally setting 1 BVT-T gives the result equation 2
Proposition A.2. Assume that L: R" + F is strictly convex and twice continuously differentiable. Assume also that for any two points wa, wh E R such that L(w,) L(w)), there exists a E G such that wa = 9. Wb. At two points w1, W2 E TR",, max VL(g wr)ll ((w2))) then [(wi) < L(w2).
Proof. Let 5(x) = (w L(w) = x} be the level sets of L, and = (L(w) : w E R*) be the image of L. Since G acts transitively the level of L, max VL(g w)ll maX V (( To simplify notation,
define function F X R. F(x) max /V(()) Since VL(w) is continuously differentiable, the directional derivative of F is defined. Additionally, since L is continuous and its domain Th" is connected, its image X also connected. This means that for any W1, W2 E R" and min(C(w1). ((w2)) S S max(C(w;)) C(w2)), there exists Wa E IR" such that L(w3) y.
Next, we show that F(-) is strictly increasing by contradiction.
Suppose that L(w)) < L(w2) and F(L(wi)) 2 F(L(w2)). By the mean value theorem. there exists a w3 such that L(w1) L(w3) L(w2) and the directional derivative of F in the direction towards L(w2) is non-positive: )[(w))-(((w)) ((((w)) < 0. Let w e
arg max max [[ be point that has the largest gradient norm in S(C(ws)). Then w... /V cannot increase along the gradient direction. However, this means
vL(w)) Ou
= V [VL(w)))
< O.
(24)
15

---

Published as a conference paper at ICLR 2024
Since assumed that L is convex and L(w3) is not minimum (C(wī) ((w1)) )), have that (((w)) & O. Therefore, equation 24 contradicts with L being strictly convex, and we have F(C(w,)) F(C(w2)).
We have shown that C(wi) < ((w2) implies F([(w1)) < F(C(w2)). Taking the contraposi- tive and switching wi and w2. F(C(w)) F(L(w2)) implies [(wi) I f(w2). Equivalently, max w I max VLig w2)ll implies that ((w,) S ((wz).
Finally, since
max VL w1) = [ max VL(g w2)ll BEG
(25)
gEG
we have L(wi) S f(w22).
B TELEPORTATION AND NEWTON'S METHOD
Lemma B.1 (One step of Newton's Method). Let f(x be p-strongly convex and L-smooth func- tion, that is, we have a global lower bound on the Hessian given by
LI = f(x) μI,
(26)
Furthermore, if the Hessian is also G-Lipschitz
- f(yill S G yll
(27)
then Newton's 's method
k+1 = x² - X f(
has mixed linear and quadratic convergence according to
I. Au
|| x111 s
|| - 2²1 + 11
2 11 - x²11
(28)
2p
2/1
Proof.
x k+1
x x² Ao V
(( (( (VS(x*)) -V/('))
v.
x² A (( / V2s s(x x ))(* x')ds
x
(Mean value theorem)
V²f(Et)-1
x1))) (ar x² dds (V2f("*) M f(a* + =sxx"
(² (v2)(x²) - f(a + +s(x" - x²
(1 f(x + s(x* - x²))) (a* x= )ds
Let 8z
| x² x1 Taking norms we have that
8R41 <
v (a
((*))
V2 f(0 + - -(="))"1))
+11 A
| V f(x* + s(x" - -I*1)]]) ōgds
equation 27 equation
26
G
$||20*
x 1|2ds 111 d _
1. (' s||2* x I|ds
1 ==0
4 s=0
G 11 2p
x=1
112 + 11
Au
11 x²11
24
The assumptions on for this proof can be relaxed, since we only require the Hessian is Lipschitz and lower bounded in # ball around x²
16

---

Published as a conference paper at ICLR 2024
Proposition
3.2
Quadratic term in convergence rate). Let C be strictly convex and let wo R..
Let
w' E arg mnax
[[V[(w)]
subject to L(w) [(wo).
(29)
WER" 2
L(W()) + then there exists Ao such that
0 20 S max ( ((wo))
and one step of gradient descent with learning rate 7 > gives
W1 w' - 7 VVL(w')
w - ->>0 12 (
(30)
Consequently, letting w' go wo. and if S No
then under the assumptions of Lemma B.
have that
G 21
L I190
||w1 w0ll
9100 wo
x²11
>Ao
wi ==
2p
Proof. The Lagrangian associated to equation 29 is given by
L(w, A) = 2
X(C(wo) L(w)).
Taking the derivative in and setting it to zero gives
V [we, Ao) = = VL(()) - VO =0
(31)
Re-arranging we have that
VL(w) = No [(
If V * 0 then from the above we have that
||V[()]] VL VVV 0.
Since V²[(w)-1 is positive definite have that VL(w) V V' 0. and conse- quently Ao > 0. Finally from equation 31 we have that Ao is an eigenvalue of V²[(w) and thus must be smaller or equal to the largest eigenvalue of 4 ((w).
C Is ONE TELEPORTATION ENOUGH TO FIND THE OPTIMAL TRAJECTORY?
This section contains proofs for the results in Section 3.3 For readability, repeat some of the definitions here.
Consider the parameter space M = R Let V R" + TR" be vector field on R", where TR" denotes the associated tangent bundle. We will write v² Bhot using the component functions
V RdD + RR and coordinates w².
Let L MM R be smooth loss function. Let G be symmetry group of L, i.e. M and G. Let & be the set of all vector fields M. Let R r² a [(w) for all ac be the reverse gradient vector field. Let x {A a" 8 where o Bur, (M() and Σ, a² (w))((()) 0, Ve E EM} be the set vector fields orthogonal to R. Gs Lie group, the infinitesimal action of its Lie algebra defines set of vector fields X, X
L(g w)
Bwis
a x a'
A gradient flow is curve IR + M where the velocity is the value of R at each point, i.e. (t))) Ry(0) for all RR The Lie bracket [A, R] defines the derivative of R with respect to A. To simplify notation, we write ([W, (()) 0 for of vector fields W C X when ([A, ())) = 0 for all A E W.
Proposition 3.4, A point w E M is optimal in a set of vector fields W if and only if [A, R]((w) = for all A W.
17

---

Published as a conference paper at ICLR 2024
Proof. Note that AC a' 0L Bu
0. We have
0C Owv
ac 2 dw 2
[A, ARC RAC A
- -A
-f-
(32)
The result then follows from Definition 3.3
Proposition 3.5 Let W S & be set of vector fields that are orthogonal to the gradient of C. If
[A, ()) O for all A E W implies that R[[A, [l = for all A E W, then the gradient flow starting optimal point in W is optimal in W.
Proof. Consider the gradient flow y that starts at an optimal point in W. The derivative of [A, R along y is
d A, R]((((()) /OO -(0)) R(4)6(0) R[A, R)((+(1))
(33)
dt
Since y(0) an optimal point, [A, R)
for A W by Proposition 3.4 By assump- tion, if [A, R ())((()))))=)))))))))))))))))))) = 0 for all A E W. then R([A, (t)) = 0 for all A E W. Therefore, both the value and the derivative of [A, RIC stay 0 along 7. Since [A, (t)) for all E R.
is optimal in W.
To help check when Proposition 3.5 is satisfied, we provide an alternative form of R[A, R] ((w) under the assumption that [A, R|, R w
We will use the following lemmas in the proof.
Lemma C.1. For two vectors v, w E RR, if v- w 0 and w + O, then there exists an anti- symmetric matrix M E Roxa such that v Mw.
Proof. Let wD = ...,,..... 0JT E R". Consider a list of n 1 anti-symmetric matrices M, E R where
1,
if and + if + and otherwise
M
1.
(34)
0,
In matrix form, the M,'s are
0 -1
0
1
0
0
1 0
0
0
0
0
0
0
0
0
0
M₁ =
0
M2 =
1
Mn
0
0
0
0
0
0
0
1
0
0
(35)
Since Mt's are anti-symmetric, M, is orthogonal to wy. The norm of M,w0 Additionally, M;wo is orthogonal to M, for i j:
C.41
1.
(M, ( = e 11+1 = dij-
(36)
Denote mn = (xc E R xF wg 0} the orthogonal complement of wn- Then M, forms basis of wp Next, we extend this to an arbitrary w E "
Let v Tw Since v has norm 1, there exists orthogonal matrix It such that v Rwo. Let
M = R.M.R. R Then Md is anti-symmetric:
(RM,RT)T- RM, RT RMMRR.
(37)
follows that M/ is orthogonal to wò The norm of M/ is (RM, = |RM,woll = |MWwwlll = 1. Additionally, M/ is orthogonal to M2 w for # j:
[
(M/w)" (MM (RM,RR Rwo) RR Rwo)
R" RM. RD R., Rd Rwo
= wW M,w
dij-
(38)
18

---

Published as a conference paper at ICLR 2024
Therefore, Mr spans w w This means that any vector w can be written linear
combination of M, w. That is, there exists "2.... kn E R, such that e = E, (a To find the anti-symmetric M that takes to v, note that
Σ K = -1001122 EKiMi w.
v
(39)
Since the sum of anti-symmetric matrices is anti- symmetremeeii,, and the product of an anti-symmetric matrix and a scalar is also anti-symmetric ||w||2 Σ, k;M, is anti-symmetric.
n
Lemma C.2. Let v E R" be nonzero vector. Then the two sets {Mu M E R "0X",M"* = -M and (w R" w² w0'h 0} are equal.
Proof. Let A = {Mu M E R " = M-13 and B = {w E R" w v (Mo)" for all anti-symmetric M, every element in B. By Lemma C.I every element in B is in A. Therefore A B.
0}. Since
Let S ((M Dar
XJ M Roxn M -M} be the of vector fields constructed by multiplying the gradient by an anti-symmetric matrix. Recall that R Jwy The the reverse
gradient vector field, and 3 = {a
a² (un) @C(w) O, Vur M) is the set of all vector we have S Dw"
Ry. Therefore, point w is an optimal point in if and only if w is an optimal point in Xi.
fields orthogonal to R. From Lemma C.2
We are now ready to prove the following proposition, which provides another way to check the condition Proposition 3.5
Proposition 3.6 If at all optimal points in S,
Oc dc L dc Mi Own Owa Sw Ows Ow
(40)
for all anti-symmetric matrix M E Ruxn, then the gradient flow starting at an optimal point in S is optimal in S.
Proof. Expanding R[A, R, we have
Oc
R[A R R
A
0)
Ow
0
Oc
r
kuo
Owg
Dwo
OC
2 Oc Ous Ow:
kk
a
Fw
Dust
Ow
0
0
Oc
OL
OL δ dc
=
a
Dw: +
Dw:
Ouri Own
Ow, Ous Ow
θ
82
dL
2
a
k
Tw,Ows Ow
Da
JLL AL
a a
Tw,wwi ow )) OC dc
J2C Oc
2
Owk Tw,dwi Owi
Dw
ac da JL Oc + +2 Oc d
2 Own Own dw,dwi Owi
(41)
Ows
Ow
Ow,Owi Ow
Assume that w is an optimal point in S. By Lemma C.2 w is also an optimal point in ET. By a Oc 8²6 Therefore, Sw, 8w X ac
Lemma C.4 in Zhao et al. (2022), ac is eigenvector of for some C. Additionally, a Mi JC. and Ba" both terms in equation 41
Md. Dw. On.
We are now ready simplify
Dwa
19

---

Published as a conference paper at ICLR 2024
For the first term in equation 41
0L dal OL ac
0L
8C
O ac Ow
Ma
Owy Ow
Sws Ow
Owk
Ow dw,dww
Ow, Owk Owy )
dc
O2C dc
Må
Ow,dwi Ow
OL
aL
Må
A1 Owa
A2 Dwo
OC
dc.
A2
owa Owv
(42)
The last equality holds because M is anti-symmetrie For the second term in equation
41 0CC ac
Oc a² 0
OC Owk
OL Ow Ow; dw
L Oc
OL
a²
+
Owk Owk Oworuus Ow
Ow,dus Ow-kuu
dL Mr,
UL
L dc
GL 82
awk
dwa @ww
Ow:
Ow,wws Ow Oc Oc
OC Oc
O Oc
MMa dwk owa
> >2 Ow, Owi
Ow
OC Oc OL Oc
Må Own owa Bwkkww;Owl Ow
(43)
In summary.
R[A, R 2M OC OC SC Oc
(44)
Own Ow Ow Ow,ww. Ow"
Since we assumed that [A, R ((w) = O, when R[A, R [(w) = for all S, the gradient flow starting at optimal point in S is optimal in S.
Proposition
C.3.
8
ac
8 Oww 3ws
ac
holds for all i,k,j,a, then
Da Ow Ow
Mí BC aC Dwn Ows Tw Oht
DC
0 holds for all anti-symmetric matrices M E Rrxx..
Oy
Proof. If
c BC
8 c Ow "m""0" am
ac Dws
for all i, k, a, then
Ow Owt Du
dL OC L
OC
M Own Owa Dw dww dw
Md. OC OL SL
OL
M OL OL 2L
OL
Σ Lkaas
Own Owa Ow Ows Ow
Owk Dwa Ow Ow;ūwi Ow
Md OL Oc 8
OC
Mn OL ac
L
OC
Σ ikaasi
+
Σ
Own Sw. Ow Ow,dri
Ow
Own dr, Ow Ow, Ow
Mi OL Oc 0L Oc
-Ma OL Oc
O Oc
= Σ i.kaa
Σ
Owp dwa du Ow, Ow
Own Ow, Ow iw Ow
Mi. ac Oc (
CC
ac
Owy Ow Own Sw Ows
Owy Ow
i,k,a
(45) where the first equality uses that the diagonal of an anti-symmetric matrix is O, the second equality swaps a and j in the second term, the third equality uses that M is anti-symmetric.
= 0,
n
20

---

Published as a conference paper at ICLR 2024
Example (Quadratic function) Consider the quadratic function [(w) w Au + +b w + c, where A RTx= is symmetric, b, R", and w E IR. Two examples of quadratic functions the ellipse Le (W,, W2) = 2 [wf + x and the Booth function C3 (W1, w2) = (wi + 2w' - 7) + (2w1 + 22 5)2. Since the third derivative of is O, one teleportation guarantees optimal trajectory.
D GROUP ACTIONS AND CURVES ON MINIMA
D.1 GROUP ACTIONS FOR MLP
Consider a multi-layer neural network with elementwise activation function o. The output of the mh layer is fimm = o Wm)m-1), where Wm E RRd Xdm the weight, hm-1 E Rdm 18k the output of the m th layer, and ho E Rdox is the data.
Assuming that (gm www.Wlm----. 2 is invertible, for 9m I (R), the following transformation is a loss-preserving group action:
{ Wno (Wm-1hm-2) a (9m Wm- 22m-m))
k= m
9m W
9m Wm
k 1=m-8
(46)
Wx
1 E m
Usually, the assumption does not hold (Zhao et al. 2023). Hence the above transformation may not preserve loss be valid group action. Nevertheless, we observe in practice that the change in the loss value is often small after such transformations on parameters. We therefore refer to equation 46) as approximate symmetry and adopt it in the teleportation algorithm. Due to the possibility that a (9m Wm hm is not invertible, we use pseudoinverses in implementations.
D.2 CURVATURE
The curvature of a curve y I. R- R" is k(t)
( b-'ll where T(t) ('))
(+) is the unit tangent
vector. The curvature can be written as a function of > and >" (Aléssio 2012 Shelekhov 202
<(t) [1 (7' - 20
(47)
1//13
D.3
THE DERIVATIVE OF CURVATURE
To compute the derivative of k(t), we first list the derivatives of a few commonly used terms:
di 11211122
d
(11
2
+ 1/2 + 7/3 + ...)= = 221n''' + 2767 + 27'5% +... = ...==//- nn
,2
dl
d
112
d
.12
+ +2
+ 13
.12
+ +..) 22'""/"" + 27/222" !M + 27/2"2"" 73 + ... 2- 7
dt
dt
d (1' -²))= (N + 2//22 + 1o. 7/7"" + 7''''' +... 117"112 2 + o/ 71
d
(48)
The derivatives of the numerator and denominator of k: are:
di [1h''P/1>>'|> - (7/ """)
d
1 [|>'/1 221--- d 1
di
[I (7' 27----0)1-
2
[I _ 11>0 12 + || d _ 111111 2('' -7") d
(r' 77"))
dt
2 [|>>/11117111 - -
-
77) -27' 2 <1' + - [m"
-| 110'' (1/ (1-""")) '00 -7"))
(49)
21

---

Published as a conference paper at ICLR 2024
and
d dt 111l
di (11> (1b>'11))= = 3 (11
d
d - (11'11))) = 31/1/(1// -"""(("'"""") (50)
3
dt
Using the derivatives above, the derivative of x: is
d de [I|''1 (r' >")) 3
1M [1|>"|1311>" (7/ //) [ [111
x'(( (t)
1
||>>/>/>>>> " " [1|'']] 711) - (7' -0") {"")] |1r'ls -11>'11311>11 112 (6'
3////11 7")
11
[11>///1>>" -7"-7") - |1
-1")) (7' -1") 1 - |1 111>"1111>~1111 - 3(7 ,1))
||''|1
(51)
D.4 THE DERIVATIVES OF CURVES ON MINIMA
Consider the curve YM IR x IR T where M E Tie(G) and
(t, w) exp (tM)) w.
(52)
In this section, derive y", and 7dh which needed to compute the curvature k(() and its derivative x² (t). We are interested in is and w' at w, or equivalently, at t = 0. To find the derivatives of y at = O, write the group action in the following form:
(t)
(53)
n!!
===0
By the uniqueness of Taylor polynomial, the derivatives are (0) = f(n). In the rest of this subsection, we expand the group action to find f(n).
Consider two consecutive layers Ux(VX) in neural network, where U R V E A are weights, X h is the output from the previous layer, and an elementwise activation function. Choosing G = (L,(R)), one group action that leaves the output of these two layers unchanged is:
9 (U, = (g U.g. V.g X) = o
(54)
Let
g = exp(tM) = Σ kt (tM)*
(55)
where M Lie(G) is in the Lie algebra of G. The action of g yields
g. gU,U, V, X) = (U
(56)
Next, we expand 7()) = g. (U, V). The Taylor expansion for g U is
V Uexp(-IM) = U Σ K! (-LM)*
k=0
= -U-UU +
UM
UM + 0(1))
(57)
21
31
22

---

Published as a conference paper at ICLR 2024
The Taylor expansion for g V is
(tM)
O(VX)) x-1
k!
((M)) x-1
1
(VVX)
(')
(o(VX)) +
((N))
x-1
k!
0Ai
A=a(VX)
L (o)1(A)
=V +
(tM))o(VX))
X
(58)
0AU
A=o(VX)
where @ denotes element-wise product: (A C B)mn Amn Bmn, and the superscript 9 denotes
elementwise power: (A47) )mn (Amn)'. The Taylor expansion of each element individually,
because o is element-wise.
Since our goal is to find the first derivatives of y, we are only interested in the terms up t Letting
1 k!
M2
M²
aaMe M
(tM)*
tM t
+ +
0(t*)
(59)
2
6
and considering only the 1,2.3 terms, we have
a- (VXx)(M))(VXX))
(LM + 12 M²
13 M3
=V+
Σ
2 +
)o(V))
x x-1 + 00(4)
6
0A
j=1
A=o(VX)
( 12 M2
+ t M3
Oa -1 (A)
+V
Jo(VX)
2
6
0A
A=o(VX)
(tM 22
22
M3 3
02
+
6 )o(VX)
2
0A2
A=o(VX)
M
M
86
(A)
+
(tM +
+
3
)o(VX))
x-1 +-1+
2
0A3
A=o(VX)
1
=V +t
(Mo(VX))
x-1
x(VX)
(MPP(VX))
d"(VX)
+
o(VX)
2(Mo(V)))
o((XJ33
x-1
+ (INPb(VX)) 2 x(VX) 6(Mo(V O (M
1
"VV
d(VX)3
+ 0 (AA
x-1
0.33
A=o(VX)
+ +O(()).
(60)
Matching terms in equation 57 and equation 60 with equation 53 7" and n" This allows us to compute the curvature and its derivative using equation 47 and equation 51
have the expressions for 7/
23

---

Published as a conference paper at ICLR 2024
E SHARPNESS, CURVATURES, AND THEIR RELATION TO GENERALIZATION
E. ALTERNATIVE DEFINITIONS OF SHARPNESS
common definition of flat minimum is based on the number of eigenvalues of the Hessian which are small. Minimizers with a large number of large eigenvalues tend to have worse generalization ability (Keskar et al. 2017). Let X (w) the 7th largest eigenvalue of the Hessian the loss
function evaluated at w. We can quantify the notion of sharpness by the number of eigenvalues larger than a threshold e RR>
Φ ur, E) ((),(H)( X, > et|
(61)
A related sharpness metric the logarithm of the product of the lt largest eigenvalues (Wu al. 2017).
22 k) log (H))
(62)
Both metrics require computing the eigenvalues of the Hessian. As result, optimizing on these metrics during teleportation is prohibitively expensive. Hence, in this paper we use the average change in loss averaged over random directions (Ø) as objective in generalization experiments.
E.2 MORE INTUITION ON CURVATURES AND GENERALIZATION
E.2.1 EXAMPLE: CURVATURE AFFECTS AVERAGE DISPLACEMENT OF MINIMA
Consider an optimization problem with two variables w1, W2 E R. Assume that the minimum is a curve 7 + R2 in the -iodeeensional parameter space. For point wy on 7 we estimate its generalization ability by computing the expected distance between and the new minimum obtained by shifting 7.
We consider the following two curves as examples:
71 :R + IR2, t + (t, k1,)
22 00, 2Tl
R2, + (log cos(O), hy sin(f) + k2),
(63)
with ki, 62 R The curve 71 a parabola with curvature Kı 2k wo (0, The curve 72 is a circle, with curvature K2 Tg
1 at Wg. Note that 71 is the only polynomial approximation with
integer power (t(t)) (t,, ke]t"),, Z*) where the eurvature at we depends k. When 1, the value of wg is undefined. When n = 1, the first derivative at we is undefined. When n > 2, k< k(w)) O.
Assume that a distribution shift in data causes y to shift by a distance r, and that the direction of the shift is chosen uniformly at random over all possible directions. Viewing from the perspective of the curve, this is equivalent to shifting wo by distance
The distance between point and curve is
dist (w, 7) min |/ w//2-
(64)
Let S, be the circle centered at the origin with radius r. The expected distance between the old solution wy and shifted curve is
Ss. dist(w, rs Jo dist((r cos 0, r rin 0), d0
Emes, [dist(w, (.))) 7)]
(65)
Ss. ds
rd0
In the limit of zero curvature, y is straight line x(t) = (t, O). In this case, the expected distance is E (dis((w,)) 7)]
12 Ir sin rd0
2r
~ 0.637r.
(66)
2xr
24

---

Published as a conference paper at ICLR 2024
Figure 7 [))(C) shows that the expected distance's dependence Using both curves " and Y2. the generalization ability of wy depends on the curvature at Wg. However, the type of dependence is affected by the type of curve used. In other words, the curvatures points around wo affect how the curvature at we affects generalization. Therefore, from these results alone, one cannot deduce whether minima with sharper curvatures generalize better worse. To find more definitive relationship between curvature and generalization, further investigation on the type of curves on the minimum required.
We emphasize that this example only serves as an intuition for connecting curvature to generaliza- tion. As a future direction, would be interesting to consider different families of parametric curves, higher dimensional parameter spaces, and deforming in addition to shifting the minima.
(a)
(b)
(c)
W2
y
0.775
disttw. Vo
0.725
0.62
r=0.1
1=1.0
0.675
r=0.02
r=0.2
0=001
W1
r=0.00
Sy
0
20
40
60
80
=m 2k1
kg
Figure 7: (a) Illustraiion of the parameter space, the minimum (7), and all shifts with distance r (S,). (b) Expected distance between we and the minimum function of is, for quadratic approximation 7 (c) Expected distance between wh and the new minimum as function of x, for constant curvature approximation 722 The expected distance is scaled by so that the curves can be plotted together.
E.2.2 HIGHER DIMENSIONS
Figure visualizes curve obtained from 2D minimum. However, it not immediately clear what curves Took like on a higher-dimensions minimum. A possible way to extend previous analysis is consider sectional curvatures.
minima
Figure 8: Left: a 2D minima in a 3D parameter space. Right: a 2D subspace of the parameter space and curve the minima (the intersection of the minima and the subspace).
E.3 COMPUTING CORRELATION TO GENERALIZATION
We generate the 100 different models used in Section 4.3 by training randomly initialized models. For all three datasets (MNIST, FashionMNIST, and CIFAR-IO) -0), we train on 50,000 samples and test different set of 10,000 samples. The labels for classification tasks belongs to of 10 classes.
For a batch of flattened input data E 2d and labels Y E I2200 the loss function is L(W,, W,, W3 X. Rhaxhı, W, E Ihixd xd are the weight matrices, and a is the LeakyReLU activation with slope coefficient 0.1. For MNIST and Fashion-MNIST, d 282. hy 16, and h 10. For CIFAR-10.
CrossEntropy (Wyo(WaooWW
Y). where Wa R 100hg W2
25

---

Published as a conference paper at ICLR 2024
d 32 3, h1 128, and h2 32. The learning rate for stochastic gradient descent is 0.01 for MNIST and Fashion-MNIST, and 0.02 for CIFAR-10. We train each model using mini- mini-batches of size 20 for 40 epochs.
When computing the sharpness Ø. we choose the displacement list that gives the highest correla- tion. The displacements used in this paper are T 0.001, 0.011, 0.021, .... 0.191 for MNIST, and = 0.001, 0.011,0.021.... 0.191 for Fashion-MNIST and CIFAR-10 10. We evaluate the change in loss over [D] 200 random directions. For curvature v,, average over 1 curves generated by random Lie algebras (invertible matrices in this case).
Figure and 10 visualizes the correlation result in Table
Each point represents model.
(b)
(c)
0.430
Carr=0.704
Corr= 0.790
1.62-
validation loss
0
0.19
validation
0.420
1.50-
0.18
1.56
1.54-
0.410
1.52-
0.405
1.50-
0.15
0.400
0.0005
0.0006
0.0007
0.00144
0.00153
0.00162
0,0057
n.0060
0.0063
Figure 9: Correlation between sharpness and validation loss MNIST (left), Sishion-MNIST (mid- dle). and CIFAR-10 (right). Sharpness and generalization are strongly correlated.
(a)
(b)
(c)
0.430
1.64
Corre-de 0.050
Corr 0.233
1.62
validation loss
0.19
о
0.420
1.58-
0.18
1 Palidation
0.415
5 1.56-
1.54
0.16
0.405
1.50
1.48
0.15
0.400
0.001
0.002
0.003
0.001
0,003
Figure 10: Correlation between curvature and validation loss on MNIST (left), Fashion-MNIS (middle), and CIFAR-10 (right). There is a weak negative correlation in all three datasets.
E.4 ADDITIONAL DETAILS FOR GENERALIZATION EXPERIMENTS
Algorithm shows an example on how to perform teleportation with an MLP
Algorithm 2 Changing curvature using teleportation
Input: loss function L(w), parameters before teleportation Wo, teleportation learning rate Meeepport, number of teleportation steps Ritelportt
Output: parameters after teleportation Wnydtgorr"
for f. 0 to nitelpprtt
do
initialize T 0axh
set w/ = (1ixh + T Wt compute grad d (
di
set T; Mellppot grad set Wt-1 = (1 + Tí) Wt end for
Return Wnseesport
26

---

AdaGrad. In AdaGrad, the rate of change in loss Published as a conference paper at ICLR 2024
On CIFAR-10, run SGD using the same three-layer architecture as in Section E.3 but with smaller hidden size h1 = 32 and h2 = 10. At epoch 20 which is close to convergence, we teleport using 5 batches of data, each of size 2000. During each teleportation for Ø, we perform 10 gradient ascent (or descent) steps on the group element. During each teleportation for th we perform gradient ascent (or descent) step the group element. The learning rate for the optimization group elements is 5 x 10-2
To investigate how teleportation affects generalization for other optimizers, repeat the same experiment but replace SGD with AdaGrad. Figure [T shows the training curve of AdaGrad on CIFAR-10. averaged across runs. 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 points with smaller curvatures increases the gap between training and validation loss.
2.0
2.0
Gra teleport(dccrese 4) teleportiincrease
AdaGrad deleport/decrease u teleporttincrease e)
Loss
Loss
1.7
1.7
1.4
1.4
20 Epoch
40
0
20 Epoch
40
Changing sharpness (left) or curvature (right) using teleportation and its effect on gener- alizability of AdaGrad solutions on CIFAR-10. Solid line represents average test loss, and dashed line represent average training loss.
Figure 11:
F INTEGRATING TELEPORTATION WITH OTHER GRADIENT-BASED ALGORITHMS
DIFFERENT METHODS OF INTEGRATING TELEPORTATION WITH MOMENTUM AND ADAGRAD
Setup. We test teleportation with various algorithms using the a 3-layer neural network and mean square error: minw;, ww.W, || Wa (W, X)) with data X R5x4, target Y E 8x4 and weight matrices W3 E _xx7, W2 E 7x6 and W1 E 16x5 The activation function is LeakyReLU with slope coefficient 0. Each element in the weight matrices is initialized uniformly random over [0, 1]. Data X, Y are randomly generated also from [0, 1].
Momentum.
We compare three strategies of integrating teleportation with momentum: teleporting both parameters and momentum, teleporting parameters but momentum, and reset momentum to after teleportation. In each run, we teleport once at epoch 5. Each strategy is repeated 5 times.
The training curves of teleporting momentum in different ways similar (Figure 12p), possibly because the momentum accumulated is small compared to the gradient right after teleportations. All methods of teleporting momentum improves convergence, which means teleportation works well with momentum.
- dC(w) aC dw dt Ow dt
(67)
where 7) E IR is the learning rate, and | is the Mahalanobis norm with A (el + diag(G (+41)) Previously, we optimize |V[l/2 in teleportation. We compare that to optimiz- ing |VL]|A Since the magnitude of A is different than 1, different learning rate for the gradient
27

---

Published as a conference paper at ICLR 2024
ascent in teleportation is required. We choose the largest learning rate (with two significant figures) that does not lead to divergence. The teleportation learning rates used are 1.2 10 for objective maxg || 2 and 7.5 10 for objective max, VC]
Teleporting using the group element that optimizes V has a slight advantage (Figure 12p). Similar to the observations in Zhao al. (2022). teleportation can be integrated into adaptive gradi-
ent descents.
(a)
(b)
Momentum
AdaGrad
10a
telepartinone)
10
AdaGrad tolepprtt
Momentum- teleport(teleport)
AdaGrad+ Ieleport_
teleport(resett)
Loss
102
Loss
10
10
10
100
10
100
200
300
100
200
300
Epoch
Epoch
Figure 12: Comparison of different methods of integrating teleportation with momentum and Ada- Grad.
F.2 ADDITIONAL DETAILS FOR EXPERIMENTS ON MNIST
We use three-layer model and cross-entropy loss for classification with minibatches of size 20. For batch of flattened input data X 282 and labels Y E R220 the loss func- tion is C(W,, W2, W,,XYY) = CrossEntropy (W)) Y), where W3 E R00x10 W2 16 W1 E 16
the weight matrices, and is the LeakyReLU activation with slope coefficient 0.1. The learning rates are 10 for AdaGrad, and 5 x 10 for SGD with momen- tum. RMSProp, and Adam. The learning rate for optimizing the group element in teleportation is 5 x 10-2, and we perform 10 gradient ascent steps when teleporting using each mini-batch. We use 50,000 samples from training set for training, and 10,000 samples in the test set for testing.
(a)
(b)
(d)
0.5
0.4
Meggrsdeiiiisor teet
Adbgrad test
Adagred Patt
Adbgrasdesslsoort tain
momentum momentum test momenntm wntpett tran patrrntumeeeeeppie. teet
Loss
0.8
RM5prup test
#MSyespennipprr trah
RMSprop
Loss
0.7
0.9
Adbmn test
Arbin sport train
Artim uport teit
Adam
0.5
0.3
0.4
0.2
0.2
10
00
70
10
10
Time
Time
Time
Time (5)
Figure 13: Runtime comparison integrating teleportation into various algorithms. Solid line represents average training loss, and dashed line represents average test loss. Shaded areas standard deviation of the test loss across runs. The plots look almost identical to Figure 5 indicating that the cost of teleportation is negligible compared to gradient descents.
28


Reviews:
[{"review_id": "BS8hzW8xDX", "reviewer": "Reviewer_bsXX", "summary": "The paper investigates how teleportation, i.e., applying a loss-invariant group action to the parameter space can improve (i) optimization speed and (ii) generalization in deep learning. For (i), the paper derives an upper bound for the gradient norm, which implies that SGD iterates converge to a basin of stationary points, from which only other stationary points are reachable via teleportation. They further show that SGD with teleportation can behave similar to Newton's method and provide a necessary condition on when one teleportation step is sufficent to accelerate optimization. They extend teleportation to commonly used optimizers and experimentally show that a teleportation step in the first epoch improves the convergence rate. Lastly, they incorporate teleportation into a meta-learned optimizer and show that learning the group element in teleportation improves the convergence rate of (meta-learned) gradient descent. For (ii), the paper introduces a novel measure for generalization based on the curvature of minima, and empirically shows that teleporting to points which decrease sharpness and increase the curvature of minima correlates with an improvement in generalization.", "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?\n2. How did you decide on 10 and 1 gradient ascent steps, respectively, for the experiments?\n3. Do you have any (preliminary) results on how teleportation affects generalization for other optimizers (e.g., AdaGrad, Adam, etc.)?\n4. 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.\n\n[1] https://openreview.net/pdf?id=VZp9X410D3", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "The paper provides novel results on exploiting parameter symmetries in the context of optimization and generalization in deep learning. For optimization, the theoretical results improve on existing work, while the paper appears to be the first to investigate teleportation with respect to generalization. The presented results also have promise to be of practical relevance, since the computational overhead of the teleportation step appears to be negligible in the experiments.", "weaknesses": "* **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.\n* **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.\n\nPlease find some minor remarks below:\n\n* p. 3: we provide theoretical analysis of teleportation -> we provide a theoretical analysis of teleportation\n* p. 3: that maximizes the magnitude of gradient -> that maximizes the magnitude of the gradient \n* p. 3: the iterates equation 4 -> the iterates in equation 4\n* p. 3, Theorem 3.1: I assume $\\theta$ should be $w$\n* Proposition 3.2: I assume $f$ should be $\\mathcal{L}$\n* p. 5: To simplify notations -> To simplify notation\n* p. 7: at the 20 epoch -> at the 20th epoch/at epoch 20\n* p. 7: teleporting to sharper point -> teleporting to sharper points\n* Lemma A.1: eq. (19) LHS: $\\xi$ seems to be missing in $\\mathcal{L}$, also 2 lines below", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oypICyl8ij", "reviewer": "Reviewer_BdD6", "summary": "Typical deep learning models have natural symmetries on their weight parameters that leave their output unchanged. The study of parameter symmetries in deep learning focuses on the effects of these symmetries in terms of optimization and generalization performance. The present work contributes to this burgeoning field by first offering an analysis of \"teleported\" SGD, an algorithm proposed by [1] which uses the symmetries to move the current iterate to one with largest (under the group action) gradient. Their first contribution is an analysis which shows that teleported SGD yields improved guarantees for SGD for smooth (possibly non-convex) functions. It was shown in [1] that teleportation (in the GD rather than SGD setting) is equivalent to a step of Newton's method, and they accordingly give an estimate of its contraction rate that demonstrates quadratic convergence as would be expected from a 2nd order method. Their final theoretical results are sufficient conditions for one teleportation step to be optimal for all times. They then provide empirical studies of the effects of teleportation for increasing/decreasing the sharpness and curvature of minima. Finally, the authors consider teleported variants of other standard optimization algorithms (Adagrad, SGD with momentenum, RMS prop, and Adam), and propose a method for meta-learning the teleportation.\n\n[1] Symmetry Teleportation for Accelerated Optimization, by Zhao, Dehmamy, Walters, and Yu 2022", "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.\n- I am a bit confused about Figure 4: the paper text says \"teleporting to points with smaller curvatures helps find a minimum with lower validation loss, while teleporting to points with larger curvatures has the opposite effect\". And this relationship, that less curvature is associated with better generalization performance, is also present in Table 1. However, Figure 4 seems to show the opposite phenomenon: the teleport(increase $\\psi$) test loss is better than the teleport(decrease $\\psi$) test loss. Am I mis-reading this plot? Could you please clarify this point, ideally in the paper text as well?\n- \"from Lemma A.1 below\" before equation 26 and on page 13 should instead read \"from Lemma A.1 above\"", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "This paper is written in a very clear and engaging style, and was a pleasure to read. The exploration of sufficient conditions for one-teleportation to be enough (section 3.3, especially Prop. 3.4) and the effects of teleporting for sharpness/curvature (section 4), plus their computationally feasible proxies $\\phi$ and $\\psi$ are, to my understanding, significant and novel. Finally, their empirical results are clear and original. For example, Figure 5 on the effects of teleportation on other first-order algorithms is particularly convincing. And their use of Pearson correlation in Table 1 to estimate the effects of curvature and sharpness is impressive.", "weaknesses": "Overall, I really like the paper. One weakness, however, is that some of their theoretical results are not especially novel. For example their results on teleported SGD (Theorem 3.1) and the Newton steps (Prop. 3.2) seems to be minor modifications of standard proof techniques. And the (very interesting) fact that teleported SGD is equivalent to a Newton iteration was already observed by previous work [1]. However, this is no way inclines me to reject the paper -- I think its clarity and other contributions are more than enough to merit acceptance.\n\n[1] Symmetry Teleportation for Accelerated Optimization, by Zhao, Dehmamy, Walters, and Yu 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ReiN7fNTHo", "reviewer": "Reviewer_Ewvi", "summary": "Teleportation is the transformation of the parameters in the parameters such that the loss is unchanged. This work theoretically shows that teleportation has accelerating effects on the convergence rate of SGD for non-convex loss. Furthermore, through experiments, they show that teleportation can also improve generalization by moving the parameters to a flatter region.", "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. \n\n- Does it matter at which epoch we perform the teleportation? If it does, how do we pick the best epoch?\n\n- 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The authors are able to prove a stronger convergence guarantee for SGD by using transportation. Instead of having a guarantee for a single stationary point like the classic SGD, they show that SGD with teleportation has convergence guarantees for a set of stationary points in group $G$. The intuition on why SGD with teleportation has accelerating effects is well explained by connecting it with second-order algorithms.\n\n- The paper shows that one teleportation might be enough which makes it feasible to do in practice.\n\n- The curvature of minima seems like a pretty interesting way to understand the generalization ability of the stationary points.\n\n- The paper is well-written overall.", "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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9k9V0OEx4b", "reviewer": "Reviewer_dBnH", "summary": "This paper analyzes the effect of teleportation technique in improving training and generalization performance of deep learning tasks. The authors first analyze the convergence property of teleportation in SGD and Newton's method. The effect of teleportation in improving generalization is justified using the concept of sharpness and curvature. Some numerical experiments verify the findings.", "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 teleportaion?\n2. 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.\n\n**References**\n\n[1] Zhao, B., Dehmamy, N., Walters, R., & Yu, R. (2022). Symmetry teleportation for accelerated optimization. *Advances in Neural Information Processing Systems*, *35*, 16679-16690.\n\n[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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The paper proposes a convergence analysis for SGD with teleportation, and shows that it can improve the generalization of the solution. This demonstrates that teleportation has the potential to become a standard tool in modern deep learning training tasks.", "weaknesses": "1. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)