remresetThe remreset package

Adaptive and Momentum Methods on Manifolds
Through Trivializations

Mario Lezcano Casado

<EMAIL_ADDRESS>Mathematical Institute. University of
Oxford

###### Abstract

Adaptive methods do not have a direct generalization to manifolds as the
adaptive term is not invariant. Momentum methods on manifolds suffer from
efficiency problems stemming from the curvature of the manifold. We introduce
a framework to generalize adaptive and momentum methods to arbitrary manifolds
by noting that for every differentiable manifold, there exists a radially
convex open set that covers almost all the manifold. Being radially convex,
this set is diffeomorphic to $\mathbb{R}^{n}$. This gives a natural
generalization of any adaptive and momentum-based algorithm to a set that
covers almost all the manifold in an arbitrary manifolds. We also show how to
extend these methods to the context of gradient descent methods with a
retraction. For its implementation, we bring an approximation to the
exponential of matrices that needs just of $5$ matrix multiplications, making
it particularly efficient on gpus. In practice, we see that this family of
algorithms closes the numerical gap created by an incorrect use of momentum
and adaptive methods on manifolds. At the same time, we see that the most
efficient algorithm of this family is given by simply pulling back the problem
to the tangent space at the initial point via the exponential map.

##  1 Introduction

Since the introduction of momentum methods by Nesterov [37], and adaptive
methods with adagrad [15], the deep learning field has seen an eruption of
methods tailored for non-convex stochastic optimization. Some of these
methods, like adagrad or amsgrad [43], offer theoretical guarantees under
standard assumptions. Others, like rmsprop [20] and adam [26], are
combinations of the former coupled with heuristics to improve their
performance.

At the same time, in the last years, optimization on manifolds has become ever
more present in the field of machine learning, in order to perform constrained
optimization. Manifolds are the natural object to treat symmetric positive
definite matrices in metric learning [52] and Bayesian statistics [42],
orthogonal matrices when performing matrix factorization and rnn
regularization [6], low-rank matrices for low-rank models [46], or invertible
matrices in normalizing flows [40].

There has been some work trying to reconcile both approaches. One of them came
by considering a surjective map from a Euclidean space onto the manifold
$\mathcal{M}$, allowing to pullback the problem from $\mathcal{M}$ to
$\mathbb{R}^{n}$ [29]. A different line of work has generalized these methods
to manifolds with a product structure [10]. However, it is still not clear how
to couple momentum and adaptive methods with Riemannian gradient descent on a
general manifold.

#### The problem

As pointed out in [[]Section $7$]lezcano2019trivializations, there is a
performance gap between gradient descent-like algorithms (dtriv$1$) and
algorithms based in pulling back the problem to a fixed tangent space
(dtriv$\infty$). This had also been noticed in the literature when using
adaptations of bfgs on manifolds. We aim to give an explanation for these
problems and offer a gradient descent-like algorithm whose performance matches
that of dtriv$\infty$ without using extra memory.

#### Summary of the contributions and main idea

We pinpoint two problems that the adaptive and momentum algorithms suffer when
considered on a Riemannian manifold

  1. 1.

Non-trivial tangent bundle. The adaptive term is not invariant under change of
coordinates

  2. 2.

Non-trivial holonomy. The momentum term depends on the optimization path

As it will be shown, the first problem completely prevents generalizing
adaptive methods to general manifolds. The second does not, but incurs a
penalty in practice. We show that these problems are intrinsic to the topology
of the manifold, so they cannot be fixed just by choosing the metric. To
circumvent them, we take a more modest approach: We will work on almost all
the manifold.

We use a deep result in Riemannian geometry that allows us to work on the
manifold as we do on $\mathbb{R}^{n}$. This holds on any connected complete
Riemannian manifold and as such, it may be implemented for any connected
differentiable manifold [39].

The second idea starts by computing all the gradients at a fixed tangent
space, as done in [29]. We then construct a family of linear maps that take
steepest descent directions to steepest descent directions to move the
gradient with the adaptive and momentum terms incorporated to the tangent
space of the current iterate. This will give us a way to move the momentum and
adaptive term without incurring in the error coming from holonomy.
Geometrically, we are taking advantage of the dense trivial contractible sub-
bundle of the tangent bundle given by the previous theorem.

Putting these two ideas together, we are able to generalize adaptive and
momentum algorithms like adam, adagrad, or rmsprop to manifolds in a way that
the practical performance of these algorithm matches that of the previous best
algorithms considered in the literature. We also give a fast implementation of
the exponential of matrices in gpus that might be of independent interest.

#### Limitations

Even though the resulting algorithm will converge to similar test accuracies,
the algorithms proposed make use of the logarithm of matrices to achieve so.
As such, even though the proposed algorithms solve the mentioned performance
gap, at the moment, regular dtriv$\infty$—pulling back the problem to a fixed
tangent space—is noticeably faster. This may change in the future when faster
gpu approximations to the logarithm of matrices become available.

##  2 Related Work

#### Riemannian gradient descent (rgd).

This algorithm was first introduced in the Riemannian setting in [33]. The
algorithm follows the steepest descent direction along a geodesic, as its
Euclidean counterpart. Following its inception, most algorithms proposed for
convex, non-convex, stochastic and non-stochastic, and constrained
optimization have found analogues in the Riemannian setting [16, 48, 1]. The
field is mostly divided in algorithms for geodesically convex optimization
[50, 54, 53], and those based on general retractions for non-convex functions
[1, 11, 12].

#### Trivializations.

A second approach to optimization on manifolds comes from considering a
surjective map from a Euclidean space onto the manifold and using it to
pullback the problem to a flat space. This approach has been used to perform
first order optimization [18, 31], gradient-free optimization [14], and
optimization of complex networks [40]. A general framework for this approach,
together with a connection to Riemannian gradient descent was proposed in
[29].

#### Adaptive and momentum methods on manifolds.

Even though most algorithms on Euclidean spaces have found a counterpart in
the Riemannian setting, many times via some bounds on the differential and
Hessian of the retraction, or through estimates given by a theorem of Rauch,
accelerated momentum methods have remained particularly elusive. The proposed
methods transport past gradients along geodesics and try to choose an
estimating sequence for the step-sizes. We will elaborate on the main idea
behind these methods in Section 3.

Current theoretical results on acceleration via momentum methods on manifolds
either require starting very close to the optimum [55], or require the
solution of a complex differential equation in every step [32], or do not
guarantee convergence to the minimum [3]. We find a similar situation when we
go to higher order. For example, there are many generalizations of bfgs to
manifolds by parallel transporting different vectors [22, 23], but their
performance in practice does not match that of their Euclidean counterparts.
Adaptive methods have recently been generalized to product manifolds [10] and
(with local convergence) to embedded matrix manifolds [25], but these methods
do not extend to general manifolds. Purely heuristic generalizations have also
been proposed, where the adaptive term is parallel-transported [44, 13].

##  3 The Issues of Momentum and Adaptive Methods on Manifolds

Consider throughout the rest of the paper a Riemannian manifold
$(\mathcal{M},\textsl{g})$ which is connected and complete.

###  3.1 Momentum Methods

Momentum methods are a generalization of Nesterov’s accelerated gradient
descent [37]. The momentum approach was first proposed in [45], and consists
of using a linear combination of previous gradients and the current gradient
as the descent direction. To be able to generalize these methods to manifolds,
we need to be able to combine gradients at different points. In order to do
so, a Riemannian manifold gives us a natural way of extending a vector to a
vector field.

######  Definition 3.1.

The adapted vector field $v^{\ast}(x)\in T_{x}\mathcal{M}$ of a vector $v\in
T_{p}\mathcal{M}$ is defined in a normal neighborhood $x\in U_{p}$ by parallel
transporting $v$ along the unique geodesic between $x$ and $p$.

