Original: Symmetry Teleportation for Accelerated Optimization
Revised: Symmetry Teleportation for Accelerated Optimization

Original: Abstract
Revised: Abstract

Original: Existing gradient-based optimization methods update parameters locally, in a direction that minimizes the loss function.
Revised: Existing gradient-based optimization methods update parameters locally, in a direction that minimizes the loss function.

Original: We study a different approach, symmetry teleportation, that allows parameters to travel a large distance on the loss level set, in order to improve the convergence speed in subsequent steps.
Revised: We study a different approach, symmetry teleportation, that allows parameters to travel a large distance on the loss level set, in order to improve the convergence speed in subsequent steps.

Original: Teleportation exploits symmetries in the loss landscape of optimization problems.
Revised: Teleportation exploits symmetries in the loss landscape of optimization problems.

Original: We derive loss-invariant group actions for test functions in optimization and multi-layer neural networks, and prove a necessary condition for teleportation to improve convergence rate.
Revised: We derive loss-invariant group actions for test functions in optimization and multi-layer neural networks, and prove a necessary condition for teleportation to improve convergence rate.

Original: We also show that our algorithm is closely related to second order methods.
Revised: We also show that our algorithm is closely related to second order methods.

Original: Experimentally, we show that teleportation improves the convergence speed of gradient descent and AdaGrad for several optimization problems including test functions, multi-layer regressions, and MNIST classification
Revised: Experimentally, we show that teleportation improves the convergence speed of gradient descent and AdaGrad for several optimization problems including test functions, multi-layer regressions, and MNIST classiﬁcation.

Revised: Our code is available at https://github.com/Rose-STL-Lab/Symmetry-Teleportation .

Original:  Introduction
Revised:  Introduction

Original: Consider the optimization problem of finding argmin w L ( w ) , where L is a loss function and w are the parameters.
Revised: Consider the optimization problem of ﬁnding argmin w L ( w ) , where L is a loss function and w are the parameters.

Original: While finding global minima of L ( w ) is hard for non-convex problems, we can use gradient descent (GD) to find local minima.
Revised: While ﬁnding global minima of L ( w ) is hard for non-convex problems, we can use gradient descent (GD) to ﬁnd local minima.

Original: In GD we apply the following update rule at every step:
Revised: In GD we apply the following update rule at every step:

Original: [Equation]
Revised: [Equation]

Original: where η is the learning rate.
Revised: where ⌘ is the learning rate.

Original: Gradient descent is a first-order method that uses only gradient information.
Revised: Gradient descent is a ﬁrst-order method that uses only gradient information.

Original: It is easy to compute but suffers from slow convergence.
Revised: It is easy to compute but suffers from slow convergence.

Original: Second-order methods such as
Revised: Second-order methods such as

Original: Newton’s method use the second derivative to account for the landscape geometry.
Revised: Newton’s method use the second derivative to account for the landscape geometry.

Original: These methods enjoy faster convergence, but calculating the second derivative (Hessian) can be computationally expensive over high dimensional spaces (Hazan, 2019).
Revised: These methods enjoy faster convergence, but calculating the second derivative (Hessian) can be computationally expensive over high dimensional spaces (Hazan, 2019).

Original: In this paper, we propose a new optimization method based on parameter space symmetries, which are group actions on the parameter space that leave the loss unchanged.
Revised: In this paper, we propose a new optimization method based on parameter space symmetries, which are group actions on the parameter space that leave the loss unchanged.

Original: Our algorithm takes advantage of higher-order landscape geometry but uses only gradient information in most steps, thus avoiding the computational cost of second-order methods.
Revised: Our algorithm takes advantage of higher-order landscape geometry but uses only gradient information in most steps, thus avoiding the computational cost of second-order methods.

Original: Specifically, we look beyond local optimization and ask: what if we allow the parameters to make a big jump once in a while?
Revised: Speciﬁcally, we look beyond local optimization and ask: what if we allow the parameters to make a big jump once in a while?

Original: As shown in Figure 1, during optimization, we teleport the parameters to another point with steeper gradients using a loss-invariant symmetry transformation.
Revised: As shown in Figure 1, during optimization, we teleport the parameters to another point with steeper gradients using a loss-invariant symmetry transformation.

Original: After teleportation, the loss stays the same, but the gradient and hence the rate of loss decay changes.
Revised: After teleportation, the loss stays the same, but the gradient and hence the rate of loss decay changes.

Original: The increased magnitude of the gradient can reduce the number of steps required to converge, leading to acceleration of the gradient descent.
Revised: The increased magnitude of the gradient can reduce the number of steps required to converge, leading to acceleration of the gradient descent.

Original: The locality of gradient descent is reflected in its formulation in terms of a proximal term; our method circumvents this locality by teleporting to new locations with the same loss.
Revised: The locality of gradient descent is reﬂected in its formulation in terms of a proximal term; our method circumvents this locality by teleporting to new locations with the same loss.

Original: A step of GD is equivalent to the following proximal mapping (Combettes and Pesquet, 2011).
Revised: A step of GD is equivalent to the following proximal mapping (Combettes and Pesquet, 2011).

Original: Let ⟨· , ·⟩ denote the inner product, we have
Revised: Let h · , · i denote the inner product, we have

Original: [Equation]
Revised: [Equation]

Original: The term 12 ∥ w − w t ∥ 2 is the proximal term that keeps w t +1 close to w t .
Revised: The term 12 k w  w t k 2 is the proximal term that keeps w t +1 close to w t .

Original: Adaptive gradient methods define the proximal term using the Mahalanobis distance ∥ w − w t ∥ 2 G − 1 / 2 to account for landscape geometry.
Revised: Adaptive gradient methods deﬁne the proximal term using the Mahalanobis distance k w  w t k 2 G  1 / 2 to account for landscape geometry.

Original: For example, in AdaGrad, G is the sum of the outer product of gradients (Duchi et al., 2011).
Revised: For example, in AdaGrad, G is the sum of the outer product of gradients (Duchi et al., 2011).

Original: Our proposed teleportation technique relaxes the requirement from the proximal term.
Revised: Our proposed teleportation technique relaxes the requirement from the proximal term.

Original: We teleport to points on the same level set of L ( w t ) , but allow w to be far from w t in Euclidean distance.
Revised: We teleport to points on the same level set of L ( w t ) , but allow w to be far from w t in Euclidean distance.

Original: In summary, our main contributions are:
Revised: In summary, our main contributions are:

Original: • We propose symmetry teleportation , an accelerated gradient-based optimization algorithm that exploits symmetries in the loss landscape.
Revised: • We propose symmetry teleportation , an accelerated gradient-based optimization algorithm that exploits symmetries in the loss landscape.

Original: • We derive loss-invariant group actions of test functions and multi-layer neural networks.
Revised: • We derive loss-invariant group actions of test functions and multi-layer neural networks.

Original: • We provide necessary conditions and examples of when teleportation improves convergence rate of the current and subsequent steps.
Revised: • We provide necessary conditions and examples of when teleportation improves convergence rate of the current and subsequent steps.

Original: • We show that our algorithm is closely related to second-order methods.
Revised: • We show that our algorithm is closely related to second-order methods.

Original: • Experimentally, we show that teleportation improves the convergence speed of gradient descent and AdaGrad in various optimization problems.
Revised: • Experimentally, we show that teleportation improves the convergence speed of gradient descent and AdaGrad in various optimization problems.

Original:  Related Work
Revised:  Related Work

Original: Continuous parameter space symmetries have been identified in neural networks with homogeneous (Badrinarayanan et al., 2015; Du et al., 2018) and radial activation functions (Ganev et al., 2021).
Revised: Continuous parameter space symmetries have been identiﬁed in neural networks with homogeneous (Badrinarayanan et al., 2015; Du et al., 2018) and radial activation functions (Ganev et al., 2021).

Original: The effect of symmetry transformations on gradients has been examined for translation, scale, and rescale symmetries in Kunin et al.
Revised: The effect of symmetry transformations on gradients has been examined for translation, scale, and rescale symmetries in Kunin et al.

Original: We contribute to this line of work by deriving loss-invariant group actions in multi-layer neural networks with invertible activation functions.
Revised: We contribute to this line of work by deriving loss-invariant group actions in multi-layer neural networks with invertible activation functions.

Original: Several works exploit symmetry to facilitate optimization.
Revised: Several works exploit symmetry to facilitate optimization.

Original: For example, Path-SGD (Neyshabur et al., 2015) improves optimization in ReLU networks by path regularization.
Revised: For example, Path-SGD (Neyshabur et al., 2015) improves optimization in ReLU networks by path regularization.

Original: G -SGD (Meng et al., 2019) performs weight updates in a scale-invariant space, which also exploits the symmetry in ReLU networks.
Revised: G -SGD (Meng et al., 2019) performs weight updates in a scale-invariant space, which also exploits the symmetry in ReLU networks.

Original: Van Laarhoven (2017) analyzes the effect of L2 regularization in batch normalization which gives scale-invariant functions.
Revised: Van Laarhoven (2017) analyzes the effect of L2 regularization in batch normalization which gives scale-invariant functions.

Original: Bamler and Mandt (2018) transforms parameters in the symmetry orbits to address slow movement along directions of weakly broken symmetry and find the optimal g ∈ G that minimizes L ( g · w ) .
Revised: Bamler and Mandt (2018) transforms parameters in the symmetry orbits to address slow movement along directions of weakly broken symmetry and ﬁnd the optimal g 2 G that minimizes L ( g · w ) .

Original: In comparison, we search within G -orbits for points which maximize | d L ( g · w ) /dt
Revised: In comparison, we search within G -orbits for points which maximize | d L ( g · w ) /dt

Original: | = ∥∇L ( g · w ) ∥ 2 .
Revised: | = kr L ( g · w ) k 2 .

Original: Natural gradient (Amari, 1998), adaptive gradient methods (Duchi et al., 2011; Kingma and Ba, 2015), and their approximations (Martens and Grosse, 2018; Gupta et al., 2018) improve the direction of parameter updates instead of directly transforming parameters.
Revised: Natural gradient (Amari, 1998), adaptive gradient methods (Duchi et al., 2011; Kingma and Ba, 2015), and their approximations (Martens and Grosse, 2018; Gupta et al., 2018) improve the direction of parameter updates instead of directly transforming parameters.

