
\section*{References}



\section{Symmetry Teleportation}
We propose \textit{symmetry teleportation}, an accelerated gradient-based optimization algorithm that exploits symmetries in the loss landscape. Symmetry teleportation searches for the best gradient descent trajectory by teleporting parameters to a point with larger gradients using a group action. The resulting algorithm requires only  gradient computations but is able to account for the global landscape geometry, leading to faster loss decay. 

Let the group $G$ be a set of symmetries which leave the loss function $\L$ invariant: $\L(g \cdot({\bm{w}}, X)) = \L({\bm{w}}, X)$, where $g \in G$. We perform gradient descent for $T$ steps. We define an index set $K \subseteq \{0, 1, 2, ..., T-1\}$ as a teleportation schedule. At epochs that are in the schedule, we transform parameters using group element $g \in G$ to the location where the gradient is largest, then continue with gradient descent steps. Algorithm \ref{alg:teleport} describes the details of this procedure.  Note that the loss does not change after teleportation (Line 2-5) since $\L$ is $G$-invariant. 

\vspace{-3mm}
\begin{algorithm}
    \SetKwInput{Input}{Input}
    \SetKwInput{Output}{Output}

    \Input{Loss function $\L({\bm{w}})$, learning rate $\eta$, number of epochs $T$, initialized parameters ${\bm{w}}_0$, symmetry group $G$, teleportation schedule $K$.}
    \Output{${\bm{w}}_{T}$.}

    \For{$t \leftarrow 0$ \KwTo $T-1$}{
        \If{$t \in K$}{
            $g \leftarrow \text{argmax}_{g \in G} \| (\nabla \L)|_{g \cdot {\bm{w}}_t} \|^2$ \\
            ${\bm{w}}_t \leftarrow g \cdot {\bm{w}}_t$
        }
       
        ${\bm{w}}_{t+1} \leftarrow {\bm{w}}_t - \eta (\nabla \L)|_{{\bm{w}}_t}$
    }

    \KwRet{${\bm{w}}_{T}$}
    \caption{Symmetry Teleportation}
    \label{alg:teleport}
\end{algorithm}
\vspace{-3mm}
Algorithm \ref{alg:teleport} can be generalized to apply teleportation to stochastic gradient descent (Appendix \ref{appendix:algorithm-sgd}). We discuss some design choices below and provide detailed analysis in the next two sections. 

When the action of $G$ is continuous, teleportation can be implemented by parameterizing and performing gradient ascent on $g$. 
For example, the $\mathrm{SO}(2)$ group can be parameterized by the rotation angle $\theta$. 
Small transformations $g\in\mathrm{GL}_d(\R)$ (general linear group) can be parameterized as $g \approx I+{\varepsilon} T$ where ${\varepsilon}\ll 1$ and $T$ are arbitrary $d\times d$ matrices. 
For discrete groups, search algorithms or random sampling can be used to find a group element that improves the magnitude of the gradient.

Although $g \cdot ({\bm{w}}, X)$ does not change $X$ in the cases we consider in this paper, Algorithm \ref{alg:teleport} can be extended to allow transformations on both parameters and data. The group actions on data during teleportations can be precomposed and applied to the input data at inference time. The  algorithm extension that allows transformation on the input can be found in Appendix \ref{appendix:algorithm-data-transform}

The teleportation schedule $K$ is a hyperparameter that determines when to perform teleportation. We define $K$ as a set to allow flexible teleportation schedules, such as with non-fixed frequencies or teleporting only at the earlier epochs.





\section{Adaptations of Algorithm \ref{alg:teleport} for different problems}
\subsection{Stochastic gradient descent}
\label{appendix:algorithm-sgd}
We extend Algorithm \ref{alg:teleport} to stochastic gradient descent (SGD). We apply the group actions using data from a mini-batch $X_i$, and repeat for $B$ mini-batches each time. The gradient we optimize, $\Tilde{\nabla} \L(X_i, g \cdot {\bm{w}}_t)$, also uses single mini-batches. Algorithm \ref{alg:teleport-sgd} provides the framework for teleportation in SGD.
\begin{algorithm}
    \SetKwInput{Input}{Input}
    \SetKwInput{Output}{Output}

    \Input{Loss function $\L({\bm{w}})$, learning rate $\eta$, number of epochs $T$, initialized parameters ${\bm{w}}_0$, symmetry group $G$, teleportation schedule $K$, number of mini-batches used to teleport $B$.}
    \Output{${\bm{w}}_{T}$.}

    \For{$t \leftarrow 0$ \KwTo $T-1$}{
        \eIf{$t \in K$}{
            \For{$X_i$ in the first $B$ mini-batches}{
                $g \leftarrow \text{argmax}_{g \in G} \| \Tilde{\nabla} \L(X_i, g \cdot {\bm{w}}_t) \|^2$ \\
                ${\bm{w}}_t \leftarrow g \cdot {\bm{w}}_t$ \\
                ${\bm{w}}_{t} \leftarrow {\bm{w}}_t - \eta \Tilde{\nabla} \L(X_i, {\bm{w}}_t)$
            }
            \For{$X_i$ in the rest mini-batches}{
                ${\bm{w}}_{t} \leftarrow {\bm{w}}_t - \eta \Tilde{\nabla} \L(X_i, {\bm{w}}_t)$
            }
        }{
            \For{all mini-batches $X_i$}{
                ${\bm{w}}_{t} \leftarrow {\bm{w}}_t - \eta \Tilde{\nabla} \L(X_i, {\bm{w}}_t)$
            }
        }
       ${\bm{w}}_{t+1} \leftarrow {\bm{w}}_t$ \\
        
    }

    \KwRet{${\bm{w}}_{T}$}
    \caption{Symmetry Teleportation (SGD)}
    \label{alg:teleport-sgd}
\end{algorithm}

\subsection{Data transformation}
\label{appendix:algorithm-data-transform}
Algorithm \ref{alg:teleport-data} here modifies Algorithm \ref{alg:teleport} to allow transformations on both parameters and data. Denote $g_X$ as the group action on data only. The group actions on data at all teleportation steps can be precomposed as a function $f$ and applied to the input data at inference time.
\begin{algorithm}
    \SetKwInput{Input}{Input}
    \SetKwInput{Output}{Output}
    \SetKwInput{Initialize}{Initialize}

    \Input{Loss function $\L({\bm{w}}, X)$, learning rate $\eta$, number of epochs $T$, initialized parameters ${\bm{w}}_0$, symmetry $G$, teleportation schedule $K$, data $X$.}
    \Output{${\bm{w}}_{T}$, data transformation $f$.}
    \Initialize{$f = $ the identity function.}
    \For{$t \leftarrow 0$ \KwTo $T-1$}{
        \If{$t \in K$}{
            $g \leftarrow \text{argmax}_{g \in G} \| \nabla_{g \cdot {\bm{w}}_t} \L (g \cdot({\bm{w}}_t, X)) \|^2$ \\
            ${\bm{w}}_t, X \leftarrow g \cdot ({\bm{w}}_t, X)$ \\
            $f \leftarrow g_X \circ f$
        }
       
        ${\bm{w}}_{t+1} \leftarrow {\bm{w}}_t - \eta \nabla_{{\bm{w}}_t} \L$
    }

    \KwRet{${\bm{w}}_{T}$, $f$}
    \caption{Symmetry Teleportation (with data transformation)}
    \label{alg:teleport-data}
\end{algorithm}

\section{Group actions}
\subsection{Continuous symmetry in test functions}
\subsubsection{Ellipse}

Consider the following loss function with $a \in \R^{\geq 0}$:
\begin{align}
    \L(x_1, x_2) &= x_1^2 + ax_2^2
\end{align}

If we change the variables to $\L(u(x_1, x_2), v(x_1, x_2)) = u^2 + v^2$, 2D rotations leave $\L$ unchanged. Therefore $SO(2)$ is a symmetry of $\L(x_1,x_2)$. Let $g_\theta \in SO(2)$, and define the group action as 
\begin{align}
    g_\theta \cdot \begin{bmatrix}
        x_1 \\
        x_2
    \end{bmatrix} = \begin{bmatrix}
        1 & 0 \\
        0 & 1/\sqrt{a}
    \end{bmatrix}
    \begin{bmatrix}
        \cos \theta & -\sin \theta \\
        \sin \theta & \cos \theta
    \end{bmatrix}
    \begin{bmatrix}
        1 & 0 \\
        0 & \sqrt{a}
    \end{bmatrix}
    \begin{bmatrix}
        x_1 \\
        x_2
    \end{bmatrix}
   
   
   
   