Local existence of a normal neighborhood follows from the inverse function
theorem. As we will see later, these neighborhoods turn out to cover almost
all the manifold (, Theorem 4.1).

Using adapted vector fields, under suitable assumptions on the iterates
falling in a normal neighborhood, there are two direct ways to generalize
momentum methods to manifolds which have been extensively considered in the
literature.

#### Transport the momentum.

Let $g_{k}\coloneqq\nabla f(x_{k})\in T_{x_{k}}\mathcal{M}$ be the gradient at
the current iterate. We may define the momentum term as $m_{k}=g_{k}+\alpha
m^{\ast}_{k-1}(x_{k})\in T_{x_{k}}\mathcal{M}$ for a mixing constant
$\alpha>0$ and then use $-\eta m_{k}$ as the descent direction, for a learning
rate $\eta>0$. In other words, we transport the momentum from the previous
tangent space to the next one, and we update it there. This method, although
direct and cheap to implement, is just an approximation to the following
variation.

#### Choose the best approximation for the gradients.

This approach defines the momentum as
$m_{k}=\sum_{t=1}^{k}\alpha^{k-t}g^{\ast}_{t}(x_{k})$. In plain words, we
parallel transport all the previous gradients along the shortest geodesic to
the current iterate. This can geometrically be interpreted as considering the
best approximation at $x_{k}$ of each of the previous gradients according to
the metric that we are considering. Of course, this method has the drawback of
not being practical, as one has to store all the gradients from previous
steps.

In summary, the second approach gives a faithful approximation of momentum
methods with respect to the metric on the manifold at the expense of having to
save all the previous iterations. This not practical, so in practice mostly
the first approximation has been studied.

The relation between these two methods comes from the holonomy of the parallel
transport system. Holonomy can be thought of as the integral of the curvature,
as formalized by the Ambrose-Singer theorem [5], and it measures how much a
vector rotates when parallel-transported around a loop. As such, the
difference between the momentum from these two algorithms is dependent on the
optimization path through the holonomy of the parallel transport system. In
$\mathbb{R}^{n}$, these two algorithms coincide, and thus, one can get the
benefits of not having to store all the previous gradients at no extra price,
but that is not the case on a curved manifold. On a general manifold, the
first algorithm is just an approximation of the second one. Furthermore, this
is a topological problem of the manifold, so it cannot be fixed by just
changing the metric. We move the exploration of this idea to Appendix B, as it
requires more advanced geometric tools. The practical implications of this
approximation were first noted in [10], where report that the holonomy breaks
the sparsity of the gradients. In practice, as we mentioned in the
introduction, one can see the efficiency problems caused by this approximation
both in [29] and in the experiments section Section 5, where we compare the
algorithms in [29], which account for an approximation transporting the
momentum with the corrected ones that we propose later in this paper.

###  3.2 Adaptive Methods

Adaptive methods were introduced by Duchi, Hazan and Singer with adagrad [15].
They add a regularization to the size of the coordinates of the gradient
coming from optimizing a certain upper bound on the regret of an online
algorithm. It accounts to forming a matrix
$H_{k}\coloneqq\sum_{t=1}^{k}g_{t}g_{t}^{\intercal}$ and regularizing the
gradient as $\hat{g}_{k}=H_{k}^{-1/2}g_{k}$, to use $-\eta\hat{g}_{t}$ as the
descent direction. In the original paper, the authors show that the matrix
$H_{t}$ can be approximated by its diagonal, being able then to effectively
store it and compute its inverse, as the product $H_{t}^{-1/2}g_{t}$ becomes a
point-wise division by the square root of the diagonal terms of $H_{t}$.

The issue with generalizing adaptive methods to manifolds does not come from
an efficiency point of view as it happened with momentum methods, but from an
intrinsic topological problem: The adaptive term cannot be defined on a
manifold without storing $\mathcal{O}(n^{2})$ parameters. The product
$g_{k}g_{k}^{\intercal}$ is invariantly described on a manifold as the
$(1,1)$-tensor given by the tensor product
$g_{k}\otimes\lparen\mathrm{d}f\rparen_{x_{k}}$. The problem comes when
considering the diagonal of this tensor. Taking the diagonal is not invariant
under changes of coordinates, and as such, the adaptive term cannot be defined
globally on a manifold without global coordinates—it may be defined only on
$\mathbb{R}^{n}$.

##  4 Adaptive Methods Through Trivializations

We will introduce the general method that solves the two problems described in
Section 3 in two steps. We will first show how one can completely avoid these
two problems by using _static trivializations_. Then, we show how these ideas
can be generalized to the setting of _dynamic trivializations_. This setting
has both static trivializations and rgd as particular cases.

###  4.1 Static trivializations

We start by recalling the static trivialization framework [29]. We state it
for case of the Riemannian exponential map to make the ideas clearer, although
it can be performed with more general maps. For a fixed $p\in\mathcal{M}$,
static trivializations pullback the problem from $\mathcal{M}$ to
$T_{p}\mathcal{M}$, going from a constrained problem on $\mathcal{M}$ to an
unconstrained problem on $T_{p}M(\cong\mathbb{R}^{n})$

$\min_{x\in\mathcal{M}}f(x)\qquad\xrightarrow{\text{trivialization}}\qquad\min_{v\in
T_{p}M}f(\exp_{p}(v)).$

As proved in [29], the two problems are equivalent to a first order in the
region in which $\exp_{p}$ is a local diffeomorphism. In this region, static
trivializations account for a change of metric in $\mathcal{M}$ and solving
the second problem via gradient descent is equivalent to solving the first one
using rgd with respect to a different metric.

On manifolds with non-positive sectional curvature, by the Cartan-Hadamard
theorem, $\exp_{p}$ is a local diffeomorphism around any point in
$T_{p}\mathcal{M}$, so $\exp_{p}$ acts globally as a change of metric. In the
general case, we cannot get the result on all the manifold, but we can still
get it on almost all of it.

######  Theorem 4.1.

Let $(\mathcal{M},\textsl{g})$ be a complete and connected Riemannian
manifold. Let $\overline{U}_{p}$ be the (closed) set of vectors $v\in
T_{p}\mathcal{M}$ for which $\gamma(t)=\exp_{p}(tv),t\in[0,1]$ is length
minimizing, and let $U_{p}$ be its interior and $C_{p}$ its boundary. Then,
$\exp_{p}$ is a diffeomorphism on $U_{p}$ and $\exp_{p}(C_{p})$ has measure
zero with respect to the Borel measure induced by g.

###### Proof.

See [41, Section 5.7.3] for the first part. The second follows from [24]. ∎

This theorem says that $\exp_{p}$ is a diffeomorphism on almost all the
manifold, and as such, $\exp_{p}$ does not create critical points on a large
enough neighborhood of $p$ diffeomorphic to $\mathbb{R}^{n}$. Furthermore, as
$T_{p}\mathcal{M}\cong\mathbb{R}^{n}$, it is clear that this approach readily
generalizes momentum and adaptive methods to almost all the manifold (or in
non-positive curvature to all the manifold) as one may use adam or adagrad to
optimize the function $f\circ\exp_{p}\colon T_{p}\mathcal{M}\to\mathbb{R}$.

###  4.2 Dynamic trivializations

We now recall the idea behind _dynamic trivializations_ [29] as a connection
between rgd and static trivializations, and we use this connection to
generalize momentum and adaptive methods to rgd.