Original: If the group acts transitively on the level set, and we teleport to the point with maximum gradient, then our update direction is the same as that in Newton’s method.
Revised: If the group acts transitively on the level set, and we teleport to the point with maximum gradient, then our update direction is the same as that in Newton’s method.

Original: We prove that our algorithm is connected to second-order optimization methods and show that teleportation can be used to improve these methods empirically.
Revised: We prove that our algorithm is connected to second-order optimization methods and show that teleportation can be used to improve these methods empirically.

Original: The concept of neural teleportation was first explored using quiver representation theory (Armenta and Jodoin, 2021; Armenta et al., 2020).
Revised: The concept of neural teleportation was ﬁrst explored using quiver representation theory (Armenta and Jodoin, 2021; Armenta et al., 2020).

Original: These works provide a way to explore level curves of the loss of neural networks and show that random teleportation speeds up gradient descent experimentally and theoretically.
Revised: These works provide a way to explore level curves of the loss of neural networks and show that random teleportation speeds up gradient descent experimentally and theoretically.

Original: Our algorithm improves neural teleportation by searching for teleportation destinations that lead to the largest improvement in the magnitude of gradient.
Revised: Our algorithm improves neural teleportation by searching for teleportation destinations that lead to the largest improvement in the magnitude of gradient.

Original: Several works study how parameter initialization affects convergence rate (Saxe et al., 2014; Tarmoun et al., 2021; Min et al., 2021).
Revised: Several works study how parameter initialization affects convergence rate (Saxe et al., 2014; Tarmoun et al., 2021; Min et al., 2021).

Original: If we apply a group transformation only at initialization, our method is similar to that of Tarmoun et al.
Revised: If we apply a group transformation only at initialization, our method is similar to that of Tarmoun et al.

Original: We do not guarantee that the transformed parameters lead to faster convergence rate throughout the entire training.
Revised: We do not guarantee that the transformed parameters lead to faster convergence rate throughout the entire training.

Original: However, we accelerate convergence at least for a short time after initialization.
Revised: However, we accelerate convergence at least for a short time after initialization.

Original: Additionally, our method is not restricted to initialization and can be applied at any time during training.
Revised: Additionally, our method is not restricted to initialization and can be applied at any time during training.

Original: Most contemporary neural networks are overparametrized.
Revised: Most contemporary neural networks are overparametrized.

Original: While this has been shown to improve generalization (Belkin et al., 2019), an important question is how overparametrization affects optimization.
Revised: While this has been shown to improve generalization (Belkin et al., 2019), an important question is how overparametrization affects optimization.

Original: A series of works starting from Arora et al.
Revised: A series of works starting from Arora et al.

Original: (2018) show that overparametrization resulting from the depth of a neural network accelerates convergence.
Revised: (2018) show that overparametrization resulting from the depth of a neural network accelerates convergence.

Original: Another view is that the symmetry created by overparametrization poses constraints on trajectories in the form of conserved quantities (Głuch and Urbanke, 2021).
Revised: Another view is that the symmetry created by overparametrization poses constraints on trajectories in the form of conserved quantities (G ł uch and Urbanke, 2021).

Original: Additionally, the symmetry generates additional trajectories.
Revised: Additionally, the symmetry generates additional trajectories.

Original: When the trajectories created by overparametrization are equivalent, model compression by removing symmetry reduces training time (Ganev et al., 2021).
Revised: When the trajectories created by overparametrization are equivalent, model compression by removing symmetry reduces training time (Ganev et al., 2021).

Original: However, when trajectories are not equivalent, gradient flows on some paths are faster than others.
Revised: However, when trajectories are not equivalent, gradient ﬂows on some paths are faster than others.

Original: We search for the better paths created by overparametrization.
Revised: We search for the better paths created by overparametrization.

Original:  Symmetry Teleportation
Revised:  Symmetry Teleportation

Original: We propose symmetry teleportation , an accelerated gradient-based optimization algorithm that exploits symmetries in the loss landscape.
Revised: We propose symmetry teleportation , an accelerated gradient-based optimization algorithm that exploits symmetries in the loss landscape.

Original: Symmetry teleportation searches for the best gradient descent trajectory by teleporting parameters to a point with larger gradients using a group action.
Revised: Symmetry teleportation searches for the best gradient descent trajectory by teleporting parameters to a point with larger gradients using a group action.

Original: The resulting algorithm requires only gradient computations but is able to account for the global landscape geometry, leading to faster loss decay.
Revised: The resulting algorithm requires only gradient computations but is able to account for the global landscape geometry, leading to faster loss decay.

Original: Let the group G be a set of symmetries which leave the loss function L invariant: L ( g · ( w , X )) =
Revised: Let the group G be a set of symmetries which leave the loss function L invariant: L ( g · ( w , X )) =

Original: L ( w , X ) , where g ∈ G .
Revised: L ( w , X ) , where g 2 G .

Original: We perform gradient descent for t max steps.
Revised: We perform gradient descent for t max steps.

Original: We define an index set
Revised: We deﬁne an index set

Original: K ⊆ { 0 , 1 , 2 , ..., t max − 1 } as a teleportation schedule.
Revised: K ✓ { 0 , 1 , 2 , ..., t max  1 } as a teleportation schedule.

Original: At epochs that are in the schedule, we transform parameters using group element g ∈ G to the location where the gradient is largest, then continue with gradient descent.
Revised: At epochs that are in the schedule, we transform parameters using group element g 2 G to the location where the gradient is largest, then continue with gradient descent.

Original: Algorithm 1 describes the details of this procedure.
Revised: Algorithm 1 describes the details of this procedure.

Original: Note that the loss does not change after teleportation (Line 2-5) since L is G -invariant.
Revised: Note that the loss does not change after teleportation (Line 2-5) since L is G -invariant.

Original: Algorithm Input: Loss function L ( w ) , learning rate η , number of epochs t max , initialized parameters w 0 , symmetry group G , teleportation schedule K .
Revised: Input: Loss function L ( w ) , learning rate ⌘ , number of epochs t max , initialized parameters w 0 , symmetry group G , teleportation schedule K .

Original: Output: w t max .
Revised: Output: w t max .

Original: 1 for t← 0 to t max − 1 do if t ∈ K then g← argmax g ∈ G ∥ ( ∇L ) | g · w t ∥ w t← g · w t end if w t +1 ← w t − η ( ∇L ) | w t 7 end for 8 return w t max
Revised: 1 for t 0 to t max  1 do if t 2 K then g argmax g 2 G k ( r L ) | g · w t k w t g · w t end if w t + w t  ⌘ ( r L ) | w t 7 end for 8 return w t max

Original: Algorithm 1 can be generalized to apply teleportation to stochastic gradient descent (Appendix A.1).
Revised: Algorithm 1 can be generalized to apply teleportation to stochastic gradient descent (Appendix A.1).

Original: We discuss some design choices below and provide detailed analysis in the next two sections.
Revised: We discuss some design choices below and provide detailed analysis in the next two sections.

Original: When the action of G is continuous, teleportation can be implemented by parameterizing and performing gradient ascent on g .
Revised: When the action of G is continuous, teleportation can be implemented by parameterizing and performing gradient ascent on g .

Original: For example, the SO(2) group can be parameterized by the rotation angle θ .
Revised: For example, the SO(2) group can be parameterized by the rotation angle ✓ .

Original: Small transformations g ∈ GL d ( R ) (general linear group) can be parameterized as g ≈ I + εT where ε ≪ 1 and T are arbitrary d × d matrices.
Revised: Small transformations g 2 GL d ( R ) (general linear group) can be parameterized as g ⇡ I + " T where " ⌧ 1 and T are arbitrary d ⇥ d matrices.

Original: For discrete groups, search algorithms or random sampling can be used to find a group element that improves the magnitude of the gradient.
Revised: For discrete groups, search algorithms or random sampling can be used to ﬁnd a group element that improves the magnitude of the gradient.

Original: Although g · ( w , X ) does not change X in the cases we consider in this paper, Algorithm 1 can be extended to allow transformations on both parameters and data.
Revised: Although g · ( w , X ) does not change X in the cases we consider in this paper, Algorithm 1 can be extended to allow transformations on both parameters and data.

Original: The group actions on data during teleportations can be precomposed and applied to the input data at inference time.
Revised: The group actions on data during teleportations can be precomposed and applied to the input data at inference time.

Original: The algorithm extension that allows transformation on the input can be found in Appendix A.2.
Revised: The algorithm extension that allows transformation on the input can be found in Appendix A.2.

Original: The teleportation schedule K is a hyperparameter that determines when to perform teleportation.
Revised: The teleportation schedule K is a hyperparameter that determines when to perform teleportation.

Original: We define K as a set to allow flexible teleportation schedules, such as with non-fixed frequencies or teleporting only at the earlier epochs.
Revised: We deﬁne K as a set to allow ﬂexible teleportation schedules, such as with non-ﬁxed frequencies or teleporting only at the earlier epochs.

Original:  Symmetry Groups of Certain Optimization Problems
Revised:  Symmetry Groups of Certain Optimization Problems

Original: We give a few practical examples to demonstrate how teleportation can be used to accelerate optimization.
Revised: We give a few practical examples to demonstrate how teleportation can be used to accelerate optimization.

Original: Specifically, we first consider two test functions which are often used to evaluate optimization algorithms (Back, 1996).
Revised: Speciﬁcally, we ﬁrst consider two test functions which are often used to evaluate optimization algorithms (Back, 1996).

Original: We then derive the symmetries of multi-layer neural networks.
Revised: We then derive the symmetries of multi-layer neural networks.

Original: 4.1 Test functions
Revised: 4.1 Test functions

Original: Rosenbrock function.
Revised: Rosenbrock function.

Original: The Rosenbrock function originally introduced by Rosenbrock (1960) has a characteristic global minimum that is inside a long, narrow, parabolicly-shaped flat valley.
Revised: The Rosenbrock function originally introduced by Rosenbrock (1960) has a characteristic global minimum that is inside a long, narrow, parabolicly-shaped ﬂat valley.

Original: Finding the valley is easy but reaching the minimum is difficult.
Revised: Finding the valley is easy but reaching the minimum is difﬁcult.

Original: On a 2-dimensional space, the Rosenbrock function has the following form:
Revised: On a 2-dimensional space, the Rosenbrock function has the following form:

Original: [Equation]
Revised: [Equation]

Original: Booth function.
Revised: Booth function.

Original: The Booth function (Jamil and Yang, 2013) is also defined on R 2 and has one global minimum at (1 , 3) where the function evaluates to 0 :
Revised: The Booth function (Jamil and Yang, 2013) is also deﬁned on R 2 and has one global minimum at (1 , 3) where the function evaluates to 0 :

Original: [Equation]
Revised: [Equation]

Original: The following proposition identifies the symmetry of these two test functions.
Revised: The following proposition identiﬁes the symmetry of these two test functions.

Original: Proposition 4.1.
Revised: Proposition 4.1.

Original: The Rosenbrock and Booth functions have rotational symmetry.
Revised: The Rosenbrock and Booth functions have rotational symmetry.

Original: In other words, there exist action maps a r , a b : SO(2) × R 2 −→ R 2 , such that for all g ∈ SO(2) ,
Revised: In other words, there exist action maps a r , a b : SO(2) ⇥ R 2 ! R 2 , such that for all g 2 SO(2) ,

Original: L r ( x 1 , x 2 ) =
Revised: L r ( x 1 , x 2 ) =

Original: L r ( a r ( g, [ x 1 , x 2 ])) and L b ( x 1 , x 2 ) =
Revised: L r ( a r ( g, [ x 1 , x 2 ])) and L b ( x 1 , x 2 ) =

Original: L b ( a b ( g, [ x 1 , x 2 ])) .
Revised: L b ( a b ( g, [ x 1 , x 2 ])) .

Original: The exact forms of the action maps are deferred to Appendix B.1.2 and B.1.3.
Revised: The exact forms of the action maps are deferred to Appendix B.1.2 and B.1.3.

Original: During the teleportation step, our goal is to maximize the gradient within a level set of the loss: max g ∈ SO(2) ∥ [Equation] · ( x 1 ,x 2 )) dt ∥ .
Revised: During the teleportation step, our goal is to maximize the gradient within a level set of the loss: max g 2 SO(2) k d L ( g · ( x 1 ,x 2 )) dt k .

Original: 4.2 Multi-layer Neural Networks
Revised: 4.2 Multi-layer Neural Networks

Original: Next, we consider feed-forward neural networks.
Revised: Next, we consider feed-forward neural networks.

Original: Denote the output of the m th layer by h m∈ R d m × n , where d m is the hidden dimension and n is the number of samples.
Revised: Denote the output of the m th layer by h m R d m ⇥ n , where d m is the hidden dimension and n is the number of samples.

Original: Denote the input by h 0 = X ∈ R d 0 × n .
Revised: Denote the input by h 0 = X 2 R d 0 ⇥ n .

Original: In terms of the previous layer output ˜ h m =
Revised: In terms of the previous layer output ˜ h m =

Original: W m h m − 1 whereW m ∈ R
Revised: W m h m  1 where

Original:  d m × d m −(we absorb biases into W m by adding an extra row of ones in ˜ h m − 1 ), the output h m is
Revised: W m 2 R d m ⇥ d m (we absorb biases into W m by adding an extra row of ones in ˜ h m  1 ), the output h m is

Original: [Equation]
Revised: [Equation]

Original: We assume the activation functions σ : R → R are bijections.
Revised: We assume the activation functions  : R ! R are bijections.

Original: For instance, Leaky-ReLU is bijective.
Revised: For instance, Leaky-ReLU is bijective.

Original: For other activation, such Sigmoid or Tanh, we analytically extend them to bijective functions (e.g. tanh( x ) + e x − N − e − x + N for N ≫ 1 ).
Revised: For other activation, such Sigmoid or Tanh, we analytically extend them to bijective functions (e.g. tanh( x ) + e x  N  e  x + N for N  1 ).

Original: For linear activations, we have linear symmetries:
Revised: For linear activations, we have linear symmetries:

Original: Proposition 4.2.
Revised: Proposition 4.2.

Original: A linear network is invariant under all groups G m ≡ GL d m acting asg · [Equation] m
Revised: A linear network is invariant under all groups G m ⌘ GL d m acting as

Original: [Equation]
Revised: [Equation]

Original: m [Equation] Similar, in the nonlinear case, we want to find g · W m that keep the outputs h k for k ̸{+
+}
Revised: Similar, in the nonlinear case, we want to ﬁnd g · W m that keep the outputs h k for k {+
+}

Original: = m − 1 invariant.
Revised: = m  1 invariant.

Original: With nonlinearity, ˜ h m
Revised: With nonlinearity, ˜ h m

Original: = W m σ ( W m − 1 h m − 2 ) .
Revised: = W m  ( W m  1 h m  2 ) .

Original: This network has a GL d m − 1 symmetry given by
Revised: This network has a GL d m  1 symmetry given by

Original: [Equation]
Revised: [Equation]

Original: Thus, we have the following proposition regarding symmetries of nonlinear networks:
Revised: Thus, we have the following proposition regarding symmetries of nonlinear networks:

Original: Proposition 4.3.
Revised: Proposition 4.3.

Original: Assume that h m − 2 is invertible.
Revised: Assume that h m  2 is invertible.

Original: A multi-layer network with bijective activation σ has a GL d m − 1 symmetry.
Revised: A multi-layer network with bijective activation  has a GL d m  1 symmetry.

Original: For g m ∈ G m = GL d m − 1 ( R ) the following group action keeps h p with p ≥ m invariant
Revised: For g m 2 G m = GL d m  1 ( R ) the following group action keeps h p with p  m invariant

Original: [Equation]
Revised: [Equation]

Original: Proofs can be found in Appendix B.2.
Revised: Proofs can be found in Appendix B.2.

Original: Note that this group action depends on the input to the network as well as the current weights of all the lower layers.
Revised: Note that this group action depends on the input to the network as well as the current weights of all the lower layers.

Original: Yet, since the action keeps the output of upper and lower layers invariant, multiple G m for different m applied at the same time still keep the network output invariant.
Revised: Yet, since the action keeps the output of upper and lower layers invariant, multiple G m for different m applied at the same time still keep the network output invariant.

Original: The proposition assumes that h m − 2 is square and full-rank.
Revised: The proposition assumes that h m  2 is square and full-rank.

Original: When n < d m − 2 and h m − 2 has rank n , (8) (with left inverse of h m − 2 ) keeps the loss invariant but does not satisfy the identity axiom of a group action.
Revised: When n < d m  2 and h m  2 has rank n , (8) (with left inverse of h m  2 ) keeps the loss invariant but does not satisfy the identity axiom of a group action.

Original:  Theoretical Analysis 5.1 What symmetries help accelerate optimization
Revised:  Theoretical Analysis 5.1 What symmetries help accelerate optimization

Original: We now discuss the conditions that need to be satisfied for teleportation to accelerate GD.
Revised: We now discuss the conditions that need to be satisﬁed for teleportation to accelerate GD.

Original: For brevity, we denote all trainable parameters (e.g. { W 1 , · · · W p } for the p -layer neural network) by a single flattened vector w ∈ R n .
Revised: For brevity, we denote all trainable parameters (e.g. { W 1 , · · · W p } for the p -layer neural network) by a single ﬂattened vector w 2 R n .

Original: Consider a symmetry G of the loss function L ( w ) , meaning for all g ∈ G ,
Revised: Consider a symmetry G of the loss function L ( w ) , meaning for all g 2 G ,

Original: L ( w ) =
Revised: L ( w ) =

Original: L ( g · w ) .
Revised: L ( g · w ) .

Original: We quantify how teleportation by G changes the rate of loss decay, given by
Revised: We quantify how teleportation by G changes the rate of loss decay, given by

Original: [Equation]
Revised: [Equation]

Original: where η is the matrix of learning rates (which must be positive semi-definite) and ∥ v ∥ 2 η = v T ηv is the Mahalanobis norm.
Revised: where ⌘ is the matrix of learning rates (which must be positive semi-deﬁnite) and k v k 2 ⌘ = v T ⌘ v is the Mahalanobis norm.

Original: A constant learning rate means η = I .
Revised: A constant learning rate means ⌘ = I .

Original: The following proposition provides the condition a symmetry needs to satisfy to accelerate optimization.
Revised: The following proposition provides the condition a symmetry needs to satisfy to accelerate optimization.

Original: (Proofs can be found in Appendix C.1.)
Revised: (Proofs can be found in Appendix C.1.)

Original: Proposition 5.1.
Revised: Proposition 5.1.

Original: Let w ′ ={+
+}
Revised: Let w {+
+}

Original:  g · w be a point we teleport to.
Revised: = g · w be a point we teleport to.

Original: Let J = ∂ w ′ /∂
Revised: Let J = @ w 0 / @

Original: w be the Jacobian.
Revised: w be the Jacobian.

Original: Symmetry teleportation using g accelerates the rate of decay in L if it satisfies
Revised: Symmetry teleportation using g accelerates the rate of decay in L if it satisﬁes

Original: [Equation]
Revised: [Equation]

Original: If the action of the symmetry group G ⊂ GL n is linear we have w ′ = g · w = g w and J = g .
Revised: If the action of the symmetry group G ⇢ GL n is linear we have w 0 = g · w = g w and J = g .

Original: It follows that if G is a subgroup of the orthogonal group, d L /dt will be invariant:
Revised: It follows that if G is a subgroup of the orthogonal group, d L /dt will be invariant:

Original: Corollary 5.2.
Revised: Corollary 5.2.

Original: Let O η denote the orthogonal group of invariances of the inverse of the learning rate, η − 1 , meaning for g ∈ O η , g T η − 1 g = η − 1 .
Revised: Let O ⌘ denote the orthogonal group of invariances of the inverse of the learning rate, ⌘  1 , meaning for g 2 O ⌘ , g T ⌘  1 g = ⌘  1 .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: In the simple case where the learning rate is a constant, O η = O ( n ) becomes the classic orthogonal group (e.g. rotations) with g T g
Revised: In the simple case where the learning rate is a constant, O ⌘ = O ( n ) becomes the classic orthogonal group (e.g. rotations) with g T g

Original: = I .
Revised: = I .

Original: In general, when g preserves the norm of the gradient, symmetry teleportation has no effect on the convergence rate.
Revised: In general, when g preserves the norm of the gradient, symmetry teleportation has no effect on the convergence rate.