\end{align}
Then 
\begin{align}
    \L(x_1, x_2) = \L(g \cdot (x_1, x_2))
\end{align}

\subsubsection{Rosenbrock function}
\label{app:rosenbrock_symmetry}
Consider the Rosenbrock function with 2 variables \cite{rosenbrock1960automatic}:
\begin{align}
    \L(x_1, x_2) = 100(x_1^2 - x_2)^2 + (x_1 - 1)^2
\end{align}

Let $u = 10(x_1^2 - x_2)$ and $v = x_1 - 1$. After changing the variables from $x$ and $y$ to $u$ and $v$, $\L$ has a rotational symmetry. Note that the function, $h: \R^2 \xrightarrow[]{} \R^2$, that maps $x_1, x_2$ to $u, v$ is bijective: 
\begin{align}
    (u, v) &= h(x_1, x_2) = (10(x_1^2 - x_2), x_1 - 1) \cr
    (x_1, x_2) &= h^{-1}(u, v) = (v + 1, (v + 1)^2 - 0.1u) \cr
    h(x_1, x_2) &= h(y_1, y_2) \Rightarrow (x_1, x_2) = (y_1, y_2)
\end{align}

Next, we show that $SO(2)$ is a symmetry of $\L(x_1,x_2)$. Let $\rho$ be a representation of $SO(2)$ acting on $\R^2$. For $g \in SO(2)$, define the following group action:
\begin{align}
    g \cdot (x_1, x_2) = 
    h^{-1} \left( \rho(g) h
    (x_1, x_2) \right)
\end{align}
Then 
\begin{align}
    \L(x_1, x_2) = \L(g \cdot (x_1, x_2))
\end{align}

For the Rosenbrock function with 2N parameters, we can construct a bijective function $h:\R^{2N} \xrightarrow[]{} \R^{2N}$ by transforming each of the $N$ pairs of variables as before, and $SO(2N)$ is a symmetry of $\L(x_1,...,x_{2N})$. However, we will only use the 2 variable version in the experiments.

\subsubsection{Booth function}
\label{app:booth_symmetry}
Consider the Booth function \cite{jamil2013literature}:
\begin{align*}
    \L(x_1, x_2) = (x_1 + 2x_2 - 7)^2 + (2x_1 + x_2 - 5)^2
\end{align*}
Similar to the Rosebrock function, a change of variables reveals a rotational symmetry of $\L$:
\begin{align}
    (u, v) &= h(x_1, x_2) = (x_1 + 2x_2 - 7, 2x_1 + x_2 - 5) \cr
    (x_1, x_2) &= h^{-1}(u, v) = (-\frac{1}{3} u + \frac{2}{3} v + 1, \frac{2}{3} u - \frac{1}{3} v + 3) \cr
\end{align}
The function $h: \R^2 \xrightarrow[]{} \R^2$ that maps $x_1, x_2$ to $u, v$ is bijective. Let $\rho$ be a representation of $SO(2)$ acting on $\R^2$. For $g \in SO(2)$, define the following group action:
\begin{align}
    g \cdot (x_1, x_2) = 
    h^{-1} \left( \rho(g) h
    (x_1, x_2) \right)
\end{align}
Then $\L(x_1,x_2)$ admits an $SO(2)$ symmetry:
\begin{align}
    \L(x_1, x_2) = \L(g \cdot (x_1, x_2))
\end{align}

\subsection{Proof of Proposition \ref{thm:mult-layer} and Proposition \ref{thm:mult-layer-nonlinear}}
\label{appendix:mult-layer}
\textbf{Proposition \ref{thm:mult-layer}}
\begin{proof} In the linear network $h_m = W_m h_{m-1}$, hence
\begin{align}
    g \cdot (W_{m}, h_{m-1})= (W_{m}g^{-1}, g h_{m-1}), \quad 
    g\cdot h_m = W_{m}g^{-1} g h_{m-1} = h_m
\end{align}
which means a $p$-layer linear network is invariant under all $G_m$ with $m\leq p$ as they keep the output $h_p$ invariant ($\forall g\in G_m$, $g\cdot h_p= h_p$). 
\eat{
For $m\geq 1$ this is a symmetry of the architecture and is independent of the  input $X$, defining 
\begin{align}
    g \cdot (W_{m}, W_{m-1}) = (W_{m}g^{-1}, g W_{m-1}),\quad  
    g\cdot h_m = W_{m}g^{-1} g W_{m-1}h_{m-2} = h_m
\end{align}
We define the group action as identity on the weights of all other layers, i.e. for  $g \in \mathrm{GL}_{d_{m-1}}$,  $g \cdot W_k = W_k$ when $k\notin \{ m, m-1\}$. 
For a $p$-layer network, we have $g\cdot h_p =h_p $, meaning the entire network is invariant under this group action.

\end{proof}

\textbf{Proposition \ref{thm:mult-layer-nonlinear}}
\begin{proof}
From (8), we want to convert $g\cdot h_{m-1}$ into a transformation on $W_{m-1}$ instead of $h_{m-1}$. 
In other words, we want to find a set of transformed weights $W'_m,W'_{m-1}$ which yields the same network output $\tilde{h}_m$: 
\begin{align}
    \tilde{h}_m&= W'_m \sigma\pa{W'_{m-1} h_{m-2}} = W_m g^{-1} g \sigma\pa{W_{m-1} h_{m-2}}  \cr 
    \Rightarrow W'_m &= W_m g^{-1}, \qquad  \sigma\pa{W'_{m-1} h_{m-2}} = g \sigma\pa{W_{m-1} h_{m-2}} 
   
   
   
    \label{eq:Wm-g-derive-0}
\end{align}
To obtain $W'_{m-1} = g\cdot W_{m-1}$ we first find a right pseudo-inverse  $h_{m-2}^{+} \in \R^{n\times d_{m-2}}$ for $h_{m-2}$ meaning $h_{m-2}^+h_{m-2} = I$ (assuming $n\ll d_{m-2}$ and that $h_{m-2}$ is full rank). 
Solving \eqref{eq:Wm-g-derive-0}
we get
\begin{align}
    W'_{m-1} &=
    \sigma^{-1}\pa{g \sigma\pa{W_{m-1} h_{m-2}} } h_{m-2}^{+}
    \label{eq:W-act-nonlin}
\end{align}

\eqref{eq:Wm-g-action-main} follows from \eqref{eq:Wm-g-derive-0} and \eqref{eq:W-act-nonlin}. 
\end{proof}


\section{Theoretical analysis of teleportation}
\subsection{Proof of Proposition \ref{prop:speedup-cond} and Corollary \ref{prop:corollary-cond}}
\label{appendix:speedup-cond-prop}
\textbf{Proposition \ref{prop:speedup-cond}}
\begin{proof}
 Let ${\bm{w}}' = g \cdot {\bm{w}}$. Denote the Jacobian as $J$, where $J_{ij} = \partial w_i' / \partial w_j$. Then the inverse of $J$ has entries $J^{-1}_{ij} = \partial w_i / \partial w_j'$. 

The gradient at ${\bm{w}}'$ is
\begin{align}
    \frac{\partial \L({\bm{w}}')}{\partial {\bm{w}}'} 
    &= \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}'} 
    = \sum_j \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}_j} \frac{\partial {\bm{w}}_j}{\partial {\bm{w}}_i'} 
    = \sum_j \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}_j} J^{-1}_{ji} 
    = \left(\left(\frac{\partial \L({\bm{w}})}{\partial {\bm{w}}}\right)^T J^{-1} \right)^T
    = (J^{-1})^T \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}}
\end{align}
The rate of change of $\L$ in gradient flow is
\begin{align}
    {d\L({\bm{w}}')\over dt}& = \bk{{\partial \L \over \partial {\bm{w}}' },{d{\bm{w}}' \over dt} } = - \left\| (J^{-1})^T \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}} \right\|_\eta^2 
   
\end{align}

Thus we will have a speedup if 
\begin{align}
    \left\| (J^{-1})^T \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}} \right\|_\eta^2  > \left\| \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}} \right\|_\eta^2
    \label{eq:speedup-cond}
\end{align}
\end{proof}