Dynamic trivializations arise naturally from the following observation: In
Section 4.1, we considered the Riemannian exponential $\exp_{p}$ at a fixed
point $p\in\mathcal{M}$ as a trivialization, but the choice of the point $p$
is completely arbitrary. Dynamic trivializations repeatedly change the
trivialization point $p$. They take $K$ steps in $T_{p_{i}}\mathcal{M}$ and
after the $K$-th iterate $v_{i,K}\in T_{p_{i}}\mathcal{M}$, they change the
trivialization point to $p_{i+1}\coloneqq\exp_{p_{i}}(v_{i,K})$ and set
$v_{i+1,0}=0$. In other words, they set the new trivialization point to the
current iterate on the manifold and start optimizing on its tangent space.

Dynamic trivializations generalize both static trivializations and rgd. For
$K=\infty$—, never changing the trivialization—dynamic trivializations are
exactly static trivializations at $\exp_{p_{0}}$. For $K=1$, dynamic
trivializations recover rgd [29].

Adaptive and momentum methods can be implemented through static
trivializations. We will now use the fact that dtriv$K$ interpolates between
rgd and static trivializations to generalize adaptive and momentum methods to
the setting of rgd.

As we argued in Section 3, the adaptive term is not invariant under changes of
coordinates, and as such, it cannot be defined after a change of basis. For
this reason, the adaptive term has to always be computed at the same tangent
space. This can be done simply as it is done, as in the setting of static
trivializations, by computing the gradient of $f\circ\exp_{p_{0}}$ and
computing the momentum and adaptive term at $p_{0}$ following the rule of adam
or rmsprop or any other adaptive optimizer. From this, we get an adapted
descent direction $\hat{g}_{t}\in T_{p_{0}}\mathcal{M}$. To transport this
vector to $T_{p_{i}}\mathcal{M}$, we shall use the map induced by the dynamic
trivialization algorithm, as per the following proposition.

######  Proposition 4.2.

Let $(\mathcal{M},\textsl{g})$ be a Riemannian manifold. Let $f$ be any
differentiable function $f$ on $\mathcal{M}$, and let
$p_{0},p_{i}\in\mathcal{M}$ be some points and $v_{i,k}\in
T_{p_{i}}\mathcal{M}$ with $x_{t}\coloneqq\exp_{p_{i}}(v_{i,k})$ such that
$p_{0},p_{i}\in U_{x_{t}}$ with $U_{x_{t}}$ as in Theorem 4.1. Then, the map

$\mathrm{d}\lparen\exp^{-1}_{p_{0}}\circ\exp_{p_{i}}\rparen^{\ast}_{v_{i,k}}\colon
T_{p_{0}}\mathcal{M}\to T_{p_{i}}\mathcal{M}$

were $(-)^{\ast}$ is the adjoint, takes directions of steepest descent of
$f\circ\exp_{p_{0}}$ to directions of steepest descent of
$f\circ\exp_{p_{i}}$. Furthermore, it is the only map between
$T_{p_{0}}\mathcal{M}$ and $T_{p_{i}}\mathcal{M}$ with this property.

###### Proof.

The hypothesis on $p_{0},p_{i}$ and Theorem 4.1 assure that the operator
$\mathrm{d}\lparen\exp^{-1}_{p_{0}}\circ\exp_{p_{i}}\rparen_{v_{i,k}}$ exists
and is an isomorphism of linear spaces. By the chain rule, letting
$w_{i,k}\coloneqq\exp^{-1}_{p_{0}}(x_{t})$ we have that

$\mathrm{d}\lparen f\circ\exp_{p_{i}}\rparen_{v_{i,k}}=\mathrm{d}\lparen
f\circ\exp_{p_{0}}\rparen_{w_{i,k}}\circ\mathrm{d}\lparen\exp_{p_{0}}^{-1}\circ\exp_{p_{i}}\rparen_{v_{i,k}}$

and taking adjoints

$\displaystyle\nabla\lparen f$
$\displaystyle\circ\exp_{p_{i}}\rparen(v_{i,k})=$
$\displaystyle\mathrm{d}\lparen\exp_{p_{0}}^{-1}\circ\exp_{p_{i}}\rparen^{\ast}_{v_{i,k}}\lparen\nabla\lparen
f\circ\exp_{p_{0}}\rparen\lparen w_{i,k}\rparen\rparen.\qed$

In other words, we compute the gradient together with the momentum and
adaptive term at $T_{p_{i}}\mathcal{M}$ and then we transport it to our
current tangent space using the linear map that preserves directions of
steepest descent. When putting together the previous discussion with this
proposition we get the algorithm atriv$K$ as described in Algorithm 1. Note
that for $K=\infty$, $\textsc{atriv}\infty$ is exactly the same as
$\textsc{dtriv}\infty$, as one would expect.

#### Comparison against the classic generalizations of rgd.

A clear geometric picture arises when comparing atriv$1$ with the methods
presented in Section 3. The second method in that section can be interpreted
as considering the flat metric from $T_{x_{t}}\mathcal{M}$ to transport the
previous gradients from $T_{x_{i}}\mathcal{M}$ to that tangent space. This
method has the problem of having to store all the previous gradients. Here, we
substitute the moving metric at $T_{x_{t}}\mathcal{M}$ by the use of a fixed
one at $T_{x_{0}}\mathcal{M}$, hence being able to accumulate the gradients.
Note also the first algorithm introduced in that section allows implementing
momentum method at the expense of suffering from the holonomy problems noted
in [10]. Adaptive trivializations do not suffer these problems since they use
a metric with trivial holonomy. Furthermore, atriv$K$ allows for being coupled
with adaptive methods, while the classic generalizations do not.

#### A note on general retractions.

The case for a general retraction is analogue and we have spared the reader
the details for concreteness. For every retraction and every
$p\in\mathcal{M}$, we know that there exists a neighborhood $U_{p}$ such that
$r_{p}$ is invertible, although in general it will not be as large as the one
given by the exponential. For this reason, we may be able to perform this
algorithm just locally. If the optimization process leaves that neighborhood
we can apply a restart process on the momentum and adaptive term. These
restarts are rather common in optimization algorithms with momentum (, [4]).
Solutions of this flavor had already been considered in the literature [47].

Algorithm 1 Adaptive dynamic trivialization

1:$K>0$, $p_{0}\in\mathcal{M}$. An adaptive optimizer (, adam, rmsprop, …)

2:for $i=0,\ldots$ do

3: $v_{i,0}=0$

4: for $k=0,\dots,K-1$ do$\triangleright$ Optimize on $T_{p_{i}}\mathcal{M}$

5: $w_{i,k}=\lparen\exp^{-1}_{p_{0}}\circ\exp_{p_{i}}\rparen(v_{i,k})$

6: $g_{t}=\nabla\lparen f\circ\exp_{p_{0}}\rparen\lparen w_{i,k}\rparen$

7: Compute $\hat{g}_{t}$ using the optimizer, aggregating the adaptive term on
$T_{p_{0}}\mathcal{M}$

8:
$\tilde{g}_{t}=\mathrm{d}\lparen\exp_{p_{0}}^{-1}\circ\exp_{p_{i}}\rparen^{\ast}_{v_{i,k}}(\hat{g}_{t})$

9: $v_{i,k+1}=v_{i,k}-\eta\tilde{g}_{t}$

10: $p_{i+1}=\exp_{p_{i}}(v_{i,K})$ $\triangleright$ Update trivialization
point

#### Convergence guarantees

In the case when the optimizer is gradient descent, we can give exact
guarantees for the algorithm by using first and second order bounds on the
Riemannian exponential.

######  Theorem 4.3.

