# Improving Convergence and Generalization Using Parameter Symmetries

Bo Zhao
University of California San Diego
<EMAIL_ADDRESS>
&Robert M. Gower
Flatiron Institute
<EMAIL_ADDRESS>
&Robin Walters
Northeastern University
<EMAIL_ADDRESS>
&Rose Yu
University of California San Diego
<EMAIL_ADDRESS>

###### Abstract

In overparametrized models, different values of the parameters may result in
the same loss value. Parameter space symmetries are transformations that
change the model parameters but leave the loss invariant. Teleportation
applies such transformations to accelerate optimization. However, the exact
mechanism behind this algorithm’s success is not well understood. In this
paper, we show that teleportation not only speeds up optimization in the
short-term, but gives overall faster time to convergence. Additionally, we
show that teleporting to minima with different curvatures improves
generalization and provide insights on the connection between the curvature of
the minima and generalization ability. Finally, we show that integrating
teleportation into a wide range of optimization algorithms and optimization-
based meta-learning improves convergence.

##  1 Introduction

Given a deep neural network architecture and a batch of data, there may be
multiple points in the parameter space that correspond to the same loss value.
Despite having the same loss, the gradients and learning dynamics that result
from starting at these points can be very different [1, 2, 3]. Parameter space
symmetries, which are transformations of the parameters that leave the loss
function invariant, allow us to _teleport_ between points in the parameter
space on the same level set of the loss function [4, 5]. In particular, we can
teleport to a steeper point in the loss landscape, which leads to faster
optimization.

Despite the empirical evidence, the exact mechanism of how teleportation
improves convergence in optimizing non-convex objectives remains elusive.
Previous work shows that gradient increases momentarily after a teleportation,
but could not show that this results in overall faster convergence [4]. In
this paper, we provide theoretical guarantees on the convergence rate. In
particular, we show that SGD (stochastic gradient descent) with teleportation
converges to a basin of stationary points, where every parameter that can be
reached by teleportation is also a stationary point. We then show that, if we
use teleportation and the loss is Polyak-Łojasiewicz (a general form of
strongly convexity [6]), then all points that can be reached by teleportation
converge to the global minima.

Previously, the application of teleportation was limited to accelerating
optimization. The second part of this paper explores a different objective –
improving generalization. We relate properties of minima to their
generalization ability and use teleportation to optimize these properties. One
property that is believed to be correlated to the model’s generalization
ability is the sharpness of the loss function [7]. We verify empirically that
some sharpness metrics are indeed correlated to generalization, and that
teleporting towards flatter regions improves the validation loss.
Additionally, we hypothesize that generalization also depends on the curvature
of the curves that lie on the minima. For fully-connected networks, we derive
an explicit expression for these curvatures and experimentally show that
teleporting towards points with larger curvatures improves the model’s
generalization ability.

The close relationship between symmetry and optimization opens up a number of
exciting opportunities. To demonstrate the wide applicability of parameter
space symmetry, we expand teleportation to other optimization algorithms
beyond SGD, including momentum, AdaGrad, RMSProp, and Adam. Experimentally,
teleportation improves the convergence rate for these standard algorithms.
Inspired by concepts from conditional programming and optimization-based meta-
learning [8], we also propose a meta-optimizer to learn where to move
parameters in symmetry directions so as to improve upon existing gradient-
based optimization methods.

In summary, our main contributions are:

  * •

theoretical guarantees that teleportation accelerates the convergence rate of
SGD,

  * •

a teleportation-based algorithm to improve generalization,

  * •

a method that approximates the curvature of a minimum and evidence that the
curvature and generalization are correlated, and

  * •

various optimization algorithms with integrated teleportation including
momentum, AdaGrad, and optimization-based meta-learning.

##  2 Related Work

##### Parameter space symmetry.

Various symmetries have been identified in neural networks. Permutation
symmetry has been linked to the structure of minima [9, 10]. Continuous
symmetries are identified in various architectures, including homogeneous
activations [11, 12], radial rescaling activations [13], and softmax and
batchnorm functions [1]. Quiver representation theory provides a more general
framework for symmetries in neural networks with pointwise [14] and rescaling
activations [15]. A new class of nonlinear and data-dependent symmetries are
identified in [16]. Since symmetry defines transformations of parameters
within a level set of the loss function, these works are the basis of the
teleportation method used in our paper.

Knowledge of parameter space symmetry motivates new optimization methods. One
line of work seeks algorithms that are invariant to symmetry transformations
[17, 18]. Others search in the orbit for parameters that can be optimized
faster [5, 4]. We build on the latter by providing theoretical analysis on the
improvement of the convergence rate and by augmenting the teleportation
objective to improve generalization.

##### Initializations and restarts.

Teleportation before training changes the initialization of parameters. The
initialization has been known to affect both optimization and generalization.
For example, imbalance between layers at initialization affects the
convergence of gradient flows in two-layer models [19]. Different
initializations, among other sources of variance, lead to different model
performance after convergence [20, 21]. The Fourier spectrum at initialization
is related to generalization because different frequency functions are learned
at different rates [22]. For shallow networks, certain initialization is
required to learn symmetric functions with generalization guarantees [23].
Teleportation during training re-initializes the parameters to a point with
the same loss. A similar effect is achieved by warm restart, [24], which
encourages parameters to move to more stable regions by periodically
increasing the learning rate. Compared to initialization methods,
teleportation allows changes in landscape multiple times throughout the
training. Compared to restarts, teleportation leads to smaller temporary
increase in loss and provides more control of where to move the parameters.

##### Sharpness of minima and generalization

The sharpness of minima has been linked to the generalization ability of
models both empirically and theoretically [25, 7, 26, 27], which motivates
optimization methods that find flatter minima [28, 29, 30, 31]. We employ
teleportation in service of this goal by searching for flatter points along
the loss level sets. The sharpness of a minimum is often defined by properties
of the Hessian of the loss function, such as the number of small eigenvalues
[7, 28, 32] or the product of the top $k$ eigenvalues [33]. Alternatively,
sharpness can be characterized by the maximum loss within a neighborhood of a
minimum [7, 29, 31] or approximated by the growth in the loss curve averaged
over random directions [34]. The sharpness of minima does not always explain
generalization [35]. Transformation of parameters that keep the function
unchanged does not affect generalization but can lead to minima with different
sharpness.

##  3 Theoretical Guarantees for Improving Optimization

In this section, we provide theoretical analysis of teleportation. We show
that with teleportation, SGD converges to a basin of stationary points.
Building on teleportation’s relation to Newton’s method, teleportation leads
to a mixture of linear and quadratic convergence. Lastly, we show that in
certain loss functions, one teleportation guarantees optimality of the entire
gradient flow trajectory.

##### Teleportation

We briefly review the teleportation algorithm [4] that exploits parameter
space symmetry to accelerate optimization. Consider the optimization problem

$w^{*}=\underset{w\in\mathbb{R}^{d}}{\operatorname{argmin}}~{}\mathcal{L}(w),\quad\mathcal{L}(w)\overset{\text{def}}{=}\mathbb{E}_{\xi\sim\cal
D}\left[\mathcal{L}(w,\xi)\right]$ (1)

where $\mathcal{D}$ is the data distribution, $\xi$ is data sampled from
$\mathcal{D},$ $\mathcal{L}$ the loss, $w$ the parameters of the model, and
$\mathbb{R}^{d}$ the parameter space. Let $\mathcal{G}$ be a group acting on
the parameter space, such that

$\displaystyle\mathcal{L}(w)\;=\;\mathcal{L}(g\cdot w),\quad\forall
g\in\mathcal{G},\;\forall w\in\mathbb{R}^{d}.$ (2)

Symmetry teleportation uses gradient ascent to find the group element $g$ that
maximizes the magnitude of gradient and applies $g$ to the parameters:

$\displaystyle w^{\prime}=$ $\displaystyle g\cdot w,\quad g=\underset{g\in
G}{\text{argmax}}\|\nabla\mathcal{L}(g\cdot w)\|^{2}.$ (3)

When each loss level set consists of a single $\mathcal{G}$-orbit, this is
equivalent to finding a point in the level set $S(x)=\\{w:\mathcal{L}(w)=x\\}$
with the largest gradient norm:

$\displaystyle w^{\prime}=\underset{v\in
S(\mathcal{L}(w))}{\text{argmax}}\|\nabla\mathcal{L}(v)\|^{2}.$ (4)

###  3.1 Teleportation and SGD

At each iteration $t\in\mathbb{N}^{+}$ in SGD, we choose a group element
$g^{t}\in\mathcal{G}$ and use teleportation before each gradient step as
follows

$w^{t+1}=g^{t}\cdot w^{t}-\eta\nabla\mathcal{L}(g^{t}\cdot w^{t},\xi^{t}).$
(5)

Here $\eta$ is a learning rate, $\nabla\mathcal{L}(w^{t},\xi^{t})$ is a
gradient of $\mathcal{L}(w^{t},\xi^{t})$ with respect to the parameters $w$,
and $\xi^{t}\sim\mathcal{D}$ is a mini-batch of data sampled i.i.d at each
iteration.

By choosing the group element that maximizes the gradient norm, we show in the
following theorem that the iterates (5) converge to a basin of stationary
points, where all points that can be reached via teleportation are also
stationary points (visualized in Figure 1).

######  Theorem 3.1.

(Smooth non-convex) Let $\mathcal{L}(w,\xi)$ be $\beta$–smooth and let

$\sigma^{2}\overset{\text{def}}{=}\mathcal{L}(w^{*})-\mathbb{E}\left[\inf_{w}\mathcal{L}(w,\xi)\right].$

Consider the iterates $w^{t}$ given by (5) where

$g^{t}\in\arg\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w^{t})\rVert^{2}.$ (6)

If $\eta=\frac{1}{\beta\sqrt{T-1}}$ then

$\displaystyle\min_{t=0,\ldots,T-1}$
$\displaystyle\mathbb{E}\left[\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w^{t})\rVert^{2}\right]$ (7) $\displaystyle\leq$
$\displaystyle\frac{2\beta}{\sqrt{T-1}}\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}(w^{*})\right]+\frac{\beta\sigma^{2}}{\sqrt{T-1}}$
(8)

where the expectation is the total expectation with respect to the data
$\xi^{t}$ for $t=0,\ldots,T-1.$

Figure 1: With teleportation, SGD converges to a basin where all points on the
level set are stationary points.

This theorem is an improvement over vanilla SGD, for which we would have
instead that

$\displaystyle\min_{t=0,\ldots,T-1}$
$\displaystyle\mathbb{E}\left[\lVert\nabla\mathcal{L}(w^{t})\rVert^{2}\right]\leq\frac{2\beta}{\sqrt{T-1}}\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}(w^{*})\right]+\frac{\beta\sigma^{2}}{\sqrt{T-1}}.$

The above only guarantees that there exists a single point $w^{t}$ for which
the gradient norm will eventually be small. In contrast, our result in (7)
guarantees that for all points over the orbit $\\{g\cdot\theta^{t}\;:\;\forall
g\in\mathcal{G}\\}$, the gradient norm will be small.

For strictly convex loss functions,
$\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot w)\rVert^{2}$ is non-
decreasing with $\mathcal{L}(w)$. In this case, the value of $\mathcal{L}$ is
smaller after $T$ steps of SGD with teleportation than vanilla SGD
(Proposition A.2). For non-convex loss functions that are smooth and are
$\mu$–PL (Polyak-Łojasiewicz), the following theorem shows that the expected
loss converges to the optimal loss for every point on the orbit of $w^{t}$.

######  Theorem 3.2.

(Error Bound) Consider the setting of Theorem 3.1. If in addition to
$\beta$–smoothness we also assume that the $\mu$-PL inequality holds, that is

$\displaystyle\mathcal{L}(w)-\inf\mathcal{L}\leq\frac{1}{2\mu}\|\nabla\mathcal{L}(w)\|^{2},\quad\forall
w\in\mathbb{R}^{s},$ (9)

then the $w^{t}$ iterates given by (5) with learning rate
$\eta=\frac{1}{\beta\sqrt{T-1}}$ converge according to

$\displaystyle\min_{t=0,\ldots,T-1}$
$\displaystyle\mathbb{E}\left[\max_{g\in\mathcal{G}}\lVert g\circ
w^{t}-w^{*}\rVert\right]\leq\frac{2\beta}{\mu}\frac{1}{\sqrt{T-1}}\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}(w^{*})\right]+\frac{\beta}{\mu}\frac{\sigma^{2}}{\sqrt{T-1}}.$
(10)

Assuming the loss function is $\mu$–PL can be reasonable for neural networks
that are overparametrized. Indeed, in [36] the authors showed that, for
overparametrized neural networks, when the task is regression, the loss
function can often be $\mu$–PL.

###  3.2 Teleportation and Newton’s method

Intuitively, teleportation can speed up optimization as it behaves similarly
to Newton’s method. The update direction in Newton’s method can be decomposed
into two components: one in the gradient direction and one orthogonal to the
gradient. Among all directions orthogonal to the gradient, this second
component matches the direction along which the gradient norm increases the
fastest (Proposition B.3). In other words, teleportation narrows the gap
between the second-order and first-order optimization methods.

After a teleportation that takes parameters to a critical point on a level
set, the gradient descent direction is the same as the Newton’s direction [4].
As a result, we can leverage the convergence of Newton’s method to derive the
convergence rate of teleportation.

######  Proposition 3.3 (Quadratic term in convergence rate).

Let $f$ be strictly convex and let $w_{0}\in\mathbb{R}^{d}$. Let

$\displaystyle
w^{\prime}\in\operatorname{argmax}_{w\in\mathbb{R}^{d}}\frac{1}{2}\|\nabla\mathcal{L}(w)\|^{2},\quad\mathrm{s.t.}\quad\mathcal{L}(w)=\mathcal{L}(w_{0}).$
(11)

Let $\nabla^{2}\mathcal{L}$ be the Hessian of $\mathcal{L}$, and
$\lambda_{\max}(\nabla^{2}\mathcal{L}(w))$ be the largest eigenvalue of
$\nabla^{2}\mathcal{L}(w)$. If $\nabla\mathcal{L}(w^{\prime})\neq 0$ then
there exists $\lambda_{0}$ such that
$0\leq\lambda_{0}\leq\lambda_{\max}(\nabla^{2}\mathcal{L}(w_{0})),$ and one
step of gradient descent after teleportation with learning rate $\gamma>0$
gives

$\displaystyle w_{1}$
$\displaystyle=w^{\prime}-\gamma\nabla\mathcal{L}(w^{\prime})=w^{\prime}-\gamma\lambda_{0}\nabla^{2}\mathcal{L}(w^{\prime})^{-1}\nabla\mathcal{L}(w^{\prime}).$
(12)

Let $w^{\prime}=g_{0}\cdot w_{0}$. If $\gamma\leq\frac{1}{\lambda_{0}}$,
$\mathcal{L}$ is a $\mu$–strongly convex $L$–smooth function, and the Hessian
is $G$–Lipschitz, then we have that

$\|w_{1}-w^{*}\|\leq\frac{G}{2\mu}\|g_{0}\cdot
w_{0}-w^{*}\|^{2}+|1-\gamma\lambda_{0}|\frac{L}{2\mu}\|g_{0}\cdot
w_{0}-w^{*}\|.$

The details of the assumptions and the proof are in Section C. Note that due
to unknown step size $\lambda_{0}$, extra care is needed in establishing this
convergence rate.

The above proposition shows that taking one step of teleportation and one
gradient step, the result is equal to taking a dampened Newton step (12).
Hence, the convergence rate in the above lemma has a quadratically contracting
term $\|g_{0}\cdot w_{0}-w^{*}\|^{2}$, which is typical of second order
methods. In particular, setting $\gamma=1/\lambda_{0}$ we would have local
quadratic convergence. In contrast, without the teleportation step and under
the same assumptions, we would have the following linear convergence

$\|w_{1}-w^{*}\|\leq\left(1-\mu\gamma\right)\|w_{0}-w^{*}\|$

for $\gamma\leq\frac{1}{L}$ using gradient descent. Thus there would be no
quadratically contracting term.

###  3.3 When is one teleportation enough

Despite the guaranteed improvement in convergence rate, teleporting at every
step in gradient descent is not computationally feasible. Generally we
teleport occasionally. In fact, for certain functions, every point on the
gradient flow has the largest gradient norm in its loss level set after one
teleportation [4]. However, this result is limited to convex quadratic
functions. In this section, we give a sufficient condition for when one
teleportation results in an optimal trajectory for general loss functions.

Let $V:\mathcal{M}\xrightarrow[]{}T\mathcal{M}$ be a vector field on the
manifold $\mathcal{M}$, where $T\mathcal{M}$ denotes the associated tangent
bundle. Here we consider the parameter space $\mathcal{M}=\mathbb{R}^{n}$, but
results in this section can be extended to optimization on other manifolds. In
this case, we may write $V=v^{i}\frac{\partial}{\partial w^{i}}$ using the
component functions $v^{i}:\mathbb{R}^{n}\xrightarrow[]{}\mathbb{R}$ and
coordinates $w^{i}$.