Original: From Theorem 3.2 in Hazan (2019), assuming that L is β -smooth and is bounded by |L| ≤ M , the gradient norm in gradient descent converges as 12 β (cid:80) Tt =1 ∥ ( ∇L ) | w t ∥ ≤ 2 M .
Revised: From Theorem 3.2 in Hazan (2019), assuming that L is  -smooth and is bounded by |L|  M , the gradient norm in gradient descent converges as 12  [Equation]

Original: Teleportation increases ( ∇L ) | w t , therefore requiring less time to reach convergence.
Revised: Teleportation increases ( r L ) | w t , therefore requiring less time to reach convergence.

Original: 5.2 Improvement of subsequent steps
Revised: 5.2 Improvement of subsequent steps

Original: Since teleportation moves the parameters to a point with a larger gradient, and subsequent GD steps are local, we would expect that teleportation improves the magnitude of the gradient for a few future steps as well.
Revised: Since teleportation moves the parameters to a point with a larger gradient, and subsequent GD steps are local, we would expect that teleportation improves the magnitude of the gradient for a few future steps as well.

Original: The following results formalize this intuition (proofs in Appendix C.2.).
Revised: The following results formalize this intuition (proofs in Appendix C.2.).

Original: Assumption 5.3 (Lipschitz Continuity) .
Revised: Assumption 5.3 (Lipschitz Continuity) .

Original: The l2 norm of the gradient is Lipschitz continuous with constant L ∈ R ≥ 0 , which is(cid:12)(cid:12)(cid:12)(cid:12)(cid:13)(cid:13)(cid:13)(cid:13) ∂ L ∂ w 1 (cid:13)(cid:13)(cid:13)(cid:13) 2 − (cid:13)(cid:13)(cid:13)(cid:13) ∂ L
Revised: The l2 norm of the gradient is Lipschitz continuous with constant L 2 R  0 , which is

Original: [Equation]
Revised: [Equation]

Original: w 2 (cid:13)(cid:13)(cid:13)(cid:13) [Equation] where w 1 , w 2 are two points in the parameter space and L is the Lipschitz constant.
Revised: (12) where w 1 , w 2 are two points in the parameter space and L is the Lipschitz constant.

Original: Proposition 5.4. Consider the gradient descent with a G -invariant loss L ( w ) and learning rate η ∈ R + .
Revised: Proposition 5.4. Consider the gradient descent with a G -invariant loss L ( w ) and learning rate ⌘ 2 R + .

Original: Let w t be the parameter at time t and w ′ t = g · w t the parameter teleported by g ∈ G .
Revised: Let w t be the parameter at time t and w 0 t = g · w t the parameter teleported by g 2 G .

Original: Let w t + T and w ′ t + T be the parameters after T more steps of gradient descent from w t and w ′ t respectively.
Revised: Let w t + T and w 0 t + T be the parameters after T more steps of gradient descent from w t and w 0 t respectively.

Original: Under Assumption 5.3, if ηL < 1 , and
Revised: Under Assumption 5.3, if ⌘ L < 1 , and

Original: [Equation]
Revised: [Equation]

Original: then
Revised: then

Original: [Equation]
Revised: [Equation]

Original: The proposition provides a sufficient condition for teleportation to improve T future steps.
Revised: The proposition provides a sufﬁcient condition for teleportation to improve T future steps.

Original: The condition is met when L is small, η is small, or the initial improvement,(cid:13)(cid:13)(cid:13)
Revised: The condition is met when L is small, ⌘ is small, or the initial improvement,

Original: [Equation][Equation]
Revised: [Equation]

Original: is large. 5.3 Convergence analysis for convex quadratic functions
Revised: 5.3 Convergence analysis for convex quadratic functions

Original: Teleportation improves the magnitude of gradient for the current step.
Revised: Teleportation improves the magnitude of gradient for the current step.

Original: We have also shown that the magnitude of gradient stays large for a few subsequent steps (Proposition 5.4).
Revised: We have also shown that the magnitude of gradient stays large for a few subsequent steps (Proposition 5.4).

Original: In this section, we analyze a class of functions where teleporting once guarantees optimality at all future times.
Revised: In this section, we analyze a class of functions where teleporting once guarantees optimality at all future times.

Original: We consider a trajectory optimal if for every point on the trajectory, the magnitude of gradient is at a local maximum in the loss level set that contains the point.
Revised: We consider a trajectory optimal if for every point on the trajectory, the magnitude of gradient is at a local maximum in the loss level set that contains the point.

Original: Consider a positive definite quadratic form L A ( w )
Revised: Consider a positive deﬁnite quadratic form L A ( w )

Original: =
Revised: =

Original: w T A w , where w ∈ R n is the parameter and
Revised: w T A w , where w 2 R n is the parameter and

Original: A ∈ R n × n is a positive definite matrix.
Revised: A 2 R n ⇥ n is a positive deﬁnite matrix.

Original: The gradient of L A is ∇L A = 

Revised: The gradient of L A is r L A =

Original: A w , and the Hessian of L A is H = 2 A .
Revised: A w , and the Hessian of L A is H = 2 A .

Original: Since A is defined to be positive definite, L A is convex.
Revised: Since A is deﬁned to be positive deﬁnite, L A is convex.

Original: The function L A has global minimum at a single point w ∗ = 0 .
Revised: The function L A has global minimum at a single point w ⇤ = 0 .

Original: Let ρ be a representation of O ( n ) acting on R n .
Revised: Let ⇢ be a representation of O ( n ) acting on R n .

Original: For g ∈ O ( n ) , we define the following group action:
Revised: For g 2 O ( n ) , we deﬁne the following group action:

Original: [Equation]
Revised: [Equation]

Original: Then it can be shown that L A ( w ) admits a O ( n ) symmetry:
Revised: Then it can be shown that L A ( w ) admits a O ( n ) symmetry:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Let S c
Revised: Let S c

Original: = { w : L A ( w ) =
Revised: = { w : L A ( w ) =

Original: c } be a level set of L A .
Revised: c } be a level set of L A .

Original: We show that after a teleportation, every point on the gradient flow trajectory is optimal in its level set (with full proof in Appendix C.3).
Revised: We show that after a teleportation, every point on the gradient ﬂow trajectory is optimal in its level set (with full proof in Appendix C.3).

Original: Proposition 5.5.
Revised: Proposition 5.5.

Original: If at point w , ∥∇L A | w ∥ 2 is at a maximum in S L A ( w ) , then for any point w ′ on the gradient flow trajectory starting from w , ∥∇L A | w ′ ∥ 2 is at a maximum in S L A ( w ′ ) .
Revised: If at point w , kr L A | w k 2 is at a maximum in S L A ( w ) , then for any point w 0 on the gradient ﬂow trajectory starting from w , kr L A | w 0 k 2 is at a maximum in S L A ( w 0 ) .

Original: Finally, we observe that for L A , teleportation moves the parameters closer to the global minimum in
Revised: Finally, we observe that for L A , teleportation moves the parameters closer to the global minimum in

Original: Euclidean distance.
Revised: Euclidean distance.

Original: In other words, maximizing the magnitude of gradient minimizes the distance to w ∗ in a loss level set.
Revised: In other words, maximizing the magnitude of gradient minimizes the distance to w ⇤ in a loss level set.

Original: Proposition 5.6.
Revised: Proposition 5.6.

Original: Consider a point w in the parameter space.
Revised: Consider a point w in the parameter space.

Original: Let g = arg max g ∈ G ∥∇L A | g · w ∥ 22 .
Revised: Let g = arg max g 2 G kr L A | g · w k 22 .

Original: Then g · w = arg min w ′ ∈ S L A ( w ) ∥ w ′ − w ∗ ∥ 22 .
Revised: Then g · w = arg min w 0 2 S L A ( w ) k w 0  w ⇤ k 22 .

Original: 5.4 Relation to second-order optimization methods
Revised: 5.4 Relation to second-order optimization methods

Original: Since our algorithm involves optimizing the gradients, symmetry teleportation is closely related to second-order optimization methods.
Revised: Since our algorithm involves optimizing the gradients, symmetry teleportation is closely related to second-order optimization methods.

Original: At the target point to teleport to, gradient descent becomes equivalent to Newton’s method (proof in Appendix C.4).
Revised: At the target point to teleport to, gradient descent becomes equivalent to Newton’s method (proof in Appendix C.4).

Original: Proposition 5.7.
Revised: Proposition 5.7.

Original: Let S L 0 = { w : L ( w )
Revised: Let S L 0 = { w : L ( w )

Original: =
Revised: =

Original: L 0 } be a level set of L .
Revised: L 0 } be a level set of L .

Original: If at a particular w ∈ S L 0 we have ∥∇L ( w ) ∥ 2 ≥ ∥∇L ( w ′ ) ∥ 2 for all w ′ in a small neighborhood of w in S L 0 , then the gradient ∇L ( w ) has the same direction as the direction from Newton’s method, H − 1 ∇L ( w ) .
Revised: If at a particular w 2 S L 0 we have kr L ( w ) k 2  kr L ( w 0 ) k 2 for all w 0 in a small neighborhood of w in S L 0 , then the gradient r L ( w ) has the same direction as the direction from Newton’s method, H  1 r L ( w ) .

Original: Proposition 5.7 provides an alternative way to interpret teleportation.
Revised: Proposition 5.7 provides an alternative way to interpret teleportation.

Original: Instead of computing the
Revised: Instead of computing the

Original: Newton’s direction, we search within the loss level set for a point where the gradient has the same[Figure] [Figure] [Figure]
Revised: Newton’s direction, we search within the loss level set for a point where the gradient has the same

Original: [Figure][Figure] [Figure] [Figure]
Revised: [Figure]

Original: [Figure] [Figure] [Figure] [Figure]
Revised: d

Original: [Figure][Figure]
Revised: [Figure]

Original: direction as Newton’s direction.
Revised: direction as Newton’s direction.

Original: However, symmetry teleportation does not require computing the full
Revised: However, symmetry teleportation does not require computing the full

Original: Hessian matrix.
Revised: Hessian matrix.