Let $(\mathcal{M},\textsl{g})$ be a connected and complete Riemannian manifold
with zeroth and first order curvature bounds $\delta\leq\sec\leq\Delta$ and
$\lVert\nabla R\rVert\leq\Lambda$. Consider a connected submanifold
$\mathcal{X}$ with $\operatorname{diam}(\mathcal{X})=r$ where
$r\leq\frac{\sqrt{2}\pi}{\sqrt{\delta+\Delta}}$ if $\Delta>-\delta$, and let
$f$ be a function on $\mathcal{X}$ bounded below by $f^{\ast}\in\mathbb{R}$
such that $\lVert\operatorname{Hess}f\rVert\leq\alpha$. Then, there exists a
constant $\widehat{\alpha}_{r}$ that does not depend on the dimension of
$\mathcal{M}$ such that $\lVert\operatorname{Hess}\lparen
f\circ\exp_{p}\rparen|_{\exp_{p}^{-1}\lparen\mathcal{X}\rparen}\rVert\leq\widehat{\alpha}_{r}$
for any $p\in\mathcal{X}$. Therefore, running atriv$K$ with gd and
$\eta=\frac{1}{\widehat{\alpha}_{r}}$, if all the iterates stay in
$\mathcal{X}$, the algorithm will find a point $v_{i,k}\in
T_{p_{i}}\mathcal{M}$ such that $\lVert\nabla\lparen
f\circ\exp_{p_{i}}\rparen(v_{i,k})\rVert<\varepsilon$ in at most
$\Big{\lceil}\mbox{$\displaystyle\frac{2\widehat{\alpha}_{r}}{\varepsilon^{2}}$}\lparen
f(p_{0})-f^{\ast}\rparen\Big{\rceil}$ steps.

###### Proof.

The result follows from [30, Proposition $5.1$] and the classical theory of
non-convex optimization [38]. ∎

The bound $\widehat{\alpha}_{r}$ is an explicit function of
$r,\delta,\Delta,\Lambda$ and takes a very simple form for most manifolds used
in optimization. For example, when $\mathcal{M}$ is the special orthogonal
group $\operatorname{SO}\lparen n\rparen$ embedded in $\mathbb{R}^{n\times
n}$, it takes the form $\widehat{\alpha}_{r}=(1+\frac{r}{3})\alpha$. In other
words, this result recovers the $\mathcal{O}(\varepsilon^{-2})$ convergence
rate of the Euclidean case modulo a curvature term and a linear size term.

The hypothesis about the iterates stay in $\mathcal{X}$, although restrictive,
is standard when dealing with manifolds [11, 53, 47, 49, 2].

#### Faster exponential of matrices.

The exponential of matrices is fundamental when dealing with the exponential
map in optimization on manifolds. On the other hand, implementing correctly
the exponential of matrices is not an easy task [35, 36]. The current standard
implementation in Scipy is based on the paper [34]. This approximation uses an
lu decomposition and the solution of a triangular system, which are sequential
algorithms.

In order to have a fast implementation in gpu of the exponential of matrices,
we implement one that is based on matrix multiplications—a polynomial
approximation. The first algorithm of this kind with numerical accuracy
guarantees was presented in the recent work by [9]. Their algorithm consists
of the scale-squaring trick with a clever evaluation of Taylor approximants.
This evaluation, requires $5$ matrix multiplications, making the exponential
of matrices a very cheap function use in the context of deep learning. We
helped PyTorch maintainers to add this algorithm to PyTorch, getting a
$14\times$ wall-clock speed improvement over the Tensorflow implementation. It
has been released in PyTorch 1.7.0.

##  5 Experiments

###  5.1 Experiment set-up

In this section, we compare the practical performance of the approach
presented in this paper with the state-of-the-art algorithms for stochastic
Riemannian optimization with orthogonal constraints for neural networks. In
particular, we compare it with dynamic trivializations [[,
dtriv]]lezcano2019trivializations, exponential trivialization [[,
exprnn]]lezcano2019cheap, Cayley trivializations [[,
scornn]]helfrich2018orthogonal and Riemannian gradient descent [[,
rgd]]wisdom2016full. We also include an lstm for reference.

#### Relations between these methods

All these methods fall under the dynamic trivialization framework. In
particular, there are some relations between them: rgd is dtriv$1$ using sgd
as the optimizer, exprnn is dtriv$\infty$ using the identity matrix as the
basis for the trivialization, and atriv$K$ is dtriv$K$ but with the momentum
and adaptive terms in the optimizer corrected. atriv$\infty$ and dtriv$\infty$
are the same algorithm. We summarize these relations in Table 1.

Table 1: Algorithms to optimize with orthogonal constraints Model | Optimizer | Retraction | Adapt. | Type
---|---|---|---|---
scornn | Any | $\operatorname*{cay}_{\mathrm{I}}$ | Yes | Static
exprnn | Any | $\exp_{\mathrm{I}}$ | Yes | Static
dtriv$\infty$ | Any | $\exp_{B_{0}}$ | Yes | Static
rgd | sgd | $\operatorname*{cay}_{B}$ | — | Dyn.
dtriv$n$ | Any | $\exp_{B}$ | No | Dyn.
atriv$n$ | Any | $\exp_{B}$ | Yes | Dyn.

In the Adapt. column in Table 1, we note which models treat adaptive methods
correctly. We have written $\varphi_{B}(A)=B\varphi(B^{-1}A)$ where
$\varphi=\exp$ or $\varphi=\operatorname*{cay}$, $A\in
T_{B}\operatorname{SO}\lparen n\rparen$, $B\in\operatorname{SO}\lparen
n\rparen$.

The model on which we perform the experiments is a vanilla rnn together with
an orthogonal constraint on the recurrent kernel. The orthogonality of $B$ is
achieved following the methods in Table 1. The non-linearity used is the
modrelu, introduced for these models in [6].

We include an implementation of the general framework in the accompanying
material. We describe in detail all the hyperparameters needed to replicate
the experiments in Appendix C.

#### Orthogonal rnns in practice

In practice, we observed that these models benefit from a lower mixing
constant in the adaptive term than the default used in adam. In particular, a
mixing constant of $\beta_{2}=0.99$ or even $\beta_{2}=0.95$ improved the
stability of the algorithms. This makes a lot of sense since, even though the
recurrent layer is orthogonal and the non-linearity is piecewise linear, so
that the recurrent gradients are bounded, the input layer might still make the
gradients unstable at times. Previous papers reported a better performance of
rmsprop over adam [6, 18]. This is explained by the fact that the mixing
constant of rmsprop is initialized by default to $0.99$, while that of adam is
initialized to $0.999$. We also observed that using these lower mixing
constants, one can almost duplicate the learning rate on the orthogonal
parameters.

We do not make use of these observations to improve our results in the
experiments section, as we are not trying to get the best result, but merely
to compare the different algorithms to optimize with orthogonal constraints.

We test these algorithms on the two standard problems that the community has
been using to test stochastic optimization with orthogonal constraints in the
context of neural networks, namely sequential mnist and timit [6, 19, 18, 31,
29].

######  Remark (Copying and adding problems).

We do not include the adding problem or the copying problem [21] which are
often considered, as these problems are too simple for recurrent models [19].
All the considered optimization methods converge to a correct answer in the
copying and adding problems.

###  5.2 Sequential mnist

This task consists on classifying the numbers from $0$ to $9$ present in the
mnist dataset [28]. This dataset is formed by $60.000$ $28\times 28$ grayscale
images. The sequential task mnist processes each of the $784$ pixels of the
image sequentially. We also consider the permuted sequential task p-mnist,
were a permutation of the bits is sampled at the beginning of the training and
fixed once and for all. This permutation is then used to permute the pixels of
all the images before flattening them.