Consider a smooth loss function
$\mathcal{L}:\mathcal{M}\xrightarrow[]{}\mathbb{R}$. Let $G$ be a symmetry
group of $\mathcal{L}$, i.e.
$\mathcal{L}(g\cdot{\bm{w}})=\mathcal{L}({\bm{w}})$ for all
${\bm{w}}\in\mathcal{M}$ and $g\in G$. Let $\mathfrak{X}$ be the set of all
vector fields on $\mathcal{M}$. Let $R=r^{i}\frac{\partial}{\partial w^{i}}$,
where $r^{i}=-\frac{\partial\mathcal{L}}{\partial w_{i}}$, be the reverse
gradient vector field. Let
$\mathfrak{X}_{\perp}=\\{A=a^{i}\frac{\partial}{\partial
w^{i}}\in\mathfrak{X}|~{}a^{i}\in C^{\infty}(\mathcal{M})\text{ and
}\sum_{i}a^{i}({\bm{w}})r^{i}({\bm{w}})=0,\forall{\bm{w}}\in\mathcal{M}\\}$ be
the set of vector fields orthogonal to $R$. If $G$ is a Lie group, the
infinitesimal action of its Lie algebra $\mathfrak{g}$ defines a set of vector
fields $\mathfrak{X}_{\mathfrak{g}}\subseteq\mathfrak{X}_{\perp}$.

A gradient flow is a curve $\gamma:\mathbb{R}\xrightarrow[]{}\mathcal{M}$
where the velocity is the value of $R$ at each point, i.e.
$\gamma^{\prime}(t)=R_{\gamma(t)}$ for all $t\in\mathbb{R}$. The Lie bracket
$[A,R]$ defines the derivative of $R$ with respect to $A$. Flows of $A$ and
$R$ commute if and only if $[A,R]=0$ (Theorem 9.44, [37]). That is,
teleportation can affect the convergence rate only if $[A,R]\mathcal{L}\neq 0$
for at least one $A\in\mathfrak{X}_{\mathfrak{g}}$. To simplify notations, we
write $([W,R]\mathcal{L})({\bm{w}})=0$ for a set of vector fields
$W\subseteq\mathfrak{X}$ when $([A,R]\mathcal{L})({\bm{w}})=0$ for all $A\in
W$.

We consider a gradient flow optimal if every point on the flow is a critical
point of the magnitude of gradient in its loss level set.

######  Definition 3.4.

Let
$f:\mathcal{M}\xrightarrow[]{}\mathbb{R},{\bm{w}}\mapsto\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}$.
A point ${\bm{w}}\in M$ is optimal with respect to a set of vector fields
$W\subseteq\mathfrak{X}_{\perp}$ if $Af({\bm{w}})=0$ for all $A\in W$. A
gradient flow $\gamma:\mathbb{R}\xrightarrow{}\mathcal{M}$ is optimal with
respect to $W$ if $\gamma(t)$ is optimal with respect to $W$ for all
$t\in\mathbb{R}$.

######  Proposition 3.5.

A point ${\bm{w}}\in\mathcal{M}$ is optimal with respect to a set of vector
fields $W$ if and only if $([W,R]\mathcal{L})({\bm{w}})=0$.

A sufficient condition for one teleportation to result in an optimal
trajectory is that whenever the function $[A,R]\mathcal{L}$ vanishes at
${\bm{w}}\in\mathcal{M}$, it vanishes along the entire gradient flow starting
at ${\bm{w}}$.

######  Proposition 3.6.

Let $W\subseteq\mathfrak{X}_{\perp}$ be a set of vector fields that are
orthogonal to $\frac{\partial\mathcal{L}}{\partial{\bm{w}}}$. Assume that for
all ${\bm{w}}\in\mathcal{M}$ such that $([W,R]\mathcal{L})({\bm{w}})=0$, we
have that $(R[W,R]\mathcal{L})({\bm{w}})=0$. Then the gradient flow starting
at any optimal point with respect to $W$ is optimal with respect to $W$.

To help check when the assumption in Proposition 3.6 is satisfied, we provide
an alternative form of $R[W,R]\mathcal{L}({\bm{w}})$ when
$[W,R]\mathcal{L}({\bm{w}})=0$.

######  Proposition 3.7.

If at all optimal points in
$S=\\{(M\frac{\partial\mathcal{L}}{\partial{\bm{w}}})^{i}\frac{\partial}{\partial
w^{i}}\in\mathfrak{X}|~{}M\in\mathbb{R}^{n\times n},M^{T}=-M\\}$ ,

$\displaystyle M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}=0$ (13)

for all anti-symmetric matrices $M\in\mathbb{R}^{n\times n}$, then the
gradient flow starting at an optimal point in $S$ is optimal in $S$.

From Proposition 3.7, we see that $R[W,R]\mathcal{L}({\bm{w}})$ is not
automatically 0 when $[W,R]\mathcal{L}({\bm{w}})=0$. Therefore, even if the
group is big enough, one teleportation does not guarantee that the gradient
flow intersects all future level sets at an optimal point. However, for loss
functions that satisfy $R[W,R]\mathcal{L}({\bm{w}})=0$ when
$[W,R]\mathcal{L}({\bm{w}})=0$, teleporting once optimizes the entire
trajectory. This is the case, for example, when
$\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w^{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{\alpha}}=\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial
w^{i}\partial w^{\alpha}}\frac{\partial\mathcal{L}}{\partial w^{j}}$ for all
$i,k,j,\alpha$ (Proposition D.3). In particular, all quadratic functions meet
this condition.

##  4 Teleportation for Improving Generalization

Teleportation was originally proposed to speedup optimization. In this
section, we explore the suitability of teleportation for improving
generalization, which is another important aspect of deep learning. We first
review definitions of the sharpness of minima. Then, we introduce a distinct
notion of curvature of the minima and discuss their implications on
generalization. By observing how sharpness and curvatures of minima are
correlated with generalization, we improve generalization by incorporating
sharpness and curvature into the objective for teleportation.

###  4.1 Sharpness of minima

Flat minima tend to generalize well [25]. A common definition of flat minimum
is based on the number of small eigenvalues of the Hessian. Although Hessian-
based sharpness metrics are known to correlate well to generalization, they
are expensive to compute and differentiate through. To use sharpness as an
objective in teleportation, we consider the change in the loss value averaged
over random directions. Let $D$ be a set of vectors drawn randomly from the
unit sphere $d_{i}\sim\\{d\in\mathbb{R}^{n}:||d||=1\\}$. Let $T$ be a list of
displacements $t_{j}\in\mathbb{R}$. Then, we have the following metric [34]:

$\displaystyle\text{Sharpness:
}\quad\phi({\bm{w}},T,D)=\frac{1}{|T||D|}\sum_{t\in T}\sum_{d\in
D}\mathcal{L}({\bm{w}}+td).\quad$ (14)

Figure 2: Gradient flow ($\mathcal{L}({\bm{w}})$) and a curve on the minimum
($\gamma$). The curvature of both curves may affect generalization.

###  4.2 Curvature of minima

At a minimum, the loss-invariant or flat directions are zero eigenvectors of
the Hessian. The curvature along these directions does not directly affect
Hessian-based sharpness metrics. However, these curvatures may affect
generalization, by themselves or by correlating to the curvature along non-
flat directions. Unlike the curvature of the loss (curve
$\mathcal{L}({\bm{w}})$ in Figure 2), the curvature of the minima (curve
$\gamma$) is less well studied. We provide a novel method to quantify the
curvature of the minima below.

Assume that the loss function $\mathcal{L}$ has a $G$ symmetry. Consider the
curve
$\gamma_{M}:\mathbb{R}\times\mathbb{R}^{n}\xrightarrow[]{}\mathbb{R}^{n}$
where $M\in\operatorname{Lie}(G)$ and
$\gamma_{M}(t,{\bm{w}})=\exp{(tM)}\cdot{\bm{w}}$. Then
$\gamma(0,{\bm{w}})={\bm{w}}$, and every point on $\gamma_{M}$ is in the
minimum if ${\bm{w}}$ is a minimum. Let $\gamma^{\prime}=\frac{d\gamma}{dt}$
be the derivative of a curve $\gamma$. The curvature of $\gamma$ is
$\kappa(\gamma,t)=\frac{\|T^{\prime}(t)\|}{\|\gamma^{\prime}(t)\|}$, where
$T(t)=\frac{\gamma^{\prime}(t)}{\|\gamma^{\prime}(t)\|}$ is the unit tangent
vector. We assume that the action map is smooth, since calculating the
curvature requires second derivatives and optimizing the curvature via
gradient descent requires third derivatives. For multi-layer network with
element-wise activations, we derive the group action, $\gamma$, and $\kappa$
in Appendix E.

Since the minimum can have more than one dimension, we measure the curvature
of a point ${\bm{w}}$ on the minimum by averaging the curvature of $k$ curves
with randomly selected $M_{i}\in\operatorname{Lie}(G)$. The resulting new
metric is

$\displaystyle\text{Curvature:
}\quad\psi({\bm{w}},k)=\frac{1}{k}\sum_{i=1}^{k}\left.\kappa(\gamma_{M_{i}}(0,{\bm{w}}),0)\right..\quad$
(15)

There are different ways to measure the curvature of a higher-dimensional
minimum, such as using the Gaussian curvature of 2D subspaces of the tangent
space. However, our method of approximating the mean curvature is easier to
compute and suitable as a differentiable objective.

(a) (b) (c) (d)

Figure 3: Illustration of the effect of sharpness (a,b) and curvature (c,d) of
minima on generalization. See Figure 2 for a 3D visualization of the curves
$\mathcal{L}({\bm{w}})$ and $\gamma$. When the loss landscape shifts due to a
change in data distribution, sharper minima have larger increase in loss. In
the example shown, minima with larger curvature moves further away from the
shifted minima.

###  4.3 Correlation with generalization

The sharpness of minima is well known to be correlated with generalization.
Figure 3(a)(b) visualizes an example of the shift in loss landscape
($\mathcal{L}({\bm{w}})$), and the change of loss $\Delta\mathcal{L}$ at a
minimizer ${\bm{w}}^{*}$ is large when the minimum is sharp. The relation
between the curvature of minimum and generalization is less well studied.
Figure 3(c)(d) shows one possible shift of the minimum ($\gamma$). Under this
shifting, the minimizer with a larger curvature becomes farther away from the
shifted minimum. The curve on the minimum can shift in other directions.
Appendix F.2 provides analytical examples of the correlation between curvature
and expected distance between the old and shifted minimum.

We verify the correlation between sharpness, curvatures, and validation loss
on MNIST [38], Fashion-MNIST [39], and CIFAR-10 [40]. On each dataset, we
train 100 three-layer neural networks with LeakyReLU using different
initializations. Details of the setup can be found in Appendix F.3.

Table 1 shows the Pearson correlation between validation loss and sharpness or
curvature (scatter plot in Figure 10 and 11 in the appendix). In all three
datasets, sharpness has a strong positive correlation with validation loss,
meaning that the average change in loss under perturbations is a good
indicator of test performance. This also confirms that wider minima are more
generalizable. For the architecture we consider, the curvatures of minima are
negatively correlated with validation loss. We observe that the magnitudes of
the curvatures are small, which suggests that the minima are relatively flat.

Table 1: Correlation with validation loss sharpness ($\phi$) | curvature ($\psi$)
---|---
MNIST | Fashion-MNIST | CIFAR-10 | MNIST | Fashion-MNIST | CIFAR-10
0.704 | 0.790 | 0.899 | -0.050 | -0.232 | -0.167

###  4.4 Teleportation for improving generalization

To improve the generalization ability of the minimizer and to gain
understanding of the curvature of minima, we teleport parameters to regions
with different sharpness and curvature. Multi-layer neural networks have
$GL(\mathbb{R})$ symmetry between layers (Appendix E.1). During teleportation,
we perform gradient descent on the group elements to change $\phi$ and $\psi$.
Results are averaged over 5 runs.

Figure 4: Changing sharpness (left) or curvature (right) using teleportation
and its effect on generalization on CIFAR-10. Solid line represents average
training loss, and dashed line represent average test loss. Teleporting to
decrease sharpness improves validation loss slightly. Teleportation changing
curvatures has a more significant impact on generalization ability.

Figure 4 shows the training curve of SGD on CIFAR-10, with one teleportation
at the 20 epoch. Teleporting to flatter points slightly improves the
validation loss, while teleporting to sharper point has no effect. Since the
group action keeps the loss invariant only on the batch of data used in
teleportation, the errors incurred in teleportation have a similar effect to a
warm restart, which makes the effect of changing sharpness less clear.

Interestingly, teleportation that changes curvature is able to affect
generalization. Teleporting to points with smaller curvatures helps find a
minimum with lower validation loss, while teleporting to points with larger
curvatures has the opposite effect. This suggests that at least locally,
curvature is correlated with generalization. Details of the experiment setup
can be found in Appendix F.4.

##  5 Applications to Other Optimization Algorithms

###  5.1 Standard Optimizers

Teleportation improves optimization not only for SGD. To show that
teleportation works well with other standard optimizers, we train a 3-layer
neural network on MNIST using different optimizers with and without
teleportation. During training, we teleport once at the first epoch, using 8
minibatches of size 200. Details can be found in Appendix G.2.

Figure 5 shows that teleportation improves the convergence rate when using
AdaGrad, SGD with momentum, RMSProp, and Adam. The runtime for a teleportation
is smaller than the time required to train one epoch, hence teleportation
improves convergence rate per epoch at almost no additional cost of time
(Figure 13 in the appendix).

Figure 5: Integrating teleportation with AdaGrad, momentum, RMSProp, and Adam
improves the convergence rate on MNIST. Solid line represents average training
loss, and dashed line average test loss. Shaded areas are 1 standard deviation
of the test loss across 5 runs.

###  5.2 Learning to Teleport

In optimization-based meta-learning, the parameter update rule or
hyperparameters are learned with a meta-optimizer [8, 41, 42, 43, 44].
Teleportation introduces an additional degree of freedom in parameter updates.
To exploit our ability to teleport without implementing optimization on
groups, we augment existing meta-learning algorithms by learning both the
local update rule and teleportation.

Let ${\bm{w}}_{t}\in\mathbb{R}^{d}$ be the parameters at time $t$, and
$\nabla_{t}=\left.\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right|_{{\bm{w}}_{t}}$
be the gradient of the loss $\mathcal{L}$. In gradient descent, the update
rule with learning rate $\eta$ is

$\displaystyle{\bm{w}}_{t+1}={\bm{w}}_{t}-\eta\nabla_{t}.$ (16)

In meta-learning [8], the update on ${\bm{w}}_{t}$ is learned using a meta-
learning optimizer $m$, which takes $\nabla_{t}$ as input. Here $m$ is an LSTM
model. Denote $h_{t}$ as the hidden state in the LSTM and $\phi$ as the
parameters in $m$. The update rule is

$\displaystyle{\bm{w}}_{t+1}$ $\displaystyle={\bm{w}}_{t}+f_{t}$ (17)
$\displaystyle\begin{bmatrix}f_{t}\\\ h_{t+1}\end{bmatrix}$
$\displaystyle=m(\nabla_{t},h_{t},\phi).$ (18)

Extending this approach beyond an additive update rule, we learn to teleport.
Let $G$ be a group whose action on the parameter space leaves $\mathcal{L}$
invariant. We use two meta-learning optimizers $m_{1},m_{2}$ to learn the
update direction $f_{t}\in\mathbb{R}^{d}$ and the group element $g_{t}\in G$:

$\displaystyle{\bm{w}}_{t+1}=g_{t}\cdot({\bm{w}}_{t}+f_{t})$ (19)
$\displaystyle\begin{bmatrix}f_{t}\\\ h_{1_{t+1}}\end{bmatrix}=m_{1}($
$\displaystyle\nabla_{t},h_{1_{t}},\phi_{1}),~{}~{}\begin{bmatrix}g_{t}\\\
h_{2_{t+1}}\end{bmatrix}=m_{2}(\nabla_{t},h_{2_{t}},\phi_{2}).$ (20)

##### Experiment setup.

We train and test on two-layer neural networks
$\mathcal{L}(W_{1},W_{2})=\|Y-W_{2}\sigma(W_{1}X)\|_{2}$, where
$W_{2},W_{1},X,Y\in\mathbb{R}^{20\times 20}$, and $\sigma$ is the LeakyReLU
function with slope coefficient 0.1. Both meta-optimizers are two-layer LSTMs
with hidden dimension 300. We train the meta-optimizers on multiple
trajectories created with different initializations, each consisting of 100
steps of gradient descent on $\mathcal{L}$ with random $X,Y$ and randomly
initialized $W$’s. We update the parameters in $m_{1}$ and $m_{2}$ by
unrolling every 10 steps. The learning rate for meta-optimizers are $10^{-4}$
for $m_{1}$ and $10^{-3}$ for $m_{2}$. We test the meta-optimizers using 5
trajectories not seen in training.