\textbf{Corollary \ref{prop:corollary-cond}}
\begin{proof}
    Since $J = g$, 
    using \eqref{eq:dLdt-0} and the l.h.s. of \eqref{eq:dL-w-cond-0} we have 
    \begin{align}
        {d\L(g\cdot {\bm{w}})\over dt }& 
        = -\nabla \L^T g^{-1}\eta \br{g^{-1}}^T \nabla \L = \nabla \L^T \br{g^T \eta^{-1} g }^{-1} \nabla \L  = \|\nabla \L\|^2_\eta 
    \end{align}
   
\end{proof}


\subsection{Proof of Proposition \ref{prop:lipschitz}}
\label{appendix:lipschitz}
\begin{proof}
From the definition of Lipschitz continuity and the update rule of gradient descent, 
\begin{align}
    \left| \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+1}} \right\|_2 - \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2 \right| \leq L \| {\bm{w}}_{t+1} - {\bm{w}}_t \|_2 = L\left\| \eta \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2
\end{align}
Equivalently,
\begin{align}
    (1 - \eta L) \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2 \leq \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+1}} \right\|_2 \leq (1 + \eta L) \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2
\end{align}
By unrolling $T$ steps, we have
\begin{align}
    (1 - \eta L)^T \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2 
    \leq \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+T}} \right\|_2 
    \leq (1 + \eta L)^T \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2
\end{align}
Similarly, for a teleported point ${\bm{w}}_t' = g \cdot {\bm{w}}_t$,
\begin{align}
    (1 - \eta L)^T \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}'} \right\|_2 
    \leq \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+T}'} \right\|_2 
    \leq (1 + \eta L)^T \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}'} \right\|_2
\end{align}
Therefore, if
\begin{align}
     (1 - \eta L)^T \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}'} \right\|_2 
     \geq (1 + \eta L)^T \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2
\end{align}
then it is guaranteed that 
\begin{align}
    \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+T}'} \right\|_2 \geq \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+T}} \right\|_2
\end{align}
\end{proof}


\subsection{Proof of Lemma \ref{lemma:newton-direction}}
\label{appendix:newton-direction-lemma}
\begin{proof}
From the definition of the directional derivative,
\begin{align}
    \partial_{\bm{v}} \left\|\frac{\partial \L}{\partial {\bm{w}}}\right\|_2^2 
    &= {\bm{v}} \cdot \frac{\partial}{\partial {\bm{w}}} \left\|\frac{\partial \L}{\partial {\bm{w}}}\right\|_2^2
\end{align}
Writing the last term in indices, 
\begin{align}
    \frac{\partial}{\partial {\bm{w}}_i} \left\|\frac{\partial \L}{\partial {\bm{w}}}\right\|_2^2
    &= \frac{\partial}{\partial {\bm{w}}_i} \sum_j \left(\frac{\partial \L}{\partial {\bm{w}}_j}\right)^2 \cr
    &= \sum_j \frac{\partial}{\partial {\bm{w}}_i} \left(\frac{\partial \L}{\partial {\bm{w}}_j}\right)^2 \cr
    &= \sum_j 2 \frac{\partial \L}{\partial {\bm{w}}_j} \frac{\partial^2 \L}{\partial {\bm{w}}_i \partial {\bm{w}}_j} \cr
    &= 2 \left(H \frac{\partial \L}{\partial {\bm{w}}}\right)_i
\end{align}
Removing the indices,
\begin{align}
    \frac{\partial}{\partial {\bm{w}}} \left\|\frac{\partial \L}{\partial {\bm{w}}}\right\|_2^2
    &= 2 H \frac{\partial \L}{\partial {\bm{w}}}
\end{align}
Substitute back and we have
\begin{align}
    \partial_{\bm{v}} \left\|\frac{\partial \L}{\partial {\bm{w}}}\right\|_2^2 
    &= {\bm{v}} \cdot \left(2 H \frac{\partial \L}{\partial {\bm{w}}}\right)
\end{align}
Since $\partial_{\bm{v}} \left\|\frac{\partial \L}{\partial {\bm{w}}}\right\|_2^2 = 0$ for all vector ${\bm{v}}$ that is orthogonal to $\frac{\partial \L}{\partial {\bm{w}}}$, ${\bm{v}} \cdot \left(2 H \frac{\partial \L}{\partial {\bm{w}}}\right) = 0$ for all vector ${\bm{v}}$ that is orthogonal to $\frac{\partial \L}{\partial {\bm{w}}}$.
In other words, $2 H \frac{\partial \L}{\partial {\bm{w}}}$ is orthogonal to all vectors that are orthogonal to $\frac{\partial \L}{\partial {\bm{w}}}$. Therefore, $2 H \frac{\partial \L}{\partial {\bm{w}}}$ has the same direction of $\frac{\partial \L}{\partial {\bm{w}}}$, and $\frac{\partial \L}{\partial {\bm{w}}}$ is an eigenvector of the Hessian of $\L$.
\end{proof}

\subsection{Proof of Proposition \ref{prop:newton-direction}}
\label{appendix:newton-direction-prop}
\begin{proof}
From Lemma \ref{lemma:newton-direction}, $\frac{dL}{dw}$ is an eigenvector of $H$. Therefore, it is also an eigenvector of $H^{-1}$. Hence $\frac{dL}{dw}$ has the same direction as $H^{-1} \frac{dL}{dw}$.
\end{proof}

\section{Experiment details and additional results}
\subsection{Test functions}
\label{appendix:exp-test-function}
\begin{figure}[t!]
\centering
\hspace{80pt} a \hspace{115pt} b \hfill ~ \\
\includegraphics[width=0.32\textwidth]{figures/test_functions/Rosenbrock_loss_vs_gradient.pdf}
\includegraphics[width=0.32\textwidth]{figures/test_functions/Booth_loss_vs_gradient.pdf}
\caption{Gradient on the trajectory of optimizing the Rosenbrock function (left) and Booth function (right). At the same loss value, the graident is larger on the trajectory with teleportation, indicating a better descent path.}
\label{fig:test_functions_supp}
\end{figure}

We compare the gradient at different loss values for gradient descent with and without teleportation. Figure \ref{fig:test_functions_supp} shows that the trajectory with teleportation has a larger $d\L/dt$ value than the trajectory without teleportation at the same loss values. Therefore, the rate of change in the loss is larger in the trajectory with teleportation, which makes it favorable. 

\subsection{Multilayer neural network}
\label{appendix:exp-multilayer}
Data $X, Y$ and initialization of parameters $W$ are set uniformly at random over $[0, 1]$. The group elements used for these transforms are found by gradient ascent on $T$ for 10 times, with learning rate $10^{-7}$ in GD and $10^{-6}$ in AdaGrad. When using teleportation, we perform symmetry transform on the parameters once every 10 epochs in the first 50 epochs. Each algorithm is run 300 steps. GD uses learning rate $10^{-4}$ and AdaGrad uses $10^{-1}$. 

To confirm that teleportation works with different learning rates, we repeat the experiment using learning rate $10^{-3}$ for GD and $10^{0}$ for AdaGrad, which is 10 times higher than the learning rate used in the main paper. Increasing learning rates further prevents convergence. Since gradient descent is expected to converge faster with higher learning rates, we adjust the teleportation schedule to once every 4 epochs for the first 20 epochs and run each algorithm 120 steps. Fig. \ref{fig:two-layer-supp} shows the results for the new learning rates. Teleportation improves both GD and AdaGrad in the similar way as using the original learning rates.

\begin{figure}[t!]
\centering
\ \ \ a\hfill b \hfill c\hfill ~ \\
\includegraphics[width=0.32\textwidth]{figures/multi_layer_loss_lr_1e-3_1e0.pdf}
\includegraphics[width=0.32\textwidth]{figures/multi_layer_loss_gradient_lr_1e-3_1e0.pdf}
\includegraphics[width=0.32\textwidth]{figures/multi_layer_loss_vs_gradient_lr_1e-3_1e0.pdf}
\caption{Multi-layer network optimization with a different learning rate. Teleportation improves both gradient descent and AdaGrad, similar to using the original learning rate. }
\label{fig:two-layer-supp}
\end{figure}