Table 2: Best test accuracy at mnist and p-mnist. Model | n | mnist | p-mnist
---|---|---|---
atriv$1$ | $170$ | $98.2$ | $95.1$
dtriv$\infty$ | $170$ | $98.1$ | $95.0$
dtriv$1$ | $170$ | $\mathbf{98.3}$ | $\mathbf{95.2}$
exprnn | $170$ | $98.0$ | $94.9$
scornn | $170$ | $97.2$ | $94.8$
scurnn | $116$ | $97.6$ | $94.9$
lstm | $128$ | $81.9$ | $79.5$
rgd | $116$ | $94.7$ | $92.5$
atriv$1$ | $360$ | $\mathbf{98.9}$ | $96.4$
dtriv$\infty$ | $360$ | $\mathbf{98.9}$ | $\mathbf{96.5}$
dtriv$1$ | $360$ | $98.4$ | $96.3$
exprnn | $360$ | $98.4$ | $96.2$
scornn | $360$ | $98.1$ | $95.9$
scurnn | $250$ | $98.3$ | $96.2$
lstm | $256$ | $88.8$ | $88.8$
rgd | $256$ | $96.1$ | $93.9$
atriv$1$ | $512$ | $\mathbf{99.0}$ | $96.7$
dtriv$\infty$ | $512$ | $\mathbf{99.0}$ | $\mathbf{96.8}$
dtriv$1$ | $512$ | $98.7$ | $96.7$
exprnn | $512$ | $98.7$ | $96.6$
scornn | $512$ | $98.2$ | $96.5$
lstm | $512$ | $91.9$ | $91.8$
rgd | $512$ | $97.3$ | $94.7$

We present the results of this experiments in Table 2, where we have denoted
by N the size of the hidden layer. This is chosen so that all the models in
each group have the same number of parameters.

Before analyzing the results, we recall that atriv$1$ is exactly the same
algorithm as dtriv$1$ but with a correct treatment of the adaptive term as per
Algorithm 1. dtriv$\infty$ and atriv$1$ both handle correctly the adaptive
term, and they both are algorithms of the same family, since dtriv$\infty$ is
the same algorithm as atriv$\infty$. In general, we observed a consistent
performance across the algorithms of the family atriv$K$ when varying $K$.

In this table we see two phenomena. First, we see that the corrected treatment
of the adaptive term in atriv$1$ gives a consistent improvement over dtriv$1$
of this algorithm making it match the performance of its static counterpart
dtriv$\infty$. Furthermore, we see that all the dynamic trivializations yield
a comparable performance when treating the adaptive term correctly, that is,
dtriv$\infty$ and atriv$1$ yield a similar perform. This hints that there is
no real improvement when it comes to dtriv$\infty$ or atriv$1$. As such, it
would be recommendable to use dtriv$\infty$, since by not requiring to change
the basis, its implementation is simpler and it is considerably faster.

###  5.3 timit speech dataset

The timit dataset is composed by real-world speech recording [17]. This is a
regression task in which, given the $n$ first points of a sequence, we want to
predict the point $n+1$. The sequences of audio are preprocessed by
downsampling them to $8$kHz. Then, they are represented in a log-scale after
applying a short-time Fourier transform. This process generates variable-
length sequences (with length between $61$ and $490$) of $129$ complex numbers
at every step. This experiment and set-up was first proposed in [51]. We
present the results of this experiment in Table 3.

Table 3: Test mse at the end of the epoch with the lowest validation mse for the timit task. Model | n | Val. mse | Test mse
---|---|---|---
atriv$1$ | $224$ | $\mathbf{4.52}$ | $\mathbf{4.50}$
dtriv$\infty$ | $224$ | $4.75$ | $4.71$
dtriv$1$ | $224$ | $6.55$ | $6.54$
exprnn | $224$ | $5.34$ | $5.30$
scornn | $224$ | $9.26$ | $8.50$
scurnn | $128$ | $9.42$ | $7.23$
lstm | $84$ | $15.42$ | $14.30$
rgd | $128$ | $15.07$ | $14.58$
atriv$1$ | $322$ | $\mathbf{3.33}$ | $\mathbf{3.29}$
dtriv$\infty$ | $322$ | $3.39$ | $3.76$
dtriv$1$ | $322$ | $4.56$ | $4.55$
exprnn | $322$ | $4.42$ | $4.38$
scornn | $322$ | $8.48$ | $7.82$
lstm | $120$ | $13.93$ | $12.95$
rgd | $192$ | $15.10$ | $14.50$
atriv$1$ | $425$ | $\mathbf{1.93}$ | $\mathbf{1.92}$
dtriv$\infty$ | $425$ | $2.00$ | $1.97$
dtriv$1$ | $425$ | $4.21$ | $4.17$
exprnn | $425$ | $5.52$ | $5.48$
scornn | $425$ | $7.97$ | $7.36$
scurnn | $258$ | $4.40$ | $3.39$
lstm | $158$ | $13.66$ | $12.62$
rgd | $256$ | $14.96$ | $14.69$

In this experiment we observe the same behavior of that present in the mnist
task. We see a consistent improvement of (atriv$1$) over the version of the
algorithm that does not correct the adaptive term (dtriv$1$). Furthermore, we
also see that the correction in the adaptive term (atriv$1$) yields a
comparable performance to that of using the point at the first tangent space
to trivialize (dtriv$\infty$).

##  6 Conclusion and Future Work

In this work we present a generalization of adaptive and momentum-based
methods to manifolds by making use of the fact that any manifold $\mathcal{M}$
has a radially convex submanifold that covers almost all $\mathcal{M}$ and
that is diffeomorphic to a subset of $\mathbb{R}^{n}$. As such, the subset
that is not on the image of this diffeomorphism has measure zero, and
restricting our optimization methods to this open subset serves as a tight
approximation to the original problem.

We see that by using the generalization of momentum and adaptive methods given
by a flat connection in the interior of the segment domain, we achieve
comparable performance in practice to that of pulling back the optimization
problem along the Riemannian exponential at a given point. In particular,
correcting the momentum and adaptive term in this fashion bridges the
performance gap created by naïvely using adaptive methods on top of dynamic
trivializations.

We expect that the map considered here can become a useful alternative to
parallel-transporting the momentum or other vectors in other optimization
methods on manifolds.

## Acknowledgements

The work of MLC was supported by the Oxford-James Martin Graduate Scholarship.