Algorithm 5.2 summarizes the training procedure. The vanilla gradient descent
baseline (“GD”) uses the largest learning rate that does not lead to
divergence ($3\times 10^{-4}$). The second baseline (“LSTM(update)”) learns
the update $f_{t}$ only and does not perform teleportation ($g_{t}=I,\forall
t$). The third baseline (“LSTM(lr,tele)”) learns the group element $g_{t}$ and
the learning rate used to perform gradient descent instead of the update
$f_{t}$. We keep training until adding more training trajectories does not
improve convergence rate. We use 700 training trajectories for our approach,
600 for the second baseline, and 30 for the third baseline.

##### Results.

By learning both the local update $f_{t}$ and nonlocal transformation $g_{t}$,
the meta-optimizers successfully learn to learn faster. Figure 6 shows the
improvement of our approach from the previous meta-learning method, which only
learns $f_{t}$. Compared to the baselines, learning the two types of updates
together (“LSTM(update,tele)”) achieves better convergence rate than learning
them separately. Additionally, learning the group element $g_{t}$ eliminates
the need for performing gradient ascent on the group manifold and reduces
hyperparameter tuning for teleportation. As an example of successful
integrations of teleportation into existing optimization algorithms, this toy
experiment demonstrates the flexibility and promising applications of
teleportation.

[H] Learning to teleport Input: Loss function $\mathcal{L}$, learning rate
$\eta$, number of epochs $T$, LSTM models $m_{1},m_{2}$ with initial
parameters $\phi_{1},\phi_{2}$, unroll step $t_{unroll}$. Output: Trained
parameters $\phi_{1}$ and $\phi_{2}$. each training initialization $t=1$ to
$T$ $f_{t},h_{1_{t+1}}=m_{1}(\nabla_{t},h_{1_{t}},\phi_{1})$
$g_{t},h_{2_{t+1}}=m_{2}(\nabla_{t},h_{2_{t}},\phi_{2})$ ${\bm{w}}\leftarrow
g_{t}\cdot({\bm{w}}+f_{t})$ $t\mod t_{unroll}=0$ update $\phi_{1},\phi_{2}$ by
back-propogation from the accumulated loss
$\sum_{i=t-t_{unroll}}^{t}\mathcal{L}({\bm{w}}_{i})$

Figure 6: Performance of the trained meta-optimizer on the test set. Learning
both local update $f_{t}$ and nonlocal transformation $g_{t}$ results in
better convergence rate than learning only local updates or learning only
teleportation.

##  6 Discussion

Teleportation provides a powerful tool to search on the loss level sets for
parameters with desired properties. We provide theoretical guarantees that
teleportation accelerates the convergence rate of SGD. Using concepts in
symmetry, we propose a distinct notion of curvature and show that
incorporating additional teleportation objectives such as changing the
curvatures can be beneficial to generalization. The close relationship between
symmetry and optimization opens up a number of exciting opportunities.
Exploring other objectives is an interesting future direction. Another
potential application is to extend the teleportation method to different
architectures, such as convolutional or graph neural networks, and to
different algorithms, such as sampling-based optimization.

The empirical results linking sharpness and curvatures to generalization are
intriguing. However, the theoretical origin of their relation remains unclear.
In particular, a more precise description of how the loss landscape changes
under distribution shifts is still an open question. More investigation of the
correlation between curvatures and generalization will help teleportation to
further improve generalization and take us a step closer to understanding the
loss landscape.

## Acknowledgements

This work was supported in part by the U.S. Department Of Energy, Office of
Science, U. S. Army Research Office under Grant W911NF-20-1-0334, Google
Faculty Award, Amazon Research Award, and NSF Grants #2134274, #2107256 and
#2134178.

## References

  * [1] Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka.  Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics.  In International Conference on Learning Representations, 2021.
  * [2] Twan Van Laarhoven.  L2 regularization versus batch and weight normalization.  Advances in Neural Information Processing Systems, 2017.
  * [3] J Elisenda Grigsby, Kathryn Lindsey, Robert Meyerhoff, and Chenxi Wu.  Functional dimension of feedforward relu neural networks.  arXiv preprint arXiv:2209.04036, 2022.
  * [4] Bo Zhao, Nima Dehmamy, Robin Walters, and Rose Yu.  Symmetry teleportation for accelerated optimization.  Advances in Neural Information Processing Systems, 2022.
  * [5] Marco Armenta, Thierry Judge, Nathan Painchaud, Youssef Skandarani, Carl Lemaire, Gabriel Gibeau Sanchez, Philippe Spino, and Pierre-Marc Jodoin.  Neural teleportation.  Mathematics, 11(2):480, 2023.
  * [6] Hamed Karimi, Julie Nutini, and Mark Schmidt.  Linear convergence of gradient and proximal-gradient methods under the polyak-lojasiewicz condition, 2016.
  * [7] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.  On large-batch training for deep learning: Generalization gap and sharp minima.  International Conference on Learning Representations, 2017.
  * [8] Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas.  Learning to learn by gradient descent by gradient descent.  Advances in Neural Information Processing Systems, 29, 2016.
  * [9] Berfin Şimşek, François Ged, Arthur Jacot, Francesco Spadaro, Clément Hongler, Wulfram Gerstner, and Johanni Brea.  Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances.  In International Conference on Machine Learning, pages 9722–9732. PMLR, 2021.
  * [10] Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur.  The role of permutation invariance in linear mode connectivity of neural networks.  International Conference on Learning Representations, 2022.
  * [11] Vijay Badrinarayanan, Bamdev Mishra, and Roberto Cipolla.  Symmetry-invariant optimization in deep networks.  arXiv preprint arXiv:1511.01754, 2015.
  * [12] Simon S Du, Wei Hu, and Jason D Lee.  Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced.  Neural Information Processing Systems, 2018.
  * [13] Iordan Ganev, Twan van Laarhoven, and Robin Walters.  Universal approximation and model compression for radial neural networks.  arXiv preprint arXiv:2107.02550v2, 2022.
  * [14] Marco Armenta and Pierre-Marc Jodoin.  The representation theory of neural networks.  Mathematics, 9(24), 2021.
  * [15] Iordan Ganev and Robin Walters.  Quiver neural networks.  arXiv preprint arXiv:2207.12773, 2022.
  * [16] Bo Zhao, Iordan Ganev, Robin Walters, Rose Yu, and Nima Dehmamy.  Symmetries, flat minima, and the conserved quantities of gradient flow.  International Conference on Learning Representations, 2023.
  * [17] Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro.  Path-SGD: Path-normalized optimization in deep neural networks.  In Advances in Neural Information Processing Systems, 2015.
  * [18] Qi Meng, Shuxin Zheng, Huishuai Zhang, Wei Chen, Zhi-Ming Ma, and Tie-Yan Liu.  $\mathcal{G}$-SGD: Optimizing relu neural networks in its positively scale-invariant space.  International Conference on Learning Representations, 2019.
  * [19] Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal.  Understanding the dynamics of gradient flow in overparameterized linear models.  In International Conference on Machine Learning, pages 10153–10161. PMLR, 2021.
  * [20] Jesse Dodge, Gabriel Ilharco, Roy Schwartz, Ali Farhadi, Hannaneh Hajishirzi, and Noah Smith.  Fine-tuning pretrained language models: Weight initializations, data orders, and early stopping.  arXiv preprint arXiv:2002.06305, 2020.
  * [21] Xavier Bouthillier, Pierre Delaunay, Mirko Bronzi, Assya Trofimov, Brennan Nichyporuk, Justin Szeto, Nazanin Mohammadi Sepahvand, Edward Raff, Kanika Madan, Vikram Voleti, et al.  Accounting for variance in machine learning benchmarks.  Proceedings of Machine Learning and Systems, 3:747–769, 2021.
  * [22] Sameera Ramasinghe, Lachlan MacDonald, Moshiur Farazi, Hemanth Sartachandran, and Simon Lucey.  How you start matters for generalization.  arXiv preprint arXiv:2206.08558, 2022.
  * [23] Ido Nachum and Amir Yehudayoff.  On symmetry and initialization for neural networks.  In Latin American Symposium on Theoretical Informatics, pages 401–412. Springer, 2021.
  * [24] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  International Conference on Learning Representations, 2017.
  * [25] Sepp Hochreiter and Jürgen Schmidhuber.  Flat minima.  Neural computation, 9(1):1–42, 1997.
  * [26] Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley.  Relative flatness and generalization.  35th Conference on Neural Information Processing Systems, 2021.
  * [27] Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al.  Towards theoretically understanding why sgd generalizes better than adam in deep learning.  Advances in Neural Information Processing Systems, 33:21285–21296, 2020.
  * [28] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina.  Entropy-sgd: Biasing gradient descent into wide valleys.  International Conference on Learning Representations, 2017.
  * [29] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.  Sharpness-aware minimization for efficiently improving generalization.  In International Conference on Learning Representations, 2021.
  * [30] Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi.  Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.  In International Conference on Machine Learning, pages 5905–5914. PMLR, 2021.
  * [31] Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales.  Fisher sam: Information geometry and sharpness aware minimisation.  In International Conference on Machine Learning, pages 11148–11161. PMLR, 2022.
  * [32] Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou.  Empirical analysis of the hessian of over-parametrized neural networks.  arXiv preprint arXiv:1706.04454, 2017.
  * [33] Lei Wu, Zhanxing Zhu, et al.  Towards understanding generalization of deep learning: Perspective of loss landscapes.  arXiv preprint arXiv:1706.10239, 2017.
  * [34] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.  Averaging weights leads to wider optima and better generalization.  Conference on Uncertainty in Artificial Intelligence, 2018.
  * [35] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio.  Sharp minima can generalize for deep nets.  In International Conference on Machine Learning, pages 1019–1028. PMLR, 2017.
  * [36] Chaoyue Liu, Libin Zhu, and Mikhail Belkin.  Loss landscapes and optimization in over-parameterized non-linear systems and neural networks.  Applied and Computational Harmonic Analysis, 59:85–116, 2022.  Special Issue on Harmonic Analysis and Machine Learning.
  * [37] John M Lee.  Introduction to Smooth Manifolds., chapter 9, page 233.  Graduate Texts in Mathematics, vol 218. Springer, New York, NY, 2013.
  * [38] Li Deng.  The mnist database of handwritten digit images for machine learning research [best of the web].  IEEE signal processing magazine, 29(6):141–142, 2012.
  * [39] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  arXiv preprint arXiv:1708.07747, 2017.
  * [40] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  Technical report, University of Toronto, 2009.
  * [41] Sachin Ravi and Hugo Larochelle.  Optimization as a model for few-shot learning.  International Conference on Learning Representations, 2017.
  * [42] Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-agnostic meta-learning for fast adaptation of deep networks.  In International conference on machine learning, pages 1126–1135. PMLR, 2017.
  * [43] Alex Nichol, Joshua Achiam, and John Schulman.  On first-order meta-learning algorithms.  arXiv preprint arXiv:1803.02999, 2018.
  * [44] Kartik Chandra, Audrey Xie, Jonathan Ragan-Kelley, and Erik Meijer.  Gradient descent: The ultimate optimizer.  In Advances in Neural Information Processing Systems, 2022.
  * [45] Sebastian U. Stich.  Unified optimal analysis of the (stochastic) gradient method.  CoRR, 2019.
  * [46] Osmar Aléssio.  Formulas for second curvature, third curvature, normal curvature, first geodesic curvature and first geodesic torsion of implicit curve in n-dimensions.  Computer Aided Geometric Design, 29(3-4):189–201, 2012.
  * [47] Aleksandr Mikhailovich Shelekhov.  On the curvatures of a curve in n-dimensional euclidean space.  Russian Mathematics, 65(11):46–58, 2021.

##  Appendix A Teleportation and SGD

######  Lemma A.1 (Descent Lemma).

Let $\mathcal{L}(w,\xi)$ be $\beta$–smooth function. It follows that

$\mathbb{E}\left[\lVert\nabla\mathcal{L}(w,\xi)\rVert^{2}\right]\leq
2\beta(\mathcal{L}(w)-\mathcal{L}(w^{*}))+2\beta(\mathcal{L}(w^{*})-\mathbb{E}\left[\inf_{w}\mathcal{L}(w,\xi)\right]).$
(21)

###### Proof.

Since $\mathcal{L}(w,\xi)$ is smooth we have that

$\displaystyle\mathcal{L}(z,\xi)-\mathcal{L}(w,\xi)$
$\displaystyle\leq\left<\nabla\mathcal{L}(w,\xi),z-w\right>+\frac{\beta}{2}\lVert
z-w\rVert^{2},\quad\forall z,w\in\mathbb{R}^{d}.$ (22)

By inserting

$z=w-\frac{1}{\beta}\nabla\mathcal{L}(w,\xi)$

into (22) we have that

$\displaystyle\mathcal{L}\big{(}w-(1/\beta)\nabla\mathcal{L}(w,\xi)\big{)}\leq\mathcal{L}(w,\xi)-\frac{1}{2\beta}\lVert\nabla\mathcal{L}(w,\xi)\rVert^{2}.$
(23)

Re-arranging we have that

$\displaystyle\mathcal{L}(w^{*},\xi)-\mathcal{L}(w,\xi)$
$\displaystyle=\mathcal{L}(w^{*},\xi)-\inf_{w}\mathcal{L}(w,\xi)+\inf_{w}\mathcal{L}(w,\xi)-\mathcal{L}(w,\xi)$
$\displaystyle\leq\mathcal{L}(w^{*},\xi)-\inf_{w}\mathcal{L}(w,\xi)+\mathcal{L}\big{(}w-(1/\beta)\nabla\mathcal{L}(w,\xi)\big{)}-\mathcal{L}(w,\xi)$
$\displaystyle\overset{\eqref{eq:smoothnessfuncstar2}}{\leq}\mathcal{L}(w^{*},\xi)-\inf_{w}\mathcal{L}(w,\xi)-\frac{1}{2\beta}\lVert\nabla\mathcal{L}(w,\xi)\rVert^{2},$

where the first inequality follows because
$\inf_{w}\mathcal{L}(w,\xi)\leq\mathcal{L}(w,\xi),\forall w.$ Re-arranging the
above and taking expectation gives

$\displaystyle\mathbb{E}\left[\lVert\nabla\mathcal{L}(w,\xi)\rVert^{2}\right]$
$\displaystyle\leq
2\mathbb{E}\left[\beta(\mathcal{L}(w^{*},\xi)-\inf_{w}\mathcal{L}(w,\xi)+\mathcal{L}(w,\xi)-\mathcal{L}(w^{*},\xi))\right]$
$\displaystyle\leq
2\beta\mathbb{E}\left[\mathcal{L}(w^{*},\xi)-\inf_{w}\mathcal{L}(w,\xi)+\mathcal{L}(w,\xi)-\mathcal{L}(w^{*},\xi)\right]$
$\displaystyle\leq
2\beta(\mathcal{L}(w)-\mathcal{L}(w^{*}))+2\beta(\mathcal{L}(w^{*})-\mathbb{E}\left[\inf_{w}\mathcal{L}(w,\xi)\right]).$

∎

######  Theorem 3.1.

Let $\mathcal{L}(w,\xi)$ be $\beta$–smooth and let

$\sigma^{2}\overset{\text{def}}{=}\mathcal{L}(w^{*})-\mathbb{E}\left[\inf_{w}\mathcal{L}(w,\xi)\right].$

Consider the iterates $w^{t}$ given by (5) where

$g^{t}\in\arg\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w^{t})\rVert^{2}.$ (24)

If $\eta=\frac{1}{\beta\sqrt{T-1}}$ then

$\displaystyle\min_{t=0,\ldots,T-1}\mathbb{E}\left[\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w^{t})\rVert^{2}\right]$
$\displaystyle\leq\frac{2\beta}{\sqrt{T-1}}\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}(w^{*})\right]+\frac{\beta\sigma^{2}}{\sqrt{T-1}}.$
(25)

###### Proof.

First note that if $\mathcal{L}(w,\xi)$ is $\beta$–smooth, then
$\mathcal{L}(w)$ is also a $\beta$–smooth function, that is

$\mathcal{L}(z)-\mathcal{L}(w)-\langle\nabla\mathcal{L}(w),z-w\rangle\leq\frac{\beta}{2}\lVert
z-w\rVert^{2}.$ (26)

Using (5) with $z=w^{t+1}$ and $w=g^{t}\cdot w^{t}$, together with (26) and
(2) we have that