\section{Discussion and Conclusions}
We proposed an accelerated optimization method that exploits the symmetry in the parameter space of  the loss function to improve convergence. It is interesting to note that optimizing $d\L/dt$ locally sometimes leads to an improved global path. One example is the ellipse function (Figure \ref{fig:sales_pic}), where teleporting once ensures that $d\L/dt$ is optimal at every $\L$ value along the trajectory.
Another example is the matrix factorization problem $\L(U, V) = \|Y-UV\|_2^2$. \cite{tarmoun2021understanding} shows that the convergence rate increases with the imbalance $U^TU - V^TV$. Consider the transformation $U, V \xrightarrow[]{} Ug, g^{-1}V$. To optimize $d\L(U,V)/dt$ locally, we would need a large $g \in GL_n$, but a large $g$ also leads to a large $U^TU - V^TV$ which is positively correlated with the overall convergence rate of the entire trajectory. Hence teleportation is guaranteed to produce a better trajectory. 

A potential future direction is to derive the exact expression for how teleportation affects the loss value at a later time in gradient flow, which may lead to a closed-form solution of the optimal teleportation destination. Additionally, inspired by the landscape view of parameter space symmetry \cite{simsek2021geometry}, teleportation using discrete (permutation) symmetries may allow us to reach a better minima. Finally, additional theory can be developed to explain the relationship between teleportation and second-order methods under the approximations we introduced, especially to quantify the improvement on the overall convergence rate, and to derive its effect on generalization bounds. Integrating teleportation with other advanced optimizers such as Adam, RMSprop would be another interesting future step.


\section{Experiments}

\subsection{Acceleration through symmetry teleportation}
We examine the effect of symmetry teleportation on optimization. 
We illustrate teleportation in the parameter space on two test functions and show the speedup in regression and classification problems using multilayer neural networks. 
For test functions, we  compared with GD for illustration purposes. 
For multi-layer neural networks, we also include AdaGrad as a more competitive baseline.
\begin{figure}[t!]
\centering
a\hfill b \hfill c\hfill d \hfill ~ \\
\includegraphics[width=0.24\textwidth, trim=35pt 20pt 0 0, clip]{figures/test_functions/Rosenbrock_level_set_GD.pdf}
\includegraphics[width=0.24\textwidth, trim=35pt 20pt 0 0, clip]{figures/test_functions/Rosenbrock_level_set_teleport.pdf}
\includegraphics[width=0.24\textwidth]{figures/test_functions/Rosenbrock_loss.pdf}
\includegraphics[width=0.24\textwidth]{figures/test_functions/Rosenbrock_loss_gradient.pdf}\\
e\hfill f \hfill g\hfill h \hfill ~ \\
\includegraphics[width=0.24\textwidth, trim=40pt 20pt 0 0, clip]{figures/test_functions/Booth_level_set_GD.pdf}
\includegraphics[width=0.24\textwidth, trim=40pt 20pt 0 0, clip]{figures/test_functions/Booth_level_set_teleport.pdf}
\includegraphics[width=0.24\textwidth]{figures/test_functions/Booth_loss.pdf}
\includegraphics[width=0.24\textwidth]{figures/test_functions/Booth_loss_gradient.pdf}
\caption{Optimization of the Rosenbrock function (top row) and Booth function (bottom row). using (a) gradient descent and (b) the proposed algorithm. Contours represent the level sets of the loss function. Loss $\L$ and convergence rate $d\L/dt$ are shown in (c) and (d). Teleportation helps move the parameters towards the target. }
\label{fig:test_functions}
\end{figure}
\eat{
\begin{figure}[t!]
\centering
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Rosenbrock_level_set_GD.pdf}}
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Rosenbrock_level_set_teleport.pdf}}
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Rosenbrock_loss.pdf}}
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Rosenbrock_loss_gradient.pdf}}
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Booth_level_set_GD.pdf}}
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Booth_level_set_teleport.pdf}}
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Booth_loss.pdf}}
\subfigure[]{\includegraphics[width=0.24\textwidth]{figures/test_functions/Booth_loss_gradient.pdf}}
\caption{Optimization of the Rosenbrock function (top row) and Booth function (bottom row). using (a) gradient descent and (b) the proposed algorithm. Contours represent the level sets of the loss function. Loss $\L$ and convergence rate $d\L/dt$ are shown in (c) and (d). Teleportation helps move the parameters towards the target. }
\label{fig:test_functions}
\vspace{-3mm}
\end{figure}

\paragraph{Rosenbrock function.} We apply symmetry teleportation to optimize the 2-variable Rosenbrock function \eqref{eqn:rosenbrock}. The parameters $x_1, x_2$ are initialized to $(-1, -1)$. Each algorithm is run 1000 steps with learning rate $10^{-3}$. We teleport the parameters every 100 steps. The group elements are found by gradient ascent on $\theta$, the parameter for the $\mathrm{SO}(2)$ group, for 10 steps with learning rate $10^{-1}$. 

The trajectory of parameters and the loss level sets are plotted in Figure \ref{fig:test_functions}a,b. The blue star denotes the final position of parameters, the green star denotes the target, and orange dots are the positions from which symmetry transforms start. While gradient descent is not able to reach the target in 1000 steps, teleportation allows large steps and reaches the target much earlier. Figure \ref{fig:test_functions}d shows that teleportation improves the norm of gradients in the following step and \ref{fig:test_functions}c shows its effect on the loss value. 
Teleportations clearly reduces the overall convergence time. 

\paragraph{Booth function.}
We also test symmetry teleportation on the Booth function defined in  Eqn. \eqref{eqn:booth}. We initialize the parameters $x_1, x_2$  to $(5, -5)$. Each algorithm is run 10 steps with learning rate $0.08$. We perform symmetry teleportation on the parameters once, before epoch 5. The group elements are found by gradient ascent on $\theta$ for 10 steps with learning rate $0.001$. $\theta$ is initialized uniformly at random over $[0, \pi)$. Similar to the Rosenbrock function, teleportation moves the parameters to a trajectory with a larger convergence rate (Figure \ref{fig:test_functions} bottom row). 

\paragraph{Multilayer neural network regression.}
We further evaluated our method on a three-layer neural network with a  regression loss   $\min_{W_1,W_2,W_3}\| Y - W_3 \sigma(W_2 \sigma(W_1 X))\|_2$. The dimension of weight matrices are $W_3 \in \R^{8 \times 7}$, $W_2 \in \R^{7 \times 6}$, and $W_1 \in \R^{6 \times 5}$. $X \in \R^{5 \times 4}$ is the data, $Y \in \R^{8 \times 4}$ is the target, and $\sigma$ is the LeakyReLU activation with slope coefficient 0.1. Discussion of the choice of learning rates and additional hyperparameter details can be found in Appendix \ref{appendix:exp-multilayer}. 

Teleportation of the $\mathrm{GL}(\R)$ group can be performed by finding an element $x$ in its Lie algebra, such that transforming ${\bm{w}}$ by the group element $g=\exp(x)$ improves the gradient $\frac{d\L}{dt}$. We use the first order approximation of the exponential map. Between each pair of weight matrices, we replace $g_m$ by $I+T$ and $g_m^{-1}$ by $I-T$ in \eqref{eq:Wm-g-action-main}, where $T \in \R^{d_m \times d_m}$ is initialized to 0. Then we perform gradient ascent steps on $T$ with objective defined in Line 3 of Algorithm \ref{alg:teleport} and update the pair of weights.

Figure \ref{fig:two-layer}a shows the training curves, where shaded area denotes one standard deviation from 5 runs. Since GD and AdaGrad use different learning rates, they are not directly comparable. However, the addition of teleportation clearly improves both algorithms. Figure \ref{fig:two-layer}b and \ref{fig:two-layer}c shows the squared norm of gradient from a single run. Teleportation increases the magnitude of gradient, and the trajectory with teleportation has a larger $d\L/dt$ value at the same loss values, which demonstrates that teleportation finds a better trajectory.
\begin{figure}[t!]
\centering
\ \ \ a\hfill b \hfill c\hfill ~ \\
\includegraphics[width=0.32\textwidth]{figures/multi_layer_loss.pdf}
\includegraphics[width=0.32\textwidth]{figures/multi_layer_loss_gradient.pdf}
\includegraphics[width=0.32\textwidth]{figures/multi_layer_loss_vs_gradient.pdf}
\caption{Multi-layer network optimization  using gradient descent with and without teleportation. Loss $\L$ and convergence rate $d\L/dt$ are shown in (a) and (b). Teleportation improves the overall convergence rate. The gradient stays large after teleportation.  }
\label{fig:two-layer}
\end{figure}