## References

  * [1] Pierre-Antoine Absil, Robert Mahony and Rodolphe Sepulchre  “Optimization algorithms on matrix manifolds”  Princeton University Press, 2009
  * [2] Kwangjun Ahn and Suvrit Sra  “From Nesterov’s Estimate Sequence to Riemannian Acceleration”  In _arXiv preprint arXiv:2001.08876_ , 2020
  * [3] Foivos Alimisis, Antonio Orvieto, Gary Bécigneul and Aurelien Lucchi  “Practical accelerated optimization on Riemannian manifolds”  In _https://arxiv.org/abs/2002.04144_ , 2020
  * [4] Zeyuan Allen-Zhu and Lorenzo Orecchia  “Linear coupling: An ultimate unification of gradient and mirror descent”  In _Innovations in Theoretical Computer Science Conference, ITCS_ , 2017, pp. 3:1–3:22
  * [5] W. Ambrose and I.. Singer  “A theorem on holonomy”  In _Transactions of the American Mathematical Society_ 75, 1953, pp. 428–443  DOI: 10.2307/1990721
  * [6] Martin Arjovsky, Amar Shah and Yoshua Bengio  “Unitary evolution recurrent neural networks”  In _International Conference on Machine Learning, ICML_ , 2016, pp. 1120–1128
  * [7] Vincent Arsigny, Olivier Commowick, Xavier Pennec and Nicholas Ayache  “A log-Euclidean framework for statistics on diffeomorphisms”  In _International Conference on Medical Image Computing and Computer-Assisted Intervention, MICCAI_ , 2006, pp. 924–931
  * [8] Vincent Arsigny, Pierre Fillard, Xavier Pennec and Nicholas Ayache  “Geometric means in a novel vector space structure on symmetric positive-definite matrices”  In _SIAM Journal on Matrix Analysis and Applications_ 29.1, 2006, pp. 328–347
  * [9] Philipp Bader, Sergio Blanes and Fernando Casas  “Computing the matrix exponential with an optimized Taylor polynomial approximation”  In _Mathematics_ 7.12, 2019
  * [10] Gary Becigneul and Octavian-Eugen Ganea  “Riemannian adaptive optimization methods”  In _International Conference on Learning Representations, ICLR_ , 2019
  * [11] Silvère Bonnabel  “Stochastic gradient descent on Riemannian manifolds”  In _IEEE Transactions on Automatic Control_ 58.9, 2013, pp. 2217–2229
  * [12] Nicolas Boumal, Pierre-Antoine Absil and Coralia Cartis  “Global rates of convergence for nonconvex optimization on manifolds”  In _IMA Journal of Numerical Analysis_ 39.1, 2019, pp. 1–33
  * [13] Minhyung Cho and Jaehyung Lee  “Riemannian approach to batch normalization”  In _Advances in Neural Information Processing Systems, NeurIPS_ , 2017, pp. 5225–5235
  * [14] David W. Dreisigmeyer  “Direct search methods on reductive homogeneous spaces”  In _Journal of Optimization Theory and Applications_ 176.3, 2018, pp. 585–604
  * [15] John Duchi, Elad Hazan and Yoram Singer  “Adaptive subgradient methods for online learning and stochastic optimization”  In _Journal of Machine Learning Research, JMLR_ 12.Jul, 2011, pp. 2121–2159
  * [16] Alan Edelman, Tomás A. Arias and Steven Thomas Smith  “The geometry of algorithms with orthogonality constraints”  In _SIAM Journal on Matrix Analysis and Applications_ 20.2, 1998, pp. 303–353
  * [17] J.. Garofolo et al.  “DARPA TIMIT acoustic phonetic continuous speech corpus”  Linguistic Data Consortium, 1993
  * [18] Kyle Helfrich, Devin Willmott and Qiang Ye  “Orthogonal recurrent neural networks with scaled Cayley transform”  In _International Conference on Machine Learning, ICML_ , 2018, pp. 1974–1983
  * [19] Mikael Henaff, Arthur Szlam and Yann LeCun  “Recurrent orthogonal networks and long-memory tasks”  In _International Conference on Machine Learning, ICML_ , 2016, pp. 2034–2042
  * [20] Geoffrey E. Hinton  “Lecture 6.5. RMSprop: Divide the gradient by a running average of its recent magnitude” Course: _Neural networks for machine learning_ , 2012  URL: http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
  * [21] Sepp Hochreiter and Jürgen Schmidhuber  “Long short-term memory”  In _Neural Computation_ 9.8, 1997, pp. 1735–1780
  * [22] Wen Huang, Pierre-Antoine Absil and Kyle A. Gallivan  “A Riemannian BFGS method for nonconvex optimization problems”  In _Numerical Mathematics and Advanced Applications, ENUMATH_ , 2015, pp. 627–634
  * [23] Wen Huang, Pierre-Antoine Absil and Kyle A. Gallivan  “A Riemannian BFGS method without differentiated retraction for nonconvex optimization problems”  In _SIAM Journal on Optimization_ 28.1, 2018, pp. 470–495  DOI: 10.1137/17M1127582
  * [24] Jin-ichi Itoh and Minoru Tanaka  “The dimension of a cut locus on a smooth Riemannian manifold”  In _Tohoku Mathematical Journal_ 50.4, 1998, pp. 571–575
  * [25] Hiroyuki Kasai, Pratik Jawanpuria and Bamdev Mishra  “Riemannian adaptive stochastic gradient algorithms on matrix manifolds”  In _International Conference on Machine Learning, ICML_ , 2019, pp. 3262–3271
  * [26] Diederik P. Kingma and Jimmy Ba  “Adam: A method for stochastic optimization”  In _International Conference on Learning Representations, ICLR_ , 2015
  * [27] Shoshichi Kobayashi and Katsumi Nomizu  “Foundations of differential geometry”  Interscience Publishers, John Wiley & Sons, 1963
  * [28] Yann LeCun, Corinna Cortes and Christopher C.J. Burges  “The MNIST database of handwritten digits” Accessed: 2019-02-06, http://yann.lecun.com/exdb/mnist/, 1998
  * [29] Mario Lezcano-Casado  “Trivializations for gradient-based optimization on manifolds”  In _Advances in Neural Information Processing Systems, NeurIPS_ , 2019, pp. 9154–9164
  * [30] Mario Lezcano-Casado  “Curvature-Dependant Global Convergence Rates for Optimization on Manifolds of Bounded Geometry”  In _https://arxiv.org/abs/2008.02517_ , 2020
  * [31] Mario Lezcano-Casado and David Martínez-Rubio  “Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group”  In _International Conference on Machine Learning, ICML_ , 2019, pp. 3794–3803
  * [32] Yuanyuan Liu et al.  “Accelerated first-order methods for geodesically convex optimization on Riemannian manifolds”  In _Advances in Neural Information Processing Systems, NeurIPS_ , 2017, pp. 4868–4877
  * [33] David G. Luenberger  “The gradient projection method along geodesics”  In _Management Science_ 18, 1972, pp. 620–631
  * [34] Awad H. Al-Mohy and Nicholas J. Higham  “A new scaling and squaring algorithm for the matrix exponential”  In _SIAM Journal on Matrix Analysis and Applications_ 31.3, 2009, pp. 970–989
  * [35] Cleve Moler and Charles Van Loan  “Nineteen dubious ways to compute the exponential of a matrix”  In _SIAM Review_ 20.4, 1978, pp. 801–836
  * [36] Cleve Moler and Charles Van Loan  “Nineteen dubious ways to compute the exponential of a matrix, twenty-five years later”  In _SIAM Review_ 45.1, 2003, pp. 3–49
  * [37] Yurii Nesterov  “A method for solving the convex programming problem with convergence rate ${O}(1/k^{2})$”  In _Soviet Mathematics Doklady_ 27, 1983, pp. 372–376
  * [38] Yurii Nesterov  “Introductory lectures on convex optimization: A basic course” 87, Applied Optimization  Kluwer Academic Publishers, 2004  DOI: 10.1007/978-1-4419-8853-9
  * [39] Katsumi Nomizu and Hideki Ozeki  “The existence of complete Riemannian metrics”  In _Proceedings of the American Mathematical Society_ 12.6, 1961, pp. 889–891
  * [40] George Papamakarios et al.  “Normalizing flows for probabilistic modeling and inference”  In _arXiv preprint arXiv:1912.02762_ , 2019
  * [41] Peter Petersen  “Riemannian geometry”, Graduate Texts in Mathematics  Springer, 2016
  * [42] Carl Edward Rasmussen and Christopher K.. Williams  “Gaussian processes for machine learning”, Adaptive Computation and Machine Learning  The MIT Press, 2005
  * [43] Sashank J. Reddi, Satyen Kale and Sanjiv Kumar  “On the convergence of Adam and beyond”  In _International Conference on Learning Representations, ICLR_ , 2018
  * [44] Soumava Kumar Roy, Zakaria Mhammedi and Mehrtash Harandi  “Geometry aware constrained optimization techniques for deep learning”  In _Computer Vision and Pattern Recognition, CVPR_ , 2018
  * [45] David E Rumelhart, Geoffrey E Hinton and Ronald J Williams  “Learning representations by back-propagating errors”  In _nature_ 323.6088 Nature Publishing Group, 1986, pp. 533–536
  * [46] Amartya Sanyal, Philip H. Torr and Puneet K. Dokania  “Stable rank normalization for improved generalization in neural networks and GANs”  In _International Conference on Learning Representations, ICLR_ , 2020
  * [47] Hiroyuki Sato, Hiroyuki Kasai and Bamdev Mishra  “Riemannian stochastic variance reduced gradient algorithm with retraction and vector transport”  In _SIAM Journal on Optimization_ 29.2, 2019, pp. 1444–1472  DOI: 10.1137/17M1116787
  * [48] Steven Thomas Smith  “Geometric optimization methods for adaptive filtering”, 1993
  * [49] Nilesh Tripuraneni, Nicolas Flammarion, Francis Bach and Michael I. Jordan  “Averaging Stochastic Gradient Descent on Riemannian Manifolds”  In _Conference on Learning Theory, COLT_ , 2018, pp. 650–687
  * [50] Constantin Udrişte  “Convex functions and optimization methods on Riemannian manifolds” 297, Mathematics and its Applications  Springer, 1994
  * [51] Scott Wisdom et al.  “Full-capacity unitary recurrent neural networks”  In _Advances in Neural Information Processing Systems, NeurIPS_ , 2016, pp. 4880–4888
  * [52] Pourya Zadeh, Reshad Hosseini and Suvrit Sra  “Geometric mean metric learning”  In _International Conference on Machine Learning, ICML_ , 2016, pp. 2464–2471
  * [53] Hongyi Zhang, Sashank J. Reddi and Suvrit Sra  “Riemannian SVRG: Fast stochastic optimization on Riemannian manifolds”  In _Advances in Neural Information Processing Systems, NeurIPS_ , 2016, pp. 4592–4600
  * [54] Hongyi Zhang and Suvrit Sra  “First-order methods for geodesically convex optimization”  In _Conference on Learning Theory, COLT_ , 2016, pp. 1617–1638
  * [55] Hongyi Zhang and Suvrit Sra  “An estimate sequence for geodesically convex optimization”  In _Conference on Learning Theory, COLT_ , 2018, pp. 1703–1723