$\displaystyle\mathcal{L}(w^{t+1})$ $\displaystyle\leq\mathcal{L}(g^{t}\cdot
w^{t})+\left<\nabla\mathcal{L}(g^{t}\cdot w^{t}),w^{t+1}-g^{t}\cdot
w^{t}\right>+\frac{\beta}{2}\lVert w^{t+1}-g^{t}\cdot w^{t}\rVert^{2}$ (27)
$\displaystyle=\mathcal{L}(w^{t})-\eta_{t}\left<\nabla\mathcal{L}(g^{t}\cdot
w^{t}),\nabla\mathcal{L}(g^{t}\cdot
w^{t},\xi^{t})\right>+\frac{\beta\eta_{t}^{2}}{2}\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t},\xi^{t})\rVert^{2}.$ (28)

Taking expectation conditioned on $w^{t}$, we have that

$\displaystyle\mathbb{E}_{t}\left[\mathcal{L}(w^{t+1})\right]$
$\displaystyle\leq\mathcal{L}(w^{t})-\eta_{t}\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t})\rVert^{2}+\frac{\beta\eta_{t}^{2}}{2}\mathbb{E}_{t}\left[\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t},\xi^{t})\rVert^{2}\right].$ (29)

Now since $\mathcal{L}(w,\xi)$ is $\beta$–smooth, from Lemma A.1 below we have
that

$\mathbb{E}\left[\lVert\nabla\mathcal{L}(w,\xi)\rVert^{2}\right]\leq
2\beta(\mathcal{L}(w)-\mathcal{L}(w^{*}))+2\beta(\mathcal{L}(w^{*})-\mathbb{E}\left[\inf_{w}\mathcal{L}(w,\xi)\right])$
(30)

Using (30) with $w=g^{t}\circ w^{t}$ we have that

$\displaystyle\mathbb{E}_{t}\left[\mathcal{L}(w^{t+1})\right]\leq\mathcal{L}(w^{t})$
$\displaystyle-\eta_{t}\lVert\nabla\mathcal{L}(g^{t}\cdot w^{t})\rVert^{2}$
(31) $\displaystyle+\beta^{2}\eta_{t}^{2}\left(\mathcal{L}(g^{t}\cdot
w^{t})-\mathcal{L}(w^{*})+\mathcal{L}(w^{*})-\mathbb{E}\left[\inf_{w}\mathcal{L}(w,\xi)\right]\right).$
(32)

Using that $\mathcal{L}(g^{t}\cdot w^{t})=\mathcal{L}(w^{t}),$ taking full
expectation and re-arranging terms gives

$\displaystyle\eta_{t}\mathbb{E}\left[\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t})\rVert^{2}\right]$
$\displaystyle\leq(1+\beta^{2}\eta_{t}^{2})\mathbb{E}\left[\mathcal{L}(w^{t})-\mathcal{L}^{*}\right]-\mathbb{E}\left[\mathcal{L}(w^{t+1})-\mathcal{L}^{*}\right]+\beta^{2}\eta_{t}^{2}\sigma^{2}.$
(33)

Now we use a re-weighting trick introduced in [45]. Let $\alpha_{t}>0$ be a
sequence such that $\alpha_{t}(1+\beta^{2}\eta_{t}^{2})=\alpha_{t-1}$.
Consequently if $\alpha_{-1}=1$ then
$\alpha_{t}=(1+\beta^{2}\eta_{t}^{2})^{-(t+1)}$ . Multiplying by both sides of
(33) by $\alpha_{t}$ thus gives

$\displaystyle\alpha_{t}\eta_{t}\mathbb{E}\left[\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t})\rVert^{2}\right]$
$\displaystyle\leq\alpha_{t-1}\mathbb{E}\left[\mathcal{L}(w^{t})-\mathcal{L}^{*}\right]-\alpha_{t}\mathbb{E}\left[\mathcal{L}(w^{t+1})-\mathcal{L}^{*}\right]+\alpha_{t}\beta^{2}\eta_{t}^{2}\sigma^{2}.$
(34)

Summing up from $t=0,\ldots,T-1$, and using telescopic cancellation, gives

$\displaystyle\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}\mathbb{E}\left[\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t})\rVert^{2}\right]$
$\displaystyle\leq\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}^{*}\right]+\beta^{2}\sigma^{2}\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}^{2}$
(35)

Let $A=\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}.$ Dividing both sides by $A$ gives

$\displaystyle\min_{t=0,\ldots,T-1}\mathbb{E}\left[\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t})\rVert^{2}\right]$
$\displaystyle\leq\frac{1}{\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}}\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t})\rVert^{2}$
$\displaystyle\leq\frac{\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}^{*}\right]+\beta^{2}\sigma^{2}\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}^{2}}{\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}}.$
(36)

Finally, if $\eta_{t}\equiv\eta$ then

$\displaystyle\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}$
$\displaystyle=\eta\sum_{t=0}^{T-1}(1+\beta^{2}\eta_{t}^{2})^{-(t+1)}=\frac{\eta}{1+\beta^{2}\eta^{2}}\frac{1-(1+\beta^{2}\eta^{2})^{-T}}{1-(1+\beta^{2}\eta^{2})^{-1}}$
(37) $\displaystyle=\frac{1-(1+\beta^{2}\eta^{2})^{-T}}{\beta^{2}\eta}$ (38)

To bound the term with the $-T$ power, we use that

$(1+\beta^{2}\eta^{2})^{-T}\leq\frac{1}{2}\quad\implies\quad\frac{\log(2)}{\log(1+\beta^{2}\eta^{2})}\leq
T.$

To simplify the above expression we can use

$\frac{x}{1+x}\leq\log(1+x)\leq x,\quad\mbox{ for }x\geq-1,$

thus

$\frac{\log(2)}{\log(1+\beta^{2}\eta^{2})}\leq
2\frac{1+\beta^{2}\eta^{2}}{\beta^{2}\eta^{2}}\leq T.$

Using the above we have that

$\displaystyle\sum_{t=0}^{T-1}\alpha_{t}\eta_{t}$
$\displaystyle\geq\frac{1}{2\beta^{2}\eta},\quad\mbox{for
}T\geq\frac{1+\beta^{2}\eta^{2}}{\beta^{2}\eta^{2}}$

Using this lower bound in (36) gives

$\displaystyle\min_{t=0,\ldots,T-1}\mathbb{E}\left[\lVert\nabla\mathcal{L}(g^{t}\cdot
w^{t})\rVert^{2}\right]$ $\displaystyle\leq
2\beta^{2}\eta\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}^{*}\right]+\eta\beta^{2}\sigma^{2},\quad\mbox{for
}T\geq\frac{1+\beta^{2}\eta^{2}}{\beta^{2}\eta^{2}}.$

Now note that

$T\geq\frac{1+\beta^{2}\eta^{2}}{\beta^{2}\eta^{2}}\Leftrightarrow\beta^{2}\eta^{2}(T-1)\geq
1\Leftrightarrow\eta\geq\frac{1}{\beta\sqrt{(T-1)}}.$

Thus finally setting $\eta=\frac{1}{\beta\sqrt{T-1}}$ gives the result (7).

∎

######  Theorem 3.2 (Error Bound).

Consider the setting of Theorem 3.1. If in addition to $\beta$–smoothness we
also assume that the $\mu$-PL (Polyak-Łojasiewicz) inequality holds, that is

$\displaystyle\mathcal{L}(w)-\inf\mathcal{L}\leq\frac{1}{2\mu}\|\nabla\mathcal{L}(w)\|^{2},$
(39)

then the $w^{t}$ iterates given by (5) with learning rate
$\eta=\frac{1}{\beta\sqrt{T-1}}$ converge according to

$\displaystyle\min_{t=0,\ldots,T-1}\mathbb{E}\left[\max_{g\in\mathcal{G}}\lVert
g\circ w^{t}-w^{*}\rVert\right]$
$\displaystyle\leq\frac{2\beta}{\mu}\frac{1}{\sqrt{T-1}}\mathbb{E}\left[\mathcal{L}(w^{0})-\mathcal{L}(w^{*})\right]+\frac{\beta}{\mu}\frac{\sigma^{2}}{\sqrt{T-1}}.$
(40)

###### Proof.

From Appendix A in [6] we have that if $\mathcal{L}$ is $\beta$–smooth, then
PL inequality is equivalent to the _error bound_ inequality given by

$\displaystyle\mu\lVert w-w^{*}\rVert\leq\|\nabla\mathcal{L}(w)\|.$ (41)

The proof now follows by directly applying the above in (7) and dividing
through by $\mu.$ ∎

Our result in (40) now shows that, for loss functions that are smooth and are
$\mu$–PL, we have that the expected loss converges to the optimal loss for
every point on the orbit of $w^{t}$. Furthermore, assuming the loss function
is $\mu$–PL can be reasonable for neural networks that are overparametrized.
Indeed, in [36] the authors showed that, for overparametrized neural networks,
when the task is regression, the loss function can often be $\mu$–PL.

######  Proposition A.2.

Assume that $\mathcal{L}:\mathbb{R}^{n}\xrightarrow[]{}\mathbb{R}$ is strictly
convex and twice continuously differentiable. Assume also that for any two
points $w_{a},w_{b}\in\mathbb{R}^{n}$ such that
$\mathcal{L}(w_{a})=\mathcal{L}(w_{b})$, there exists a $g\in G$ such that
$w_{a}=g\cdot w_{b}$. At two points $w_{1},w_{2}\in\mathbb{R}^{n}$, if
$\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w_{1})\rVert^{2}=\lVert\nabla\mathcal{L}(w_{2})\rVert^{2}$, then
$\mathcal{L}(w_{1})\leq\mathcal{L}(w_{2})$.

###### Proof.

Let $S(x)=\\{w:\mathcal{L}(w)=x\\}$ be the level sets of $\mathcal{L}$, and
$X=\\{\mathcal{L}(w):w\in\mathbb{R}^{n}\\}$ be the image of $\mathcal{L}$.
Since $G$ acts transitively on the level sets of $\mathcal{L}$,
$\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot w)\rVert^{2}=\max_{w\in
S(x)}\lVert\nabla\mathcal{L}(w)\rVert^{2}$. To simplify notation, we define a
function $F:X\xrightarrow[]{}\mathbb{R}$, $F(x)=\max_{w\in
S(x)}\lVert\nabla\mathcal{L}(w)\rVert^{2}$. Since $\nabla\mathcal{L}(w)$ is
continuously differentiable, the directional derivative of $F$ is defined.
Additionally, since $\mathcal{L}$ is continuous and its domain
$\mathbb{R}^{n}$ is connected, its image $X$ is also connected. This means
that for any $w_{1},w_{2}\in\mathbb{R}^{n}$ and
$\min(\mathcal{L}(w_{1}),\mathcal{L}(w_{2}))\leq
y\leq\max(\mathcal{L}(w_{1}),\mathcal{L}(w_{2}))$, there exists a
$w_{3}\in\mathbb{R}^{n}$ such that $\mathcal{L}(w_{3})=y$.

Next, we show that $F(\cdot)$ is strictly increasing by contradiction.

Suppose that $\mathcal{L}(w_{1})<\mathcal{L}(w_{2})$ and
$F(\mathcal{L}(w_{1}))\geq F(\mathcal{L}(w_{2}))$. By the mean value theorem,
there exists a $w_{3}$ such that
$\mathcal{L}(w_{1})<\mathcal{L}(w_{3})<\mathcal{L}(w_{2})$ and the directional
derivative of $F$ in the direction towards $\mathcal{L}(w_{2})$ is non-
positive:
$\partial_{\mathcal{L}(w_{2})-\mathcal{L}(w_{3})}F(\mathcal{L}(w_{3}))\leq 0$.
Let $w_{3}^{*}\in\operatorname{argmax}_{w\in
S(\mathcal{L}(w_{3}))}\lVert\nabla\mathcal{L}(w)\rVert^{2}$ be a point that
has the largest gradient norm in $S(\mathcal{L}(w_{3}))$. Then at $w_{3}^{*}$,
$\lVert\nabla\mathcal{L}\rVert^{2}$ cannot increase along the gradient
direction. However, this means

$\displaystyle\nabla\mathcal{L}(w_{3}^{*})\cdot\frac{\partial}{\partial
w}\lVert\nabla\mathcal{L}(w_{3}^{*})\rVert^{2}=\nabla\mathcal{L}(w_{3}^{*})^{T}H\nabla\mathcal{L}(w_{3}^{*})\leq
0.$ (42)

Since we assumed that $\mathcal{L}$ is convex and $\mathcal{L}(w_{3}^{*})$ is
not a minimum ($\mathcal{L}(w_{3}^{*})>\mathcal{L}(w_{1})$), we have that
$\nabla\mathcal{L}(w_{3}^{*})\neq 0$. Therefore, (42) contradicts with
$\mathcal{L}$ being strictly convex, and we have
$F(\mathcal{L}(w_{1}))<F(\mathcal{L}(w_{2}))$.

We have shown that $\mathcal{L}(w_{1})<\mathcal{L}(w_{2})$ implies
$F(\mathcal{L}(w_{1}))<F(\mathcal{L}(w_{2}))$. Taking the contrapositive and
switching $w_{1}$ and $w_{2}$, $F(\mathcal{L}(w_{1}))\leq
F(\mathcal{L}(w_{2}))$ implies $\mathcal{L}(w_{1})\leq\mathcal{L}(w_{2})$.
Equivalently, $\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w_{1})\rVert^{2}\leq\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w_{2})\rVert^{2}$ implies that $\mathcal{L}(w_{1})\leq\mathcal{L}(w_{2})$.

Finally, since

$\displaystyle\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w_{1})\rVert^{2}=\lVert\nabla\mathcal{L}(w_{2})\rVert^{2}\leq\max_{g\in\mathcal{G}}\lVert\nabla\mathcal{L}(g\cdot
w_{2})\rVert^{2},$ (43)

we have $\mathcal{L}(w_{1})\leq\mathcal{L}(w_{2})$. ∎

##  Appendix B Increase of gradient norm in second-order methods

The update direction in Newton’s method can be decomposed into two components:
one in the gradient direction and another one orthogonal to the gradient. We
refer to these two components as the gradient component and the symmetry
component, respectively.

Next, we show that for convex functions, the symmetry component moves the
parameters towards points with larger gradient norms, which coincides with
teleportation’s direction. Our results help explain why teleportation improves
convergence - second-order methods have a component that resembles
teleportation. In other words, teleportation brings first-order updates closer
to second-order updates.

Consider a twice differentiable convex function $\mathcal{L}({\bm{w}})$ that
has an invertible Hessian. Let ${\bm{v}}_{1}=-\nabla\mathcal{L}$ be the
gradient descent update direction, and ${\bm{v}}_{2}=-H^{-1}\nabla\mathcal{L}$
be the update direction in Newton’s method. The projection of ${\bm{v}}_{2}$
on ${\bm{v}}_{1}$ is given by

${\bm{v}}_{\parallel}=\frac{{\bm{v}}_{2}\cdot{\bm{v}}_{1}}{\|{\bm{v}}_{1}\|^{2}}{\bm{v}}_{1},$
(44)

and the component of ${\bm{v}}_{2}$ that is orthogonal to ${\bm{v}}_{1}$ is

${\bm{v}}_{\bot}={\bm{v}}_{2}-{\bm{v}}_{\parallel}.$ (45)

Teleportation moves parameters in the symmetry direction that increases
$\|\nabla\mathcal{L}\|$. The symmetry component of Newton’s method has a
similar role.

######  Lemma B.1.

For vector ${\bm{w}}\in\mathbb{R}^{n}$, positive definite matrix
$A\in\mathbb{R}^{n\times n}$, and $\alpha,\beta\in\mathbb{Z}$,

$\displaystyle({\bm{w}}^{T}A^{\alpha}{\bm{w}})^{2}\leq({\bm{w}}^{T}A^{\alpha+\beta}{\bm{w}})({\bm{w}}^{T}A^{\alpha-\beta}{\bm{w}}).$
(46)

###### Proof.

Since $A$ is positive definite, there exists orthogonal matrix
$P\in\mathbb{R}^{n\times n}$ and diagonal matrix $D\in\mathbb{R}^{n\times n}$
such that $A$ can be decomposed into $A=PDP^{-1}$. Then
$A^{-1}=(PDP^{-1})^{-1}=PDP^{-1}$, and $A^{2}=PDP^{-1}PDP^{-1}=PD^{2}P^{-1}$.
Similarly $A^{\gamma}=PD^{\gamma}P^{-1}$ for all integer $\gamma$.

Let ${\bm{x}}=P^{T}{\bm{w}}\in\mathbb{R}^{n}$. Substitute ${\bm{w}}=P{\bm{x}}$
into the left and right hand side of (46) and let $d_{i}$ be the $i^{th}$
diagonal element in $D$, we have