\paragraph{MNIST classification.}
We apply symmetry teleportation on the MNIST classification task \cite{deng2012mnist}. We split the training set into 48,000 for training and 12,000 for validation. The input data has dimension $28 \times 28$ and is flattened into a vector. The output of the neural network has dimension $10$ corresponding to the 10 digit classes.  We used a three-layer neural network with hidden dimension $[512, 512]$, LeakyReLU activations, and cross-entropy loss. Learning rate is $2 \times 10^{-3}$, and learning rate for teleportation is $10^{-3}$. Each optimization algorithm is run 80 epochs with batch size of 20. Immediately after the first epoch, we apply teleportation using data from one mini-batch, and repeat for 4 different mini-batches. For each mini-batch, 10 gradient ascent steps are used to optimize $g_m$. 



Figure \ref{fig:mnist}a shows the effect of teleportation on training and validation loss, and Figure \ref{fig:mnist}b,c shows the norm of the gradient in training. 
While teleportation significantly accelerates the decrease of the training loss, its effect on the validation loss is limited for SGD and detrimental for AdaGrad. 
Therefore, teleportation on MNIST makes training faster at the beginning but leads to earlier overfit and slightly worse validation accuracy. A possible reason is that regions with large gradients have sharp minimum that do not generalize well. 

\begin{figure}[t!]
\centering
\ \ \ a\hfill b \hfill c\hfill ~ \\
\includegraphics[width=0.32\textwidth]{figures/mnist_loss.pdf}
\includegraphics[width=0.32\textwidth]{figures/mnist_loss_gradient.pdf}
\includegraphics[width=0.32\textwidth]{figures/mnist_loss_vs_gradient.pdf}
\caption{MNIST classification using gradient descent with and without teleportation. Solid lines are training loss and dashed lines are validation loss.}
\label{fig:mnist}
\vspace{-3mm}
\end{figure}

\subsection{Teleportation schedule}
The effect of teleportation varies depending on its time and frequency. Figure \ref{fig:schedule} shows the result of teleportation on MNIST with different hyperparameters. In all experiments, we use SGD with batch size 20, learning rate $2 \times 10^{-3}$, and 10 gradient ascent steps for each teleportation. 

In Figure \ref{fig:schedule}a, we randomly select 4 different nimi-batches and apply teleportation on each of them individually, but at different epochs. Teleportation before training has the worst performance. After epoch 0, the effect of teleportation is stronger when it is applied earlier. In Figure \ref{fig:schedule}b, we again apply teleportations on 4 randomly selected mini-batches, but repeat this with 5 different teleportation schedules (hyperparameter $K$ in Algorithm \ref{alg:teleport-sgd}) as shown in the legend. All schedules has the same number of teleportations. Smaller intervals between teleportation accelerate convergence more significantly. In Figure \ref{fig:schedule}c, we apply teleportation immediately after the first epoch, but use different numbers of mini-batches (hyperparameter $B$ in Algorithm \ref{alg:teleport-sgd}) and teleport using each of them individually. Using more mini-batches to teleport leads to faster decrease in training loss but is also more prone to overfitting.




\begin{figure}[t!]
\centering
\ \ \ a\hfill b \hfill c\hfill ~ \\
\includegraphics[width=0.32\textwidth]{figures/MNIST_SGD_tele_epoch_batch4_telestep10_telelr1e-3_plot.pdf}
\includegraphics[width=0.32\textwidth]{figures/MNIST_SGD_tele_schedule_nbatch4_telestep10_telelr1e-4_plot.pdf}
\includegraphics[width=0.32\textwidth]{figures/MNIST_SGD_tele_batch_epoch1_telestep10_telelr1e-3_plot.pdf}
\caption{Teleportation (a) once at different epoch, (b) 5 times with different teleportation schedules, and (c) using different number of mini-batches. The lists in the legend of (a) and (c) denote the epoch numbers in teleportation schedule where teleportation happens.}
\label{fig:schedule}
\vspace{-3mm}
\end{figure}

\subsection{Runtime analysis}
The additional amount of computation introduced by teleportation depends on the implementation of Line 2-5 of Algorithm \ref{alg:teleport}. We discuss our implementation for teleporting multi-layer neural networks as an example. Teleportating each pair of adjacent weight matrices requires computing the inverse of the output from the previous layer. 
Denote the batch size as $n$, the largest dimension of all layers as $d_{max}$, and the number of layers as $l$. Assume that $d_{max} > n$.  Computing the pseudoinverse of the output of each layer has complexity $O(d_{max}^2 n)$, and computing the pseudoinverse for all layers has complexity $O(d_{max}^2 nl)$. Note that all matrices we invert has dimensions at most $d_{max} \times n$. 

For one gradient ascent step on $g$, the forward and backward pass both have complexity $O(d_{max}^2 nl)$. This is the same as the forward and backward pass of gradient descents on ${\bm{w}}$ because the architecture is the same except with approximately twice as many layers.
We perform $t$ gradient ascent steps on $g$. 
Therefore, the computation cost for one teleportation is $O(d_{max}^2 nlt)$. 

We show empirically that the runtime for teleportation scales polynomially with matrix dimensions and linearly with the number of layers. We record the runtime of gradient descent on a Leaky-ReLU neural network for 300 epochs, with a 10-step teleportation every 10 epochs. Each pair of adjacent weight matrices is transformed by a group action during teleportation. Figure \ref{fig:multi-layer-time}(a) shows the runtime for a 3-layer network using square weights and data matrices with different dimensions. Figure \ref{fig:multi-layer-time}(b) shows the runtime for 128-by-128 weight and data matrices with different number of layers. 

Although the teleportation step has the same complexity as a gradient descent step, the runtime is dominated by teleportation due to larger constants in the complexity analysis. The trade-off between teleportation time and convergence rate depends on specific problems. In our experiments, the convergence rate is improved by a small number of teleportation steps which does not add significant computational overhead.

\begin{figure}[t!]
\centering
~\hfill  \ \ a\hfill ~\hfill \ \ b \hfill ~ \hfill~ \hfill ~ \\
\includegraphics[width=0.32\textwidth]{figures/multi_layer_time-dim.pdf}
\includegraphics[width=0.32\textwidth]{figures/multi_layer_time-layer.pdf}
\caption{Gradient descent training time on a Leaky-ReLU neural network for 300 epochs, with a 10-step teleportation every 10 epochs. (a) 3-layer network using square weights and data matrices, with different dimensions. (b) Weights and data are all 128 by 128 matrices, but the network has different number of layers. }
\label{fig:multi-layer-time}
\vspace{-3mm}
\end{figure}
\section{Symmetry Groups of Certain Optimization Problems}
We give a few practical examples to demonstrate how teleportation can be used to accelerate optimization. 
Specifically, we first consider two test functions
which are often used to evaluate optimization algorithms \cite{back1996evolutionary}. 
We then derive the symmetries of   multi-layer  neural networks. 

\subsection{Test functions}
\paragraph{Rosenbrock function.}
The Rosenbrock function originally introduced by  \cite{rosenbrock1960automatic} has a characteristic global minimum that is inside a long, narrow, parabolicly-shaped flat valley. Finding the valley is easy but reaching the minimum is difficult. On a 2-dimensional space, the Rosenbrock function has the following form: 
\begin{align}
    \L_{r}(x_1, x_2) = 100(x_1^2 - x_2)^2 + (x_1 - 1)^2.
\label{eqn:rosenbrock}
\end{align}

\paragraph{Booth function.}
The Booth function \cite{jamil2013literature} is also defined on $\mathbb{R}^2$ and has one global minimum at $(1,3)$ where the function evaluates to $0$:
\begin{align}
    \L_{b}(x_1, x_2) = (x_1 + 2x_2 - 7)^2 + (2x_1 + x_2 - 5)^2.
\label{eqn:booth}
\end{align}
The following proposition identifies the symmetry of these two test functions. 
\begin{proposition}
The Rosenbrock and Booth functions have rotational symmetry. In other words, there exist action maps $a_{r}, a_{b}: \mathrm{SO}(2) \times \R^2 \xrightarrow[]{} \R^2$, such that for all $g \in \mathrm{SO}(2)$, 
\[\L_{r} (x_1, x_2) = \L_{r} (a_{r}(g, [x_1, x_2])) \quad \text{and} \quad  \L_{b} (x_1, x_2) = \L_{b} (a_{b}(g, [x_1, x_2])).\]
\end{proposition}
The exact form of the action maps are deferred to Appendix \ref{app:rosenbrock_symmetry} and \ref{app:booth_symmetry}.
During the teleportation step, our goal is to maximize the gradient within a level set of the loss:
$
    \max_{ g \in \mathrm{SO}(2)} \| \frac{d \L(g \cdot (x_1, x_2))}{dt} \| 