Original: The second derivative required for optimizing over continuous groups is obtained by taking derivatives with respect to parameters and then with respect to the group element, as opposed to taking the derivative with respect to parameters twice.
Revised: The second derivative required for optimizing over continuous groups is obtained by taking derivatives with respect to parameters and then with respect to the group element, as opposed to taking the derivative with respect to parameters twice.

Original: This makes the computation significantly more feasible than Newton’s method on neural networks with large number of parameters.
Revised: This makes the computation signiﬁcantly more feasible than Newton’s method on neural networks with large number of parameters.

Original: In practice, however, the group action used for teleportation is usually not transitive.
Revised: In practice, however, the group action used for teleportation is usually not transitive.

Original: Additionally, we do not teleport using the optimal group element since it can be unbounded.
Revised: Additionally, we do not teleport using the optimal group element since it can be unbounded.

Original: We also do not apply teleportation in every gradient descent step.
Revised: We also do not apply teleportation in every gradient descent step.

Original: Therefore, Proposition 5.7 serves as an intuition instead of an exact formulation of how teleportation works.
Revised: Therefore, Proposition 5.7 serves as an intuition instead of an exact formulation of how teleportation works.

Original: We provide empirical evidence in Sectionthat these approximations do not erase the benefits of teleportation completely, and leave theoretical investigations of the connection to second-order methods under approximations as future work.
Revised: We provide empirical evidence in Sectionthat these approximations do not erase the beneﬁts of teleportation completely, and leave theoretical investigations of the connection to second-order methods under approximations as future work.

Original:  Experiments 6.1 Acceleration through symmetry teleportation
Revised:  Experiments 6.1 Acceleration through symmetry teleportation

Original: We examine the effect of symmetry teleportation on optimization.
Revised: We examine the effect of symmetry teleportation on optimization.

Original: We illustrate teleportation in the parameter space on two test functions and show a speedup in regression and classification problems using multilayer neural networks.
Revised: We illustrate teleportation in the parameter space on two test functions and show a speedup in regression and classiﬁcation problems using multilayer neural networks.

Original: For test functions, we compared with GD for illustration purposes.
Revised: For test functions, we compared with GD for illustration purposes.

Original: For multi-layer neural networks, we also include AdaGrad as a more competitive baseline.
Revised: For multi-layer neural networks, we also include AdaGrad as a more competitive baseline.

Original: Rosenbrock function.
Revised: Rosenbrock function.

Original: We apply symmetry teleportation to optimize the 2-variable Rosenbrock function (3).
Revised: We apply symmetry teleportation to optimize the 2-variable Rosenbrock function (3).

Original: The parameters x 1 , x 2 are initialized to ( − 1 , − 1) .
Revised: The parameters x 1 , x 2 are initialized to (  1 ,  1) .

Original: Each algorithm is run 1000 steps with learning rate 10 − 3 .
Revised: Each algorithm is run 1000 steps with learning rate 10  3 .

Original: We teleport the parameters every 100 steps.
Revised: We teleport the parameters every 100 steps.

Original: The group elements are found by gradient ascent on θ , the parameter for the SO(2) group, for 10 steps with learning rate 10 − 1 .
Revised: The group elements are found by gradient ascent on ✓ , the parameter for the SO(2) group, for 10 steps with learning rate 10  1 .

Original: The trajectory of parameters and the loss level sets are plotted in Figure 2a,b.
Revised: The trajectory of parameters and the loss level sets are plotted in Figure 2a,b.

Original: The blue star denotes the final position of parameters, the green star denotes the target, and orange dots are the positions from which symmetry transforms start.
Revised: The blue star denotes the ﬁnal position of parameters, the green star denotes the target, and orange dots are the positions from which symmetry transforms start.

Original: While gradient descent is not able to reach the target in 1000 steps, teleportation allows large steps and reaches the target much earlier.
Revised: While gradient descent is not able to reach the target in 1000 steps, teleportation allows large steps and reaches the target much earlier.

Original: Figure 2d shows that teleportation improves the norm of gradients in the following step, and 2c shows its effect on the loss value.
Revised: Figure 2d shows that teleportation improves the norm of gradients in the following step, and 2c shows its effect on the loss value.

Original: Teleportations clearly reduce the number of steps needed for convergence.
Revised: Teleportations clearly reduce the number of steps needed for convergence.

Original: Booth function.
Revised: Booth function.

Original: We also test symmetry teleportation on the Booth function defined in Eqn.
Revised: We also test symmetry teleportation on the Booth function deﬁned in Eqn.

Original: We initialize the parameters x 1 , x 2 to (5 , − 5) .
Revised: We initialize the parameters x 1 , x 2 to (5 ,  5) .

Original: Each algorithm is run 10 steps with learning rate 0 .
Revised: Each algorithm is run 10 steps with learning rate 0 .

Original: We perform symmetry teleportation on the parameters once, before epoch 5.
Revised: We perform symmetry teleportation on the parameters once, before epoch 5.

Original: The group elements are found by gradient ascent on θ for 10 steps with learning rate 0 .
Revised: The group elements are found by gradient ascent on ✓ for 10 steps with learning rate 0 .

Original: θ is initialized uniformly
Revised: ✓ is initialized uniformly

Original: [Figure]
Revised: [Figure]

Revised: at random over [0 , ⇡ ) .

Revised: Similar to the Rosenbrock function, teleportation moves the parameters to a trajectory with a larger convergence rate (Figure 2 bottom row).

Original: Multilayer neural network regression.
Revised: Multilayer neural network regression.

Original: We further evaluated our method on a three-layer neural network with a regression loss min W 1 ,W 2 ,W 3 ∥ Y − W 3 σ ( W 2 σ ( W 1 X ))
Revised: We further evaluated our method on a three-layer neural network with a regression loss min W 1 ,W 2 ,W 3 k Y  W 3  ( W 2  ( W 1 X ))

Original: ∥ 2 .
Revised: k 2 .

Original: The dimension of weight matrices are W 3 ∈ R 8 × 7 , W 2 ∈ R 7 × 6 , and W 1 ∈ R 6 × 5 .
Revised: The dimension of weight matrices are W 3 2 R 8 ⇥ 7 , W 2 2 R 7 ⇥ 6 , and W 1 2 R 6 ⇥ 5 .

Original: X ∈ R 5 × 4 is the data, Y ∈ R 8 × 4 is the target, and σ is the LeakyReLU activation with slope coefficient 0.1.
Revised: X 2 R 5 ⇥ 4 is the data, Y 2 R 8 ⇥ 4 is the target, and  is the LeakyReLU activation with slope coefﬁcient 0.1.

Original: Additional hyperparameter details can be found in Appendix D.2.
Revised: Additional hyperparameter details can be found in Appendix D.2.

Original: Teleportation of the GL( R ) group can be performed by finding an element x in its Lie algebra, such that transforming w by the group element g = exp( x ) improves the gradient d L dt .
Revised: Teleportation of the GL( R ) group can be performed by ﬁnding an element x in its Lie algebra, such that transforming w by the group element g = exp ( x ) improves the gradient d L dt .

Original: We use the first order approximation of the exponential map.
Revised: We use the ﬁrst order approximation of the exponential map.

Original: Between each pair of weight matrices, we replace g m by I + T and g − 1 m by I − T in (8), where T ∈ R d m × d m is initialized to 0.
Revised: Between each pair of weight matrices, we replace g m by I + T and g  1 m by I  T in (8), where T 2 R d m ⇥ d m is initialized to 0.

Original: Then we perform gradient ascent steps on T with objective defined in Line 3 of Algorithm 1 and update the pair of weights.
Revised: Then we perform gradient ascent steps on T with objective deﬁned in Line 3 of Algorithm 1 and update the pair of weights.

Original: Figure 3a and 3b show the training curves plotted against epochs and time.
Revised: Figure 3a and 3b show the training curves plotted against epochs and time.

Original: Shaded area denotes one standard deviation from 5 runs.
Revised: Shaded area denotes one standard deviation from 5 runs.

Original: Since GD and AdaGrad use different learning rates, they are not directly comparable.
Revised: Since GD and AdaGrad use different learning rates, they are not directly comparable.

Original: However, the addition of teleportation clearly improves both algorithms.
Revised: However, the addition of teleportation clearly improves both algorithms.

Original: Figure 3c and 3d shows the squared norm of gradient from a single run.
Revised: Figure 3c and 3d shows the squared norm of gradient from a single run.

Original: Teleportation increases the magnitude of gradient, and the trajectory with teleportation has a larger d L /dt value at the same loss values, which demonstrates that teleportation finds a better trajectory.
Revised: Teleportation increases the magnitude of gradient, and the trajectory with teleportation has a larger d L /dt value at the same loss values, which demonstrates that teleportation ﬁnds a better trajectory.

Original: MNIST classification.
Revised: MNIST classiﬁcation.

Original: We apply symmetry teleportation on the MNIST classification task (Deng, 2012).
Revised: We apply symmetry teleportation on the MNIST classiﬁcation task (Deng, 2012).

Original: We split the training set into 48,000 for training and 12,000 for validation.
Revised: We split the training set into 48,000 for training and 12,000 for validation.

Original: The input data has dimension 28 × 28 and is flattened into a vector.
Revised: The input data has dimension 28 ⇥ 28 and is ﬂattened into a vector.

Original: The output of the neural network has dimensioncorresponding to the 10 digit classes.
Revised: The output of the neural network has dimensioncorresponding to the 10 digit classes.

Original: We used a three-layer neural network with hidden dimension [512 , 512] , LeakyReLU activations, and cross-entropy loss.
Revised: We used a three-layer neural network with hidden dimension [512 , 512] , LeakyReLU activations, and cross-entropy loss.

Original: Learning rate is 2 × 10 − 3 , and learning rate for teleportation is 10 − 3 .
Revised: Learning rate is 2 ⇥ 10  3 , and learning rate for teleportation is 10  3 .

Original: Each optimization algorithm is run 80 epochs with batch size of 20.
Revised: Each optimization algorithm is run 80 epochs with batch size of 20.

Original: Immediately after the first epoch, we apply teleportation using data from one mini-batch, and repeat for 4 different mini-batches.
Revised: Immediately after the ﬁrst epoch, we apply teleportation using data from one mini-batch, and repeat for 4 different mini-batches.

Original: For each mini-batch, 10 gradient ascent steps are used to optimize g m .
Revised: For each mini-batch, 10 gradient ascent steps are used to optimize g m .