$\displaystyle({\bm{w}}^{T}A^{\alpha}{\bm{w}})^{2}$
$\displaystyle=\left({\bm{x}}^{T}P^{T}(PD^{\alpha}P^{-1})P{\bm{x}}\right)^{2}$
(47) $\displaystyle=({\bm{x}}^{T}D^{\alpha}{\bm{x}})^{2}$ (48)
$\displaystyle=\left(\sum_{i}^{n}d_{i}^{\alpha}x_{i}^{2}\right)\left(\sum_{j}^{n}d_{j}^{\alpha}x_{j}^{2}\right)$
(49) $\displaystyle=\sum_{i\leq
j}(2d_{i}^{\alpha}d_{j}^{\alpha})x_{i}^{2}x_{j}^{2}$ (50)

and

$\displaystyle({\bm{w}}^{T}A^{\alpha+\beta}{\bm{w}})({\bm{w}}^{T}A^{\alpha-\beta}{\bm{w}})$
$\displaystyle=({\bm{x}}^{T}P^{T}(PD^{\alpha+\beta}P^{-1})P{\bm{x}})({\bm{x}}^{T}P^{T}(PD^{\alpha-\beta}P^{-1})P{\bm{x}})$
(51)
$\displaystyle=({\bm{x}}^{T}D^{\alpha+\beta}{\bm{x}})({\bm{x}}^{T}D^{\alpha-\beta}{\bm{x}})$
(52)
$\displaystyle=\left(\sum_{i}^{n}d_{i}^{\alpha+\beta}x_{i}^{2}\right)\left(\sum_{j}^{n}d_{j}^{\alpha-\beta}x_{j}^{2}\right)$
(53) $\displaystyle=\sum_{i\leq
j}(d_{i}^{\alpha+\beta}d_{j}^{\alpha-\beta}+d_{i}^{\alpha-\beta}d_{j}^{\alpha+\beta})x_{i}^{2}x_{j}^{2}.$
(54)

To prove (50) $\leq$ (54), it suffices to show that
$2d_{i}^{\alpha}d_{j}^{\alpha}\leq
d_{i}^{\alpha+\beta}d_{j}^{\alpha-\beta}+d_{i}^{\alpha-\beta}d_{j}^{\alpha+\beta}$
for all $(i,j)$. Since $d_{i}$ are the eigenvalues of $A$, all $d_{i}$’s are
positive. By the inequality of arithmetic and geometric means,

$\displaystyle(d_{i}d_{j}^{-1})^{\beta}+(d_{i}^{-1}d_{j})^{\beta}\geq
2\sqrt{(d_{i}d_{j}^{-1})^{\beta}(d_{i}^{-1}d_{j})^{\beta}}=2.$ (55)

Therefore,

$\displaystyle
d_{i}^{\alpha+\beta}d_{j}^{\alpha-\beta}+d_{i}^{\alpha-\beta}d_{j}^{\alpha+\beta}=\left[(d_{i}d_{j}^{-1})^{\beta}+(d_{i}^{-1}d_{j})^{\beta}\right](d_{i}d_{j})^{\alpha}\geq
2d_{i}^{\alpha}d_{j}^{\alpha}.$ (56)

∎

The following proposition states that the gradient norm increases along the
symmetry component for convex functions.

######  Proposition B.2.

For convex function $\mathcal{L}$, the directional derivative of
$\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}$ along
the direction of ${\bm{v}}_{\bot}$ is non-negative. That is,

$\displaystyle{\bm{v}}_{\bot}\cdot\frac{\partial}{\partial{\bm{w}}}\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}\geq
0.$ (57)

###### Proof.

Note that
$\frac{\partial}{\partial{\bm{w}}}\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}=2H\nabla\mathcal{L}$.
Then we have

$\displaystyle{\bm{v}}_{\bot}\cdot\frac{\partial}{\partial{\bm{w}}}\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}$
$\displaystyle=2(H\nabla\mathcal{L})^{T}{\bm{v}}_{\bot}$ (58)
$\displaystyle=-2\nabla\mathcal{L}^{T}H^{T}H^{-1}\nabla\mathcal{L}+2\frac{\nabla\mathcal{L}^{T}H^{-1}\nabla\mathcal{L}}{\nabla\mathcal{L}^{T}\nabla\mathcal{L}}\nabla\mathcal{L}^{T}H^{T}\nabla\mathcal{L}$
(59)
$\displaystyle=2\left(-\nabla\mathcal{L}^{T}\nabla\mathcal{L}+\frac{\nabla\mathcal{L}^{T}H^{-1}\nabla\mathcal{L}}{\nabla\mathcal{L}^{T}\nabla\mathcal{L}}\nabla\mathcal{L}^{T}H^{T}\nabla\mathcal{L}\right).$
(60)

Setting $\alpha=0$ and $\beta=1$ in Lemma B.1 and substitute
$\nabla\mathcal{L},H$ for ${\bm{w}},A$, we have

$\displaystyle-\nabla\mathcal{L}^{T}\nabla\mathcal{L}+\frac{\nabla\mathcal{L}^{T}H^{-1}\nabla\mathcal{L}}{\nabla\mathcal{L}^{T}\nabla\mathcal{L}}\nabla\mathcal{L}^{T}H^{T}\nabla\mathcal{L}\geq
0.$ (61)

Therefore,

$\displaystyle{\bm{v}}_{\bot}\cdot\frac{\partial}{\partial{\bm{w}}}\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}=2\left(-\nabla\mathcal{L}^{T}\nabla\mathcal{L}+\frac{\nabla\mathcal{L}^{T}H^{-1}\nabla\mathcal{L}}{\nabla\mathcal{L}^{T}\nabla\mathcal{L}}\nabla\mathcal{L}^{T}H^{T}\nabla\mathcal{L}\right)\geq
0.$ (62)

∎

Figure 7: The gradient norm increases along ${\bm{v}}_{\bot}$, the component
of Newton’s direction (${\bm{v}}_{2}$) that is orthogonal to the gradient
(${\bm{v}}_{1}$).

Figure 7 visualizes the vectors used in the Proposition B.2 for a convex
quadratic function with two parameters. While any vector can be decomposed
into a gradient component ${\bm{v}}_{\parallel}$ and a symmetry component
${\bm{v}}_{\bot}$, the gradient norm does not necessarily increase along
${\bm{v}}_{\bot}$. Therefore, Proposition B.2 implies that the symmetry
component of Newton’s method resembles teleportation.

Note that Lemma B.1 does not hold for non-convex functions. For example,

$\displaystyle\alpha=0,\beta=1,{\bm{w}}=\begin{bmatrix}1\\\
3\end{bmatrix},A=\begin{bmatrix}1&0\\\
0&-2\end{bmatrix},A^{-1}=\begin{bmatrix}1&0\\\ 0&-\frac{1}{2}\end{bmatrix}.$
(63)
$\displaystyle({\bm{w}}^{T}{\bm{w}})^{2}=100>(-3.5)\times(-17)=({\bm{w}}^{T}A^{-1}{\bm{w}})({\bm{w}}^{T}A{\bm{w}}).$
(64)

As a result, Proposition B.2 also does not hold.

The following Proposition gives a stronger result, namely that among all
directions orthogonal to the gradient, ${\bm{v}}_{\bot}$ is the direction
along which the gradient norm increases the fastest.

######  Proposition B.3.

Let $\hat{{\bm{v}}}_{\bot}=\frac{{\bm{v}}_{\bot}}{\|{\bm{v}}_{\bot}\|}$. Then

$\displaystyle\hat{{\bm{v}}}_{\bot}$
$\displaystyle\in\operatorname{argmax}_{\hat{{\bm{v}}}\in\mathbb{R}^{n}}\hat{{\bm{v}}}\cdot\frac{\partial}{\partial{\bm{w}}}\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}$
(65) $\displaystyle\mbox{subject
to}\quad\hat{{\bm{v}}}\cdot\nabla\mathcal{L}=0\text{ and
}\hat{{\bm{v}}}^{T}\hat{{\bm{v}}}=1.$ (66)

###### Proof.

Recall that
$\frac{\partial}{\partial{\bm{w}}}\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}=2H\nabla\mathcal{L}$.
The Lagrange multiplier of this constrained optimization is

$\displaystyle
L(\hat{{\bm{v}}},\lambda_{1},\lambda_{2})=2\hat{{\bm{v}}}^{T}H\nabla\mathcal{L}-\lambda_{1}\hat{{\bm{v}}}^{T}\nabla\mathcal{L}-\lambda_{2}(\hat{{\bm{v}}}^{T}\hat{{\bm{v}}}-1)$
(67)

The solution can be found by solving $\frac{\partial
L}{\partial\hat{{\bm{v}}}}=0$, $\frac{\partial L}{\partial\lambda_{1}}=0$, and
$\frac{\partial L}{\partial\lambda_{2}}=0$:

$\displaystyle\frac{\partial
L}{\partial\hat{{\bm{v}}}}=2H\nabla\mathcal{L}-\lambda_{1}\nabla\mathcal{L}-2\lambda_{2}\hat{{\bm{v}}}=0$
$\displaystyle\quad\Rightarrow\quad\hat{{\bm{v}}}=\frac{1}{2\lambda_{2}}\left(H\nabla\mathcal{L}-\frac{\lambda_{1}}{2}\nabla\mathcal{L}\right)$
(68) $\displaystyle\frac{\partial
L}{\partial\lambda_{1}}=-\hat{{\bm{v}}}^{T}\nabla\mathcal{L}=-\nabla\mathcal{L}^{T}H\nabla\mathcal{L}+\frac{\lambda_{1}}{2}\nabla\mathcal{L}^{T}\nabla\mathcal{L}=0$
$\displaystyle\quad\Rightarrow\quad\lambda_{1}=\frac{2\nabla\mathcal{L}^{T}H\nabla\mathcal{L}}{\nabla\mathcal{L}^{T}\nabla\mathcal{L}}$
(69)

Substitute $\lambda_{1}$ into $\hat{{\bm{v}}}$, we have

$\displaystyle\hat{{\bm{v}}}=\frac{1}{2\lambda_{2}}\left(H\nabla\mathcal{L}-\frac{2\nabla\mathcal{L}^{T}H\nabla\mathcal{L}}{\nabla\mathcal{L}^{T}\nabla\mathcal{L}}\nabla\mathcal{L}\right).$
(70)

Since $\lambda_{2}$ is a scalar, $\hat{{\bm{v}}}$ has the same direction as
$\hat{{\bm{v}}}_{\bot}$, which means $\hat{{\bm{v}}}_{\bot}$ is a solution to
(66). ∎

We have shown that a component in second-order method update is in the
teleportation direction. Viewing from the other direction, these results also
suggest that small teleportation brings the subsequent first-order updates
closer to second-order updates, at least for convex functions. This advances
the results in [4], which only says that teleportation brings subsequent
first-order updates equal to second-order updates at critical points of
$\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}$.

The equality in (57) holds at critical points of
$\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}$ in the
level set. Proposition B.2 also holds for nonstrictly convex $\mathcal{L}$.
The proposition does not hold for non-convex functions. However, since
Newton’s method may not move in the right direction in nonconvex problems,
this is actually good. For smooth functions, it may be possible to find a
modified version of (57) that depends on the smoothness constant.

It has been observed that the solutions found by adaptive gradient descent
have worse generalization ability compared to stochastic gradient descent
[27]. Depending on what other classes of functions the proposition can
generalize to, our results may support this observation. If second-order
methods prefer steeper trajectories, then the minima they find are likely to
be sharper.

##  Appendix C Teleportation and Newton’s method

######  Lemma C.1 (One step of Newton’s Method).

Let $f(x)$ be a $\mu$–strongly convex and $L$–smooth function, that is, we
have a global lower bound on the Hessian given by

$LI\succeq\nabla^{2}f(x)\succeq\mu I,\quad\forall x\in\mathbb{R}^{n}.$ (71)

Furthermore, if the Hessian is also $G$–Lipschitz

$\|\nabla^{2}f(x)-\nabla^{2}f(y)\|\leq G\|x-y\|$ (72)

then Newton’s method

$x^{k+1}=x^{k}-\lambda_{k}\nabla^{2}f(x^{k})^{-1}\nabla f(x^{k})$

has a mixed linear and quadratic convergence according to

$\|x^{k+1}-x^{*}\|\leq\frac{G}{2\mu}\|x^{k}-x^{*}\|^{2}+|1-\lambda_{k}|\frac{L}{2\mu}\|x^{k}-x^{*}\|.$
(73)

###### Proof.

$\displaystyle x^{k+1}-x^{*}$
$\displaystyle=x^{k}-x^{*}-\lambda_{k}\nabla^{2}f(x^{k})^{-1}\left(\nabla
f(x^{k})-\nabla f(x^{*})\right)$
$\displaystyle=x^{k}-x^{*}-\lambda_{k}\nabla^{2}f(x^{k})^{-1}\int_{s=0}^{1}\nabla^{2}f(x^{k}+s(x^{*}-x^{k}))(x^{k}-x^{*})ds\quad\mbox{(Mean
value theorem)}$
$\displaystyle=\nabla^{2}f(x^{k})^{-1}\int_{s=0}^{1}\left(\nabla^{2}f(x^{k})-\lambda_{k}\nabla^{2}f(x^{k}+s(x^{*}-x^{k}))\right)(x^{k}-x^{*})ds$
$\displaystyle=\nabla^{2}f(x^{k})^{-1}\int_{s=0}^{1}\left(\nabla^{2}f(x^{k})-\nabla^{2}f(x^{k}+s(x^{*}-x^{k}))\right.$
$\displaystyle\hskip
90.0pt\left.+(1-\lambda_{k})\nabla^{2}f(x^{k}+s(x^{*}-x^{k}))\right)(x^{k}-x^{*})ds$

Let $\delta_{k}:=\|x^{k+1}-x^{*}\|$. Taking norms we have that

$\displaystyle\delta_{k+1}$
$\displaystyle\leq\|\nabla^{2}f(x^{k})^{-1}\|\int_{s=0}^{1}\left(\|\nabla^{2}f(x^{k})-\nabla^{2}f(x^{k}+s(x^{*}-x^{k}))\|\right.$
$\displaystyle\hskip
100.0pt\left.+|1-\lambda_{k}\|\|\nabla^{2}f(x^{k}+s(x^{*}-x^{k}))\|\right)\delta_{k}ds$
$\displaystyle\overset{\eqref{eq:Hesslip}+\eqref{eq:strconv}}{\leq}\frac{G}{\mu}\int_{s=0}^{1}s\|x^{k}-x^{*}\|^{2}ds+|1-\lambda_{k}|\frac{L}{\mu}\int_{s=0}^{1}s\|x^{k}-x^{*}\|ds$
$\displaystyle=\frac{G}{2\mu}\|x^{k}-x^{*}\|^{2}+|1-\lambda_{k}|\frac{L}{2\mu}\|x^{k}-x^{*}\|.$

∎

The assumptions on for this proof can be relaxed, since we only require the
Hessian is Lipschitz and lower bounded in a $\frac{\mu}{2L}$–ball around
$x^{*}$.

######  Proposition 3.3 (Quadratic term in convergence rate).

Let $f$ be strictly convex and let $w_{0}\in\mathbb{R}^{d}$. Let

$w^{\prime}\in\operatorname{argmax}_{w\in\mathbb{R}^{d}}\frac{1}{2}\|\nabla\mathcal{L}(w)\|^{2}\quad\mbox{subject
to}\quad\mathcal{L}(w)=\mathcal{L}(w_{0}).$ (74)

If $\nabla\mathcal{L}(w^{\prime})\neq 0$ then there exists $\lambda_{0}$ such
that

$0\leq\lambda_{0}\leq\lambda_{\max}(\nabla^{2}\mathcal{L}(w_{0}))$

and one step of gradient descent with learning rate $\gamma>0$ gives

$\displaystyle w_{1}$
$\displaystyle=w^{\prime}-\gamma\nabla\mathcal{L}(w^{\prime})$
$\displaystyle=w^{\prime}-\gamma\lambda_{0}\nabla^{2}\mathcal{L}(w^{\prime})^{-1}\nabla\mathcal{L}(w^{\prime}).$
(75)

Consequently, letting $w^{\prime}=g_{0}\circ w_{0}$, and if
$\gamma\leq\frac{1}{\lambda_{0}}$ then under the assumptions of Lemma C.1 we
have that

$\|w_{1}-w^{*}\|\leq\frac{G}{2\mu}\|g_{0}\circ
w_{0}-x^{*}\|^{2}+|1-\gamma\lambda_{0}|\frac{L}{2\mu}\|g_{0}\circ
w_{0}-w^{*}\|.$

###### Proof.

The Lagrangian associated to (74) is given by

$L(w,\lambda)=\frac{1}{2}\|\nabla\mathcal{L}(w)\|^{2}+\lambda(\mathcal{L}(w_{0})-\mathcal{L}(w)).$

Taking the derivative in $w$ and setting it to zero gives