$.





\eat{
\subsection{Two-layer neural network}
\bz{To be removed.}

Given a two-layer neural network with Leaky ReLU as the activation function.  $X \in \R^{d \times n}, Y \in \R^{d \times n}$ are the input and output. The objective function for regression is
\begin{align}
    \L(X, Y, U, V) = \| Y - U \sigma(VX) \|_2^2
    \label{eq:two-layer-obj}
\end{align}
where $U \in \R^{d \times k}, V \in \R^{k\times d}$, and $\sigma: \R^{k\times d} \xrightarrow[]{} \R^{k\times d}$ is the element-wise Leaky ReLU. As shown in section 2 in this document, $\L$ has a $GL_8(\R)$ symmetry. For $g \in G$, the following action leaves $\L$ unchanged:
\begin{align}
    g \cdot (U, V) = (g \cdot U, g \cdot V) = (Ug^{-1}, \sigma^{-1} (g \sigma (VX)) X^{-1})
    \label{eq:two-layer-group-action}
\end{align}

During the symmetry transform step, our goal is to maximize the learning rate within a level set of the loss:
$
    \max_{g \in GL(\R)} \frac{d \L(g \cdot (U, V))}{dt} 
    \label{eq:two-layer-group-obj}
$
One approach is to find an element $x$ in the Lie algebra of $G$, such that transforming $U,V$ by the group element $g=\exp(x)$ improves the gradient $\frac{d\L}{dt}$. 
}


\subsection{Multi-layer Neural Networks}
Next, we consider feed-forward neural networks.
Denote the output of the $m$th layer by $h_m \in \R^{d_m \times n}$, where $d_m$ is the hidden dimension and $n$ is the number of samples. 
Denote the input by $h_0 = X \in \R^{d_0\times n}$. 
Each layer first performs a linear transformation $\tilde{h}_m = W_m h_{m-1}$ where $W_m\in \R^{d_{m}\times d_{m-1}}$ (we absorb biases into $W_m$ by adding an extra row of ones in $\tilde{h}_{m-1}$). The output $h_m$ is
\begin{align}
   
    h_m &= \sigma(\tilde{h}_m) = \sigma \pa{W_m h_{m-1} }.
   
\end{align}

We assume the activation functions $\sigma: \R\to \R $ are bijections. 
For instance, Leaky-ReLU is bijective. 
For other activation, such Sigmoid or Tanh, we analytically extend them to  bijective functions (e.g. $\tanh(x) + e^{x-N} - e^{-x+N}$ for $N\gg 1$). 
For linear activations, we have linear symmetries: 

\begin{proposition}
\label{thm:mult-layer}
A linear network is invariant under all groups $G_m \equiv \mathrm{GL}_{d_m}$ acting as 
\begin{align}
    g \cdot (W_{m}, W_{m-1}) &= (W_{m}g^{-1}, g W_{m-1}), &
    g \cdot W_{k} = W_k, \quad \forall k \notin \{m,m-1\}. 
\end{align}
\end{proposition}



Similar, in the nonlinear case, we want to find $g\cdot W_m$ that keep the outputs $h_k$ for $k\ne m-1$ invariant. 
With nonlinearity, $\tilde{h}_m = W_m \sigma(W_{m-1} h_{m-2})$. 
This network has a
$\mathrm{GL}_{d_{m-1}}$ symmetry 
given by
\begin{align}
    g \cdot (W_{m}, h_{m-1}) = (W_{m}g^{-1}, g h_{m-1}), \quad
    g\cdot h_m = W_{m}g^{-1} g h_{m-1} = h_m.
    \label{eq:g-act-W-h-0}