##  Appendix A Review of Vector Bundles and Connections

In this section we summarize some standard results on theory of connections on
vector bundles and get up to parallel transport systems and their relation
with the usual definition of a covariant derivative. We choose to later use
this framework rather than that of Riemannian metrics since all the problems
that we encounter can already be seen in this more abstract framework, and one
may see more geometrically the geometric obstructions to some constructions.

These results and definitions are all classic. A good reference for all of
them is the book [[]Chapters II, III]kobayashi1963foundations.

#### Vector bundles.

A _vector bundle $E$ of rank $k$ over $M$_ consists of a pair of manifolds
$E,M$ together with a submersion $\pi\colon E\to M$ such that for every $x\in
M$, the fibre $E_{x}\coloneqq\pi^{-1}(\\{x\\})$ has a vector space structure,
and for every $x\in M$ there exists a neighborhood $U\subseteq M$ and a linear
isomorphism $t_{x}$ called a _local trivialization_ such that
$(\pi,t_{x})\colon\pi^{-1}(U)\to U\times\mathbb{R}^{k}$ is an isomorphism.
This isomorphism indicates how to locally glue the fibers together.
Intuitively, $E$ is a collection of vector spaces which represents the idea of
having glued a vector space of dimension $k$ at every point of $M$. $E$ has
locally the structure of product of manifolds. The tangent bundle $TM$ of a
manifold $M$ of dimension $n$ is an example of a vector space of rank $n$ over
$M$. The cylinder $S^{1}\times\mathbb{R}$ and the Möbius bundle (a Möbius band
of infinite width) are two rank $1$ vector bundles over the circle $S^{1}$.
The cylinder is a trivial bundle, as it is globally diffeomorphic to a product
of spaces, while the Möbius bundle is a non-trivial bundle, as it is just
locally diffeomorphic to $U\times\mathbb{R}$, $U\subseteq\mathbb{R}$, but not
globally. A _section_ of a vector bundle is a smooth map $s\colon M\to E$ such
that $\pi\circ s=\mathrm{Id}_{M}$. Intuitively, a section is a smooth choice
of vectors on each vector space. When $E=TM$, a section is just a vector
field. We denote by $\Gamma(E)$ the space of all sections on $E$.

#### Abstract and linear connections.

The tangent bundle $TE$ of a vector bundle has a distinguished sub-bundle
$VE\coloneqq\ker\mathrm{d}\pi$. In the same fashion that for vector spaces
there is no natural inclusion of a quotient vector space onto the original
vector space (it requires a choice of basis), here there is no natural
complement to $VE$ in $TE$. We say that a sub-bundle $\mathcal{H}\subseteq TE$
is an _Ehresmann connection on $E$_ if $VE\oplus\mathcal{H}=TE$, in other
words, for every $p\in E$, $V_{p}E$ and $\mathcal{H}_{p}$ have zero
intersection and generate the whole $T_{p}E$. As a consequence, we have that
$\mathrm{d}\pi(\mathcal{H}_{p})=T_{\pi(p)}M$ is a linear isomorphism. We say
that an Ehresmann connection is a connection if it is linear on the fibers,
that is, given the scalar multiplication on the fibers $\mu_{a}(p)=ap$, then
$\lparen\mathrm{d}\mu_{a}\rparen_{p}\lparen\mathcal{H}_{p}\rparen=\mathcal{H}_{ap}$.

#### Parallel transport systems.

An equivalent way to talk about connections is through parallel transport
systems. A _parallel transport system_ $\mathbb{P}$ is an assignment for every
curve $\gamma\colon[a,b]\to M$ and point $p\in E_{\gamma(a)}$ of a curve
$\mathbb{P}_{\gamma}(p)$ on $E$ with initial point $p$ such that
$\pi\circ\mathbb{P}_{\gamma}(p)=\gamma$ such that certain linearity and
regularity conditions hold. These conditions are: $1)$ For every curve
$\gamma$, $\widehat{\mathbb{P}}_{\gamma}(p)\coloneqq\mathbb{P}_{\gamma}(p)(b)$
gives a linear isomorphism between $E_{\gamma(a)}$ and $E_{\gamma(b)}$ with
inverse
$\widehat{\mathbb{P}}^{-1}_{\gamma}=\widehat{\mathbb{P}}_{\gamma^{-}}$, where
$\gamma^{-}$ is the curve walked in the reverse direction. $2)$ It is
independent of the parametrization of $\gamma$. $3)$ Depends smoothly on
$\gamma$ and $p$, and the initial tangent to the lifted curve only depends on
the initial tangent vector to the curve on $M$.

When $\nabla$ is the Levi-Civita connection on the tangent bundle $\pi\colon
TM\to M$ induced by a metric, the associated parallel transport system comes
from parallel transporting vectors along curves. The trivial parallel
transport on a trivial vector bundle $E\cong M\times\mathbb{R}^{k}$ for a
point $(x,v)\in E$ and a curve such that $\gamma(0)=x$, is given by
$\mathbb{P}_{\gamma}(x,v)(t)=(\gamma(t),v)\in M\times\mathbb{R}^{k}$. This
gives the parallel transport system on $\mathbb{R}^{n}$ with the canonical
metric, since $T\mathbb{R}^{n}\cong\mathbb{R}^{n}\times\mathbb{R}^{n}$. One
can prove that defining a parallel transport system is equivalent to defining
a connection on a vector bundle. This makes explicit the idea that a
connection is a way of identifying nearby tangent spaces on a manifold.