$\nabla_{w}L(w,\lambda_{0})=0\;\implies\;\nabla^{2}\mathcal{L}(w)\nabla\mathcal{L}(w)-\lambda_{0}\nabla\mathcal{L}(w)=0.$
(76)

Re-arranging we have that

$\nabla\mathcal{L}(w)=\lambda_{0}\nabla^{2}\mathcal{L}(w)^{-1}\nabla\mathcal{L}(w).$

If $\nabla\mathcal{L}(w^{\prime})\neq 0$ then from the above we have that

$\|\nabla\mathcal{L}(w)\|^{2}=\lambda_{0}\nabla\mathcal{L}(w)^{\top}\nabla^{2}\mathcal{L}(w)^{-1}\nabla\mathcal{L}(w)>0.$

Since $\nabla^{2}\mathcal{L}(w)^{-1}$ is positive definite we have that
$\nabla\mathcal{L}(w)^{\top}\nabla^{2}\mathcal{L}(w)^{-1}\nabla\mathcal{L}(w)\geq
0$, and consequently $\lambda_{0}>0.$ Finally from (76) we have that
$\lambda_{0}$ is an eigenvalue of $\nabla^{2}\mathcal{L}(w)$ and thus it must
be smaller or equal to the largest eigenvalue of $\nabla^{2}\mathcal{L}(w)$.

∎

##  Appendix D Is one teleportation enough to find the optimal trajectory?

######  Proposition 3.5.

A point ${\bm{w}}\in M$ is optimal in a set of vector fields $W$ if and only
if $[A,R]\mathcal{L}({\bm{w}})=0$ for all $A\in W$.

###### Proof.

Note that $A\mathcal{L}=a^{i}\frac{\partial\mathcal{L}}{\partial w^{i}}=0$. We
have

$\displaystyle[A,R]\mathcal{L}$
$\displaystyle=AR\mathcal{L}-RA\mathcal{L}=A\left(r^{i}\frac{\partial\mathcal{L}}{\partial
w_{i}}\right)-0=-A\left\|\frac{\partial\mathcal{L}}{\partial{\bm{w}}}\right\|_{2}^{2}=-Af.$
(77)

The result then follows from Definition 3.4. ∎

######  Proposition 3.6.

Let $W\subseteq\mathfrak{X}_{\perp}$ be a set of vector fields that are
orthogonal to the gradient of $\mathcal{L}$. If $[A,R]\mathcal{L}({\bm{w}})=0$
for all $A\in W$ implies that $R([A,R]\mathcal{L})({\bm{w}})=0$ for all $A\in
W$, then the gradient flow starting at an optimal point in $W$ is optimal in
$W$.

###### Proof.

Consider the gradient flow $\gamma$ that starts at an optimal point in $W$.
The derivative of $[A,R]\mathcal{L}$ along $\gamma$ is

$\displaystyle\frac{d}{dt}[A,R]\mathcal{L}(\gamma(t))=\gamma^{\prime}(t)([A,R]\mathcal{L})(\gamma(t))=-R[A,R]\mathcal{L}(\gamma(t)).$
(78)

Since $\gamma(0)$ is an optimal point, $[A,R]\mathcal{L}(\gamma(0))=0$ for all
$A\in W$ by Proposition 3.5. By assumption, if $[A,R]\mathcal{L}(\gamma(t))=0$
for all $A\in W$, then $R([A,R]\mathcal{L})(\gamma(t))=0$ for all $A\in W$.
Therefore, both the value and the derivative of $[A,R]\mathcal{L}$ stay 0
along $\gamma$. Since $[A,R]\mathcal{L}(\gamma(t))=0$ for all
$t\in\mathbb{R}$, $\gamma$ is optimal in $W$. ∎

To help check when Proposition 3.6 is satisfied, we provide an alternative
form of $R[A,R]\mathcal{L}({\bm{w}})$ under the assumption that
$[A,R]\mathcal{L}({\bm{w}})=0$. We will use the following lemmas in the proof.

######  Lemma D.1.

For two vectors ${\bm{v}},{\bm{w}}\in\mathbb{R}^{n}$, if
${\bm{v}}^{T}{\bm{w}}=0$ and ${\bm{w}}\neq\mathbf{0}$, then there exists an
anti-symmetric matrix $M\in\mathbb{R}^{n\times n}$ such that
${\bm{v}}=M{\bm{w}}$.

###### Proof.

Let ${\bm{w}}_{0}=[1,0,...,0]^{T}\in\mathbb{R}^{n}$. Consider a list of $n-1$
anti-symmetric matrices $M_{i}\in\mathbb{R}^{n\times n}$, where

$\displaystyle M_{ij}^{~{}k}=\begin{cases}-1,&\text{if $j=1$ and $k=i+1$}\\\
1,&\text{if $j=i+1$ and $k=1$}\\\ 0,&\text{otherwise}\end{cases}$ (79)

In matrix form, the $M_{i}$’s are

$\displaystyle M_{1}=\begin{bmatrix}0&-1&0&...&0\\\ 1&0&0&...&0\\\
0&0&0&...&0\\\ &&...&&\\\
0&0&0&...&0\end{bmatrix},M_{2}=\begin{bmatrix}0&0&-1&...&0\\\ 0&0&0&...&0\\\
1&0&0&...&0\\\ &&...&&\\\
0&0&0&...&0\end{bmatrix},...,M_{n-1}=\begin{bmatrix}0&0&0&...&-1\\\
0&0&0&...&0\\\ 0&0&0&...&0\\\ &&...&&\\\ 1&0&0&...&0\end{bmatrix}.$ (80)

Since $M_{i}$’s are anti-symmetric, $M_{i}{\bm{w}}_{0}$ is orthogonal to
${\bm{w}}_{0}$. The norm of $M_{i}{\bm{w}}_{0}=\mathbf{e}_{i+1}$ is 1.
Additionally, $M_{i}{\bm{w}}_{0}$ is orthogonal to $M_{j}{\bm{w}}_{0}$ for
$i\neq j$:

$\displaystyle(M_{i}{\bm{w}}_{0})^{T}(M_{j}{\bm{w}}_{0})=\mathbf{e}_{i+1}^{T}\mathbf{e}_{j+1}=\delta_{ij}.$
(81)

Denote
${\bm{w}}_{0}^{\bot}=\\{{\bm{x}}\in\mathbb{R}^{n}:{\bm{x}}^{T}{\bm{w}}_{0}=0\\}$
as the orthogonal complement of ${\bm{w}}_{0}$. Then $M_{i}{\bm{w}}_{0}$ forms
a basis of ${\bm{w}}_{0}^{\bot}$. Next, we extend this to an arbitrary
${\bm{w}}\in\mathbb{R}^{n}$.

Let $\hat{{\bm{w}}}=\frac{{\bm{w}}}{\|{\bm{w}}\|_{2}}$. Since $\hat{{\bm{w}}}$
has norm 1, there exists an orthogonal matrix $R$ such that
$\hat{{\bm{w}}}=R{\bm{w}}_{0}$. Let $M_{i}^{\prime}=RM_{i}R^{T}$. Then
$M_{i}^{\prime}$ is anti-symmetric:

$\displaystyle(RM_{i}R^{T})^{T}=RM^{T}_{i}R^{T}=-RM_{i}R^{T}.$ (82)

It follows that $M_{i}^{\prime}\hat{{\bm{w}}}$ is orthogonal to
$\hat{{\bm{w}}}$. The norm of $M_{i}^{\prime}\hat{{\bm{w}}}$ is
$\|(RM_{i}R^{T})(R{\bm{w}}_{0})\|=\|RM_{i}{\bm{w}}_{0}\|=\|M_{i}{\bm{w}}_{0}\|=1$.
Additionally, $M_{i}^{\prime}\hat{{\bm{w}}}$ is orthogonal to
$M_{j}^{\prime}\hat{{\bm{w}}}$ for $i\neq j$:

$\displaystyle(M_{i}^{\prime}\hat{{\bm{w}}})^{T}(M_{j}^{\prime}\hat{{\bm{w}}})$
$\displaystyle=(RM_{i}R^{T}R{\bm{w}}_{0})^{T}(RM_{j}R^{T}R{\bm{w}}_{0})$ (83)
$\displaystyle={\bm{w}}_{0}^{T}R^{T}RM_{i}^{T}R^{T}RM_{j}R^{T}R{\bm{w}}_{0}$
(84) $\displaystyle={\bm{w}}_{0}^{T}M_{i}^{T}M_{j}{\bm{w}}_{0}$ (85)
$\displaystyle=\delta_{ij}.$ (86)

Therefore, $M_{i}^{\prime}\hat{{\bm{w}}}$ spans
$\hat{{\bm{w}}}^{\bot}={\bm{w}}^{\bot}$. This means that any vector
${\bm{v}}\in{\bm{w}}^{\bot}$ can be written as a linear combination of
$M_{i}^{\prime}\hat{{\bm{w}}}$. That is, there exists
$k_{1},...,k_{n}\in\mathbb{R}$, such that
${\bm{v}}=\sum_{i}k_{i}(M_{i}^{\prime}\hat{{\bm{w}}})$. To find the anti-
symmetric $M$ that takes ${\bm{w}}$ to ${\bm{v}}$, note that

$\displaystyle{\bm{v}}=\left(\sum_{i}k_{i}M_{i}^{\prime}\right)\hat{{\bm{w}}}=\left(\|{\bm{w}}\|_{2}^{-1}\sum_{i}k_{i}M_{i}^{\prime}\right){\bm{w}}.$
(87)

Since the sum of anti-symmetric matrices is anti-symmetric, and the product of
an anti-symmetric matrix and a scalar is also anti-symmetric,
$\|{\bm{w}}\|_{2}^{-1}\sum_{i}k_{i}M_{i}^{\prime}$ is anti-symmetric. ∎

######  Lemma D.2.

Let ${\bm{v}}\in\mathbb{R}^{n}$ be a nonzero vector. Then the two sets
$\\{M{\bm{v}}:M\in\mathbb{R}^{n\times n},M^{T}=-M\\}$ and
$\\{{\bm{w}}\in\mathbb{R}^{n}:{\bm{w}}^{T}{\bm{v}}=0\\}$ are equal.

###### Proof.

Let $A=\\{M{\bm{v}}:M\in\mathbb{R}^{n\times n},M^{T}=M^{-1}\\}$ and
$B=\\{{\bm{w}}\in\mathbb{R}^{n}:{\bm{w}}^{T}{\bm{v}}=0\\}$. Since
$(M{\bm{v}})^{T}{\bm{v}}=0$ for all anti-symmetric $M$, every element in $A$
is in $B$. By Lemma D.1, every element in $B$ is in $A$. Therefore $A=B$. ∎

Let
$S=\\{(M\frac{\partial\mathcal{L}}{\partial{\bm{w}}})^{i}\frac{\partial}{\partial
w^{i}}\in\mathfrak{X}|~{}M\in\mathbb{R}^{n\times n},M^{T}=-M\\}$ be the set of
vector fields constructed by multiplying the gradient by an anti-symmetric
matrix. Recall that $R=-\frac{\partial\mathcal{L}}{\partial
w_{i}}\frac{\partial}{\partial w^{i}}$ is the reverse gradient vector field,
and $\mathfrak{X}_{\perp}=\\{a^{i}\frac{\partial}{\partial
w^{i}}|~{}\sum_{i}a^{i}({\bm{w}})\frac{\partial\mathcal{L}({\bm{w}})}{\partial
w^{i}}=0,\forall{\bm{w}}\in\mathcal{M}\\}$ is the set of all vector fields
orthogonal to $R$. From Lemma D.2, we have $S=\mathfrak{X}_{\perp}$.
Therefore, a point ${\bm{w}}$ is an optimal point in $S$ if and only if
${\bm{w}}$ is an optimal point in $\mathfrak{X}_{\perp}$.

We are now ready to prove the following proposition, which provides another
way to check the condition in Proposition 3.6.

######  Proposition 3.7.

If at all optimal points in $S$,

$\displaystyle M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}=0$ (88)

for all anti-symmetric matrix $M\in\mathbb{R}^{n\times n}$, then the gradient
flow starting at an optimal point in $S$ is optimal in $S$.

###### Proof.

Expanding $R[A,R]\mathcal{L}$, we have

$\displaystyle R[A,R]\mathcal{L}$
$\displaystyle=R\left(A\left(r^{i}\frac{\partial\mathcal{L}}{\partial
w^{i}}\right)-0\right)$ (89) $\displaystyle=r^{k}\frac{\partial}{\partial
w^{k}}\left(a^{j}\frac{\partial}{\partial
w_{j}}\left(r^{i}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)\right)$
(90) $\displaystyle=r^{k}\frac{\partial}{\partial
w^{k}}\left(a^{j}\left(\frac{\partial r^{i}}{\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}+r^{i}\frac{\partial}{\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)\right)$ (91)
$\displaystyle=-r^{k}\frac{\partial}{\partial
w^{k}}\left(a^{j}\left(\left(\frac{\partial}{\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w_{i}}\right)\frac{\partial\mathcal{L}}{\partial
w^{i}}+\frac{\partial\mathcal{L}}{\partial w_{i}}\frac{\partial}{\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)\right)$ (92)
$\displaystyle=-2r^{k}\frac{\partial}{\partial
w^{k}}\left(a^{j}\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)$ (93)
$\displaystyle=-2r^{k}\left(\frac{\partial a^{j}}{\partial
w^{k}}\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}+a^{j}\frac{\partial}{\partial
w^{k}}\left(\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)\right)$ (94)
$\displaystyle=2\frac{\partial\mathcal{L}}{\partial w_{k}}\frac{\partial
a^{j}}{\partial w^{k}}\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}+2\frac{\partial\mathcal{L}}{\partial
w_{k}}a^{j}\frac{\partial}{\partial
w^{k}}\left(\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)$ (95)

Assume that ${\bm{w}}$ is an optimal point in $S$. By Lemma D.2, ${\bm{w}}$ is
also an optimal point in $\mathfrak{X}_{\perp}$. By Lemma C.4 in [4],
$\frac{\partial\mathcal{L}}{\partial{\bm{w}}}$ is an eigenvector of
$\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial w^{j}}$. Therefore,
$\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}=\lambda\frac{\partial\mathcal{L}}{\partial w^{j}}$ for some
$\lambda\in\mathbb{C}$. Additionally,
$a^{j}=M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial w_{\alpha}}$ and
$\frac{\partial a^{j}}{\partial
w^{k}}=M_{\alpha}^{j}\frac{\partial^{2}\mathcal{L}}{\partial
w_{\alpha}\partial w^{k}}$. We are now ready to simplify both terms in (95).

For the first term in (95),

$\displaystyle\frac{\partial\mathcal{L}}{\partial w_{k}}\frac{\partial
a^{j}}{\partial w^{k}}\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}$
$\displaystyle=\frac{\partial\mathcal{L}}{\partial
w_{k}}M_{\alpha}^{j}\frac{\partial^{2}\mathcal{L}}{\partial w_{\alpha}\partial
w^{k}}\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}$ (96)
$\displaystyle=M_{\alpha}^{j}\left(\frac{\partial^{2}\mathcal{L}}{\partial
w_{\alpha}\partial w^{k}}\frac{\partial\mathcal{L}}{\partial
w_{k}}\right)\left(\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)$ (97)
$\displaystyle=M_{\alpha}^{j}\left(\lambda_{1}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\right)\left(\lambda_{2}\frac{\partial\mathcal{L}}{\partial
w^{j}}\right)$ (98)
$\displaystyle=\lambda_{1}\lambda_{2}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial\mathcal{L}}{\partial w^{j}}$ (99) $\displaystyle=0$
(100)

The last equality holds because $M$ is anti-symmetric.

For the second term in (95),

$\displaystyle\frac{\partial\mathcal{L}}{\partial
w_{k}}a^{j}\frac{\partial}{\partial
w^{k}}\left(\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}\right)$
$\displaystyle=\frac{\partial\mathcal{L}}{\partial
w_{k}}a^{j}\left(\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial
w_{i}\partial w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}+\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial^{2}\mathcal{L}}{\partial w^{k}\partial w^{i}}\right)$
(101) $\displaystyle=\frac{\partial\mathcal{L}}{\partial
w_{k}}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\left(\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial
w_{i}\partial w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}+\frac{\partial^{2}\mathcal{L}}{\partial w_{i}\partial
w^{j}}\frac{\partial^{2}\mathcal{L}}{\partial w^{k}\partial w^{i}}\right)$
(102) $\displaystyle=M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}+\lambda_{1}\lambda_{2}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial\mathcal{L}}{\partial w^{j}}$ (103)
$\displaystyle=M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}$ (104)

In summary,

$\displaystyle
R[A,R]\mathcal{L}=2M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}.$ (105)

Since we assumed that $[A,R]\mathcal{L}({\bm{w}})=0$, when
$R[A,R]\mathcal{L}({\bm{w}})=0$ for all $A\in S$, the gradient flow starting
at an optimal point in $S$ is optimal in $S$. ∎