\end{align}
\eat{
First, we note that \eqref{eq:g-act-W-h-0} yields a nonlinear action on $\tilde{h}_{m-1}=W_{m-1}h_{m-2}$
\begin{align}
    g \cdot \{W_{m}, \tilde{h}_{m-1}\}&= \left\{W_{m}g^{-1}, \sigma^{-1}\pa{g \sigma( \tilde{h}_{m-1})}\right\}.
\end{align}


Thus, we have the following proposition regarding symmetries of nonlinear networks: 
\begin{proposition}
\label{thm:mult-layer-nonlinear}
A multi-layer network with bijective activation has a $\mathrm{GL}_{d_{m-1}}$ symmetry.
For $g_m \in G_m = \mathrm{GL}_{d_{m-1}}(\R)$ the following group action keeps $h_p$ with $p\geq m$ invariant 
\begin{align}
    g_m \cdot W_k = \left\{
    \begin{array}{lc}
        W_m g_m^{-1} & k = m \\
        \sigma^{-1}\pa{g \sigma\pa{W_{m-1} h_{m-2}} } h_{m-2}^{+} & k = m-1 \\
        W_k & k \not\in \{m, m-1\}
    \end{array}
    \right.
    \label{eq:Wm-g-action-main}
\end{align}
where $h_{m-2}^{+}$ is the right pseudo-inverse for $h_{m-2}$ (assuming $n\ll d_{m-2}$ and that $h_{m-2}$ is full rank). 
\end{proposition}


Proofs can be found in Appendix \ref{appendix:mult-layer}. Note that this group action depends on the input to the network as well as the current weights of all the lower layers. 
Yet, since this action keeps the output of upper and lower layers invariant, multiple $G_m$ for different $m$ applied at the same time still keep the network output invariant. 
That is, \eqref{eq:Wm-g-action-main} gives a well-defined non-linear group action of $G = GL_{d_1} \times \ldots G_{d_{p-1}}$ on $(X,W_1,\ldots,W_p)$ which depends on but fixes $X$ and keeps the loss function invariant. 

\section{Introduction}
Consider the optimization problem of finding $\text{argmin}_{\bm{w}} \L({\bm{w}})$, where $\L$ is a loss function and ${\bm{w}} $
are the parameters. 
While finding global minima of $\L({\bm{w}})$ is hard for non-convex problems, we can use gradient descent (GD) to find local minima. 
In GD we apply the following update rule at every step:
\begin{align}
    {\bm{w}}_{t+1} \leftarrow {\bm{w}}_t - \eta \nabla \L,
\end{align}
where $\eta$ is the learning rate. Gradient-descent is a first-order method that uses only gradient information. 
It is easy to compute but suffers from slow convergence. 
Second-order methods such as Newton's method use the second derivative to account for the landscape geometry. 
These methods enjoy faster convergence, but calculating the second derivative (Hessian) can be computationally expensive in high dimensions \cite{hazan2019lecture}. 


\begin{figure}[ht!]
\centering
\subfigure{\label{fig:GD-first}\includegraphics[width=0.32\textwidth]{figures/GD_first.pdf}}
\subfigure{\label{fig:GD-second}\includegraphics[width=0.32\textwidth]{figures/GD_second.pdf}}
\subfigure{\label{fig:GD-proposed}\includegraphics[width=0.32\textwidth]{figures/GD_proposed.pdf}}
\caption{Left to right: gradient descent, second-order methods, proposed method.}
\label{fig:sales_pic}
\end{figure}

In this paper, we propose a new optimization method based on parameter space symmetries. Our algorithm takes advantage of higher-order landscape geometry but uses only gradient information in most steps and thus avoids the computational cost of second-order methods.

Specifically, we look beyond local optimization and ask: \textit{what if we allow the parameters to make a big jump once in a while?} 
As shown in Figure \ref{fig:sales_pic}, during optimization, we teleport the parameters to another point with steeper gradients using a loss-invariant symmetry transformation.
After teleportation, the loss stays the same, but the gradient and hence the rate of loss decay may change. 
These transformations move the parameters to a better trajectory under gradient descent, which improves the convergence speed of subsequent steps.


The locality of gradient descent is reflected in its formulation in terms of a proximal term; our method circumvents this locality by teleporting to new locations with the same loss.
For GD, an equivalent  problem  uses the following proximal mapping \cite{combettes2011proximal}. Let $\bk{\cdot, \cdot}$ denote the inner product, we have
\begin{align}
    {\bm{w}}_{t+1} = \text{argmin}_{{\bm{w}}} \left\{\eta \langle (\nabla \L)|_{{\bm{w}}_t}, {\bm{w}} \rangle + \frac{1}{2}\|{\bm{w}} - {\bm{w}}_t\|_2^2 \right\}.
\end{align}
The term $\frac{1}{2}\|{\bm{w}} - {\bm{w}}_t\|^2$ is the proximal term that keeps ${\bm{w}}_{t+1}$ close to ${\bm{w}}_t$. 
Adaptive gradient methods, such as Adagrad, define the proximal term using the Mahalanobis distance $\|{\bm{w}} - {\bm{w}}_t\|_{G^{-1/2}}^2$ \citep{duchi2011adaptive} to account for landscape geometry. 
Our proposed teleportation technique  relaxes the requirement from the proximal term. 
We teleport to points on the same level set of $\L({\bm{w}}_t)$, but allow ${\bm{w}}$  to be far from ${\bm{w}}_t$ in Euclidean distance. 
\eat{
For example, the update rule in AdaGrad replaces the Euclidean distance by the Mahalanobis distance $\|{\bm{w}} - {\bm{w}}_t\|_{G^{-1/2}}^2$ \cite{duchi2011adaptive}.
In our case,  we define the  $NS$ as the quotient space \ry{$NS =V/S$ ?? } 
\nd{This requires more careful definition of $S$ and whether $NS$ is a metric space. It may be too much to discuss here.}
\begin{align}
    {\bm{w}}_{t+1} = \text{argmin}_{{\bm{w}} \in NS } \left\{\eta \langle (\nabla \L)|_{{\bm{w}}_t}, {\bm{w}} \rangle + \frac{1}{2}\|{\bm{w}} - {\bm{w}}_t\|_2^2 \right\}.
\end{align}


In summary, our main contributions are:
\begin{itemize}
    \item We propose \textit{symmetry teleportation}, an accelerated gradient-based optimization algorithm that exploits symmetries in the loss landscape.
    \item We derive loss-invariant group actions of test functions and multi-layer neural networks.
    \item We provide necessary conditions and examples of when teleportation improves convergence rate of the current and subsequent steps.
    \item We show that our algorithm is closely related to second order methods. 
    \item Experimentally, we show that teleportation improves the convergence speed of gradient descent and AdaGrad in different optimization problems. 
\end{itemize}
\section{Related Work}


Continuous parameter space  symmetries have been identified in neural networks with homogeneous \citep{badrinarayanan2015symmetry,du2018algorithmic} and radial activation functions \citep{ganev2021radial}.
The effect of symmetry transformations on gradients has been examined for translation, scale, and rescale symmetries in \cite{kunin2021neural}.  We contribute to this line of work by deriving loss-invariant group actions in multi-layer neural networks with invertible activation functions.

Several works exploit symmetry to facilitate optimization. For example, Path-SGD \citep{neyshabur2015path-sgd} improves optimization in ReLU networks by path regularization. $\mathcal{G}$-SGD \citep{meng2019mathcal} performs weight updates in a scale-invariant space, which also exploits the symmetry in ReLU networks.
\cite{van2017l2} analyzes the effect of L2 regularization in batch normalization which gives scale-invariant functions. 
\citet{bamler2018improving} transforms parameters in the symmetry orbits to address slow movement along directions of weakly broken symmetry and find the optimal $g \in G$ that minimizes $\L(g \cdot {\bm{w}})$. In comparison, we search within $G$-orbits for a point that has the largest gradient, i.e., we find the optimal $g$ to maximize $d\L(g \cdot {\bm{w}})/dt$.

Natural gradient \cite{amari1998natural}, adaptive gradient methods \cite{duchi2011adaptive, kingma2014adam}, and their approximations \cite{martens2018kfac, gupta2018shampoo} improve the direction of parameter updates instead of directly transforming parameters. If the group acts transitively on the level set, and we teleport to the point with maximum gradient, then our update direction is the same as that in Newton's method. 
We prove that our algorithm is connected to second-order optimization methods and show that teleportation can be used to improve these methods empirically. 

 The concept of neural teleportation was first explored using quiver representation theory \cite{armenta2021representation, armenta2020neural}. These works provide a way to explore level curves of the loss of neural networks and show that random teleportation speeds up gradient descent experimentally and theoretically. Our algorithm improves neural teleportation by searching for teleportation destinations that lead to the largest improvement in the magnitude of gradient. 


Several works study how parameter initialization affects convergence rate \cite{saxe2014exact, tarmoun2021understanding, min2021explicit}. If we apply a group transformation only at initialization, our method is similar to that of \cite{tarmoun2021understanding}. We do not guarantee that the transformed parameters lead to faster convergence rate throughout the entire training. However, we accelerate convergence at least for a short time after initialization. Additionally, our method is not restricted to initialization and can be applied at any time during training. 



Most contemporary neural networks are overparametrized.  While this has been shown to improve generalization \citep{belkin2019reconciling}, an important question is how overparametrization affects optimization. A series of works starting from \cite{arora2018optimization} show that overparametrization resulting from the depth of a neural network accelerates convergence. Another view is that the symmetry created by overparametrization poses constraints on trajectories in the form of conserved quantities \citep{gluch2021noether}. Additionally, the symmetry generates additional trajectories. When the trajectories created by overparametrization are equivalent, model compression by removing symmetry reduces training time \citep{ganev2021radial}. However, when trajectories are not equivalent, gradient flow on some paths are faster than others. We search for the better paths created by overparametrization. 



\section{Theoretical Analysis}
\subsection{What symmetries help accelerate optimization}
We now discuss the conditions that need to be satisfied for teleportation to accelerate GD. 
For brevity, we denote all trainable parameters (e.g. $ \left\{W_1,\cdots W_p \right\}$ for the $p$-layer neural network) by a single flattened vector ${\bm{w}} \in \R^n$. 
Consider a symmetry $G$ of the loss function $\L({\bm{w}})$, meaning
$\forall g \in G$, $\L({\bm{w}}) = \L(g \cdot {\bm{w}})$. 
We quantify how teleportation by $G$ changes the rate of loss decay, given by 
\begin{align}
    {d\L({\bm{w}})\over dt}& = \bk{{\partial \L \over \partial {\bm{w}} },{d{\bm{w}} \over dt} } = - \br{\nabla \L}^T\eta \nabla\L = - \|\nabla\L\|^2_\eta 
    \label{eq:dLdt-0}
\end{align}
where $\eta$ is the matrix of learning rates (which must be positive semi-definite) and $\|v \|_\eta^2=v^T\eta v$ is the Mahalanobis norm. 

The following proposition states the condition a symmetry needs to satisfy to accelerate optimization. 
(Proofs can be found in Appendix \ref{appendix:speedup-cond-prop}.) 


\begin{proposition}
\label{prop:speedup-cond}
Let ${\bm{w}}'=g\cdot {\bm{w}}$ be a point we teleport to. 
Let $J= \partial {\bm{w}}' /\partial {\bm{w}} $ be the Jacobian. 
Symmetry teleportation using $g$ accelerates the rate of decay in $\L$ if it satisfies
\begin{align}
    \left\|\br{J^{-1}}^T\nabla \L({\bm{w}})\right\|^2_\eta > \left\|\nabla \L({\bm{w}})\right\|^2_\eta
    \label{eq:dL-w-cond-0}
\end{align}

\end{proposition}

If the action of the symmetry group $G\subset GL_n$ is linear we have ${\bm{w}}' = g \cdot {\bm{w}} = g {\bm{w}}$ and $J=g$.
It follows that if $G$ is a subgroup of the certain orthogonal groups, $d\L/dt$ will be invariant: 
\begin{corollary}
\label{prop:corollary-cond}
Let $O_\eta $ denote the orthogonal group of invariances of the inverse of the learning rate, $\eta^{-1}$, meaning for $g\in O_\eta$, $ g^T\eta^{-1} g= \eta^{-1}$. 
Then
\begin{align}
    \forall g &\in O_\eta, \qquad {d\L (g\cdot {\bm{w}})\over dt } = {\L({\bm{w}})\over dt } 
\end{align}
\end{corollary}

In the simple case where the learning rate is a constant, $O_\eta = O(n)$ becomes the classic orthogonal group (e.g. rotations)
with $g^Tg = I $. 
In general, when $g$ preserves the norm of the gradient, symmetry teleportation has no effect on the convergence rate. 



\eat{
\begin{proposition}
Let   $ M&\equiv {\partial \L\over \partial {\bm{w}} } \br{{\partial \L\over \partial {\bm{w}} }}^T$, and $J = \partial (g\cdot {\bm{w}}) / \partial {\bm{w}}$ and $J^{-1} = \partial {\bm{w}} / \partial (g\cdots {\bm{w}})$, then symmetry teleportation accelerates optimization if it satisfies
\[ \Tr\br{ M [J^{-1}]^TJ^{-1} } > \Tr[M] \]
\bz{equivalently, \[\left\| J^{-1} \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}} \right\|^2 > \left\| \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}} \right\|^2\]
which means the Jacobian of the group action increases the gradient norm.} \ry{simplify and explain}
\end{proposition}

