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

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AffiliationAddress email

Original: Abstract
Revised: Abstract

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

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

Original: Teleportationexploits parameter space symmetries of the optimization problem and transformsparameters while keeping the loss invariant.
Revised: Teleportationexploits parameter space symmetries of the optimization problem and transformsparameters while keeping the loss invariant.

Original: We derive the loss-invariant groupactions for test functions and multi-layer neural networks, and prove a necessarycondition of when teleportation improves convergence rate.
Revised: We derive the loss-invariant groupactions for test functions and multi-layer neural networks, and prove a necessarycondition of when teleportation improves convergence rate.

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

Original: Experimentally, we show thatteleportation improves the convergence speed of gradient descent and AdaGradfor several optimization problems including test functions, multi-layer regressions,and MNIST classification.
Revised: Experimentally, we show thatteleportation improves the convergence speed of gradient descent and AdaGradfor several optimization problems including test functions, multi-layer regressions,and MNIST classification.

Original:  Introduction
Revised:  Introduction

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

Original: While finding global minima of L ( w ) is hard for non-convex problems, we can usegradient descent (GD) to find local minima.
Revised: While finding global minima of L ( w ) is hard for non-convex problems, we can usegradient descent (GD) to find 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 first-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 methodsenjoy faster convergence, but calculating the second derivative (Hessian) can be computationallyexpensive in high dimensions Hazan (2019).
Revised: These methodsenjoy faster convergence, but calculating the second derivative (Hessian) can be computationallyexpensive in high dimensions Hazan (2019).

Original: In this paper, we propose a new optimization method based on parameter space symmetries.
Revised: In this paper, we propose a new optimization method based on parameter space symmetries.

Original: Ouralgorithm takes advantage of higher-order landscape geometry but uses only gradient information inmost steps and thus avoids the computational cost of second-order methods.
Revised: Ouralgorithm takes advantage of higher-order landscape geometry but uses only gradient information inmost steps and thus avoids the computational cost of second-order methods.

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

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

Original: Afterteleportation, the loss stays the same, but the gradient and hence the rate of loss decay may change.
Revised: Afterteleportation, the loss stays the same, but the gradient and hence the rate of loss decay may change.

Original: The increase in the magnitude of gradient can reduce the time required to converge, leading toacceleration of the gradient flow.
Revised: The increase in the magnitude of gradient can reduce the time required to converge, leading toacceleration of the gradient flow.

Original: These transformations move the parameters to a better trajectoryunder gradient descent, which improves the convergence speed of subsequent steps.
Revised: These transformations move the parameters to a better trajectoryunder gradient descent, which improves the convergence speed of subsequent steps.

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

Original: For GD, an equivalentproblem uses the following proximal mapping Combettes and Pesquet (2011).
Revised: For GD, an equivalentproblem uses the following proximal mapping Combettes and Pesquet (2011).

Original: Let ⟨· , ·⟩ denote theinner product, we have
Revised: Let ⟨· , ·⟩ denote theinner 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 ∥ w − w t ∥ 2 is the proximal term that keeps w t +1 close to w t .

Original: Adaptive gradient methods,such as Adagrad, define the proximal term using the Mahalanobis distance ∥ w − w t ∥ 2 G − 1 / 2 (Duchiet al., 2011) to account for landscape geometry.
Revised: Adaptive gradient methods,such as Adagrad, define the proximal term using the Mahalanobis distance ∥ w − w t ∥ 2 G − 1 / 2 (Duchiet al., 2011) to account for landscape geometry.

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

Original: We teleport to points on the same level set of L ( w t ) , but alloww to be far from w t in Euclidean distance.
Revised: We teleport to points on the same level set of L ( w t ) , but alloww 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 algorithmthat exploits symmetries in the loss landscape.
Revised: • We propose symmetry teleportation , an accelerated gradient-based optimization algorithmthat 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 convergencerate of the current and subsequent steps.
Revised: We provide necessary conditions and examples of when teleportation improves convergencerate 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
Revised: Experimentally, we show that teleportation improves the convergence speed of gradient

Original: descent and AdaGrad in different optimization problems.
Revised: descent and AdaGrad in different 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 and Walters,2021).
Revised: 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 and Walters,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 derivingloss-invariant group actions in multi-layer neural networks with invertible activation functions.
Revised: We contribute to this line of work by derivingloss-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 (Neyshaburet al., 2015) improves optimization in ReLU networks by path regularization.
Revised: For example, Path-SGD (Neyshaburet al., 2015) improves optimization in ReLU networks by path regularization.

Original: G -SGD (Meng et al.
Revised: G -SGD (Meng et al.

Original: ,2019) performs weight updates in a scale-invariant space, which also exploits the symmetry in ReLUnetworks.
Revised: ,2019) performs weight updates in a scale-invariant space, which also exploits the symmetry in ReLUnetworks.

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

Original: Bamler and Mandt (2018) transforms parameters in the symmetryorbits to address slow movement along directions of weakly broken symmetry and find the optimalg ∈ G that minimizes L ( g · w ) .
Revised: Bamler and Mandt (2018) transforms parameters in the symmetryorbits to address slow movement along directions of weakly broken symmetry and find the optimalg ∈ 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: ∥∇L ( g · w ) ∥ 2 .

Original: Natural gradient Amari (1998), adaptive gradient methods Duchi et al.
Revised: Natural gradient Amari (1998), adaptive gradient methods Duchi et al.

Original: (2011); Kingma and Ba (2014),and their approximations Martens and Grosse (2018)
Revised: (2011); Kingma and Ba (2014),and their approximations Martens and Grosse (2018)

Original: ; Gupta et al.
Revised: ; Gupta et al.

Original: (2018) improve the direction ofparameter updates instead of directly transforming parameters.
Revised: (2018) improve the direction ofparameter updates instead of directly transforming parameters.

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

Original: We prove that our algorithm is connected to second-order optimizationmethods and show that teleportation can be used to improve these methods empirically.
Revised: We prove that our algorithm is connected to second-order optimizationmethods 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
Revised: The concept of neural teleportation was first explored using quiver representation theory Armenta and

Original: Jodoin (2021); Armenta et al. (2020).
Revised: Jodoin (2021); Armenta et al. (2020).

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

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

Original: Several works study how parameter initialization affects convergence
Revised: Several works study how parameter initialization affects convergence

Original: [Table]
Revised: [Table]

Original: is similar to that of Tarmoun et al.
Revised: is similar to that of Tarmoun et al.

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

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