######  Proposition D.3.

If $\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w^{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{\alpha}}=\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial
w^{i}\partial w^{\alpha}}\frac{\partial\mathcal{L}}{\partial w^{j}}$ holds for
all $i,k,j,\alpha$, then $M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}=0$ holds for all anti-
symmetric matrices $M\in\mathbb{R}^{n\times n}$.

###### Proof.

If $\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w^{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{\alpha}}=\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial
w^{i}\partial w^{\alpha}}\frac{\partial\mathcal{L}}{\partial w^{j}}$ for all
$i,k,j,\alpha$, then

$\displaystyle\quad\quad M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}$
$\displaystyle=\sum_{i,k,\alpha<j}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}+\sum_{i,k,\alpha>j}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial w^{i}}$
$\displaystyle=\sum_{i,k,\alpha<j}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}+\sum_{i,k,j>\alpha}M_{j}^{\alpha}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{j}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{\alpha}}\frac{\partial\mathcal{L}}{\partial w^{i}}$
$\displaystyle=\sum_{i,k,\alpha<j}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}\frac{\partial\mathcal{L}}{\partial
w^{i}}+\sum_{i,k,j>\alpha}-M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w_{j}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{\alpha}}\frac{\partial\mathcal{L}}{\partial w^{i}}$
$\displaystyle=\sum_{i,k,\alpha<j}M_{\alpha}^{j}\frac{\partial\mathcal{L}}{\partial
w_{k}}\frac{\partial\mathcal{L}}{\partial
w^{i}}\left(\frac{\partial\mathcal{L}}{\partial
w_{\alpha}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{j}}-\frac{\partial\mathcal{L}}{\partial
w_{j}}\frac{\partial^{3}\mathcal{L}}{\partial w^{k}\partial w_{i}\partial
w^{\alpha}}\right)$ $\displaystyle=0,$

where the first equality uses that the diagonal of an anti-symmetric matrix is
0, the second equality swaps $\alpha$ and $j$ in the second term, the third
equality uses that $M$ is anti-symmetric. ∎

##### Example (Quadratic function)

Consider the quadratic function
$\mathcal{L}({\bm{w}})=\frac{1}{2}{\bm{w}}^{T}A{\bm{w}}+\mathbf{b}^{T}{\bm{w}}+\mathbf{c}$,
where $A\in\mathbb{R}^{n\times n}$ is symmetric,
$\mathbf{b},\mathbf{c}\in\mathbb{R}^{n}$, and ${\bm{w}}\in\mathbb{R}^{n}$. Two
examples of quadratic functions are the ellipse
$\mathcal{L}_{e}(w_{1},w_{2})=\frac{1}{2}(w_{1}^{2}+\lambda^{2}w_{2}^{2})$ and
the Booth function
$\mathcal{L}_{b}(w_{1},w_{2})=(w_{1}+2w_{2}-7)^{2}+(2w_{1}+w_{2}-5)^{2}$.
Since the third derivative of $\mathcal{L}$ is 0, one teleportation guarantees
optimal trajectory.

##  Appendix E Group actions and curves on minima

###  E.1 Group actions for MLP

Consider a multi-layer neural network with elementwise activation function
$\sigma$. The output of the $m^{th}$ layer is $h_{m}=\sigma(W_{m}h_{m-1})$,
where $W_{m}\in\mathbb{R}^{d_{m}\times d_{m-1}}$ is the weight,
$h_{m-1}\in\mathbb{R}^{d_{m-1}\times k}$ is the output of the $m-1^{th}$
layer, and $h_{0}\in\mathbb{R}^{d_{0}\times k}$ is the data. There are at
least two ways to define a $\mathrm{GL}_{d_{m-1}}(\mathbb{R})$ symmetry acting
on $W_{m}$ and $W_{m-1}$. Unless stated otherwise, we use the second group
action since it does not require $\sigma$ to be invertible. We use
pseudoinverses in experiments.

##### Group action 1

[4]. Assume that $h_{m-2}$ is invertible and $\sigma$ is bijective. For
$g_{m}\in\mathrm{GL}_{d_{m-1}}(\mathbb{R})$,

$\displaystyle g_{m}\cdot
W_{k}=\left\\{\begin{array}[]{lc}W_{m}g_{m}^{-1}&k=m\\\
\sigma^{-1}\left(g_{m}\sigma\left(W_{m-1}h_{m-2}\right)\right)h_{m-2}^{-1}&k=m-1\\\
W_{k}&k\not\in\\{m,m-1\\}\end{array}\right.$ (109)

##### Group action 2

[16]. Assume that $g_{m}\sigma\left(W_{m-1}h_{m-2}\right)$ is invertible. For
$g_{m}\in\mathrm{GL}_{d_{m-1}}(\mathbb{R})$,

$\displaystyle g_{m}\cdot
W_{k}=\left\\{\begin{array}[]{lc}W_{m}\sigma\left(W_{m-1}h_{m-2}\right)\sigma\left(g_{m}W_{m-1}h_{m-2}\right)^{-1}&k=m\\\
g_{m}W_{m-1}&k=m-1\\\ W_{k}&k\not\in\\{m,m-1\\}\end{array}\right.$ (113)

###  E.2 Curvature

The curvature of a curve $\gamma:\mathbb{R}\xrightarrow[]{}\mathbb{R}^{n}$ is
$\kappa(t)=\frac{\|T^{\prime}(t)\|}{\|\gamma^{\prime}(t)\|}$, where
$T(t)=\frac{\gamma^{\prime}(t)}{\|\gamma^{\prime}(t)\|}$ is the unit tangent
vector. The curvature can be written as a function of $\gamma^{\prime}$ and
$\gamma^{\prime\prime}$ [46, 47]:

$\displaystyle\kappa(t)=\frac{\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}}{\|\gamma^{\prime}\|^{3}}.$
(114)

###  E.3 The derivative of curvature

To compute the derivative of $\kappa(t)$, we first list the derivatives of a
few commonly used terms:

$\displaystyle\frac{d}{dt}\|\gamma^{\prime}\|^{2}$
$\displaystyle=\frac{d}{dt}({\gamma^{\prime}_{1}}^{2}+{\gamma^{\prime}_{2}}^{2}+{\gamma^{\prime}_{3}}^{2}+...)=2\gamma^{\prime}_{1}\gamma^{\prime\prime}_{1}+2\gamma^{\prime}_{2}\gamma^{\prime\prime}_{2}+2\gamma^{\prime}_{3}\gamma^{\prime\prime}_{3}+...=2\gamma^{\prime}\cdot\gamma^{\prime\prime}$
(115) $\displaystyle\frac{d}{dt}\|\gamma^{\prime\prime}\|^{2}$
$\displaystyle=\frac{d}{dt}({\gamma^{\prime\prime}_{1}}^{2}+{\gamma^{\prime\prime}_{2}}^{2}+{\gamma^{\prime\prime}_{3}}^{2}+...)=2\gamma^{\prime\prime}_{1}\gamma^{\prime\prime\prime}_{1}+2\gamma^{\prime\prime}_{2}\gamma^{\prime\prime\prime}_{2}+2\gamma^{\prime\prime}_{3}\gamma^{\prime\prime\prime}_{3}+...=2\gamma^{\prime\prime}\cdot\gamma^{\prime\prime\prime}$
(116) $\displaystyle\frac{d}{dt}(\gamma^{\prime}\cdot\gamma^{\prime\prime})$
$\displaystyle=\frac{d}{dt}(\gamma^{\prime}_{1}\gamma^{\prime\prime}_{1}+\gamma^{\prime}_{2}\gamma^{\prime\prime}_{2}+\gamma^{\prime}_{3}\gamma^{\prime\prime}_{3}...)=\gamma^{\prime}_{1}\gamma^{\prime\prime\prime}_{1}+\gamma^{\prime\prime}_{1}\gamma^{\prime\prime}_{1}+...=\|\gamma^{\prime\prime}\|^{2}+\gamma^{\prime}\cdot\gamma^{\prime\prime\prime}$
(117)

The derivatives of the numerator and denominator of $\kappa$ are:

$\displaystyle\frac{d}{dt}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}$
$\displaystyle=\frac{1}{2}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}\frac{d}{dt}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)$
(118)
$\displaystyle=\frac{1}{2}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}$
(119)
$\displaystyle\quad\quad\quad\left(\|\gamma^{\prime}\|^{2}\frac{d}{dt}\|\gamma^{\prime\prime}\|^{2}+\|\gamma^{\prime\prime}\|^{2}\frac{d}{dt}\|\gamma^{\prime}\|^{2}-2(\gamma^{\prime}\cdot\gamma^{\prime\prime})\frac{d}{dt}(\gamma^{\prime}\cdot\gamma^{\prime\prime})\right)$
(120)
$\displaystyle=\frac{1}{2}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}$
(121)
$\displaystyle\quad\quad\quad\left(2\|\gamma^{\prime}\|^{2}(\gamma^{\prime\prime}\cdot\gamma^{\prime\prime\prime})+2\|\gamma^{\prime\prime}\|^{2}(\gamma^{\prime}\cdot\gamma^{\prime\prime})-2(\gamma^{\prime}\cdot\gamma^{\prime\prime})(\|\gamma^{\prime\prime}\|^{2}+\gamma^{\prime}\cdot\gamma^{\prime\prime\prime})\right)$
(122)
$\displaystyle=\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}\left(\|\gamma^{\prime}\|^{2}(\gamma^{\prime\prime}\cdot\gamma^{\prime\prime\prime})-(\gamma^{\prime}\cdot\gamma^{\prime\prime})(\gamma^{\prime}\cdot\gamma^{\prime\prime\prime})\right),$
(123)

and

$\displaystyle\frac{d}{dt}\|\gamma^{\prime}\|^{3}=\frac{d}{dt}(\|\gamma^{\prime}\|^{2})^{\frac{3}{2}}=\frac{3}{2}(\|\gamma^{\prime}\|^{2})^{\frac{1}{2}}\frac{d}{dt}\|\gamma^{\prime}\|^{2}=\frac{3}{2}(\|\gamma^{\prime}\|^{2})^{\frac{1}{2}}(2\gamma^{\prime}\cdot\gamma^{\prime\prime})=3\|\gamma^{\prime}\|(\gamma^{\prime}\cdot\gamma^{\prime\prime}).$
(124)

Using the derivatives above, the derivative of $\kappa$ is

$\displaystyle\kappa^{\prime}(t)$
$\displaystyle=\frac{\left[\frac{d}{dt}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}\right]\|\gamma^{\prime}\|^{3}-\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}\left[\frac{d}{dt}\|\gamma^{\prime}\|^{3}\right]}{\|\gamma^{\prime}\|^{6}}$
(125)
$\displaystyle=\frac{\begin{multlined}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}\left(\|\gamma^{\prime}\|^{2}(\gamma^{\prime\prime}\cdot\gamma^{\prime\prime\prime})-(\gamma^{\prime}\cdot\gamma^{\prime\prime})(\gamma^{\prime}\cdot\gamma^{\prime\prime\prime})\right)\|\gamma^{\prime}\|^{3}\\\
-\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}3\|\gamma^{\prime}\|(\gamma^{\prime}\cdot\gamma^{\prime\prime})\end{multlined}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}\left(\|\gamma^{\prime}\|^{2}(\gamma^{\prime\prime}\cdot\gamma^{\prime\prime\prime})-(\gamma^{\prime}\cdot\gamma^{\prime\prime})(\gamma^{\prime}\cdot\gamma^{\prime\prime\prime})\right)\|\gamma^{\prime}\|^{3}\\\
-\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}3\|\gamma^{\prime}\|(\gamma^{\prime}\cdot\gamma^{\prime\prime})}{\|\gamma^{\prime}\|^{6}}$
(128)
$\displaystyle=\frac{\begin{multlined}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}\left(\|\gamma^{\prime}\|^{2}(\gamma^{\prime\prime}\cdot\gamma^{\prime\prime\prime})-(\gamma^{\prime}\cdot\gamma^{\prime\prime})(\gamma^{\prime}\cdot\gamma^{\prime\prime\prime})\right)\|\gamma^{\prime}\|^{2}\\\
-\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}3(\gamma^{\prime}\cdot\gamma^{\prime\prime})\end{multlined}\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{-\frac{1}{2}}\left(\|\gamma^{\prime}\|^{2}(\gamma^{\prime\prime}\cdot\gamma^{\prime\prime\prime})-(\gamma^{\prime}\cdot\gamma^{\prime\prime})(\gamma^{\prime}\cdot\gamma^{\prime\prime\prime})\right)\|\gamma^{\prime}\|^{2}\\\
-\left(\|\gamma^{\prime}\|^{2}\|\gamma^{\prime\prime}\|^{2}-(\gamma^{\prime}\cdot\gamma^{\prime\prime})^{2}\right)^{\frac{1}{2}}3(\gamma^{\prime}\cdot\gamma^{\prime\prime})}{\|\gamma^{\prime}\|^{5}}.$
(131)

###  E.4 The derivatives of curves on minima

Consider the curve
$\gamma_{M}:\mathbb{R}\times\mathbb{R}^{n}\xrightarrow[]{}\mathbb{R}^{n}$
where $M\in\operatorname{Lie}(G)$ and

$\displaystyle\gamma_{M}(t,{\bm{w}})=\exp{(tM)}\cdot{\bm{w}}.$ (133)

In this section, we derive $\gamma^{\prime}$, $\gamma^{\prime\prime}$, and
$\gamma^{\prime\prime\prime}$, which are needed to compute the curvature
$\kappa(t)$ and its derivative $\kappa^{\prime}(t)$. We are interested in
$\kappa$ and $\kappa^{\prime}$ at ${\bm{w}}$, or equivalently, at $t=0$. To
find the derivatives of $\gamma$ at $t=0$, we write the group action in the
following form:

$\displaystyle\gamma(t)=\sum_{n=0}^{\infty}\frac{f(n)}{n!}t^{n}.$ (134)

By the uniqueness of Taylor polynomial, the derivatives are
$\gamma^{(n)}(0)=f(n)$.

Consider two consecutive layers $U\sigma(VX)$ in a neural network, where
$U\in\mathbb{R}^{m\times h},V\in\mathbb{R}^{h\times n}$ are weights,
$X\in\mathbb{R}^{h\times k}$ is the output from the previous layer, and
$\sigma$ is an elementwise activation function. Choosing
$G=GL_{h}(\mathbb{R})$, one group action that leaves the output of these two
layers unchanged is:

$\displaystyle g\cdot(U,V,X)=(g\cdot U,g\cdot V,g\cdot
X)=(Ug^{-1},\sigma^{-1}(g\sigma(VX))X^{-1},X).$ (135)

Let

$\displaystyle g=\exp(tM)=\sum_{k=0}^{\infty}\frac{1}{k!}(tM)^{k},$ (136)

where $M\in\operatorname{Lie}(G)$ is in the Lie algebra of $G$. The action of
$g$ yields

$\displaystyle g\cdot(U,V,X)$
$\displaystyle=(U\exp(-tM),\sigma^{-1}(\exp(tM)\sigma(VX))X^{-1},X).$ (137)

Next, we expand $\gamma(t)=g\cdot(U,V)$. The Taylor expansion for $g\cdot U$
is

$\displaystyle U\exp(-tM)$
$\displaystyle=U\sum_{k=0}^{\infty}\frac{1}{k!}(-tM)^{k}$ (138)
$\displaystyle=U-tUM+\frac{t^{2}}{2!}UM^{2}-\frac{t^{3}}{3!}UM^{3}+O(t^{4}).$
(139)

The Taylor expansion for $g\cdot V$ is

$\displaystyle\sigma^{-1}(\exp(tM)\sigma(VX))X^{-1}$ (140) $\displaystyle=$
$\displaystyle\sigma^{-1}\left(\left(\sum_{k=0}^{\infty}\frac{1}{k!}(tM)^{k}\right)\sigma(VX)\right)X^{-1}$
(141) $\displaystyle=$
$\displaystyle\sigma^{-1}\left(\sigma(VX)+\sum_{k=1}^{\infty}\frac{1}{k!}(tM)^{k}\sigma(VX)\right)X^{-1}$
(142) $\displaystyle=$
$\displaystyle\left[\sigma^{-1}(\sigma(VX))+\sum_{j=1}^{\infty}\left(\sum_{k=1}^{\infty}\frac{1}{k!}(tM)^{k}\sigma(VX)\right)^{\odot
j}\odot\left.\frac{\partial^{j}\sigma^{-1}(A)}{\partial
A^{j}}\right|_{A=\sigma(VX)}\right]X^{-1}$ (143) $\displaystyle=$
$\displaystyle
V+\left[\sum_{j=1}^{\infty}\left(\sum_{k=1}^{\infty}\frac{1}{k!}(tM)^{k}\sigma(VX)\right)^{\odot
j}\odot\left.\frac{\partial^{j}\sigma^{-1}(A)}{\partial
A^{j}}\right|_{A=\sigma(VX)}\right]X^{-1},$ (144)