#### Covariant derivatives.

A _covariant derivative operator on $E$_ is an application $\nabla$ that lets
us derive sections $s\in\Gamma(E)$, on a given direction $X\in TM$. We ask for
this operator $\nabla_{X}s$ to be linear in $X$ and Leibnitz in $s$, meaning
that $\nabla_{X}(fs)=X(f)s+f\nabla_{X}s$. An example of a covariant derivative
on $TM$ is the usual Levi-Civita connection given by a metric on $M$. Every a
parallel transport system (or equivalently, every connection on $E$) gives
raise to a covariant derivative operator on $E$ by the formula
$\nabla_{\dot{\gamma}(0)}s(\gamma(0))\coloneqq\frac{\mathrm{d}}{\mathrm{d}t}|_{t=0}\widehat{\mathbb{P}}^{-1}_{\gamma}\lparen
s\lparen\gamma(t)\rparen\rparen$. Note that this notion coincides with that of
the Levi-Civita connection, where $\nabla_{v}X$ at a point $x\in M$ may be
computed as the derivative of the parallel transport of $X(\gamma(t))$ back to
$x$ where $\gamma(t)$ is the geodesic emanating from $x$ with
$\gamma^{\prime}(0)=v$. Covariant derivatives are a convenient way to perform
computations, but they tend to hide the geometric intuition behind the ideas
in differential geometry.

##  Appendix B Topological obstructions to admitting a connection with trivial
holonomy

In this section we give a survey of some topological obstructions to admitting
a flat connection and a flat metric. This comes to justify the assertion that
the holonomy problems cannot be solved just by choosing a metric, as these
problems are intrinsic to the topology of the manifold. We work in the context
of vector bundles, as these are much more suited for this kind of questions
than simply working on the tangent bundle. We give a short summary of these
concepts in Appendix A.

Let $\pi\colon E\to M$ be a vector bundle over $M$ together with a connection
$\nabla$. Associated to this connection, we have a parallel transport system
$\mathbb{P}$. Given that
$\widehat{\mathbb{P}}^{-1}_{\gamma}=\widehat{\mathbb{P}}_{\gamma^{-}}$, we
define the _holonomy group_ at a point $p$ as

$\mathrm{Hol}(p)\coloneqq\\{\widehat{\mathbb{P}}_{\gamma}\>|\>\gamma(a)=\gamma(b)=0\\}.$

One may also define the _restricted holonomy group_ as

$\mathrm{Hol}_{0}(p)\coloneqq\\{\widehat{\mathbb{P}}_{\gamma}\>|\>\gamma(a)=\gamma(b)=0,\gamma\text{
is null-homotopic}\\}.$

That is, $\mathrm{Hol}_{0}(p)$ is the group composed just by the isomorphisms
given by curves that are contractible to a point.

One can prove that $\mathrm{Hol}(p)$ is a Lie group and $\mathrm{Hol}_{0}(p)$
is the connected component containing the identity. More importantly, we have
that if the group is trivial, that is, $\mathrm{Hol}(p)=\\{\mathrm{Id}\\}$,
then the parallel transport is also trivial. A trivial parallel transport
system gives a global linear isomorphism between the fibers of the vector
bundle, meaning that the vector bundle is globally trivial, that is,
$E=M\times\mathbb{R}^{k}$. This imposes a topological restriction on the type
of connections that a bundle admits. What this means for us is that most
manifolds do not admit a connection with trivial holonomy on their tangent
bundle. For example, it is easy to see that the sphere $S^{2}$ does not have a
trivial bundle, by the hairy ball theorem, as it does not admit a constant
non-zero vector field. In general, the tangent bundle of $S^{n}$ is non-
trivial for $n\neq 1,3,7$. This is a very strong obstruction given by the
differentiable structure of the manifold to finding connections with trivial
holonomy.

Even if the manifold has a trivial tangent bundle, this may still not be
enough for our purposes. For example, it is not difficult to show using left-
invariant vector fields that the tangent bundle of a Lie group $G$ is trivial,
that is $TG=G\times\mathfrak{g}$. As such, one may put a trivial connection on
this bundle following this trivialization. This connection will of course have
trivial holonomy. On the other hand, this connection has torsion, unless $G$
is Abelian. In particular, a direct computation shows that this connection has
torsion $T(X,Y)=-[X,Y]$ for left-invariant vector fields. In other words,
whenever $G$ is not Abelian this connection does not come from any metric.
Remember that a Lie group (in particular a matrix Lie group) is Abelian if and
only if it is the product of a Euclidean space and a torus. As such, for most
interesting Lie groups like $\operatorname{GL}\lparen
n\rparen,\operatorname{SO}\lparen n\rparen,\operatorname{U}\lparen
n\rparen,\operatorname{SU}\lparen n\rparen,\operatorname{SL}\lparen
n\rparen,\ldots$ this trivialization does not induce a parallel transport
system without holonomy that comes from a metric when working on dimensions
larger than $2$.

More generally, if the manifold admits a metric of non-positive curvature, and
it is simply connected, then it is diffeomorphic to $\mathbb{R}^{n}$ through
the Riemannian exponential map. These are called Hadamard manifolds. On these
manifolds, one may clearly put a flat metric—and hence with trivial holonomy
given that it is simply connected—simply by pushing forward the flat metric
from $\mathbb{R}^{n}$ along the diffeomorphism. This was explored for the case
of symmetric positive definite matrices in the context of medical imaging and
numerical analysis in [7, 8].

We use a similar idea to this one in our approach to define momentum and
adaptive methods on manifolds. In particular, we use that for every manifold,
the exponential map is a diffeomorphism from a radially convex subset of
$\mathbb{R}^{n}$ onto almost all the manifold, therefore being able to make
almost all the tangent bundle trivial, and being able to put a flat metric on
almost all the manifold. That is the reason why this approach is called a
trivialization. Note that if our manifold were the symmetric positive definite
matrices, the trivialization map would be a translation by left-multiplication
by a matrix of that considered in [8], so trivializations can be seen as a
generalization of this paper to arbitrary manifolds.

##  Appendix C Experiment Set-Up and Hyperparameters

The experiments tried to replicate faithfully the results of the previous
papers. For that reason, we used the implementation of [[]Section
F]lezcano2019trivializations as a base.

The initialization, size, optimizers, and mixing constants of the adaptive
terms in the algorithms, were the same as those described in previous papers
across all the experiments. The only hyperparameters that were tuned were the
learning rates of the orthogonal and the non-orthogonal parameters for the
comparison to be fair.

The learning rates used were as follows:

Table 4: Hyperparameters for atriv$1$. Dataset | Size | Optimizer | Learning Rate | Orth. Optimizer | Orth. Learning Rate
---|---|---|---|---|---
mnist | $170$ | rmsprop | $7\cdot 10^{-4}$ | rmsprop | $10^{-4}$
$360$ | $5\cdot 10^{-4}$ | $10^{-4}$
$512$ | $3\cdot 10^{-4}$ | $5\cdot 10^{-5}$
p-mnist | $170$ | $7\cdot 10^{-4}$ | $10^{-4}$
$360$ | $7\cdot 10^{-4}$ | $7\cdot 10^{-5}$
$512$ | $5\cdot 10^{-4}$ | $7\cdot 10^{-5}$
timit | $224$ | adam | $10^{-3}$ | rmsprop | $2\cdot 10^{-4}$
$322$ | $10^{-3}$ | $2\cdot 10^{-4}$
$425$ | $10^{-3}$ | $2\cdot 10^{-4}$