Original: Additionally, our method is not restricted to initialization and canbe applied at any time during training.
Revised: Additionally, our method is not restricted to initialization and canbe 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 improvegeneralization (Belkin et al., 2019), an important question is how overparametrization affects optimization.
Revised: While this has been shown to improvegeneralization (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 resultingfrom the depth of a neural network accelerates convergence.
Revised: (2018) show that overparametrization resultingfrom the depth of a neural network accelerates convergence.

Original: Another view is that the symmetrycreated by overparametrization poses constraints on trajectories in the form of conserved quantities(Głuch and Urbanke, 2021).
Revised: Another view is that the symmetrycreated 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: Whenthe trajectories created by overparametrization are equivalent, model compression by removingsymmetry reduces training time (Ganev and Walters, 2021).
Revised: Whenthe trajectories created by overparametrization are equivalent, model compression by removingsymmetry reduces training time (Ganev and Walters, 2021).

Original: However, when trajectories are notequivalent, gradient flows on some paths are faster than others.
Revised: However, when trajectories are notequivalent, gradient flows on some paths are faster than others.

Original: We search for the better paths createdby overparametrization.
Revised: We search for the better paths createdby overparametrization.

Original:  Symmetry Teleportation
Revised:  Symmetry Teleportation

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

Original: Symmetry teleportation searches for the best gradientdescent trajectory by teleporting parameters to a point with larger gradients using a group action.
Revised: Symmetry teleportation searches for the best gradientdescent 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 globallandscape geometry, leading to faster loss decay.
Revised: The resulting algorithm requires only gradient computations but is able to account for the globallandscape 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 ∈ 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 define 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, wetransform parameters using group element g ∈ G to the location where the gradient is largest, thencontinue with gradient descent steps.
Revised: At epochs that are in the schedule, wetransform parameters using group element g ∈ G to the location where the gradient is largest, thencontinue with gradient descent steps.

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

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

Original: t −
Revised: t −

Original: 7 end
Revised: 7 end

Original: 8 return
Revised: 8 return

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 andperforming gradient ascent on g .
Revised: When the action of G is continuous, teleportation can be implemented by parameterizing andperforming gradient ascent on g .

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

Original: Small transformations g ∈ GL d ( R ) (general linear group) can be parameterized asg
Revised: Small transformations g ∈ GL d ( R ) (general linear group) can be parameterized asg

Original: ≈ I + εT where ε ≪ 1 and T are arbitrary d × d matrices.
Revised: ≈ I + εT where ε ≪ 1 and T are arbitrary d × d matrices.

Original: For discrete groups, search algorithmsor random sampling can be used to find a group element that improves the magnitude of the gradient.
Revised: For discrete groups, search algorithmsor random sampling can be used to find 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 beextended 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 beextended to allow transformations on both parameters and data.

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

Original: The algorithmextension that allows transformation on the input can be found in Appendix A.2.
Revised: The algorithmextension 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 orteleporting only at the earlier epochs.
Revised: We define K as a set to allow flexible teleportation schedules, such as with non-fixed frequencies orteleporting 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 optimizationalgorithms Back (1996).
Revised: Specifically, we first consider two test functions which are often used to evaluate optimizationalgorithms 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 acharacteristic global minimum that is inside a long, narrow, parabolicly-shaped flat valley.
Revised: The Rosenbrock function originally introduced by Rosenbrock (1960) has acharacteristic global minimum that is inside a long, narrow, parabolicly-shaped flat valley.

Original: Findingthe valley is easy but reaching the minimum is difficult.
Revised: Findingthe valley is easy but reaching the minimum is difficult.

Original: On a 2-dimensional space, the Rosenbrockfunction has the following form:
Revised: On a 2-dimensional space, the Rosenbrockfunction 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 oneglobal minimum at (1 , 3) where the function evaluates to 0 :
Revised: The Booth function Jamil and Yang (2013) is also defined on R 2 and has oneglobal 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 identifies 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 ∈ 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 teleportationstep, our goal is to maximize the gradient within a level set of the loss: max g ∈ SO(2)
Revised: During the teleportationstep, our goal is to maximize the gradient within a level set of the loss: max g ∈ SO(2)

Original: ∥
Revised: ∥

Original: d L ( g · ( x 1 ,x 2 )) dt ∥ .
Revised: d L ( g · ( x 1 ,x 2 )) dt ∥ .

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 ∈
Revised: Denote the input by h 0 = X ∈

Original: R d 0 × n .
Revised: R d 0 × n .

Original: Each layer first performs a linear transformation ˜
Revised: Each layer first performs a linear transformation ˜

Original: h m
Revised: h m

Original: =
Revised: =

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

Original: d m × d m − 1(we absorb biases into W m by adding an extra row of ones in ˜ h m − 1 ).
Revised: d m × d m − 1(we absorb biases into W m by adding an extra row of ones in ˜ h m − 1 ).

Original: The output h m is
Revised: 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
Revised: A linear network is invariant under all groups

Original: G m ≡ GL d m acting as
Revised: G m ≡ GL d m acting as

Original: [Equation]
Revised: [Equation]

Original: 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 find 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: 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: Forg m ∈ G m =
Revised: Forg m ∈ G m =

Original: GL d m − 1 ( R ) the following group action keeps h p with p ≥ m invariant
Revised: GL d m − 1 ( R ) the following group action keeps h p with p ≥ m invariant

Original: [Equation]
Revised: [Equation]

Original: where h + m − 2 is the right pseudo-inverse for h m − 2 (assuming n ≪ d m − 2 and that h m − 2 is full rank).
Revised: where h + m − 2 is the right pseudo-inverse for h m − 2 (assuming n ≪ d m − 2 and that h m − 2 is full rank).

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 networkas well as the current weights of all the lower layers.
Revised: Note that this group action depends on the input to the networkas well as the current weights of all the lower layers.

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

Original: That is, (8) gives a well-defined non-linear group action of G = GL d 1 × . . .
Revised: That is, (8) gives a well-defined non-linear group action of G = GL d 1 × . . .

Original: GL d
Revised: GL d

Original: [Equation]
Revised: [Equation]

Original: − 1on ( X, W 1 , . . .
Revised: − 1on ( X, W 1 , . . .

Original: , W p ) which depends on but fixes X and keeps the loss function invariant.
Revised: , W p ) which depends on but fixes X and keeps the loss function invariant.

Original:  Theoretical Analysis
Revised:  Theoretical Analysis

Original: 5.1 What symmetries help accelerate optimization
Revised: 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 satisfied 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 singleflattened 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 singleflattened vector w ∈ R n .

Original: Consider a symmetry G of the loss function L ( w ) , meaning ∀ g ∈ G ,
Revised: Consider a symmetry G of the loss function L ( w ) , meaning ∀ g ∈ 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 isthe Mahalanobis norm.
Revised: where η is the matrix of learning rates (which must be positive semi-definite) and ∥ v ∥ 2 η = v T ηv isthe Mahalanobis norm.

Original: For constant learning rate, η = I .
Revised: For constant learning rate, η = I .

Original: The following proposition states the condition a symmetry needs to satisfy to accelerate optimization.(Proofs can be found in Appendix C.1.)
Revised: The following proposition states the condition a symmetry needs to satisfy to accelerate optimization.(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 ′ /∂

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 satisfies

Original: [Equation]
Revised: [Equation]

Original: If the action of the symmetry group G ⊂ GL n is linear we have w ′ =
Revised: If the action of the symmetry group G ⊂ GL n is linear we have w ′ =

Original: g · w = g w and J = g .
Revised: g · w = g w and J = g .

Original: Itfollows that if G is a subgroup of the certain orthogonal groups, d L /dt will be invariant:
Revised: Itfollows that if G is a subgroup of the certain orthogonal groups, 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 ∈
Revised: Let O η denote the orthogonal group of invariances of the inverse of the learning rate,η − 1 , meaning for g ∈

Original: O η , g T η − 1 g
Revised: O η , g T η − 1 g

Original: =
Revised: =

Original: η − 1 .
Revised: η − 1 .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: η ,
Revised: η ,

Original: [Equation]
Revised: [Equation]

Original: dt (11)
Revised: dt (11)

Original: In the simple case where the learning rate is a constant, O η
Revised: In the simple case where the learning rate is a constant, O η

Original: =
Revised: =

Original: O
Revised: O

Original: ( n ) becomes the classic orthogonalgroup (e.g. rotations) with g T g
Revised: ( n ) becomes the classic orthogonalgroup (e.g. rotations) with g T g

Original: = I .
Revised: = I .

Original: In general, when g preserves the norm of the gradient, symmetryteleportation has no effect on the convergence rate.
Revised: In general, when g preserves the norm of the gradient, symmetryteleportation 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 , thegradient 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 , thegradient norm in gradient descent converges as 12 β (cid:80) Tt =1 ∥ ( ∇L ) | w t ∥ ≤ 2 M .

Original: Teleportation increases( ∇L ) | w t , therefore requiring less time to reach convergence.
Revised: Teleportation increases( ∇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 stepsare local, we would expect that teleportation improves the magnitude of gradient for a few futuresteps as well.
Revised: Since teleportation moves the parameters to a point with a larger gradient, and subsequent GD stepsare local, we would expect that teleportation improves the magnitude of gradient for a few futuresteps 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 withconstant L ∈ R ≥ 0 , which is
Revised: The l2 norm of the gradient is Lipschitz continuous withconstant L ∈ R ≥ 0 , which is

Original: [Equation]
Revised: [Equation]

Original: (12) 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η ∈ R + .

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

Original: = g · w t the parameter teleported by g ∈ G .
Revised: = 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 ′ trespectively.
Revised: Let w t + T and w ′ t + T be the parameters after T more steps of gradient descent from w t and w ′ trespectively.

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] [Equation]
Revised: [Equation] [Equation]

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

Original: Thecondition is met when L is small, η is small, or the initial improvement, (cid:13)(cid:13)(cid:13) ∂
Revised: Thecondition is met when L is small, η is small, or the initial improvement, (cid:13)(cid:13)(cid:13) ∂

Original: L ∂ w ′ t (cid:13)(cid:13)(cid:13) 2 / (cid:13)(cid:13)(cid:13) ∂
Revised: L ∂ w ′ t (cid:13)(cid:13)(cid:13) 2 / (cid:13)(cid:13)(cid:13) ∂

Original: L ∂ w t (cid:13)(cid:13)(cid:13) 2 , is large.
Revised: L ∂ w t (cid:13)(cid:13)(cid:13) 2 , is large.

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

Original: Symmetry teleportation is closely related to second-order optimization methods.
Revised: Symmetry teleportation is closely related to second-order optimization methods.

Original: We show that at thetarget point to teleport to, gradient descent becomes equivalent to Newton’s method.
Revised: We show that at thetarget point to teleport to, gradient descent becomes equivalent to Newton’s method.

Original: Full proofs ofthe results in this section can be found in Appendix C.3 and C.4.
Revised: Full proofs ofthe results in this section can be found in Appendix C.3 and C.4.

Original: Lemma 5.5.
Revised: Lemma 5.5.

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 aneigenvector of the Hessian of L .
Revised: 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 aneigenvector of the Hessian of L .

Original: The following proposition is a direct consequence of Lemma 5.5.
Revised: The following proposition is a direct consequence of Lemma 5.5.

Original: Proposition 5.6.
Revised: Proposition 5.6.

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 wehave ∥∇L ( w ) ∥ 2 ≥ ∥∇L ( w ′ )
Revised: If at a particular w ∈ S L 0 wehave ∥∇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 ) .
Revised: ∥ 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: Proposition 5.6 provides an alternative way to interpret teleportation.
Revised: Proposition 5.6 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 samedirection as Newton’s direction.
Revised: Newton’s direction, we search within the loss level set for a point where the gradient has the samedirection as Newton’s direction.

Original: However, symmetry teleportation does not require computing thefull Hessian matrix.
Revised: However, symmetry teleportation does not require computing thefull Hessian matrix.

Original: The second derivative required for optimizing continuous groups are obtained bytaking derivatives with respect to parameters and then with respect to the group element, as opposedto taking derivatives with respect to parameters twice.
Revised: The second derivative required for optimizing continuous groups are obtained bytaking derivatives with respect to parameters and then with respect to the group element, as opposedto taking derivatives with respect to parameters twice.

Original: This makes the computation significantly morefeasible than Newton’s method on neural networks with large number of parameters.
Revised: This makes the computation significantly morefeasible 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 applyteleportation in every gradient descent step.
Revised: We also do not applyteleportation in every gradient descent step.

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

Original: We provide empirical evidence in Section 6that
Revised: We provide empirical evidence in Section 6that

Original: these approximations do not erase the benefits of teleportation completely, and leave theoreticalinvestigations of the connection to second-order methods under approximations as future work.
Revised: these approximations do not erase the benefits of teleportation completely, and leave theoreticalinvestigations of the connection to second-order methods under approximations as future work.

Original:  Experiments
Revised:  Experiments

Original: 6.1 Acceleration through symmetry teleportation
Revised: 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 theparameter space on two test functions and show the speedup in regression and classification problemsusing multilayer neural networks.
Revised: We illustrate teleportation in theparameter space on two test functions and show the speedup in regression and classification problemsusing 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 Rosenbrockfunction (3).
Revised: We apply symmetry teleportation to optimize the 2-variable Rosenbrockfunction (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 stepswith learning rate 10 − 3 .
Revised: Each algorithm is run 1000 stepswith 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 foundby gradient ascent on θ , the parameter for the SO(2) group, for 10 steps with learning rate 10 − 1 .
Revised: The group elements are foundby 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 denotesthe final position of parameters, the green star denotes the target, and orange dots are the positionsfrom which symmetry transforms start.
Revised: The blue star denotesthe final position of parameters, the green star denotes the target, and orange dots are the positionsfrom which symmetry transforms start.

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

Original: Figure 2d shows thatteleportation improves the norm of gradients in the following step and 2c shows its effect on the lossvalue.
Revised: Figure 2d shows thatteleportation improves the norm of gradients in the following step and 2c shows its effect on the lossvalue.

Original: Teleportations clearly reduce the overall convergence time.
Revised: Teleportations clearly reduce the overall convergence time.

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 defined 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 elementsare found by gradient ascent on θ for 10 steps with learning rate 0 .
Revised: The group elementsare found by gradient ascent on θ for 10 steps with learning rate 0 .

Original: θ is initialized uniformlyat random over [0 , π ) .
Revised: θ is initialized uniformlyat random over [0 , π ) .

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

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

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

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

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

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

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

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

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

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

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

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

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

Original: We further evaluated our method on a three-layer neuralnetwork 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 neuralnetwork with a regression loss min W 1 ,W 2 ,W 3 ∥ Y − W 3 σ ( W 2 σ ( W 1 X ))

Original: ∥ 2 .
Revised: ∥ 2 .

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

Original: X ∈ R 5 × 4 is the data, Y ∈ R 8 × 4 is thetarget, and σ is the LeakyReLU activation with slope coefficient 0.1.
Revised: X ∈ R 5 × 4 is the data, Y ∈ R 8 × 4 is thetarget, and σ is the LeakyReLU activation with slope coefficient 0.1.

Original: Discussion of the choice oflearning rates and additional hyperparameter details can be found in Appendix E.2.
Revised: Discussion of the choice oflearning rates and additional hyperparameter details can be found in Appendix E.2.

Original: Teleportation of the GL( R ) group can be performed by finding an element x in its Lie algebra, suchthat 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 finding an element x in its Lie algebra, suchthat transforming w by the group element g = exp( x ) improves the gradient d L dt .

Original: We use the firstorder approximation of the exponential map.
Revised: We use the firstorder approximation of the exponential map.

Original: Between each pair of weight matrices, we replace g mby I + T and g − 1 m by I − T in (8),
Revised: Between each pair of weight matrices, we replace g mby I + T and g − 1 m by I − T in (8),

Original: where T ∈ R d m × d m is initialized to 0.
Revised: where T ∈ R d m × d m is initialized to 0.

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

Original: Figure 3a shows the training curves, where shaded area denotes one standard deviation from 5 runs.
Revised: Figure 3a shows the training curves, where 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 3b and 3c shows the squarednorm of gradient from a single run.
Revised: Figure 3b and 3c shows the squarednorm of gradient from a single run.

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

Original: MNIST classification.
Revised: MNIST classification.

Original: We apply symmetry teleportation on the MNIST classification task Deng(2012).
Revised: We apply symmetry teleportation on the MNIST classification 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 hasdimension 28 × 28 and is flattened into a vector.
Revised: The input data hasdimension 28 × 28 and is flattened into a vector.

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

Original: We used a three-layer neural network with hidden dimension
Revised: We used a three-layer neural network with hidden dimension

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

Original: [512 , 512] , LeakyReLU activations, and cross-entropy loss.
Revised: [512 , 512] , LeakyReLU activations, and cross-entropy loss.

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

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

Original: Immediately after the first epoch, we apply teleportation using data from one mini-batch, and repeatfor 4 different mini-batches.
Revised: Immediately after the first epoch, we apply teleportation using data from one mini-batch, and repeatfor 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,d shows the effect of teleportation on training and validation loss, and Figure 4b,c showsthe norm of the gradient in training.
Revised: Figure 4a,d shows the effect of teleportation on training and validation loss, and Figure 4b,c showsthe norm of the gradient in training.

Original: While teleportation significantly accelerates the decrease ofthe training loss in SGD, its effect on the validation loss is limited and detrimental for AdaGrad.
Revised: While teleportation significantly accelerates the decrease ofthe 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 overfitand slightly worse validation accuracy.
Revised: Therefore, teleportation on MNIST makes training faster at the beginning but leads to earlier overfitand slightly worse validation accuracy.

Original: A possible reason is that regions with large gradients havesharp minima that do not generalize well.
Revised: A possible reason is that regions with large gradients havesharp 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 ofteleportation on MNIST with different hyperparameters.
Revised: Figure 5 shows the result ofteleportation on MNIST with different hyperparameters.

Original: In all experiments, we use SGD with batchsize 20, learning rate 2 × 10 − 3 , and 10 gradient ascent steps for each teleportation.
Revised: In all experiments, we use SGD with batchsize 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 themindividually, but at different epochs.
Revised: In Figure 5a, we randomly select 4 different mini-batches and apply teleportation on each of themindividually, but at different epochs.

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

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

Original: In Figure 5b, we again applyteleportations on 4 randomly selected mini-batches, but repeat this with 5 different teleportationschedules (hyperparameter K in Algorithm 2) as shown in the legend.
Revised: In Figure 5b, we again applyteleportations on 4 randomly selected mini-batches, but repeat this with 5 different teleportationschedules (hyperparameter K in Algorithm 2) as shown in the legend.

Original: All schedules have thesame number of teleportations.
Revised: All schedules have thesame number of teleportations.

Original: Smaller intervals between teleportation accelerate convergencemore significantly.
Revised: Smaller intervals between teleportation accelerate convergencemore significantly.

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

Original: Using more mini-batches to teleport leads to faster decrease in training loss but isalso more prone to overfitting.
Revised: Using more mini-batches to teleport leads to faster decrease in training loss but isalso more prone to overfitting.

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 networksas an example.
Revised: We discuss our implementation for teleporting multi-layer neural networksas an example.

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

Original: Denote the batch size as n , the largest dimension of all layers as
Revised: Denote the batch size as n , the largest dimension of all layers as

Original: [Figure]
Revised: [Figure]

Original: d max , and the number of layers as l .
Revised: d max , and the number of layers as l .

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

Original: Computing the pseudoinverse of theoutput of each layer has complexity O ( d 2 max n ) , and computing the pseudoinverse for all layers hascomplexity O ( d 2 max nl ) .
Revised: Computing the pseudoinverse of theoutput of each layer has complexity O ( d 2 max n ) , and computing the pseudoinverse for all layers hascomplexity 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 descents on w because the architectureis the same except with approximately twice as many layers.
Revised: This is the same as the forward and backward pass of gradient descents on w because the architectureis 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 dimensionsand linearly with the number of layers.
Revised: We show empirically that the runtime for teleportation scales polynomially with matrix dimensionsand linearly with the number of layers.

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

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

Original: Figure 6(a) shows the runtimefor a 3-layer network using square weights and data matrices with different dimensions.
Revised: Figure 6(a) shows the runtimefor 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 isdominated 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 isdominated by teleportation due to larger constants in the complexity analysis.

Original: The trade-off betweenteleportation time and convergence rate depends on specific problems.
Revised: The trade-off betweenteleportation time and convergence rate depends on specific problems.

Original: In our experiments, theconvergence rate is improved by a small number of teleportation steps which does not add significantcomputational overhead.
Revised: In our experiments, theconvergence rate is improved by a small number of teleportation steps which does not add significantcomputational overhead.

Original:  Discussion and Conclusions
Revised:  Discussion and Conclusions

Original: We proposed an accelerated optimization method that exploits the symmetry in the parameter spaceof the loss function to improve convergence.
Revised: We proposed an accelerated optimization method that exploits the symmetry in the parameter spaceof the loss function to improve convergence.

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

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

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

Original: ∥
Revised: ∥

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

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

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

Original: Consider the transformation
Revised: Consider the transformation

Original: U, V −→ Ug, g − 1 V .
Revised: 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 largeg also leads to a large U T U − V T V which is positively correlated with the overall convergence rateof the entire trajectory.
Revised: /dt locally, we would need a large g ∈ GL n , but a largeg also leads to a large U T U − V T V which is positively correlated with the overall convergence rateof 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 lossvalue at a later time in gradient flow, which may lead to a closed-form solution of the optimalteleportation destination.
Revised: A potential future direction is to derive the exact expression for how teleportation affects the lossvalue at a later time in gradient flow, which may lead to a closed-form solution of the optimalteleportation 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 reacha better minimum.
Revised: (2021), teleportation using discrete (permutation) symmetries may allow us to reacha better minimum.

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

Original: Integrating teleportation with other advanced optimizers such as Adam, RMSprop would beanother interesting future step.
Revised: Integrating teleportation with other advanced optimizers such as Adam, RMSprop would beanother interesting future step.

Original: Checklist
Revised: Checklist

Original: The checklist follows the references.
Revised: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.
Revised: Please read the checklist guidelines carefully for information onhow to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
Revised: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
Revised: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:
Revised: For example:

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

Original: [Yes] See Section ?? .
Revised: [Yes] See Section ?? .

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

Original: [No] The code and the data areproprietary.•
Revised: [No] The code and the data areproprietary.•

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

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

Original: Please do not modify the questions and only use the provided macros for your answers.
Revised: Please do not modify the questions and only use the provided macros for your answers.

Original: Note that the
Revised: Note that the

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

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

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?

Original: [Yes](b) Did you describe the limitations of your work?
Revised: [Yes](b) Did you describe the limitations of your work?

Original: [Yes](c)
Revised: [Yes](c)

Original: Did you discuss any potential negative societal impacts of your work?
Revised: 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 tothem?
Revised: [No](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

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

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (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: (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]
Revised: [Yes]

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

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

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

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

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

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

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?

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

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiableinformation or offensive content?
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiableinformation or offensive content?

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

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

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

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

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

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?
Revised: spent on participant compensation?

Original: [N
Revised: [N

Original: /A]
Revised: /A]

Original: A Adaptations of Algorithm 1 for different problems
Revised: A Adaptations of Algorithm 1 for different problems

Original: A.1 Stochastic gradient descent
Revised: A.1 Stochastic gradient descent

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

Original: We apply the group actions usingdata from a mini-batch X i , and repeat for B mini-batches each time.
Revised: We apply the group actions usingdata 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.
Revised: The gradient we optimize,˜ ∇L ( X i , g · w t ) , also uses single mini-batches.

Original: Algorithm 2 provides the framework for teleportationin SGD.
Revised: Algorithm 2 provides the framework for teleportationin SGD.

Original: Algorithm
Revised: Algorithm

Original: Input: Loss function L ( w ) , learning rate η , number of epochs t max , initialized parameters w 0 , symmetry group G , teleportation schedule K , number of mini-batches used to teleport B .
Revised: Input: Loss function L ( w ) , learning rate η , number of epochs t max , initialized parameters w 0 , symmetry group G , teleportation schedule K , number of mini-batches used to teleport B .

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

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 ∥
Revised: 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 )
Revised: ˜ ∇L ( X i , g · w t )

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

Original: ∇L
Revised: ∇L

Original: ( X i , w t ) end for for X i in the rest mini-batches do w t ← w t − η ˜ ∇L
Revised: ( 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 − η ˜
Revised: ( X i , w t ) end for else for all mini-batches X i do w t ← w t − η ˜

Original: ∇L
Revised: ∇L

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

Original: A.2 Data transformation
Revised: A.2 Data transformation

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

Original: Denoteg X as the group action on data only.
Revised: Denoteg X as the group action on data only.

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

Original: Algorithm 3:
Revised: Algorithm 3:

Original: Teleportation transformation
Revised: Teleportation transformation

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

Original: Output: w t max , data transformation f .
Revised: Output: w t max , data transformation f .

Original: Initialize: f = the identity function.
Revised: Initialize: f = the identity function.

Original: 1 for t ← 0 to t max − 1 do if t ∈ K then g ← argmax g ∈ G ∥∇ g · w t L ( g · ( w t , X ))
Revised: 1 for t ← 0 to t max − 1 do if t ∈ 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
Revised: ∥ 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
Revised: B Group actions

Original: In this section, we derive the group actions for the test functions and multi-layer neural networks.
Revised: 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).
Revised: More details about group theory can be found in textbooks such as Lang (2002).

Original: B.1 Continuous symmetry in test functionsB.1.1 Ellipse
Revised: B.1 Continuous symmetry in test functionsB.1.1 Ellipse

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

Original: [Equation]
Revised: [Equation]

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

Original: =
Revised: =

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

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

Original: Let g θ ∈
Revised: Let g θ ∈

Original: SO
Revised: SO

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

Original: [Equation]
Revised: [Equation]

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: B.1.2 Rosenbrock function
Revised: B.1.2 Rosenbrock function

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

Original: [Equation]
Revised: [Equation]

Original: Let u = 10( x 21 − x 2 ) and v = x 1 − 1 .
Revised: 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 arotational symmetry.
Revised: After changing the variables from x and y to u and v , L has arotational symmetry.

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

Original: [Equation]
Revised: [Equation]

Original: Next, we show that SO (2) is a symmetry of L ( x 1 , x 2 ) .
Revised: 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 461 R. For g ∈ SO (2) , define the following group action:
Revised: Let ρ be a representation of SO (2) acting on 461 R. For g ∈ SO (2) , define the following group action:

Original: [Equation]
Revised: [Equation]

Original: Then
Revised: Then

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

Original: For the Rosenbrock function with 2N parameters, we can construct a bijective function h : R 2 N −→ 464 R 2 N by transforming each of the N pairs of variables as before, and SO (2 N ) is a symmetry of
Revised: For the Rosenbrock function with 2N parameters, we can construct a bijective function h : R 2 N −→ 464 R 2 N by transforming each of the N pairs of variables as before, and SO (2 N ) is a symmetry of

Original: L ( x 1 , ..., x 2 N ) .
Revised: L ( x 1 , ..., x 2 N ) .

Original: However, we will only use the 2 variable version in the experiments.
Revised: However, we will only use the 2 variable version in the experiments.

Original: B.1.3 Booth function
Revised: B.1.3 Booth function

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: The function h : R 2 −→ R 2 that maps x 1 , x 2 to u, v is bijective.
Revised: 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 .
Revised: Let ρ be a representation of SO (2)acting on R 2 .

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

Original: [Equation]
Revised: [Equation]

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

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

Original: B. 
Revised: B.

Original: Proof of Proposition 4.2 and Proposition 4.

Revised: Proof of Proposition 4.2 and Proposition 4.

Original: Proposition 4.

Revised: Proposition 4.

Original: Proof.
Revised: Proof.

Original: In the linear network h m =
Revised: In the linear network h m =

Original: W m h m − 1 , hence
Revised: W m h m − 1 , hence

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

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

Original: Proposition 4.

Revised: Proposition 4.

Original: Proof.
Revised: Proof.

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

Original: Inother words, we want to find a set of transformed weights W ′ m , W ′ m − 1 which yields the same networkoutput ˜ h m :
Revised: Inother words, we want to find a set of transformed weights W ′ m , W ′ m − 1 which yields the same networkoutput ˜ h m :

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: To obtain W ′ m − 1 = g · W m − 1 we first find a right pseudo-inverse h + m − 2 ∈ R n × d m − 2 for h m − 2meaning h + m − 2 h m − 

Revised: To obtain W ′ m − 1 = g · W m − 1 we first find a right pseudo-inverse h + m − 2 ∈ R n × d m − 2 for h m − 2meaning h + m − 2 h m −

Original: =
Revised: =

Original: I (assuming n ≪
Revised: I (assuming n ≪

Original: d m − 2 and that h m − 2 is full rank).
Revised: d m − 2 and that h m − 2 is full rank).

Original: Solving (26) we get
Revised: Solving (26) we get

Original: [Equation]
Revised: [Equation]

Original: gσ
Revised: gσ

Original: [Equation]
Revised: [Equation]

Original: (27) (8) follows from (26) and (27).
Revised: (27) (8) follows from (26) and (27).

Original: C Theoretical analysis of teleportation
Revised: C Theoretical analysis of teleportation

Original: C.1 Proof of Proposition 5.1 and Corollary
Revised: C.1 Proof of Proposition 5.1 and Corollary

Original: Proposition 5.

Revised: Proposition 5.

Original: Proof.
Revised: Proof.

Original: Let w ′ =
Revised: Let w ′ =

Original: g · w .
Revised: g · w .

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

Original: i /∂w j .
Revised: i /∂w j .

Original: Then the inverse of J hasentries J − 1 ij
Revised: Then the inverse of J hasentries J − 1 ij

Original: =
Revised: =

Original: ∂w
Revised: ∂w

Original: i /∂w ′
Revised: i /∂w ′

Original: j .
Revised: j .

Original: The gradient at w ′ is
Revised: The gradient at w ′ is

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: Thus we will have a speedup if
Revised: Thus we will have a speedup if

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

Original: Corollary 5.

Revised: Corollary 5.

Original: Proof.
Revised: Proof.

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

Original: of (10) we have
Revised: of (10) we have

Original: [Equation]
Revised: [Equation]

Original: C.2 Proof of Proposition 5. 
Revised: C.2 Proof of Proposition 5.

Original: Proof.
Revised: Proof.

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

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

Original: Equivalently,
Revised: Equivalently,

Original: [Equation]
Revised: [Equation]

Original: By unrolling T steps, we have
Revised: By unrolling T steps, we have

Original: [Equation]
Revised: [Equation]

Original: ≤
Revised: ≤

Original: [Equation]
Revised: [Equation]

Original: 1 +
Revised: 1 +

Original: [Equation]
Revised: [Equation]

Original: (cid:13)(cid:13)(cid:13)(cid:13) ∂
Revised: (cid:13)(cid:13)(cid:13)(cid:13) ∂

Original: [Equation]
Revised: [Equation]

Original: (cid:13)(cid:13)(cid:13)(cid:13)(34)
Revised: (cid:13)(cid:13)(cid:13)(cid:13)(34)

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

Original: [Equation]
Revised: [Equation]

Original: Therefore, if
Revised: Therefore, if

Original: [Equation]
Revised: [Equation]

Original: (36) then it is guaranteed that
Revised: (36) then it is guaranteed that

Original: [Equation]
Revised: [Equation]

Original: C.3 Proof of Lemma
Revised: C.3 Proof of Lemma

Original: Proof.
Revised: Proof.

Original: From the definition of the directional derivative,∂ v (cid:13)(cid:13)(cid:13)(cid:13) ∂ L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  = v · ∂ ∂ w (cid:13)(cid:13)(cid:13)(cid:13) ∂
Revised: From the definition of the directional derivative,∂ v (cid:13)(cid:13)(cid:13)(cid:13) ∂ L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  = v · ∂ ∂ w (cid:13)(cid:13)(cid:13)(cid:13) ∂

Original: L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  (38)
Revised: L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  (38)

Original: Writing the last term in indices,∂ ∂ w i (cid:13)(cid:13)(cid:13)(cid:13)
Revised: Writing the last term in indices,∂ ∂ w i (cid:13)(cid:13)(cid:13)(cid:13)

Original: ∂
Revised: ∂

Original: L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  = ∂
Revised: L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  = ∂

Original: ∂ w i (cid:88) j (cid:18) ∂
Revised: ∂ w i (cid:88) j (cid:18) ∂

Original: L ∂ w j (cid:19)= (cid:88) j ∂
Revised: L ∂ w j (cid:19)= (cid:88) j ∂

Original: ∂ w i (cid:18) ∂
Revised: ∂ w i (cid:18) ∂

Original: L ∂ w j (cid:19)= (cid:88)
Revised: L ∂ w j (cid:19)= (cid:88)

Original: j 2 ∂ L ∂ w j ∂ 2 L ∂ w i ∂ w j =(cid:18)
Revised: j 2 ∂ L ∂ w j ∂ 2 L ∂ w i ∂ w j =(cid:18)

Original: H ∂ L ∂ w (cid:19) i (39)
Revised: H ∂ L ∂ w (cid:19) i (39)

Original: Removing the indices,∂ ∂ w (cid:13)(cid:13)(cid:13)(cid:13)
Revised: Removing the indices,∂ ∂ w (cid:13)(cid:13)(cid:13)(cid:13)

Original: ∂
Revised: ∂

Original: L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  = 2 H ∂ L ∂ w (40)
Revised: L ∂ w (cid:13)(cid:13)(cid:13)(cid:13)  = 2 H ∂ L ∂ w (40)

Original: Substitute back and we have∂ v (cid:13) (cid:13)(cid:13)(cid:13) ∂
Revised: Substitute back and we have∂ v (cid:13) (cid:13)(cid:13)(cid:13) ∂

Original: L ∂ w (cid:13) (cid:13)(cid:13)(cid:13)  = v · (cid:18) 2 H ∂ L ∂ w (cid:19) (41)
Revised: L ∂ w (cid:13) (cid:13)(cid:13)(cid:13)  = v · (cid:18) 2 H ∂ L ∂ w (cid:19) (41)

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

Original: ∂
Revised: ∂

Original: L ∂ w (cid:13)(cid:13) is orthogonal to all vectors that are orthogonal to
Revised: L ∂ w (cid:13)(cid:13) is orthogonal to all vectors that are orthogonal to

Original: Therefore
Revised: Therefore

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

Original: C.4 Proof of Proposition 5.6
Revised: C.4 Proof of Proposition 5.6

Original: Proof.
Revised: Proof.

Original: From Lemma 5.5, dLdw is an eigenvector of H .
Revised: From Lemma 5.5, dLdw is an eigenvector of H .

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

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

Original: D Global convergence rate
Revised: C.5 Teleportation using mini-batches

Revised: In this section, we provide a preliminary analysis on whether teleportation on SGD is expected toimprove convergence rate over SGD. Let b be the number of samples in a mini-batch, and n be the number of features for each sample. Denote X b ∈ R b × n as the data in a mini-batch, and X ib ∈ R n ( i ∈ { 1 , ..., b } ) as the i th sample in themini-batch.

Revised: We assume that the loss of a set of data points is the average of loss values evaluated oneach data point.

Revised: This is a reasonable assumption because the loss of neural networks usually shouldnot depend on the order or amount of samples.

Revised: Concretely, we assume that

Revised: [Equation]

Revised: [Equation]

Revised: Consider a group element g found to maximize the convergence gain for a mini-batch X b :

Revised: [Equation]

Revised: Then, the improvement on the gradient norm evaluated on X b is

Revised: [Equation]

Revised: w ′ =

Revised: g · w (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) − (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 b b (cid:88) i =

Revised: ∂ L ( X ib , w ) ∂ w (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) .

Revised: Let X ∈ R m × n be the set of all training data, where m is the total number of samples.

Revised: Using the gfound using the mini-batch X b , the expected improvement in the gradient norm on all data is

Revised: [Equation]

Revised: (cid:12)(cid:12)(cid:12)(cid:12)

Revised: [Equation]

Revised: (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Revised: [Equation]

Revised: (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Revised: [Equation]

Revised: (cid:88)

Revised: [Equation]

Revised: (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Revised: [Equation]

Revised: What remains to be done is to show that ∆ m ≥ 0 with high probability, which would imply that ifteleporting using a minibatch improves | d L /dt | over that minibatch, then with high probability, thisteleportation also improves | d L /dt | evaluated using another mini-batch or the entire dataset.

Revised: As astep forward, we make the following observations: first, as the batch size approaches m , the amountof improvement on the minibatch approaches the improvement evaluated using all data.

Revised: Second,the gap between ∆ b and ∆ m depends on the variance of gradient with respect to data points X i .

Revised: When the variance is small, we expect the teleportation using a mini-batch to have a similar effects asteleportation using all data.

Revised: D Global convergence rate

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

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

Original: However, whether oneteleportation is sufficient to find the optimal gradient flow trajectory remains unknown.
Revised: However, whether oneteleportation is sufficient to find the optimal gradient flow trajectory remains unknown.

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

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

Original: We then develop a general condition forwhen one teleportation gives the optimal trajectory.
Revised: We then develop a general condition forwhen one teleportation gives the optimal trajectory.

Original: D.1 Convex quadratic functions
Revised: D.1 Convex quadratic functions

Original: Consider a positive definite quadratic form L A ( w )
Revised: Consider a positive definite 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 ∈ R n is the parameter and

Original: A ∈ R n × n is a positive definite matrix.
Revised: A ∈ R n × n is a positive definite matrix.

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

Original: = 

Revised: =

Original: A w , and the Hessian of
Revised: A w , and the Hessian of

Original: L A is H
Revised: L A is H

Original: =
Revised: =

Original: A .
Revised: A .

Original: Since A is defined to be positive definite, L A is convex.
Revised: Since A is defined to be positive definite, L A is convex.

Original: The function L A has globalminimum at a single point w ∗ = 0 .
Revised: The function L A has globalminimum 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 ∈ O ( n ) , we define 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:

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 first show that starting from a point in S c , allother points in S c can be reached with one teleportation.
Revised: We first show that starting from a point in S c , allother points in S c can be reached with one teleportation.

Original: Proposition D.1.
Revised: Proposition D.1.

Original: S c contains a single orbit.
Revised: S c contains a single orbit.

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

Original: S c for all w ∈ S c .
Revised: S c for all w ∈ S c .

Original: Proof.
Revised: Proof.

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

Original: Then w T 1 A w 1 =
Revised: Then w T 1 A w 1 =

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

Original: =
Revised: =

Original: c and w T 2 A w 2 =(
Revised: c and w T 2 A w 2 =(

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

Original: =
Revised: =

Original: c .
Revised: c .

Original: Let v 

Revised: Let v

Original: =
Revised: =

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

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

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

Original: Since ∥ v 1 ∥ =
Revised: Since ∥ v 1 ∥ =

Original: ∥ v 2 ∥ =1 , there exists g 1 , g 2 ∈
Revised: ∥ 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 .
Revised: O ( n ) , such that g 1 e 1 = v 1 and g 2 e 1 = v 2 .

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

Original: Letg = g 2 g − 1 1 .
Revised: Letg = g 2 g − 1 1 .

Original: Then v 

Revised: Then v

Original: = g
Revised: = g

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

Original: A − 12 gA 12 w 1 = g · w 1 .
Revised: 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 .
Revised: 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.
Revised: Therefore,the group action of G on S c is transitive.

Original: Equivalently, S c contains a single orbit.
Revised: Equivalently, S c contains a single orbit.

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

Original: [Equation]
Revised: [Equation]

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

Original: In other words, (44) is equivalent to the followingoptimization problem:
Revised: In other words, (48) is equivalent to the followingoptimization problem:

Original: [Equation]
Revised: [Equation]

Original: Let c =
Revised: Let c =

Original: L A ( w ) .
Revised: L A ( w ) .

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

Original: [Equation]
Revised: [Equation]

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

Original: Proposition D.2.
Revised: Proposition D.2.

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

Original: Proof.
Revised: Proof.

Original: We solve (46) using the method of Lagrangian multipliers.
Revised: We solve (50) using the method of Lagrangian multipliers.

Original: The Lagrangian of (46) is
Revised: The Lagrangian of (50) is

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: 48) which gives
Revised: (52) which gives

Original: [Equation]
Revised: [Equation]

Original: Since A is positive definite, A
Revised: Since A is positive definite, A

Original: =
Revised: =

Original: A T and A is invertible.
Revised: A T and A is invertible.

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

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

Original: Then, the constraint is w ′ T
Revised: Then, the constraint is w ′ T

Original: A w ′ =
Revised: A w ′ =

Original: λ
Revised: λ

Original: ∥ w ′ ∥ 2 = c , and theobjective becomes max w ′ λ 2 ∥ w ′ ∥ 2 = max w ′ cλ .
Revised: ∥ w ′ ∥ 2 = c , and theobjective becomes max [Equation] ∥ 2 = max w ′ cλ .

Original: Therefore, we want λ to be the largest eigenvalueof A .
Revised: Therefore, we want λ to be the largest eigenvalueof A .

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

Original: Proposition D.3.
Revised: Proposition D.3.

Original: If at point w , ∥∇ L A
Revised: If at point w , ∥∇ L A

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

Original: Proof.
Revised: Proof.

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

Original: Thenthe gradient of L A is A w = λ w .
Revised: Thenthe gradient of L A is A w = λ w .

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

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

Original: Finally, we make the observation that for L A , teleportation moves the parameters closer to the optimalsolution in Euclidean distance.
Revised: Finally, we make the observation that for L A , teleportation moves the parameters closer to the optimalsolution in Euclidean distance.

Original: The following proposition states that maximizing the magnitude ofgradient is equivalent to minimizing the distance to w ∗ in a loss level set.
Revised: The following proposition states that maximizing the magnitude ofgradient is equivalent to minimizing the distance to w ∗ in a loss level set.

Original: Proposition D.4.
Revised: Proposition D.4.

Original: The solution to the following optimization problem is the same as the solution to(46) :
Revised: The solution to the following optimization problem is the same as the solution to(50) :

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Similar to Proposition D.2, we solve this optimization problem using the method of Lagrangianmultipliers.
Revised: Similar to Proposition D.2, we solve this optimization problem using the method of Lagrangianmultipliers.

Original: Note that w ∗ = 0 .
Revised: Note that w ∗ = 0 .

Original: The Lagrangian is
Revised: The Lagrangian is

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: 53) which gives
Revised: (57) which gives

Original: [Equation]
Revised: [Equation]

Original: [Equation] so the solution to (51) is an eigenvector of A .
Revised: (58) so the solution to (55) is an eigenvector of A .

Original: Then, the constraint is w ′ T
Revised: Then, the constraint is w ′ T

Original: A w ′ =
Revised: A w ′ =

Original: λ
Revised: λ

Original: ∥ w ′ ∥ 2 = c , andthe objective becomes min w ′ ∥
Revised: ∥ w ′ ∥ 2 = c , andthe objective becomes min w ′ ∥

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

Original: Therefore, we want λ to be the largest eigenvalueof A .
Revised: Therefore, we want λ to be the largest eigenvalueof A .

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

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

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

Original: The point with largest ∥∇ L A ∥ 2 on a level set is in the eigendirectionof
Revised: The point with largest ∥∇ L A ∥ 2 on a level set is in the eigendirectionof

Original: A corresponding to its largest eigenvalue, or equivalently, a point on the smallest semi-axes of theellipsoid.
Revised: A corresponding to its largest eigenvalue, or equivalently, a point on the smallest semi-axes of theellipsoid.

Original: Note that this point has the smallest distance to the global minimum w ∗ = 0 among allpoints in the same level set.
Revised: Note that this point has the smallest distance to the global minimum w ∗ = 0 among allpoints in the same level set.

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

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

Original: D.2 General smooth functions
Revised: D.2 General smooth functions

Original: The previous section shows that for convex quadratic loss functions, one teleportation ensures thatevery point on the future gradient flow trajectory has the optimal gradient norm in its loss level set.
Revised: The previous section shows that for convex quadratic loss functions, one teleportation ensures thatevery point on the future gradient flow trajectory has the optimal gradient norm in its loss level set.

Original: Inthis section, we provide a necessary condition for one teleportation to result in a globally optimalgradient flow trajectory.
Revised: Inthis section, we provide a necessary condition for one teleportation to result in a globally optimalgradient flow trajectory.

Original: We denote a vector field on a manifold M as V : M −→ TM , where TM is the tangent bundle of M .
Revised: We denote a vector field on a manifold M as V : M −→ TM , where TM is the tangent bundle of M .

Original: Here we consider the parameter space M = R n .
Revised: Here we consider the parameter space M = R n .

Original: We can also write V = v
Revised: We can also write V = v

Original: i ∂∂x
Revised: i ∂∂x

Original: i using the componentfunctions v i : R n −→ R .
Revised: i using the componentfunctions v i : R n −→ R .

Original: Consider a smooth loss function L : R n −→ R .
Revised: Consider a smooth loss function L : R n −→ R .

Original: Let R =
Revised: Let R =

Original: r i ∂∂w i be the vector field of the gradientflow direction at every point on the parameter manifold.
Revised: r i ∂∂w i be the vector field of the gradientflow direction at every point on the parameter manifold.

Original: Then r i = − ∂ L ∂w i .
Revised: Then r i = − ∂ L ∂w i .

Original: Let A = a i ∂∂w
Revised: Let A = a i ∂∂w

Original: i be avector field orthogonal to R , i.e. r i δ ij a j = 0 at every point w , where δ ij is the Kronecker delta.
Revised: i be avector field orthogonal to R , i.e. r i δ ij a j = 0 at every point w , where δ ij is the Kronecker delta.

Original: The Lie bracket [ A, R ] gives the difference between flowing a short time in the symmetry directionthen in the gradient direction, and flowing in the gradient direction then in the symmetry direction.
Revised: The Lie bracket [ A, R ] gives the difference between flowing a short time in the symmetry direction then in the gradient direction, and flowing in the gradient direction then in the symmetry direction.

Original: Teleportation has an effect on the convergence rate if and only if the two vector fields do not commute,i.e. [ A, R ] L ̸ = 0 .
Revised: Teleportation has an effect on the convergence rate if and only if the two vector fields do not commute,i.e. [ A, R ] L ̸ = 0 .

Original: Even when A and R do not commute, there may exist some w such that [ A, R ] L ( w ) = 0 .
Revised: Even when A and R do not commute, there may exist some w such that [ A, R ] L ( w ) = 0 .

Original: Applying[ A, R ] to L , we have
Revised: Applying[ A, R ] to L , we have

Original: [Equation]
Revised: [Equation]

Original: At the target teleportation destination w , (cid:13)(cid:13)
Revised: At the target teleportation destination w , (cid:13)(cid:13)

Original: ∂
Revised: ∂

Original: L ∂ w (cid:13)(cid:13) 2 2 is at a local maximum on a loss level set, hence
Revised: L ∂ w (cid:13)(cid:13) 2 2 is at a local maximum on a loss level set, hence

Original: [ A, R ] L ( w ) = 0 .
Revised: [ A, R ] L ( w ) = 0 .

Original: If [ A, R ] L ( w ) = 0 for all A , then w is a local extremum.
Revised: If [ A, R ] L ( w ) = 0 for all A , then w is a local extremum.

Original: Therefore, a necessarycondition for one teleportation to result in an optimal trajectory is that if [ A, R ] L ( w ) = 0 , then itstays 0 along the gradient flow starting at w .
Revised: Therefore, a necessarycondition for one teleportation to result in an optimal trajectory is that if [ A, R ] L ( w ) = 0 , then itstays 0 along the gradient flow starting at w .

Original: In other words, we need that [ A, R ] L ( w ) = 0 implies
Revised: In other words, we need that [ A, R ] L ( w ) = 0 implies

Original: R [ A, R ] L ( w ) = 0 .
Revised: R [ A, R ] L ( w ) = 0 .

Original: To see when this condition is satisfied, we first expand R
Revised: To see when this condition is satisfied, we first expand R

Original: [ A, R ] L :
Revised: [ A, R ] L :

Original: [Equation]

Original: ∂w

Original: [Equation]
Revised: [Equation]

Original: From Lemma 5.5, r = − ∂ L ∂ w is an eigenvector of ∂ 2 L ∂w i ∂w j
Revised: From Lemma 5.5, r = − ∂ L ∂ w is an eigenvector of ∂ 2 L ∂w i ∂w j

Original: when A (cid:13)(cid:13) ∂
Revised: when A (cid:13)(cid:13) ∂

Original: L ∂ w (cid:13)(cid:13) 2 

Revised: L ∂ w (cid:13)(cid:13) 2

Original: =
Revised: =

Original: for all symmetrydirection A .
Revised: for all symmetrydirection A .

Original: Let λ 1 , λ 2 be eigenvalues of the Hessian.
Revised: Let λ 1 , λ 2 be eigenvalues of the Hessian.

Original: Then we have
Revised: Then we haver k a j

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Hence, only the first term of R
Revised: Hence, only the first term of R

Original: [ A, R ]
Revised: [ A, R ]

Original: L is left:
Revised: L is left:

Original: [Equation]
Revised: [Equation]

Original: [Figure]

Original: where H denotes the Hessian of L and J ( H a ) denotes the Jacobian of H a .
Revised: where H denotes the Hessian of L and J ( H a ) denotes the Jacobian of H a .

Original: This term is notautomatically 0 when [ A, R ] V ( w ) = 0 .
Revised: This term is notautomatically 0 when [ A, R ] V ( w ) = 0 .

Original: Therefore, in general, one teleportation does not guaranteethat the trajectory intersects all future level sets at an optimal point.
Revised: Therefore, in general, one teleportation does not guaranteethat the trajectory intersects all future level sets at an optimal point.

Original: However, for loss functions thatsatisfy R [ A, R ] V ( w ) = 0 when [ A, R ] V ( w ) = 

Revised: However, for loss functions thatsatisfy R [ A, R ] V ( w ) = 0 when [ A, R ] V ( w ) =

Original: (for example, when J ( H a ) is skew-symmetric or0), teleporting once optimizes the entire trajectory.
Revised: (for example, when J ( H a ) is skew-symmetric or0), teleporting once optimizes the entire trajectory.

Original: E Experiment details and additional results
Revised: E Experiment details and additional results

Original: E.1 Test functions
Revised: E.1 Test functions

Original: We compare the gradient at different loss values for gradient descent with and without teleportation.
Revised: 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 trajectorywithout teleportation at the same loss values.
Revised: Figure 7 shows that the trajectory with teleportation has a larger d L /dt value than the trajectorywithout teleportation at the same loss values.

Original: Therefore, the rate of change in the loss is larger in thetrajectory with teleportation, which makes it favorable. 
Revised: Therefore, the rate of change in the loss is larger in thetrajectory with teleportation, which makes it favorable. [Figure]

Original: E.2 Multilayer neural network
Revised: E.2 Multilayer neural network

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

Original: The group elements used for these transforms are found by gradient ascent on Tfor 10 times, with learning rate 10 − 7 in GD and 10 − 6 in AdaGrad.
Revised: The group elements used for these transforms are found by gradient ascent on Tfor 10 times, with learning rate 10 − 7 in GD and 10 − 6 in AdaGrad.

Original: When using teleportation, weperform symmetry transform on the parameters once every 10 epochs in the first 50 epochs.
Revised: When using teleportation, weperform symmetry transform on the parameters once every 10 epochs in the first 50 epochs.

Original: Eachalgorithm is run 300 steps.
Revised: Eachalgorithm is run 300 steps.

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

Original: Learning rates To confirm that teleportation works with different learning rates, we repeat theexperiment using learning rate 10 − 3 for GD and 10 0 for AdaGrad, which is 10 times higher than thelearning rate used in the main paper.
Revised: Learning rates To confirm that teleportation works with different learning rates, we repeat theexperiment using learning rate 10 − 3 for GD and 10 0 for AdaGrad, which is 10 times higher than thelearning rate used in the main paper.

Original: Increasing learning rates further prevents convergence.
Revised: Increasing learning rates further prevents convergence.

Original: Sincegradient descent is expected to converge faster with higher learning rates, we adjust the teleportationschedule to once every 4 epochs for the first 20 epochs and run each algorithm 120 steps.
Revised: Sincegradient descent is expected to converge faster with higher learning rates, we adjust the teleportationschedule to once every 4 epochs for the first 20 epochs and run each algorithm 120 steps.

Original: Fig.
Revised: Fig.

Original: 8shows the results for the new learning rates.
Revised: 8shows the results for the new learning rates.

Original: Teleportation improves both GD and AdaGrad in thesimilar way as using the original learning rates.
Revised: Teleportation improves both GD and AdaGrad in thesimilar way as using the original learning rates.

Original: Hyperparameter tuning To observe the effect of hyperparameters on the speedup in computationtime, we did a hyperparameter sweep on the number of steps and the learning rate used in eachteleportation.
Revised: Hyperparameter tuning To observe the effect of hyperparameters on the speedup in computationtime, we did a hyperparameter sweep on the number of steps and the learning rate used in eachteleportation.

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

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

Original: This experiment is run on one CPU.
Revised: This experiment is run on one CPU.

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

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

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

Original: Omittedpoints in the figure indicate that the gradient descent fails to converge within 2000 steps or diverges.[Figure] [Figure] of steps
Revised: Omittedpoints in the figure indicate that the gradient descent fails to converge within 2000 steps or diverges.

Original: and the learning rate used to find the optimal group element in teleportation.

Original: The wall-clock speedup of applying teleportation is shown separately for gradient descent (a)(b) and AdaGrad (c)(d).

Original: [Figure]

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

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

Original: Increasing the number of steps used to optimize teleportationtarget allows us to find a better point in the parameter space but increases the cost of one teleportation.
Revised: Increasing the number of steps used to optimize teleportationtarget 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 likelyto lead to divergence since ∥ ∂ L /∂ w ∥ can become too large for the gradient descent learning rate.
Revised: Increasing the learning rate of optimizing the group element improves ∥ ∂ L /∂ w ∥ but is more likelyto lead to divergence since ∥ ∂ L /∂ w ∥ can become too large for the gradient descent learning rate.

Original: Using the set of hyperparameter that gives the maximum speedup, we plot the training curves withand without teleportation in Figure 10.
Revised: Using the set of hyperparameter that gives the maximum speedup, we plot the training curves withand without teleportation in Figure 10.

Original: As shown in Figure 10(a) and (b), teleportation reduces boththe number of epochs and the total computational time required to reach convergence. 
Revised: As shown in Figure 10(a) and (b), teleportation reduces boththe number of epochs and the total computational time required to reach convergence. [Figure] [Figure]

Revised: steps and the learning rate used to find the optimal group element in teleportation.

Revised: The wall-clock speedup of applying teleportation is shown separately for gradient descent (a)(b) and AdaGrad (c)(d).

Revised: [Figure]