Original: Figure 4a,b shows the effect of teleportation on training and validation loss, and Figure 4c,d shows the norm of the gradient in training.
Revised: Figure 4a,b shows the effect of teleportation on training and validation loss, and Figure 4c,d shows the norm of the gradient in training.

Original: While teleportation significantly accelerates the decrease of the training loss in SGD, its effect on the validation loss is limited and detrimental for AdaGrad.
Revised: While teleportation signiﬁcantly accelerates the decrease of the training loss in SGD, its effect on the validation loss is limited and detrimental for AdaGrad.

Original: Therefore, teleportation on MNIST makes training faster at the beginning but leads to earlier overfit and slightly worse validation accuracy.
Revised: Therefore, teleportation on MNIST makes training faster at the beginning but leads to earlier overﬁt and slightly worse validation accuracy.

Original: A possible reason is that regions with large gradients have sharp minima that do not generalize well.
Revised: A possible reason is that regions with large gradients have sharp minima that do not generalize well.

Original: 6.2 Teleportation schedule
Revised: 6.2 Teleportation schedule

Original: The effect of teleportation varies depending on its time and frequency.
Revised: The effect of teleportation varies depending on its time and frequency.

Original: Figure 5 shows the result of teleportation on MNIST with different hyperparameters.
Revised: Figure 5 shows the result of teleportation on MNIST with different hyperparameters.

Original: In all experiments, we use SGD with batch size 20, learning rate 2 × 10 − 3 , and 10 gradient ascent steps for each teleportation.
Revised: In all experiments, we use SGD with batch size 20, learning rate 2 ⇥ 10  3 , and 10 gradient ascent steps for each teleportation.

Original: In Figure 5a, we randomly select 4 different mini-batches and apply teleportation on each of them individually, but at different epochs.
Revised: In Figure 5a, we randomly select 4 different mini-batches and apply teleportation on each of them individually, but at different epochs.

Original: Teleportation before training has the worst performance.
Revised: Teleportation before training has the worst performance.

Original: After epoch 0, the effect of teleportation is stronger when it is applied earlier.
Revised: After epoch 0, the effect of teleportation is stronger when it is applied earlier.

Original: In Figure 5b, we again apply teleportations on 4 randomly selected mini-batches, but repeat this with 5 different teleportation
Revised: In Figure 5b, we again apply teleportations on 4 randomly selected mini-batches, but repeat this with 5 different teleportation

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: schedules (hyperparameter K in Algorithm 2) as shown in the legend.
Revised: schedules (hyperparameter K in Algorithm 2) as shown in the legend.

Original: All schedules have the same number of teleportations.
Revised: All schedules have the same number of teleportations.

Original: Smaller intervals between teleportations accelerate convergence more significantly.
Revised: Smaller intervals between teleportations accelerate convergence more signiﬁcantly.

Original: In Figure 5c, we apply teleportation immediately after the first epoch, but use different numbers of mini-batches (hyperparameter B in Algorithm 2) and teleport using each of them individually.
Revised: In Figure 5c, we apply teleportation immediately after the ﬁrst epoch, but use different numbers of mini-batches (hyperparameter B in Algorithm 2) and teleport using each of them individually.

Original: Using more mini-batches to teleport leads to faster decrease in training loss but is also more prone to overfitting.
Revised: Using more mini-batches to teleport leads to faster decrease in training loss but is also more prone to overﬁtting.

Original: 6.3 Runtime analysis
Revised: 6.3 Runtime analysis

Original: The additional amount of computation introduced by teleportation depends on the implementation of
Revised: The additional amount of computation introduced by teleportation depends on the implementation of

Original: Line 2-5 of Algorithm 1.
Revised: Line 2-5 of Algorithm 1.

Original: We discuss our implementation for teleporting multi-layer neural networks as an example.
Revised: We discuss our implementation for teleporting multi-layer neural networks as an example.

Original: Teleportating each pair of adjacent weight matrices requires computing the inverse of the output from the previous layer.
Revised: Teleportating each pair of adjacent weight matrices requires computing the inverse of the output from the previous layer.

Original: Denote the batch size as n , the largest dimension of all layers as d max , and the number of layers as l .
Revised: Denote the batch size as n , the largest dimension of all layers as d max , and the number of layers as l .

Original: Assume that d max > n .
Revised: Assume that d max > n .

Original: Computing the inverse of the output of each layer has complexity O ( d 2 max n ) , and computing the pseudoinverse for all layers has complexity
Revised: Computing the inverse of the output of each layer has complexity O ( d 2 max n ) , and computing the pseudoinverse for all layers has complexity

Original: O ( d 2 max nl ) .
Revised: O ( d 2 max nl ) .

Original: Note that all matrices we invert have dimensions at most d max × n .
Revised: Note that all matrices we invert have dimensions at most d max ⇥ n .

Original: For one gradient ascent step on g , the forward and backward pass both have complexity O
Revised: For one gradient ascent step on g , the forward and backward pass both have complexity O

Original: ( d 2 max nl ) .
Revised: ( d 2 max nl ) .

Original: This is the same as the forward and backward pass of gradient descent on w because the architecture is the same except with approximately twice as many layers.
Revised: This is the same as the forward and backward pass of gradient descent on w because the architecture is the same except with approximately twice as many layers.

Original: We perform t gradient ascent steps on g .
Revised: We perform t gradient ascent steps on g .

Original: Therefore, the computation cost for one teleportation is O ( d 2 max nlt ) .
Revised: Therefore, the computation cost for one teleportation is O ( d 2 max nlt ) .

Original: We show empirically that the runtime for teleportation scales polynomially with matrix dimensions and linearly with the number of layers.
Revised: We show empirically that the runtime for teleportation scales polynomially with matrix dimensions and linearly with the number of layers.

Original: We record the runtime of gradient descent on a Leaky-ReLU neural network for 300 epochs, with a 10-step teleportation every 10 epochs.
Revised: We record the runtime of gradient descent on a Leaky-ReLU neural network for 300 epochs, with a 10-step teleportation every 10 epochs.

Original: Each pair of adjacent weight matrices is transformed by a group action during teleportation.
Revised: Each pair of adjacent weight matrices is transformed by a group action during teleportation.

Original: Figure 6(a) shows the runtime for a 3-layer network using square weights and data matrices with different dimensions.
Revised: Figure 6(a) shows the runtime for a 3-layer network using square weights and data matrices with different dimensions.

Original: Figure 6(b) shows the runtime for 128-by-128 weight and data matrices with different number of layers.
Revised: Figure 6(b) shows the runtime for 128-by-128 weight and data matrices with different number of layers.

Original: Although the teleportation step has the same complexity as a gradient descent step, the runtime is dominated by teleportation due to larger constants in the complexity analysis.
Revised: Although the teleportation step has the same complexity as a gradient descent step, the runtime is dominated by teleportation due to larger constants in the complexity analysis.

Original: The trade-off between teleportation time and convergence rate depends on specific problems.
Revised: The trade-off between teleportation time and convergence rate depends on speciﬁc problems.

Original: In our experiments, the convergence rate is improved by a small number of teleportation steps which does not add significant computational overhead.
Revised: In our experiments, the convergence rate is improved by a small number of teleportation steps which does not add signiﬁcant computational overhead.

Original: [Figure]
Revised: [Figure]

Original:  Discussion and Conclusions
Revised:  Discussion and Conclusions

Original: We proposed a new optimization algorithm that exploits the symmetry in the loss landscape to improve convergence.
Revised: We proposed a new optimization algorithm that exploits the symmetry in the loss landscape to improve convergence.

Original: It is interesting to note that optimizing d L /dt locally sometimes leads to an improved global path.
Revised: It is interesting to note that optimizing d L /dt locally sometimes leads to an improved global path.

Original: One example is the ellipse function (Figure 1), where teleporting once ensures that d L /dt is optimal at every L value along the trajectory.
Revised: One example is the ellipse function (Figure 1), where teleporting once ensures that d L /dt is optimal at every L value along the trajectory.

Original: Another example is the matrix factorization problem L ( U, V ) =
Revised: Another example is the matrix factorization problem L ( U, V ) =

Original: ∥ Y − UV ∥ 22 .
Revised: k Y  UV k 22 .

Original: Tarmoun et al.
Revised: Tarmoun et al.

Original: (2021) shows that the convergence rate increases with the imbalance U T U − V T V .
Revised: (2021) shows that the convergence rate increases with the imbalance U T U  V T V .

Original: Consider the transformation U, V −→ Ug, g − 1 V .
Revised: Consider the transformation U, V ! Ug, g  1 V .

Original: To optimize d L ( U, V )
Revised: To optimize d L ( U, V )

Original: /dt locally, we would need a large g ∈ GL n , but a large g also leads to a large
Revised: /dt locally, we would need a large g 2 GL n , but a large g also leads to a large

Original: U T U − V T V which is positively correlated with the overall convergence rate of the entire trajectory.
Revised: U T U  V T V which is positively correlated with the overall convergence rate of the entire trajectory.

Original: Hence teleportation is guaranteed to produce a better trajectory.
Revised: Hence teleportation is guaranteed to produce a better trajectory.

Original: A potential future direction is to derive the exact expression for how teleportation affects the loss value at a later time in gradient flow, which may lead to a closed-form solution of the optimal teleportation destination.
Revised: A potential future direction is to derive the exact expression for how teleportation affects the loss value at a later time in gradient ﬂow, which may lead to a closed-form solution of the optimal teleportation destination.

Original: Additionally, inspired by the landscape view of parameter space symmetry ¸Sim¸sek et al.
Revised: Additionally, inspired by the landscape view of parameter space symmetry ¸Sim¸sek et al.

Original: (2021), teleportation using discrete (permutation) symmetries may allow us to reach a better minimum.
Revised: (2021), teleportation using discrete (permutation) symmetries may allow us to reach a better minimum.

Original: Finally, additional theory can be developed to explain the relationship between teleportation and second-order methods under the approximations we introduced, especially to quantify the improvement on the overall convergence rate, and to derive its effect on generalization bounds.
Revised: Finally, additional theory can be developed to explain the relationship between teleportation and second-order methods under the approximations we introduced, especially to quantify the improvement on the overall convergence rate, and to derive its effect on generalization bounds.