\begin{proof}
Let ${\bm{w}}' = g \cdot {\bm{w}}$. 
Denote the Jacobian as $J$, where $J_i^j = \partial {\bm{w}}_i' / \partial {\bm{w}}_j$. 
Then the inverse of $J$ has entries $[J^{-1}]_i^j = \partial {\bm{w}}_i / \partial {\bm{w}}_j'$.
The gradient at ${\bm{w}}'$ is
\begin{align}
    \frac{\partial \L({\bm{w}}')}{\partial {\bm{w}}'_i} 
    &= \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}'_i} 
   
   
   
    = {\partial {\bm{w}}_j \over {\bm{w}}'_i} \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}_j} \cr 
    & = [J^{-1}]_j^i\frac{\partial \L({\bm{w}})}{\partial {\bm{w}}_j} 
    = \br{J^{-1} \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}} }^i
\end{align}
The rate of change of $\L$ during gradient flow $d{\bm{w}} /dt = -\eta \partial \L/\partial {\bm{w}} $ is 
\begin{align}
    \frac{d \L({\bm{w}}')}{dt} 
    &= \sum_i \frac{\partial \L({\bm{w}}')}{\partial {\bm{w}}_i'} \frac{d {\bm{w}}_i'}{d t}
    = -\eta \sum_i \frac{\partial \L({\bm{w}}')}{\partial {\bm{w}}_i'} \frac{\partial \L({\bm{w}}')}{\partial {\bm{w}}_i'} \cr 
    &= -\eta \left\| J^{-1} \frac{\partial \L({\bm{w}})}{\partial {\bm{w}}} \right\|^2 
    = -\eta \Tr\br{ M [J^{-1}]^TJ^{-1} }
    \cr
   
    M&\equiv {\partial \L\over \partial {\bm{w}} } \br{{\partial \L\over \partial {\bm{w}} }}^T
\end{align}
\nd{ derive the full AdaGrad version }
Thus we will have a speedup if 
\begin{align}
    \mbox{if }\Tr\br{ M [J^{-1}]^TJ^{-1} } > \Tr[M]
    \label{eq:MJ-sppedup-cond}
\end{align}
\end{proof}
}


From Theorem 3.2 in \cite{hazan2019lecture}, assuming that $\L$ is $\beta$-smooth and is bounded by $|\L| \leq M$, the gradient norm in gradient descent converges as
$\frac{1}{2\beta} \sum_{t=1}^T \| (\nabla \L)|_{{\bm{w}}_t} \| \leq 2M$. Teleportation increases $(\nabla \L)|_{{\bm{w}}_t}$, therefore requiring less time to reach convergence.

\subsection{Improvement of subsequent steps }
Since teleportation moves the parameters to a point with a larger  gradient, and subsequent GD steps are local, we would expect that teleportation improves the magnitude of gradient for a few future steps as well. The following results formalize this intuition (proofs in Appendix \ref{appendix:lipschitz}.).
\begin{assumption}[Lipschitz Continuity]
The l2 norm of the gradient is Lipschitz continuous with constant $L \in \R^{\geq 0}$ , which is
\begin{align}
    \left| \left\| \frac{\partial \L}{\partial {\bm{w}}_{1}} \right\|_2 - \left\| \frac{\partial \L}{\partial {\bm{w}}_{2}} \right\|_2 \right| \leq L \| {\bm{w}}_1 - {\bm{w}}_2 \|_2
\end{align}
where ${\bm{w}}_1, {\bm{w}}_2$ are two points in the parameter space and $L$ is the Lipschitz constant.
\label{thm:lipschitz}
\end{assumption}
\begin{proposition}
Consider the gradient descent with a $G$-invariant loss $\L({\bm{w}})$ and learning rate $\eta \in \R^+$. Let ${\bm{w}}_t$ be the parameter at time $t$ and ${\bm{w}}_t' = g \cdot {\bm{w}}_t$ the parameter teleported by $g\in G$. Let ${\bm{w}}_{t+T}$ and ${\bm{w}}_{t+T}'$ be the parameters after $T$ more steps of gradient descent from ${\bm{w}}_t$ and ${\bm{w}}_t'$ respectively. Under Assumption \ref{thm:lipschitz}, if $\eta L < 1$, and 
\begin{align}
    \frac{\left\| \partial \L / \partial {\bm{w}}_{t}' \right\|_2}{\left\| \partial \L / \partial {\bm{w}}_{t} \right\|_2} \geq \frac{(1 + \eta L)^T}{(1 - \eta L)^T} 
\end{align}
then 
\begin{align}
    \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+T}'} \right\|_2 \geq \left\| \frac{\partial \L}{\partial {\bm{w}}_{t+T}} \right\|_2
\end{align}
\label{prop:lipschitz}
\end{proposition}
The proposition provides a sufficient condition for teleportation to improve $T$ future steps. The condition is met when $L$ is small, $\eta$ is small, or the initial improvement, $\left\| \frac{\partial \L}{\partial {\bm{w}}_{t}'} \right\|_2 / \left\| \frac{\partial \L}{\partial {\bm{w}}_{t}} \right\|_2$, is large.


\subsection{Relation to second-order optimization methods}
Symmetry teleportation is closely related to second-order optimization methods. We show that at the target point to teleport to, gradient descent becomes equivalent to Newton's method. 
Full proofs of the results in this section can be found in Appendix \ref{appendix:newton-direction-lemma} and \ref{appendix:newton-direction-prop}.
\begin{lemma}
If $\partial_{\bm{v}} \left\|\frac{\partial \L}{\partial {\bm{w}}}\right\|_2^2 = 0$ for all unit vector ${\bm{v}}$ that is orthogonal to $\frac{\partial \L}{\partial {\bm{w}}}$, then $\frac{\partial \L}{\partial {\bm{w}}}$ is an eigenvector of the Hessian of $\L$.
\label{lemma:newton-direction}
\end{lemma}

The following proposition is a direct consequence of Lemma \ref{lemma:newton-direction}.
\begin{proposition}
Let $S_{\L_0} = \{{\bm{w}}: \L({\bm{w}}) = \L_0\}$ be a level set of $\L$. 
If at a particular ${\bm{w}} \in S_{\L_0}$ we have $ \|\nabla\L({\bm{w}})\|_2\geq  \|\nabla\L({\bm{w}}')\|_2 $  for all ${\bm{w}}'$ in a small neighborhood of ${\bm{w}}$ in $S_{\L_0}$, then the gradient $\nabla \L({\bm{w}})$ has the same direction as the Newton's direction $H^{-1} \nabla\L({\bm{w}}) $.
\label{prop:newton-direction}
\end{proposition}


Proposition \ref{prop:newton-direction} provides an alternative way to interpret teleportation. Instead of computing the Newton's direction, we search within the loss level set for a point where the gradient has the same direction as Newton's direction. 
However, symmetry teleportation does not require computing the full Hessian matrix. The second derivative required for optimizing continuous groups are obtained by taking derivatives with respect to parameters and then with respect to the group element, as opposed to taking derivatives with respect to parameters twice. This makes the computation significantly more feasible than Newton's method on neural networks with large number of parameters.

In practice, however, the group action used for teleportation is usually not transitive. Additionally, we do not teleport using the optimal group element since it can be unbounded. We also do not apply teleportation in every gradient descent step. Therefore, Proposition \ref{prop:newton-direction} serves as an intuition instead of an exact formulation of how teleportation works. We provide empirical evidence in Section 6 that these approximations do not erase the benefits of teleportation completely, and leave theoretical investigations of the connection to second-order methods under approximations as future work. 


