Practical Structured Riemannian Optimization with Momentum by using Generalized Normal Coordinates Sophia Sanborn Christian Shewmake
Simone
Abstract
Adding momentum into Riemannian optimization is computationally challenging due to the intractable ODEs needed to define the exponential and parallel transport maps.
We address these issues for Gaussian Fisher-Rao manifolds by proposing new local coordinates to exploit sparse structures and efficiently approximate the ODEs, which results in a numerically stable update scheme.
Our approach extends the structured natural-gradient descent method of Lin et al.
(2021a)
by incorporating momentum into it and scaling the method for large-scale applications arising in numerical optimization and deep learning.
1. Introduction and Background
Riemannian gradient methods are useful for optimization on manifolds.
Momentum (Goodfellow et al., 2016) is a common approach to improve gradient-based methods.
Unfortunately, adding momentum into Riemannian optimization is computationally intensive due to the intractable ODEs needed to define the exponential and parallel transport map.
Moreover, Riemannian methods could perform poorly and be numerically unstable in singleprecision floating-point settings such as training large-scale neural networks.
Straightforward approximations of the ODEs in global coordinates involve computing nonsparse and high-order derivatives which limits the large-scale usage of Riemannian methods.
Many well-known algorithms in machine learning can be derived from Gaussian FisherRao manifolds, for example, see Khan and Rue (2021); Lin et al.
(2021a).
In this work, we propose special local coordinates for Gaussian Fisher-Rao manifolds to relax the requirement of the use of global coordinates, and the exact exponential and Riemannian transport maps required by previous methods, and exploit sparse Lie-group structures to obtain numerically stable updates. It is challenging to add momentum in local coordinates, which involves the computation of the metric, the exponential and transport maps, and the change of local coordinates. Our Riemannian momentum update can be viewed as Euclidean GD with momentum in a special coordinate. We use the approach to extend the structured naturalgradient descent (NGD) method (Lin et al., 2021a) by adding momentum to the method, and scaling it to large-scale applications (see Fig.1).
We will begin with the basics, introduce the method, and conclude with an example.
Additional results can be found in appendices.
1.1. Riemannian Gradient Descent and Natural-gradient Descent
Consider a geodesically complete manifold M with the Fisher-Rao F metric.
The metric is known as the Fisher information matrix and is defined in A.2.
Under a global parametrization τ , Riemannian gradient descent (Bonnabel, 2013) (RGD) is defined as
[Equation]
where v :=
F − 1 τ ( τ ) g ( τ ) is a Riemannian gradient, g ( τ ) is a Euclidean gradient (1-form),
F τ is the metric, RExp( τ , v ) is the Riemannian exponential map (defined in A.4) induced by the geodesic ODE at the current τ with Riemannian direction v , and α is a step-size.
The exponential map often does not admit a closed-form expression (see Table 1).
For an approximation of the exponential map, we have to evaluate the Christoffel symbols Γ acb ( τ ) defined at A.3, which involves computing partial derivatives of the Fisher matrix.
Natural-gradient descent (NGD) is a special case of RGD in (1) that uses a first-order
Taylor approximation of the exponential map to avoid evaluating the Christoffel symbols,
[Equation]
This approximation, which is known as the Euclidean retraction, is defined as,
[Equation]
where r ( t ) is the solution of the geodesic ODE with initial values r (0)
= x and ˙ r (0)
= ν .
1.2. Standard Riemannian Normal Parametrization
We now discuss updates using a special local parametrization that will allow us to reduce computational costs, and we will generalize this later.
Let τ be the global coordinate, and η be a normal (local) coordinate at the current point τ (cur) so that τ =
RExp( τ (cur) , F − 1 / 2 τ ( τ (cur) ) η ), where τ (cur) and F − 1 / 2 τ ( τ (cur) ) are constants in this coordinate.
In the coordinate η , the origin η 0 := 0 represents point τ (cur) .
The following properties hold at point τ (cur) in local coordinate η at η 0 = 0 .
[Equation]
The update of RGD in Eq.
1 can be written as a NGD update in local coordinate η : η 1 ← η 0 − α g ( η 0 ) , τ (new)
← RExp( τ (cur) , F − 1 / 2 τ
[Equation]
≡ [Equation] (4) where η 0 = 0 , F η ( η 0 ) =
I and g ( η 0 )
=
F − 1 / 2 τ ( τ (cur) )
g ( τ (cur) ) is a Euclidean gradient.
This update remains on the manifold while NGD in global coordinate τ (see Eq.2) does not, especially when τ is in a constraint set (e.g., the positive-definite constraint of S in Eq.10).
[Figure]
2. Natural Gradient Descent with Momentum 2.1. Momentum in Riemannian Gradient Descent
To incorporate momentum in RGD, we have to use a parallel transport map to move a (gradient) vector at the current point τ (cur) to a new point τ (new) on a curved space.
This map is induced by the transport ODE (defined in Eq.15 and Eq.17 in the appendix).
We use the transport map T τ (cur) → τ (new) (defined in A.6) to transport a Euclidean gradient at τ (cur) to a Euclidean gradient at τ (new) .
We consider the following update with momentum weight β under global coordinate τ .
[Equation]
where w is a Euclidean gradient and m is a momentum for Euclidean gradients.
This update is equivalent to the update of Alimisis et al.
(2020), shown in A.8, if the exact exponential and the exact transport are used.
Many Riemannian momentum methods including Alimisis et al.
(2020) use the exact transport map ˆ T for Riemannian gradients (see A.5) while we use the transport map for Euclidean gradients.
In most Gaussian cases, the exact transport maps (see Table 1) are unknown.
We choose to approximate the Euclidean transport map T since it is more computationally efficient than the Riemannian map ˆ T as we will discuss in Sec. 2.3. Thus, our update is distinct from Alimisis et al.
2.2. Generalized Riemannian Normal Parametrizations
Our idea is to use a generalized Riemannian normal parametrization to simplify the computation.
The standard Riemannian normal parametrization is computationally challenging since it involves the computation of the Riemannian exponential map.
To get rid of that dependence on the map, we define a generalized Riemannian normal coordinate η associated to point τ (cur) via a diffeomorphic map τ
= ϕ τ (cur) ( η ) such that the following conditions hold, where the Christoffel symbols Γ acb need not vanish at η 0 = 0 .
[Equation]
We recover the standard normal coordinate if the exponential map is used as the diffeomorphic map.
However, we can instead choose a more efficient map.
Moreover, if necessary, we can efficiently compute the Christoffel symbols at η 0 for Gaussian manifolds.
2.3. NGD with Momentum using Generalized Normal Coordinates
We propose an approximation of Eq. 5 by local coordinate η associated to global τ (cur)
[Equation]
where F − 1 η ( η 0 )
=
I , η 0 = 0 , and, η 0 and η 1 represent τ (cur) and τ (new) , respectively.
To use normal coordinate ξ associated to τ (new) at the next iteration, we have to transform the Euclidean gradient w ( η 1 ) in coordinate η into coordinate ξ via the chain rule:
[Equation]
where ξ 0 = 0 , J ( ξ ) is the Jacobian, and both ξ 0 and η 1 represent τ (new) .
We can compute the Jacobian as J (
[Equation]
ξ since by construction we have τ = ϕ
τ (cur) ( η ) =
ϕ
τ (new) ( ξ ).
In practice, we could also directly compute w ( ξ 0 ) by automaticdifferentiation.
As mentioned in Sec.
2.1, we use the transport map for Euclidean gradients so that the Jacobian can be easy to compute at ξ 0 = 0 .
If the transport map for Riemannian gradients is used, we have to compute the inverse of the Jacobian.
We now discuss approximations of the transport map for Euclidean gradients.
2.4. Approximations of the Transport Map
To avoid evaluating the Christoffel symbols, we could use a zero-order Taylor expansion to approximate the parallel transport map as shown below,
[Equation]
(7) where ω ( t ) is the solution of the transport ODE with initial value ω (0)
= m ( η 0 ) .
A more accurate approximation is a first-order approximation as in A.7, where we have to compute the Christoffel symbols at η 0 = 0 to compute the first-order term ˙ ω (0).
For Gaussian manifolds, we can simplify the computation since η is a generalized normal coordinate.
If we use the standard normal coordinate, the zero-order approximation becomes the first-order approximation since the Christoffel symbols vanish.
Our update with the zeroorder approximation is exactly the Euclidean momentum update in the local coordinate.
Our goal is to replace the standard normal coordinate with a generalized one.
In the next section, we give an example of Gaussian manifolds, where the zero-order approximation is also a first-order approximation since the first-order term ˙ ω (0) vanishes. In structured cases, the zero-order approximation works very well for many machine-learning problems as supported by the experiments even when the first order term does not vanish. The firstorder approximation does not improve the performance while increasing the computational cost. The first-order term is O ( α 2 ) (See Appendix B), which is small when step-size α is small, especially in noisy, large-scale, or low-precision floating-point settings.
3. Gaussian Fisher-Rao Manifolds
We will now illustrate our update described in Sec 2.3 by considering an example of Gaussian
Fisher-Rao manifolds.
Gaussian Fisher-Rao manifolds play an important role in machine learning.
Consider the following unconstrained optimization problem.
[Equation] where ℓ is a real-valued loss function to be minimized. Khan et al. ; Lin et al.
(2021a,b) reformulate the problem in Eq.
8 as a new optimization problem over a Gaussian (surrogate) distribution q ( z | τ ) with mean µ and inverse covariance S ,
[Equation]
with global coordinate τ = { µ , S } , and where ℓ ( z ) is the loss function defined in Eq. 8.
Khan et al. (2017, 2018) show a connection between NGD (see Eq.2) in global coordinate τ for problem 8 and Newton’s method for problem 9 at steps-size α = 1 as [Equation] This update recovers Newton’s method when the expectations are approximated at the mean µ .
However, the parameter constraint in S is not satisfied in Eq.10.
Lin et al.
(2021a,b) propose NGD updates via local coordinates for structured S and handle the constraint in
S , which enables usage of these Newton-like methods in non-convex and stochastic settings.
We will extend the local coordinates in the method of Lin et al.
(2021a,b) by using the generalized normal coordinates.
Our extension allows the method to incorporate momentum and improves its performance for large-scale applications.
Due to the space limit, we only consider full Gaussians with a constant mean.
A structured case is given in Appendix B.
Consider a d -dimension Gaussian family with zero mean q ( z | τ ), where the global coordinate τ = S is the inverse covariance, which is symmetric positive-definite.
For efficiency, we propose a new re-parameterization instead of the one in Lin et al.
(2021a).
We introduce an auxiliary parameterization B such that τ = B T B , where B ∈ GL
d × d is an invertible matrix.
The auxiliary parameter B has an additional Lie group structure since it is a member of the general linear group.
We define a local parameterization η associated to B (cur) as B = h ( 1 √ 2 η ) B (cur) , where h ( N ) = Expm( N ) :
= (cid:80) ∞ n =0 N n n ! is the matrix exponential (the Lie-group exponential map) and η is a symmetric matrix.
In practice, we approximate Expm() by a truncation since the exact map is numerically unstable and computationally intensive in noisy, large-scale, and low precision floating-point settings.
Local parameter η lives in a subspace of the Lie algebra of B since η is symmetric. Note that m is also symmetric in our update.
In this example, η is a generalized Riemannian normal coordinate associated to τ (cur) via diffeomorphic map τ = ϕ τ (cur) ( η ) :=
( B (cur) )
T h T ( 1 √ 2 η ) h ( 1 √ 2 η ) B (cur) , where τ (cur) =
( B (cur) )
T B (cur) =
ϕ τ (cur) ( η 0 ) and F η ( η 0 ) =
I at η 0 = 0 .
The Christoffel symbols at η 0 do not vanish. Thus, η is not the standard normal coordinate.
We can write w as w = w ( ξ 0 ) =
w ( η 1 ) since w ( ξ 0 ) = J ( ξ 0 ) w ( η 1 ) = w ( η 1 ) as shown in B.
Our proposed update can be simplified as below, where we can efficiently compute the Euclidean gradient g ( η 0 ) and use the zero-order approximation for the transport map.
Although the Christoffel symbols do not vanish, the first-order term of the transport map does vanish as shown in A.7. Thus, this approximation is also a first-order approximation.
[Equation]
where in the normal coordinate, our update is exactly the Euclidean GD with momentum.
The intermediate step in B allows us to include structures in B by considering sparse subgroup structures in B as suggested by Lin et al.
(2021a).
This gives rise to structured
Gaussian (sub)manifolds as we will discuss in B. In structured cases, our update also enjoys
Lie subgroup invariance (Lin et al., 2021b). Appendix A. Backgound A.1. Exact Riemannian exponential map and Riemannian transport map Manifold Riemannian exponential map
Riemannian transport map Gaussian (constant mean) (see Fig.1( a )) RExp( τ , v ) = τ Expm( τ − 1 v ); symmetric ˆ T τ 0 → τ 1 ( v ) = EvE T ; E = ( τ 1 τ − 1 0 )
Structured Gaussian (constant mean) Unknown Unknown (Structured) Gaussian (see Fig.1( b ),1( c ) ) Unknown
Unknown A.2. Fisher Information Matrix Under parametrization τ , the Fisher information matrix is defined as
[Equation]
(11) where q ( z | τ ) is a probabilistic distribution parameterized by τ such as the Gaussian distribution with τ = { µ , S } , where µ is the mean and S is the inverse covariance matrix.
A.3. Christoffel Symbols
The Christoffel symbols of the first kind of the Fisher metric connection are defined as
[Equation] [Equation]
where F bd ( τ ) denotes the ( b, d ) entry of metric F τ and ∂ b denotes the partial derivative w.r.t.
the b -th entry of τ .
The Christoffel symbols of the second kind are defined as Γ cab ( τ ) : = (cid:80) [Equation] F cd
[Equation] ab [Equation] τ ) where F cd ( τ ) denotes the (c,d) entry of F − 1 τ .
Observe that the Christoffel symbols involves computing all partial derivatives of the Fisher matrix F τ .
A.4. Riemannian Exponential Map
The exponential map is defined via a geodesic.
The geodesic generalizes the notion of a straight line to a manifold.
The geodesic r ( t ) satisifies the following second-order nonlinear system of ODEs with initial values r (0) =
x and ˙ r (0) =
ν , where x denotes a point on the manifold and ν is a Riemannian gradient,
[Equation]
where r c ( t ) denotes the c -th entry and Γ cab is the Christoffel symbol of the second kind.
The Riemannian exponential map is defined as [Equation]
[Equation]
where x denotes an initial point and ν is an initial Riemannian gradient so that r (0)
=
x and ˙ r (0) =
ν .
A.5. Parallel Transport Map for Riemannian Gradients
In a curved space, the transport map along a given curve generalizes the notion of the parallel transport.
In Riemannian optimization, we consider the transport map along a geodesic curve. Given a geodesic curve r ( t ), denote by v ( t ) a smooth Riemannian gradient field that satisfies the following first-order linear system of ODEs with initial value v (0) =
ν ,
[Equation]
The transport map ˆ T τ (old) → τ (old) ( ν ) transports the Riemannian gradient ν at τ (old) to τ (new) as follows,
[Equation]
where r (0) = τ (cur) , r (1) =
τ (new) , and v (0)
=
ν .
It can be computationally challenging to solve this linear ODE due to the Christoffel symbol.
A.6. Parallel Transport Map for Euclidean Gradients
Given a geodesic curve r ( t ), denote by ω ( t ) a smooth Euclidean gradient field on manifold
M that satisfies the following first-order linear system of ODEs with initial value ω (0) = m ,
[Equation]
[Equation]
The transport map transports the Euclidean gradient g at τ (cur) to τ (new) as shown below,
[Equation]
cur [Equation] 18) where r (0) = τ (cur) , r (1) =
τ (new) , and ω (0) = g .
These two transport maps are related as follows,
[Equation]
where g is a Euclidean gradient evaluated at τ (cur) .
A.7. A First-order Approximation of the Transport Map
We can also consider a first-order approximation of the transport map for Euclidean gradients as follows,
[Equation]
where we have to evaluate the Christoffel symbols as discussed below.
By the transport ODE in Eq.
17, we can compute ˙ ω (0) as follows,
[Equation]
where r (0) =
η 0 is the current point and ˙ r (0) is the Riemannian gradient so that η 1 =
Exp( η 0 , ˙ r (0)).
In our case, ˙ r (0) =
− F − 1 η ( η 0 ) m ( η 0 ) and ω (0) = m ( η 0 ) .
Note that the Fisher matrix and Christoffel symbols are evaluated at
a generalized Riemannian normal coordinate η 0 = 0 . The computation can be simplified due to the identities: F [Equation]
[Equation]
I and
[Equation]
d F cb ( η 0 )], where F ad ( η 0 ) =
δ ad .
Thus, we have the following simplification.
[Equation]
In the Gaussian case considered in Sec. 3, we can show the first-order term vanishes since m is symmetric.
[Equation]
where < m , m T > :
= mm T − m T m is the Lie bracket.
A.8. Update 5 is Equivalent to Alimisis et al. (2020)
The update of Alimisis et al.
(2020) is defined as
[Equation]
21) where ν is a Riemannian gradient, and ˆ T τ (cur) → τ (new) is the parallel transport map for Riemannian gradients.
Due to Eq.
19, update 21 and update 5 are equivalent by the observation that m (cur) =
F ( τ (cur) )
ν (cur) and w (new) =
F ( τ (new) )
z (new) .
Appendix B. Structured Gaussian with Constant Mean
Consider a d -dimensional Gaussian family with zero mean q ( z | τ ), where the global parametrization τ is the inverse covariance, which is symmetric positive-definite.
We introduce an auxiliary parameterization B such that τ = B T B and B = (cid:20) B A
B C
B D (cid:21) , where B A is an invertible upper-triangular matrix and B D is an invertible diagonal matrix.
The auxiliary parameter B lives in a subgroup of the upper-triangular group illustrated in
Figure 2.
We first define a constant matrix D = (cid:34)
Tril( 1 A ) √ 1 B 1 √
I D (cid:35) , where 1 is the matrix of all ones, I D is the identity matrix, Tril() extracts the upper-triangular half.
We define a local parameterization η = (cid:20)
η
A η C η D (cid:21) associated to B (cur) as B =
h
( D ⊙ η ) B (cur) , where ⊙ is the element-wise product, h ( N ) is the matrix exponential, η A is an upper-triangular matrix, η D is a diagonal matrix.
In this case, the local parameter η lives in the Lie algebra of B .
We can show that η is a generalized Riemannian normal coordinate associated to τ (cur) as discussed in the main text.
Moreover, F η ( η 0 ) = I where η
= 0 .
Since m has the same structure of η , m is not symmetric. The first-order approximation of the transport map at η 0 do not vanish since m is not symmetric.
However, in practice, the first-order approximation of the transport does not improve the performance since step-size α is small. Recall that m is O ( α ) in our proposed update. The first-order approximation is O ( α 2 ) as shown below. ˙
[Equation] , where < A, B > := AB − BA is the Lie bracket.
Notice that τ =
( B (cur) )
T h T ( D ⊙ η ) h ( D ⊙ η ) B (cur) =
( B (new) )
T h T ( D ⊙ ξ ) h ( D ⊙ ξ ) B (new) and B (new) =
h ( D ⊙ η 1 ) B (cur) .
We have h T ( D ⊙ η ) h ( D ⊙ η ) =
h T ( D ⊙ η 1 ) h T ( D ⊙ ξ ) h ( D ⊙ ξ ) h ( D ⊙ η 1 ).
Due to the structure restriction on η and ξ , h ( D ⊙ η ) and h ( D ⊙ ξ ) cannot be an orthogonal matrix and change the sign of the determinant.
Thus, h ( D ⊙ η ) = h ( D ⊙ ξ ) h ( D ⊙ η 1 ).
Since h is the matrix exponential, we have η = D − 1 ⊙ Expm − 1 (cid:0) Expm( D ⊙ ξ )Expm( D ⊙ η 1 ) (cid:1) .
By the Baker–Campbell–Hausdorff formula, we have η = ξ + η 1 + (2 D ) − 1 ⊙ < D ⊙ ξ , D ⊙ η 1 > + O ( ξ 2 )+ O ( < η 1 , < η 1 , ξ > ) , where < A, B > := AB − BA is the Lie bracket. We consider the zero-order approximation of the transport map.
In this case, w = m and η 1 = − m .
Recall that m =
O ( α ), where α is the step size. The Jacobian-vector product at ξ 0 = 0 can be simplified as [Equation] >
[Equation]
In the Gaussian with constant mean cases, the high-order terms in the Jacobian-vector product all vanish since m is symmetric and D = 1 √ 2 1 , where 1 is a matrix of all ones.
In structured Gaussia cases, the high-order terms in the Jacobian-vector product do not vanish.
However, in practice, we can ignore the high-order terms since the step-size α is very small and the high-order terms do not improve performance.
We can similarly obtain the proposed update.
Appendix C. Numerical Results C.1. Optimization on Symmetric Positive-definite Matrices
We consider manifold optimization on positive-definite matrices S d × d
++ , where the Riemannian exponential and transport maps admit a closed-form expression as shown in Table.1.
However, our approach can be extended to structured cases, where these maps do not have a closed-form expression.
We evaluate the proposed method in an example of the metric nearness problem considered in Lin et al.
(2021a).
The objective function is ℓ ( Z ) = 12 N (cid:80) Ni =1 ∥ ZQx i − x i ∥ 22 , where x i ∈ R d ,
Q ∈ S
d × d ++ and Z ∈ S
d × d ++ .
The optimal is Q − 1 .
We randomly generate x i and Q with d = 100, N train = 5 × 10 5 for training and N test = 10 5 for testing.
All methods are trained using mini-batches, where the size of the mini-batch is 250.
We also consider existing Riemannian momentum methods such as Ahn and Sra (2020), Alimisis et al. (2021), and Alimisis et al.
(2020) as baselines.
From Figure 1( a ), we can see that the proposed method using a local coordinate performs as well as existing Riemannian momentum methods using a global coordinate with the exact exponential and transport maps.
C.2. Unconstrained Optimization
We consider the following valley-shaped function suggested by in Lin et al.
(2021a).
Rosenbrock : ℓ ( z ) =
1 d (cid:80) d − 1 i =1 (cid:2) 100( z i +1 − z i )
+ ( z i − 1) 2 (cid:3) , where d = 10 3 .
In this example, we consider a structured Gaussian manifold with a non-constant mean, where these exponential and transport maps do not have a closed-form expression.
Our baselines are the BFGS provided by SciPy and the method of Lin et al.
(2021a).
Our method and the method of Lin et al.
(2021a) use a structured Hessian information, which can be efficiently computed via a Hessian-vector product.
We consider a block Heisenberg structure as suggested by Lin et al.
(2021a), where our update introduce momentum into the method of Lin et al.
(2021a).
From Figure 1( b ), we can see that the momentum can improve the performance.
C.3. Deep Learning
We consider a CNN model with the VGG-16 neural network architecture.
We train the model with matrix Gaussian (Lin et al., 2021a) for each layer-wise matrix weights on dataset“CIFAR-10”.
In this example, we consider a structured Gaussian manifold with a non-constant mean, where these exponential and transport maps do not have a closedform expression.
For each Gaussian, we use a Kronecker product of two dense groups.
The iteration cost and the structure are similar to the KFAC approximation (Martens and
Grosse, 2015).
We also use a similar approximation as KFAC to compute natural gradients.
We include the KFAC method as our baseline method.
We train the model with mini-batch size 256.
From Figure 1( c ), we can see that the momentum can improve the performance and our method performs similarly to KFAC.
However, our approach enables us to introduce more sparse structures along with the
Kronecker structure such as a Kronecker product of two sparse groups, which can further reduce the iteration cost of our method.
We are currently investigating sparse updates that can further reduce the iteration cost while maintaining the performance.