Original: Integrating teleportation with other advanced optimizers such as Adam and RMSprop would be another interesting future step.
Revised: Integrating teleportation with other advanced optimizers such as Adam and RMSprop would be another interesting future step.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: This work was supported in part by U S. Department Of Energy, Office of Science, U S. Army
Revised: This work was supported in part by U.S. Department Of Energy, Ofﬁce of Science, U. S. Army

Original: Research Office under Grant W911NF-20-1-0334, Google Faculty Award, Amazon Research Award, and NSF Grants #2134274, #2107256 and #2134178.
Revised: Research Ofﬁce under Grant W911NF-20-1-0334, Google Faculty Award, Amazon Research Award, and NSF Grants #2134274, #2107256 and #2134178.

Original: R. Walters is supported by the Roux Institute and the Harold Alfond Foundation.
Revised: R. Walters is supported by the Roux Institute and the Harold Alfond Foundation.

Original: We are grateful to Iordan Ganev for many helpful discussions
Revised: We are grateful to Iordan Ganev for many helpful discussions

Original: Checklist
Revised: Checklist

Original: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information on how to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] or [N/A] . You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description. For example: • Did you include the license to the code and datasets?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Original: [Yes] See Section ?? .

Original: • Did you include the license to the code and datasets?

Original: [No] The code and the data are proprietary.

Original: • Did you include the license to the code and datasets? [N/A] Please do not modify the questions and only use the provided macros for your answers. Note that the
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: Checklist section does not count towards the page limit.

Original: In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes](b) Did you describe the limitations of your work? [Yes] (c) Did you discuss any potential negative societal impacts of your work?
Revised: [Yes] (c) Did you discuss any potential negative societal impacts of your work?

Original: [No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: [No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 2.
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] (b) Did you include complete proofs of all theoretical results?
Revised: [Yes] (b) Did you include complete proofs of all theoretical results?

Original: [Yes] 3.
Revised: [Yes] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] 4.
Revised: [Yes] 4.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [N/A] (b) Did you mention the license of the assets?
Revised: [N/A] (b) Did you mention the license of the assets?

Original: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A] (d) Did you discuss whether
Revised: [N/A] (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: and how consent was obtained from people whose data you’re using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

Original: [N/A] 5.
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: A Adaptations of Algorithm 1 for different problems A.1 Stochastic gradient descent

Original: We extend Algorithm 1 to stochastic gradient descent (SGD).

Original: We apply the group actions using data from a mini-batch X i , and repeat for B mini-batches each time.

Original: The gradient we optimize, ˜ ∇L ( X i , g · w t ) , also uses single mini-batches.

Original: Algorithm 2 provides the framework for teleportation in SGD.

Original: Input batches teleport

Original: Output: w t

Original: 1 for t ← 0 to t max − 1 do if t ∈ K then for X i in the first B mini-batches do g ← argmax g ∈ G ∥

Original: ˜ ∇L ( X i , g · w t )

Original: ∥ w t ← g · w t w t ← w t − η ˜

Original: ∇L

Original: ( X i , w t ) end for for X i in the rest mini-batches do w t ← w t − η ˜ ∇L

Original: ( X i , w t ) end for else for all mini-batches X i do w t ← w t − η ˜

Original: X i , w t ) end for end if w t +1 ← t 17 end for 18 return w t max

Original: A.2 Data transformation

Original: Algorithm 3 here modifies Algorithm 1 to allow transformations on both parameters and data.

Original: Denote g X as the group action on data only.

Original: The group actions on data at all teleportation steps can be precomposed as a function f and applied to the input data at inference time.

Original: Algorithm 3: Symmetry Teleportation (with data transformation)

Original: Initialize: f = the identity function.

Original: 1 for t ← 0 to t max − 1 do if t

Original: K then g ← argmax g ∈ G ∥∇ g · w t L ( g · ( w t , X ))

Original: ∥ w t , X ← g · ( w t , X ) f ← g X ◦ f end if w t +1 ← w t − η ∇ w t L 8 end for 9 return w t max f

Original: B Group actions

Original: In this section, we derive the group actions for the test functions and multi-layer neural networks.

Original: More details about group theory can be found in textbooks such as Lang (2002).

Original: B.1 Continuous symmetry in test functions B.1.1 Ellipse

Original: Consider the following loss function with a ∈ R ≥ 0 :

Original: [Equation]

Original: If we change the variables to L ( u ( x 1 , x 2 ) , v ( x 1 , x 2 ))

Original: =

Original: u 2 + v 2 , 2D rotations leave L unchanged.

Original: Therefore SO (2) is a symmetry of L ( x 1 , x 2 ) .

Original: Let g θ ∈

Original: SO

Original: (2) , and define the group action as

Original: [Equation]

Original: Then

Original: [Equation]

Original: B.1.2 Rosenbrock function

Original: Consider the Rosenbrock function with 2 variables Rosenbrock (1960):

Original: [Equation]

Original: Let u = 10( x 21 − x 2 ) and v = x 1 − 1 .

Original: After changing the variables from x and y to u and v , L has a rotational symmetry.

Original: Note that the function, h : R 2 −→ R 2 , that maps x 1 , x 2 to u, v is bijective:

Original: [Equation]

Original: Next, we show that SO (2) is a symmetry of L ( x 1 , x 2 ) .

Original: Let ρ be a representation of SO (2) acting on R. For g ∈ SO (2) , define the following group action:

Original: [Equation]

Original: Then

Original: L

Original: [Equation] For the Rosenbrock function with 2N parameters, we can construct a bijective function h : R 2 N −→ R 2 N by transforming each of the N pairs of variables as before, and SO (2 N ) is a symmetry of L ( x 1 , ..., x 2 N ) . However, we will only use the 2 variable version in the experiments.
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: B.1.3 Booth function

Original: Consider the Booth function Jamil and Yang (2013):

Original: [Equation]

Original: Similar to the Rosebrock function, a change of variables reveals a rotational symmetry of L :

Original: [Equation]

Original: [Equation]

Original: The function h : R 2 −→ R 2 that maps x 1 , x 2 to u, v is bijective.

Original: Let ρ be a representation of SO (2) acting on R 2 .

Original: For g ∈ SO (2) , define the following group action:

Original: [Equation]

Original: Then L ( x 1 , x 2 ) admits an SO (2) symmetry:

Original: [Equation]

Original: B.2 Continuous symmetry in multi-layer neural Networks

Original: In this and the following sections, we provide proofs for the theoretical results.

Original: We restate the propositions from the main text for readability.

Original: Proposition 4.2.

Original: A linear network is invariant under all groups

Original: G m ≡ GL d m acting as

Original: [Equation]

Original: Proof.

Original: In the linear network h m =

Original: W m h m − 1 , hence

Original: [Equation]

Original: which means a p -layer linear network is invariant under all G m with m ≤ p as they keep the output h p invariant ( ∀ g ∈ G m , g · h p = h p ).

Original: Proposition 4.3.

Original: Assume that h m − 2 is invertible.

Original: A multi-layer network with bijective activation σ has a GL d m − 1 symmetry. For g m ∈ G m = GL d m − 1 ( R ) the following group action keeps h p with p ≥ m invariant [Equation]
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: Proof.

Original: From (7), we want to convert g · h m − 1 into a transformation on W m − 1 instead of h m − 1 .

Original: In other words, we want to find a set of transformed weights W ′ m , W ′ m − 1 which yields the same network output ˜ h m :

Original: [Equation]

Original: Solving (29) we get

Original: [Equation]

Original: 30) (28) follows from (29) and (30).

Original: To verify that (28) is a valid group action,

Original: [Equation]

Original: and

Original: [Equation]

Original: C Theoretical analysis of teleportation C.1 What symmetries help accelerate optimization

Original: Proposition 5.1.

Original: Let w ′ =

Original: g · w be a point we teleport to.

Original: Let J = ∂ w ′ /∂

Original: w be the Jacobian.

Original: Symmetry teleportation using g accelerates the rate of decay in L if it satisfies

Original: [Equation]

Original: Proof.

Original: Let w ′ =

Original: g · w .

Original: Denote the Jacobian as J , where J ij = ∂w ′

Original: i /∂w j .

Original: Then the inverse of J has entries

Original: [Equation]

Original: The gradient at w ′ is

Original: [Equation]

Original: The rate of change of L in gradient flow is

Original: [Equation]

Original: Thus we will have a speedup if

Original: [Equation]

Original: Proof of Corollary 5.

Original: Proof.

Original: Since J = g , using (9) and the l.h.s.

Original: of (10) we have

Original: [Equation]

Original: C.2 Improvement of subsequent steps

Original: Proposition 5.4. Consider the gradient descent with a G -invariant loss L ( w ) and learning rate η ∈ R + .

Original: Let w t be the parameter at time t and w ′ t

Original: = g · w t the parameter teleported by g ∈ G .

Original: Let w t + T and w ′ t + T be the parameters after T more steps of gradient descent from w t and w ′ t respectively.

Original: Under Assumption 5.3, if ηL < 1 , and

Original: [Equation]

Original: then

Original: [Equation]

Original: Proof.

Original: From the definition of Lipschitz continuity and the update rule of gradient descent,

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Equation]

Original: Equivalently,

Original: [Equation]

Original: By unrolling T steps

Original: [Equation]

Original: we have

Original: [Equation]

Original: [Equation]

Original: Similarly, for a teleported point w ′ t = g · w t ,

Original: [Equation]

Original: Therefore, if

Original: [Equation]

Original: then it is guaranteed that

Original: [Equation]

Original: C.3 Convergence analysis for convex quadratic functions

Original: We first show that starting from a point in S c , all other points in S c can be reached with one teleportation.

Original: Proposition C.1.

Original: S c contains a single orbit.

Original: That is, G · w ≡ { g · w : g ∈ G } =

Original: S c for all w ∈ S c .

Original: Proof.

Original: Consider two points w 1 , w 2 ∈ S c .

Original: Then w T 1 A w 1 =

Original: ( A 12 w 1 ) T ( A 12 w 1 )

Original: =

Original: c and w T 2 A w 2 =

Original: ( A 12 w 2 ) T ( A 12 w 2 )

Original: =

Original: c .

Original: [-Let v 
-]
Revised: [-Let v 
-]

Original: =

Original: A 12 w 1 √ c , v 2 =

Original: A 12 w 2 √ c and e 1 =