where $\odot$ denotes element-wise product: $(A\odot B)_{mn}=A_{mn}B_{mn}$,
and the superscript ⊙ denotes elementwise power: $(A^{\odot
j})_{mn}=(A_{mn})^{j}$. The Taylor expansion is of each element individually,
because $\sigma$ is element-wise.

Since our goal is to find the first 3 derivatives of $\gamma$, we are only
interested in the terms up to $t^{3}$. Letting

$\displaystyle\sum_{k=1}^{\infty}\frac{1}{k!}(tM)^{k}=tM+t^{2}\frac{M^{2}}{2}+t^{3}\frac{M^{3}}{6}+O(t^{4})$
(145)

and considering only the $j=1,2,3$ terms, we have

$\displaystyle\sigma^{-1}(\exp(tM)\sigma(VX))X^{-1}$ (146) $\displaystyle=$
$\displaystyle
V+\left[\sum_{j=1}^{\infty}\left((tM+t^{2}\frac{M^{2}}{2}+t^{3}\frac{M^{3}}{6})\sigma(VX)\right)^{\odot
j}\odot\left.\frac{\partial^{j}\sigma^{-1}(A)}{\partial
A^{j}}\right|_{A=\sigma(VX)}\right]X^{-1}+O(t^{4})$ (147) $\displaystyle=$
$\displaystyle
V+\left[\left((tM+t^{2}\frac{M^{2}}{2}+t^{3}\frac{M^{3}}{6})\sigma(VX)\right)\odot\left.\frac{\partial\sigma^{-1}(A)}{\partial
A}\right|_{A=\sigma(VX)}\right.$ (148)
$\displaystyle\quad\quad\quad+\left((tM+t^{2}\frac{M^{2}}{2}+t^{3}\frac{M^{3}}{6})\sigma(VX)\right)^{\odot
2}\odot\left.\frac{\partial^{2}\sigma^{-1}(A)}{\partial
A^{2}}\right|_{A=\sigma(VX)}$ (149)
$\displaystyle\quad\quad\quad+\left.\left((tM+t^{2}\frac{M^{2}}{2}+t^{3}\frac{M^{3}}{6})\sigma(VX)\right)^{\odot
3}\odot\left.\frac{\partial^{3}\sigma^{-1}(A)}{\partial
A^{3}}\right|_{A=\sigma(VX)}\right]X^{-1}+O(t^{4})$ (150) $\displaystyle=$
$\displaystyle
V+t\left(\left(M\sigma(VX)\right)\odot\frac{1}{\sigma^{\prime}(VX)}\right)X^{-1}$
(151)
$\displaystyle\quad\quad+\frac{t^{2}}{2}\left(\left(M^{2}\sigma(VX)\right)\odot\frac{1}{\sigma^{\prime}(VX)}-2(M\sigma(VX))^{\odot
2}\odot\frac{\sigma^{\prime\prime}(VX)}{\sigma^{\prime}(VX)^{3}}\right)X^{-1}$
(152)
$\displaystyle\quad\quad+\frac{t^{3}}{6}\left(\left(M^{3}\sigma(VX)\right)\odot\frac{1}{\sigma^{\prime}(VX)}-6(M\sigma(VX))\odot(M^{2}\sigma(VX))\odot\frac{\sigma^{\prime\prime}(VX)}{\sigma^{\prime}(VX)^{3}}\right.$
(153) $\displaystyle\quad\quad\quad\quad\quad\quad\left.+6(M\sigma(VX))^{\odot
3}\odot\left.\frac{\partial^{3}\sigma^{-1}(A)}{\partial
A^{3}}\right|_{A=\sigma(VX)}\right)X^{-1}$ (154)
$\displaystyle\quad\quad+O(t^{4}).$ (155)

Matching terms in (138) and (146) with (134), we have the expressions for
$\gamma^{\prime}$, $\gamma^{\prime\prime}$, and $\gamma^{\prime\prime\prime}$.
This allows us to compute the curvature and its derivative using (114) and
(125).

##  Appendix F Sharpness, Curvatures, and Their Relation to Generalization

###  F.1 Alternative definitions of sharpness

A common definition of flat minimum is based on the number of eigenvalues of
the Hessian which are small. Minimizers with a large number of large
eigenvalues tend to have worse generalization ability [7]. Let
$\lambda_{i}(H)({\bm{w}})$ be the $i^{th}$ largest eigenvalue of the Hessian
of the loss function evaluated at ${\bm{w}}$. We can quantify the notion of
sharpness by the number of eigenvalues larger than a threshold
${\varepsilon}\in\mathbb{R}^{>0}$:

$\displaystyle\phi_{1}({\bm{w}},{\varepsilon})=\left|\\{\lambda_{i}(H)({\bm{w}}):\lambda_{i}>{\varepsilon}\\}\right|.$
(156)

A related sharpness metric uses the logarithm of the product of the $k$
largest eigenvalues [33],

$\displaystyle\phi_{2}({\bm{w}},k)=\sum_{i=1}^{k}\log\lambda_{i}(H)({\bm{w}}).$
(157)

Note that both metrics require computing the eigenvalues of the Hessian.
Optimizing on these metrics during teleportation is prohibitively expensive.
Hence, in this paper we use the average change in loss averaged over random
directions ($\phi$) as objective in generalization experiments.

###  F.2 More intuition on curvatures and generalization

####  F.2.1 Example: curvature affects average displacement of minima

Consider an optimization problem with two variables
$w_{1},w_{2}\in\mathbb{R}$. Assume that the minimum is a one-dimensional curve
$\gamma:\mathbb{R}\to\mathbb{R}^{2}$ in the two-dimensional parameter space.
For a point ${\bm{w}}_{0}$ on $\gamma$, we estimate its generalization ability
by computing the expected distance between ${\bm{w}}_{0}$ and the new minimum
obtained by shifting $\gamma$.

We consider the following two curves:

$\displaystyle\gamma_{1}:$
$\displaystyle\mathbb{R}\to\mathbb{R}^{2},t\mapsto(t,kt^{2})$ (158)
$\displaystyle\gamma_{2}:$
$\displaystyle[0,2\pi]\to\mathbb{R}^{2},\theta\mapsto(k\cos(\theta),k\sin(\theta).$
(159)

The curve $\gamma_{1}$ is a parabola, and the curvature at
${\bm{w}}_{0}=(0,0)$ is $\kappa_{1}=2k$. The curve $\gamma_{2}$ is a circle,
and the curvature at ${\bm{w}}_{0}=(0,0)$ is $\kappa_{2}=\frac{1}{k}$. Note
that $\gamma_{1}$ is the only polynomial approximation with integer power
($\gamma(t)=(t,k|t|^{n})$) where the curvature at ${\bm{w}}_{0}$ depends on
$k$. When $n<1$, the value of ${\bm{w}}_{0}$ is undefined. When $n=1$, the
first derivative at ${\bm{w}}_{0}$ is undefined. When $n>2$,
$\kappa({\bm{w}}_{0})=0$.

Assume that a distribution shift in data causes $\gamma$ to shift by a
distance $r$, and that the direction of the shift is chosen uniformly at
random over all possible directions. Viewing from the perspective of the
curve, this is equivalent to shifting ${\bm{w}}_{0}$ by distance $r$.

The distance between a point ${\bm{w}}$ and a curve $\gamma$ is

$\displaystyle
dist({\bm{w}},\gamma)=\min_{{\bm{w}}^{\prime}\in\gamma_{2}}\|{\bm{w}}^{\prime}-{\bm{w}}\|_{2}.$
(160)

Let $S_{r}$ be the circle centered at the origin with radius $r$. Figure
8(b)(c) shows that the expected distance’s dependence on $\kappa$. Using both
curves $\gamma_{1}$ and $\gamma_{2}$, the generalization ability of
${\bm{w}}_{0}$ depends on the curvature at ${\bm{w}}_{0}$. However, the type
of dependence is affected by the type of curve used. In other words, the
curvatures at points around ${\bm{w}}_{0}$ affect how the curvature at
${\bm{w}}_{0}$ affects generalization. Therefore, from these results alone,
one cannot deduce whether minima with sharper curvatures generalize better or
worse. To find a more definitive relationship between curvature and
generalization, further investigation on the type of curves on the minimum is
required.

We emphasize that this example only serves as an intuition for connecting
curvature to generalization. As a future direction, it would be interesting to
consider different families of parametric curves, higher dimensional parameter
spaces, and deforming in addition to shifting the minima.

(a) (b) (c)

Figure 8: (a) Illustration of the parameter space, the minimum ($\gamma$), and
all shifts with distance $r$ ($S_{r}$). (b) Expected distance between
${\bm{w}}_{0}$ and the new minimum as a function of $\kappa$, for quadratic
approximation $\gamma_{1}$. (c) Expected distance between ${\bm{w}}_{0}$ and
the new minimum as a function of $\kappa$, for constant curvature
approximation $\gamma_{2}$. The expected distance is scaled by $r^{-2}$
because the arc length of $S_{r}$ is proportional to $r$, and the average
distance at each point on $S_{r}$ is also roughly proportional to $r$.

####  F.2.2 Higher dimensions

Figure 9 visualizes a curve obtained from a 2D minima. However, it is not
immediately clear what curves look like on a higher-dimensional minimum. A
possible way to extend previous analysis is to consider sectional curvatures.

Figure 9: Left: a 2D minima in a 3D parameter space. Right: a 2D subspace of
the parameter space and a curve on the minima (the intersection of the minima
and the subspace).

###  F.3 Computing correlation to generalization

We generate the 100 different models used in Section 4.3 by training randomly
initialized models. For all three datasets (MNIST, FashionMNIST, and
CIFAR-10), we train on 50,000 samples and test on a different set of 10,000
samples. The labels for classification tasks belongs to 1 of 10 classes.

For a batch of flattened input data $X\in\mathbb{R}^{d\times 20}$ and labels
$Y\in\mathbb{R}^{20}$, the loss function is
$\mathcal{L}(W_{1},W_{2},W_{3},X,Y)=\text{CrossEntropy}\left(W_{3}\sigma(W_{2}\sigma(W_{1}X)),Y\right)$,
where $W_{3}\in\mathbb{R}^{10\times h_{2}}$, $W_{2}\in\mathbb{R}^{h_{2}\times
h_{1}}$, $W_{1}\in\mathbb{R}^{h_{1}\times d}$ are the weight matrices, and
$\sigma$ is the LeakyReLU activation with slope coefficient 0.1. For MNIST and
Fashion-MNIST, $d=28^{2}$, $h_{1}=16$, and $h_{2}=10$. For CIFAR-10,
$d=32^{3}\times 3$, $h_{1}=128$, and $h_{2}=32$. The learning rate for
stochastic gradient descent is $0.01$ for MNIST and Fashion-MNIST, and $0.02$
for CIFAR-10. We train each model using mini-batches of size 20 for 40 epochs.

When computing the sharpness $\phi$, we choose the displacement list T that
gives the highest correlation. The displacements used in this paper are
$T={0.001,0.011,0.021,...,0.191}$ for MNIST, and
$T={0.001,0.011,0.021,...,0.191}$ for Fashion-MNIST and CIFAR-10. We evaluate
the change in loss over $|D|=200$ random directions. For curvature $\psi$, we
average over $k=1$ curves generated by random Lie algebras (invertible
matrices in this case).

Figure 10 and 11 visualizes the correlation result in Table 1. Each point
represents one model.

(a) (b) (c)

Figure 10: Correlation between sharpness and validation loss on MNIST (left),
Fashion-MNIST (middle), and CIFAR-10 (right). Sharpness and generalization are
strongly correlated.

(a) (b) (c)

Figure 11: Correlation between curvature and validation loss on MNIST (left),
Fashion-MNIST (middle), and CIFAR-10 (right). There is a weak negative
correlation in all three datasets.

###  F.4 Additional details for generalization experiments

On CIFAR-10, we run SGD using the same three-layer architecture as in Section
F.3, but with a smaller hidden size $h_{1}=32$ and $h_{2}=10$. At epoch 20
which is close to convergence, we teleport using 5 batches of data, each of
size 2000. During each teleportation for $\phi$, we perform 10 gradient ascent
(or descent) steps on the group element. During each teleportation for $\psi$,
we perform 1 gradient ascent (or descent) step on the group element. The
learning rate for the optimization on group elements is $5\times 10^{-2}$.

##  Appendix G Integrating teleportation with other gradient-based algorithms

###  G.1 Different methods of integrating teleportation with momentum and
AdaGrad

##### Setup.

We test teleportation with various algorithms using the a 3-layer neural
network and mean square error:
$\min_{W_{1},W_{2},W_{3}}\|Y-W_{3}\sigma(W_{2}\sigma(W_{1}X))\|_{2}$, with
data $X\in\mathbb{R}^{5\times 4}$, target $Y\in\mathbb{R}^{8\times 4}$, and
weight matrices $W_{3}\in\mathbb{R}^{8\times 7}$, $W_{2}\in\mathbb{R}^{7\times
6}$, and $W_{1}\in\mathbb{R}^{6\times 5}$. The activation function $\sigma$ is
LeakyReLU with slope coefficient 0.1. Each element in the weight matrices is
initialized uniformly at random over $[0,1]$. Data $X,Y$ are randomly
generated also from $[0,1]$.

##### Momentum.

We compare three strategies of integrating teleportation with momentum:
teleporting both parameters and momentum, teleporting parameters but not
momentum, and reset momentum to 0 after a teleportation. In each run, we
teleport once at epoch 5. Each strategy is repeated 5 times.

The training curves of teleporting momentum in different ways are similar
(Figure 12a), possibly because the momentum accumulated is small compared to
the gradient right after teleportations. All methods of teleporting momentum
improves convergence, which means teleportation works well with momentum.

##### AdaGrad.

In AdaGrad, the rate of change in loss is

$\displaystyle{d\mathcal{L}({\bm{w}})\over
dt}={\partial\mathcal{L}\over\partial{\bm{w}}}^{T}{d{\bm{w}}\over
dt}=-\eta\|\nabla\mathcal{L}\|_{A},$ (161)

where $\eta\in\mathbb{R}$ is the learning rate, and
$\|\nabla\mathcal{L}\|_{A}$ is the Mahalanobis norm with
$A=({\varepsilon}I+diag(G_{t+1}))^{-\frac{1}{2}}$. Previously, we optimize
$\|\nabla\mathcal{L}\|_{2}$ in teleportation. We compare that to optimizing
$\|\nabla\mathcal{L}\|_{A}$. Since the magnitude of $A$ is different than 1, a
different learning rate for the gradient ascent in teleportation is required.
We choose the largest learning rate (with two significant figures) that does
not lead to divergence. The teleportation learning rates used are $1.2\times
10^{-5}$ for objective $\max_{g}\|\nabla\mathcal{L}\|_{2}$ and $7.5\times
10^{-3}$ for objective $\max_{g}\|\nabla\mathcal{L}\|_{A}$.

Teleporting using the group element that optimizes $\|\nabla\mathcal{L}\|_{A}$
has a slight advantage (Figure 12b). Similar to the observations in [4],
teleportation can be integrated into adaptive gradient descents.

(a) (b)

Figure 12: Comparison of different method of integrating teleportation with
momentum and AdaGrad.

###  G.2 Additional details for experiments on MNIST

We use a three-layer model and cross-entropy loss for classification with
minibatches of size 20. For a batch of flattened input data
$X\in\mathbb{R}^{28^{2}\times 20}$ and labels $Y\in\mathbb{R}^{20}$, the loss
function is
$\mathcal{L}(W_{1},W_{2},W_{3},X,Y)=\text{CrossEntropy}\left(W_{3}\sigma(W_{2}\sigma(W_{1}X)),Y\right)$,
where $W_{3}\in\mathbb{R}^{10\times 10}$, $W_{2}\in\mathbb{R}^{10\times 16}$,
$W_{1}\in\mathbb{R}^{16\times 28^{2}}$ are the weight matrices, and $\sigma$
is the LeakyReLU activation with slope coefficient 0.1. The learning rates are
$10^{-4}$ for AdaGrad, and $5\times 10^{-2}$ for SGD with momentum, RMSProp,
and Adam. The learning rate for optimizing the group element in teleportation
is $5\times 10^{-2}$, and we perform 10 gradient ascent steps when teleporting
using each mini-batch. We use 50,000 samples from training set for training,
and 10,000 samples in the test set for testing.

(a) (b) (c) (d)

Figure 13: Runtime comparison for integrating teleportation into various
algorithms. Solid line represents average training loss, and dashed line
represents average test loss. Shaded areas are 1 standard deviation of the
test loss across 5 runs. The plots look almost identical to Figure 5,
indicating that the cost of teleportation is negligible compared to gradient
descents.