Original: [1 , 0 , ..., 0] T .

Original: Since ∥ v 1 ∥ =

Original: ∥ v 2 ∥ = 1 , there exists g 1 , g 2 ∈

Original: O ( n ) , such that g 1 e 1 = v 1 and g 2 e 1 = v 2 .

Original: One way to construct such gis let the first column be equal to v 1 and other columns be the rest of the orthonormal basis.

Original: Let g

Original: = g 2 g − 1 1 .

Original: [-Then v 
-]
Revised: [-Then v 
-]

Original: = g

Original: v 1 , A 12 w 2 = gA 12 w 1 , and w 2 =

Original: A − 12 gA 12 w 1 = g · w 1 .

Original: We have shown that for any w 1 , w 2 ∈ S c , there exists a g ∈ G such that w 2 = g · w 1 .

Original: Therefore, the group action of G on S c is transitive.

Original: Equivalently, S c contains a single orbit.

Original: The objective of teleportation is transforming parameters using a group element to maximize the norm of gradient:

Original: [Equation]
Revised: [N/A]

Original: Since all points on the level set are reachable, the target teleportation destination is the point with the largest gradient norm on the same level set.

Original: In other words, (43) is equivalent to the following optimization problem:

Original: [Equation]

Original: Let c =

Original: L A ( w ) .

Original: Substitute in L A and ∇L A , we have the following equivalent formulation:

Original: [Equation]

Original: Next, we solve this optimization problem and show that the gradient norm is maximized on the gradient flow trajectory starting from its solution.

Original: Proposition C.2.

Original: The solution to (45) is an eigenvector of A corresponding to its largest eigenvalue.

Original: Proof.

Original: We solve (45) using the method of Lagrangian multipliers.

Original: The Lagrangian of (45) is

Original: [Equation]

Original: Setting the derivative with respect of w ′ to 0, we have

Original: [Equation]

Original: which gives

Original: [Equation]

Original: Since A is positive definite, A

Original: =

Original: A T and A is invertible.

Original: Therefore,

Original: [Equation]

Original: so the solution to (45) is an eigenvector of A .

Original: Then, the constraint is w ′ T

Original: A w ′ =

Original: λ

Original: ∥ w ′ ∥ 2 = c , and the objective becomes max w ′ λ 2 ∥ w ′ ∥ 2 = max w ′ cλ .

Original: Therefore, we want λ to be the largest eigenvalue of A .

Original: Hence the optimal w ′ is an eigenvector of A corresponding to its largest eigenvalue.

Original: Proposition 5.5.

Original: If at point w , ∥∇L A | w ∥ 2 is at a maximum in S L A ( w ) , then for any point w ′ on the gradient flow trajectory starting from w , ∥∇L A | w ′ ∥ 2 is at a maximum in S L A ( w ′ ) .

Original: Proof.

Original: From Proposition C.2, w is an eigenvector of A corresponding to its largest eigenvalue.

Original: Then the gradient of L A is A w = λ w .

Original: Therefore, on the gradient flow trajectory starting from w , every point has the same direction as w , meaning that the points are all eigenvectors of A corresponding to its largest eigenvalue.

Original: Therefore, ∥∇L A ∥ 2 is always at a maximum on the loss level sets.

Original: Finally, we show that maximizing the magnitude of gradient is equivalent to minimizing the distance to w ∗ in a loss level set (Proposition 5.6).

Original: Proposition C.3.

Original: The solution to the following optimization problem is the same as the solution to (45) :

Original: [Equation]

Original: Proof.

Original: Similar to Proposition C.2, we solve this optimization problem using the method of Lagrangian multipliers.

Original: Note that w ∗ = 0 .

Original: The Lagrangian is

Original: [Equation]

Original: Setting the derivative with respect of w ′ to 0, we have

Original: [Equation]

Original: which gives

Original: [Equation]

Original: so the solution to (50) is an eigenvector of A .

Original: Then, the constraint is w ′ T

Original: A w ′ =

Original: λ

Original: ∥ w ′ ∥ 2 = c , and the objective becomes min w ′ ∥

Original: w ′ ∥ 2 = min w ′ cλ .

Original: Therefore, we want λ to be the largest eigenvalue of A .

Original: Hence the optimal w ′ is an eigenvector of A corresponding to its largest eigenvalue, which is the same as the solution to (45).

Original: For a more concrete example, consider a diagonal matrix A with positive diagonal elements.

Original: Then the level sets of L A are n -dimensional ellipsoids centered at the origin 0, with axes in the same directions as the standard basis.

Original: The point with largest ∥∇L A ∥ 2 on a level set is in the eigendirection of A corresponding to its largest eigenvalue, or equivalently, a point on the smallest semi-axes of the ellipsoid.

Original: [-Note that this point has the smallest distance to the global minimum w ∗ = 
-]
Revised: [-Note that this point has the smallest distance to the global minimum w ∗ = 
-]

Original: among all points in the same level set.

Original: In addition, the gradient flow trajectory from this point always points to w ∗ .

Original: Therefore, like the 2D ellipse function, one teleportation on the n -dimensional ellipsoid also guarantees optimal gradient norm at all points on the trajectory.

Original: C.4 Relation to second-order optimization methods

Original: To prove Proposition 5.7, we first note that when the norm of the gradient is at a critical point on the level set of the loss function, the gradient is an eigenvector of the Hession.

Original: Lemma C.4.

Original: If ∂ v (cid:13)(cid:13) ∂ L ∂ w (cid:13)(cid:13) 2 2 = 0 for all unit vector v that is orthogonal to ∂ L ∂ w , then ∂ L ∂ w is an eigenvector of the Hessian of L .

Original: Proof.

Original: From the definition of the directional derivative,

Original: [Equation]

Original: Writing the last term in indices,

Original: [Equation]

Original: (cid:19)

Original: [Equation]

Original: Removing the indices,

Original: [Equation]

Original: (cid:13)(cid:13)(cid:13)(cid:13)

Original: [Equation]

Original: Substitute back and we have

Original: [Equation]

Original: Since ∂ v (cid:13)(cid:13)

Original: ∂

Original: [-L ∂ w (cid:13)(cid:13) 2 2 = 0 for all vector v that is orthogonal to ∂ L ∂ w , v · (cid:0) 2 H ∂ L ∂ w (cid:1) = 
-]
Revised: [-L ∂ w (cid:13)(cid:13) 2 2 = 0 for all vector v that is orthogonal to ∂ L ∂ w , v · (cid:0) 2 H ∂ L ∂ w (cid:1) = 
-]

Original: for all vector v that is orthogonal to ∂ L ∂ w .

Original: In other words, 2 H ∂ L ∂ w is orthogonal to all vectors that are orthogonal to ∂ L ∂ w .

Original: Therefore, 2 H ∂ L ∂ w has the same direction of ∂ L ∂ w , and ∂ L ∂ w is an eigenvector of the Hessian of L .

Original: Proposition 5.7 is a direct consequence of Lemma C.4.

Original: Proposition 5.7.

Original: Let S L 0 = { w : L ( w )

Original: =

Original: L 0 } be a level set of L .

Original: If at a particular w ∈ S L 0 we have ∥∇L ( w ) ∥ 2 ≥ ∥∇L ( w ′ )

Original: ∥ 2 for all w ′ in a small neighborhood of w in S L 0 , then the gradient ∇L ( w ) has the same direction as the Newton’s direction H − 1 ∇L ( w ) .

Original: Proof.

Original: From Lemma C.4, dLdw is an eigenvector of H .

Original: Therefore, it is also an eigenvector of H − 1 .

Original: Hence dLdw has the same direction as H − 1 dLdw .

Original: D Experiment details and additional results D.1 Test functions

Original: We compare the gradient at different loss values for gradient descent with and without teleportation.

Original: Figure 7 shows that the trajectory with teleportation has a larger d L /dt value than the trajectory without teleportation at the same loss values.

Original: Therefore, the rate of change in the loss is larger in the trajectory with teleportation, which makes it favorable.

Original: [Figure]

Original: [Figure]

Original: D.2 Multilayer neural network

Original: Additional training details Data X, Y and initialization of parameters W are set uniformly at random over [0 , 1] .

Original: GD uses learning rate 10 − 4 and AdaGrad uses 10 − 1 .

Original: Each algorithm is runsteps.

Original: When using teleportation, we perform symmetry transform on the parameters once at epoch 5.

Original: In GD, the group elements used for these transforms are found by gradient ascent on T for 8 steps, with learning rate 10 − 7 .

Original: In AdaGrad, the group elements are found by gradient ascent for 2 steps, with learning rate 10 − 5 .

Original: The choice of hyperparameters comes from a grid search described in the next section.

Original: Hyperparameter tuning To observe the effect of hyperparameters on the speedup in computation time, we did a hyperparameter sweep on the number of steps and the learning rate used in each teleportation.

Original: The speedup of teleportation on SGD and AdaGrad is defined by t sgd /t sgd + teleport and t adagrad /t adagrad + teleport respectively, where t sgd , t adagrad are the wall-clock time required to reach convergence using SGD or AdaGrad, and t sgd + teleport , t adagrad + teleport are convergence time with teleportation.

Original: We consider the optimization algorithm converged if the difference between the loss of two consecutive steps is less than 10 − 3 .

Original: This experiment is run on one CPU.

Original: Figure 8 shows the speedup of the same multilayer neural network regression problem defined in

Original: Section 6.1, but teleporting only once at epoch 5.

Original: We did a grid search for teleportation learning rates in [10 − 9 , 10 − 8 , 10 − 7 , 10 − 6 , 10 − 5 ] and number of teleportation steps in [1 , 2 , 4 , 8 , 16 , 32] .

Original: Omitted points in the figure indicate that the gradient descent fails to converge within 2000 steps or diverges.

Original: When converged, most hyperparameter combinations improve the convergence speed in wall-clock time (speedup > 1 ).

Original: There are trade-offs in both the number of steps used for each teleportation and the teleportation learning rate.

Original: Increasing the number of steps used to optimize teleportation target allows us to find a better point in the parameter space but increases the cost of one teleportation.

Original: Increasing the learning rate of optimizing the group element improves ∥ ∂ L /∂ w ∥ but is more likely to lead to divergence since ∥ ∂ L /∂ w ∥ can become too large for the gradient descent learning rate.