Title: A Framework for Bilevel Optimization on Riemannian Manifolds
Venue: NeurIPS 2024
Accepted: yes
Keywords: Bilevel optimization

Riemannian optimization

Riemannian hypergradient
Abstract: Bilevel optimization has gained prominence in various applications. In this study, we introduce a framework for solving bilevel optimization problems, where the variables in both the lower and upper levels are constrained on Riemannian manifolds. We present several hypergradient estimation strategies on manifolds and analyze their estimation errors. Furthermore, we provide comprehensive convergence and complexity analyses for the proposed hypergradient descent algorithm on manifolds. We also extend our framework to encompass stochastic bilevel optimization and incorporate the use of general retraction. The efficacy of the proposed framework is demonstrated through several applications.

A Framework for Bilevel Optimization on Riemannian Manifolds

Abstract

Bilevel optimization has gained prominence in various applications. In this study, we introduce a framework for solving bilevel optimization problems, where the variables in both the lower and upper levels are constrained on Riemannian manifolds. We present several hypergradient estimation strategies on manifolds and analyze their estimation errors. Furthermore, we provide comprehensive convergence and complexity analyses for the proposed hypergradient descent algorithm on manifolds. We also extend our framework to encompass stochastic bilevel optimization and incorporate the use of general retraction. The efficacy of the proposed framework is demonstrated through several applications.

1 Introduction

Bilevel optimization is a hierarchical optimization problem where the upper-level problem depends on the solution of the lower-level, i.e.,

Applications involving bilevel optimization include meta learning [16], hyperparameter optimization [18], and neural architecture search (NAS) [53], to name a few. The lower-level problem is usually assumed to be strongly convex.

Common strategies for solving such problem can be classified into two categories: single-level reformulation [29, 60] and approximate hypergradient descent [19, 40]. The former aims to reformulate the bilevel optimization problem into a single-level one using the optimality conditions of the lower-level problem as constraints. However, this may impose a large number of constraints for machine learning applications. The latter scheme directly solves the bilevel problem through iteratively updating the lower and upper-level parameters and, hence, is usually more efficient. Nevertheless, existing works have mostly focused on unconstrained bilevel optimization [19, 32, 40, 11, 52, 45, 14].

In this work, we study bilevel optimization problems where x and y are on Riemannian manifolds M x and M y , respectively. We focus on the setup where the lower-level function g ( x, y ) is geodesic strongly convex (a generalized notion of convexity on manifolds, defined in Section 2) in y . This ensures the lower-level problem has a unique solution y ∗ ( x ) given x . The upper-level function f can be nonconvex on M x × M y . Because the unconstrained bilevel optimization is a special case of our formulation on manifolds, such a formulation includes a wider class of applications. Examples of Riemannian bilevel optimization include Riemannian meta learning [64] and NAS over SPD networks [62]. Moreover, there has been a surge of interest of min-max optimization over Riemannian manifolds [37, 41, 73, 27, 25, 67, 35], which also gets subsumed in the framework of bilevel optimization with g = -f .

Contributions. (i) We derive intrinsic Riemannian hypergradient via the implicit function theorem and propose four strategies for estimating the hypergradient, i.e., through Hessian inverse, conjugate gradient, truncated Neumann series, and automatic differentiation. We then provide hypergradient estimation error bounds for all the proposed strategies. (ii) We introduce the Riemannian hypergradient descent algorithm to solve bilevel optimization problems on manifolds and provide convergence guarantees. We also generalize the framework to the stochastic setting and to allow the use of retraction. (iii) The efficacy of the proposed modeling is shown on several problem instances including hyper-representation over SPD matrices, Riemannian meta learning, and unsupervised domain adaptation. The proofs, extensions, and experimental details are deferred to the appendix sections.

Related works in unconstrained setting. Unconstrained bilevel optimization where the lowerlevel problem is strongly convex has been widely studied [19, 32, 40, 11, 52, 45, 14]. A crucial ingredient is the notion of hypergradient in bilevel optimization problems and its computation. There exist strategies for approximating the hypergradient, e.g., using conjugate gradient [40], Neumann series [19], iterative differentiation [21], and Nyström method [31]. While bilevel optimization with constraints is relatively unexplored, a few works exists that impose constraints only for the upper level problem [32, 10]. Recently, linearly lower-level constrained bilevel optimization has been explored in [65, 68], where a projected gradient method is employed for the lower-level problem.

Related works on manifolds. There has been limited work on bilevel optimization problems on manifolds. [7] studies semivectorial bilevel optimization on Riemannian manifolds where the upperlevel is a scalar optimization problem while the lower-level is a multiobjective problem under greatest coalition. [50, 49] reformulate bilevel problems on manifolds into a single-level problem based on the KKT conditions on manifolds. However, for all those works, it is unclear whether there exists an algorithm that efficiently solves the problem in large-scale settings. In contrast, we aim to provide a general framework for solving bilevel optimization on Riemannian manifolds. [47] is a contemporary work that also proposes gradient-based algorithms for bilevel optimization on Riemannian manifolds. The main differences of our work with respect to [47] are as follows: (1) We provide an analysis for various hypergradient estimators while [47] focuses on conjugate gradient for deterministic setting and Neumann series for stochastic setting; (2) We provide an analysis for retraction which is more computationally efficient than exponential map and parallel transport employed in [47]; and (3) We explore the utility of Riemannian bilevel optimization in various machine learning applications, which is not the case with [47].

2 Preliminaries and notations

A Riemannian manifold M is a smooth manifold equipped with a smooth inner product structure (a Riemannian metric) ⟨· , ·⟩ p : T z M× T z M→ R for any z ∈ M and its tangent space T z M . The induced norm is thus ∥ u ∥ z = √ ⟨ u, u ⟩ z for any u ∈ T z M . A geodesic c : [0 , 1] →M generalizes the line segment in the Euclidean space as the locally shortest path on manifolds. The exponential map on a manifold is defined as Exp z ( u ) = c (1) for a geodesic c that satisfies c (0) = z, c ′ (0) = u . In a totally normal neighbourhood U where exponential map has a smooth inverse, the Riemannian distance d ( x, y ) = ∥ Exp -1 x ( y ) ∥ x = ∥ Exp -1 y ( x ) ∥ y . The parallel transport operation Γ z 2 z 1 : T z 1 M→ T z 2 M is a linear map which preserves the inner product, i.e., ⟨ u, v ⟩ z 1 = ⟨ Γ z 2 z 1 u, Γ z 2 z 1 v ⟩ z 2 , ∀ u, v ∈ T z 1 M . The (Cartesian) product of Riemannian manifolds M x ×M y is also a Riemannian manifold.

For a differentiable function f : M→ R , the Riemannian gradient G f ( z ) ∈ T z M is the tangent vector that satisfies ⟨G f ( z ) , u ⟩ z = D f ( z )[ u ] for all u ∈ T z M . Here D is the differential operator and D f ( z )[ u ] represents the directional derivative of f at z along u . For a twice differentiable function f , Riemannian Hessian H f ( z ) is defined as the covariant derivative of Riemannian gradient.

Geodesic convexity extends the convexity notion in the Euclidean space to Riemannian manifolds. A geodesic convex set Z ⊆ M is where any two points can be joined by a geodesic. A function f : M→ R is said to be geodesic (strongly) convex if for all geodesics c : [0 , 1] → Z , f ( c ( t )) is (strongly) convex in t ∈ [0 , 1] . If the function is smooth, then f is called µ -geodesic strongly convex if and only if f (Exp z ( tu )) ≥ f ( z ) + t ⟨G f ( z ) , u ⟩ z + t 2 µ 2 ∥ u ∥ 2 z , ∀ t ∈ [0 , 1] . An equivalent second-order characterization is H ( z ) ⪰ µ id , where we denote id as the identity operator.

For a bifunction ϕ : M x ×M y → R , we denote G x ϕ ( x, y ) , G y ϕ ( x, y ) as the Riemannian (partial) gradient and H x ϕ ( x, y ) , H y ϕ ( x, y ) as the Riemannian Hessian. The Riemannian cross-derivatives

are linear operators G 2 xy ϕ ( x, y ) : T y M y → T x M x , G 2 yx ϕ ( x, y ) : T x M x → T y M y defined as G 2 xy ϕ ( x, y )[ v ] = D y G x ϕ ( x, y )[ v ] for any v ∈ T y M y (with D representing the differential operator) and similarly for G 2 yx ϕ ( x, y ) . For a linear operator T : T x M x → T y M y , the adjoint operator, denoted as T † is defined with respect to the Riemannian metric, i.e., ⟨ T [ u ] , v ⟩ y = ⟨ T † [ v ] , u ⟩ x for any u ∈ T x M x , v ∈ T y M y . The operator norm of T is defined as ∥ T ∥ y := sup u ∈ T x M x : ∥ u ∥ x =1 ∥ T [ u ] ∥ y .

3 Proposed Riemannian hypergradient algorithm

In this work, we consider the constrained bilevel optimization problem

where M x , M y are two Riemannian manifolds and f, g : M x ×M y → R are real-valued jointly smooth functions. We focus on the setting where the lower-level function g ( x, y ) is geodesic strongly convex. This ensures the lower-level problem has a unique solution y ∗ ( x ) for a given x . The upper-level function f can be nonconvex on M x ×M y .

We propose to minimize F ( x ) directly within the Riemannian optimization framework. To this end, we need the notion of the Riemannian gradient of F ( x ) := f ( x, y ∗ ( x )) , which we call the Riemannian hypergradient.

Proposition 1. The differential of y ∗ ( x ) and the Riemannian hypergradient of F ( x ) are given by

The above proposition crucially relies on the implicit function theorem on manifolds [25] and requires the invertibility of the Hessian of the lower level function f with respect to y . This is guaranteed in our setup as f is geodesic strongly convex in y . Hence, there exists a unique differentiable function y ∗ ( x ) that maps x to the lower-level solution. We show the Riemannian hypergradient descent (RHGD) algorithm for (1) in Algorithm 1.

Algorithm 1 Riemannian hypergradient descent (RHGD)

1: Initialize x 0 ∈ M x , y 0 ∈ M y . 2: for k = 0 , ..., K -1 do 3: y 0 k = y k . 4: for s = 0 , ..., S -1 do 5: y s +1 k = Exp y s k ( -η y G y g ( x k , y s k )) . 6: end for 7: Set y k +1 = y S k . 8: Compute approximated hypergradient ̂ G F ( x k ) . 9: Update x k +1 = Exp x k ( -η x ̂ G F ( x k )) . 10: end for

Steps 3 to 7 solve the lower-level problem using the Riemannian gradient descent method. Since computing the optimal solution y ∗ ( x ) is computationally challenging, we obtain an approximate solution y k +1 .

Step 8 involves computing the Riemannian hypergradient G F ( x ) of F ( x ) . For computational efficiency, we compute an approximation ̂ G F ( x k ) .

Step 9 is the usual exponential map to find the updated point x k +1 .

We highlight that Step 8 of Algorithm 1 approximates the Riemannian hypergradient. In the rest of the section, we discuss various computationally efficient ways to estimate the Riemannian hypergradient and discuss the corresponding theoretical guarantees for RHGD. The error of hypergradient approximation comes from the inaccuracies of y k +1 to y ∗ ( x k ) and also from the Hessian inverse.

3.1 Hypergradient estimation

When the inverse Hessian of the lower-level problem can be computed efficiently, we can estimate the hypergradient directly by evaluating the Hessian inverse ( HINV ) at y k +1 , i.e., ̂ G hinv F ( x k ) = G x f ( x k , y k +1 ) - -G 2 xy g ( x k , y k +1 ) [ H -1 y g ( x k , y k +1 )[ G y f ( x k , y k +1 )] ] . However, computing the inverse Hessian is computationally expensive in many scenarios. We now discuss three practical strategies for estimating the Riemannian hypergradient when y k +1 is given.

Conjugate gradient approach (CG). When evaluating the Hessian inverse is difficult, we can solve the linear system H y g ( x k , y k +1 )[ u ] = G y f ( x k , y k +1 ) for some u ∈ T y k +1 M y . To this

end, we employ the tangent space conjugate gradient algorithm (Appendix F, Algorithm 3) that solves the linear system on the tangent space T y k +1 M y with only access to Hessian-vector products, i.e., ̂ G cg F ( x k ) = G x f ( x k , y k +1 ) - G 2 xy g ( x k , y k +1 )[ˆ v T k ] , where ˆ v T k is computed as a solution to H y g ( x k , y k +1 )[ˆ v T k ] = G y f ( x k , y k +1 ) , where T is the number of iterations of the tangent space conjugate gradient algorithm.

Truncated Neumann series approach (NC). The Neumann series states for an invertible operator H such that ∥ H ∥ ≤ 1 , its inverse H -1 = ∑ ∞ i =0 (id -H ) i , where id is the identity operator. An alternative approach to estimate the Hessian inverse is to use a truncated Neumann series, which leads to the following approximated hypergradient, ̂ G ns F ( x k ) = G x f ( x k , y k +1 ) -G 2 xy g ( x k , y k +1 )[ γ ∑ T -1 i =0 (id -γ H y g ( x k , y k +1 )) i [ G y f ( x k , y k +1 )]] , where γ is chosen such that (id -γ H y g ( x k , y k +1 )) ≻ 0 . γ can be set as γ = 1 L , where the gradient operator is L -Lipschitz (discussed later in Definition 1). Empirically, we observe that this approach is faster than the conjugate gradient approach. However, it requires estimating T and L beforehand.

Automatic differentiation approach (AD). Another hypergradient estimation strategy follows the idea of iterative differentiation by backpropagation. After running several iterations of gradient update to obtain y k +1 (which is a function of x k ), we can use automatic differentiation to compute directly the Riemannian gradient of f ( x k , y k +1 ( x k )) with respect to x k . We can compute the Riemannian hypergradient from the differential in the direction of arbitrary u ∈ T x k M x using basic chain rules.

3.2 Theoretical analysis

This section provides theoretical analysis for the proposed hypergradient estimators as well as the Riemannian hypergradient descent. First, we require the notion of Lipschitzness of functions and operators defined on Riemannian manifolds. Below, we introduce the definition in terms of bi-functions and bi-operators and state the assumptions that are required for the analysis.

Definition 1 (Lipschitzness) . (1) For a bifunction f : M x ×M y → R , we say f has L Lipschitz Riemannian gradient in U x × U y ⊆ M x × M y if it satisfies for any x, x 1 , x 2 ∈ U x , y, y 1 , y 2 ∈ U y , ∥ Γ y 2 y 1 G y f ( x, y 1 ) - G y f ( x, y 2 ) ∥ y 2 ≤ Ld ( y 1 , y 2 ) , ∥G x f ( x, y 1 ) - G x f ( x, y 2 ) ∥ x ≤ Ld ( y 1 , y 2 ) , ∥ Γ x 2 x 1 G x f ( x 1 , y ) -G x f ( x 2 , y ) ∥ x 2 ≤ Ld ( x 1 , x 2 ) and ∥G y f ( x 1 , y ) -G y f ( x 2 , y ) ∥ y ≤ Ld ( x 1 , x 2 ) .

(2) For an operator G ( x, y ) : T y M y → T x M x , we say G ( x, y ) is ρ -Lipschitz if it satisfies, ∥ Γ x 2 x 1 G ( x 1 , y ) -G ( x 2 , y ) ∥ x 2 ≤ ρd ( x 1 , x 2 ) and ∥G ( x, y 1 ) -G ( x, y 2 )Γ y 2 y 1 ∥ x ≤ ρd ( y 1 , y 2 ) .

(3) For an operator H ( x, y ) : T y M y → T y M y , we say H ( x, y ) is ρ -Lipschitz if it satisfies, ∥ Γ y 2 y 1 H ( x, y 1 )Γ y 1 y 2 -H ( x, y 2 ) ∥ y 2 ≤ ρd ( y 1 , y 2 ) and ∥H ( x 1 , y ) -H ( x 2 , y ) ∥ y ≤ ρd ( x 1 , x 2 ) .

It is worth mentioning that Definition 1 implies the joint Lipschitzness over the product manifold M x ×M y , which is verified in Appendix C.2. Due to the possible nonconvexity for the upper level problem, the optimality is measured in terms of the Riemannian gradient norm of F ( x ) .

Definition 2 ( ϵ -stationary point) . We call x ∈ M x an ϵ -stationary point of bilevel optimization (1) if it satisfies ∥G F ( x ) ∥ 2 x ≤ ϵ .

Assumption 1. All the iterates in the lower level problem are bounded in a compact subset that contains the optimal solution, i.e., there exists a constants D k > 0 , for all k such that d ( y s k , y ∗ ( x k )) ≤ D k for all s . Such a neighbourhood has unique geodesic. We take ¯ D := max k { D 1 , ..., D k } .

Assumption 2. Function f ( x, y ) has bounded Riemannian gradients, i.e., ∥G y f ( x, y ) ∥ y ≤ M , ∥G x f ( x, y ) ∥ x ≤ M for all ( x, y ) ∈ U and the Riemannian gradients are L -Lipschitz in U .

Assumption 3. Function g ( x, y ) is µ -geodesic strongly convex in y ∈ U y for any x ∈ U x and has L Lipschitz Riemannian gradient G x g ( x, y ) , G y g ( x, y ) in U . Further, the Riemannian Hessian H y g ( x, y ) , cross derivatives G 2 xy g ( x, y ) , G 2 yx g ( x, y ) are ρ -Lipschitz in U .

Assumption 1 is standard in Riemannian optimization literature by properly bounding the domain of variables, which allows to express Riemannian distance in terms of (inverse) Exponential map. Also, the boundedness of the domain implies the bound on curvature, as is required for analyzing convergence for geodesic strongly convex lower-level problems [41, 71]. Assumptions 2 and 3 are common regularity conditions imposed on f and g in the bilevel optimization literature. This translates into the smoothness of the function F and D y ∗ ( x ) (discussed in Appendix C.3).

We first bound the estimation error of the proposed schemes of approximated hypergradient as follows. For the hypergradient computed by automatic differentiation, we highlight that due to the presence of exponential map in the chain of differentiation, it is non-trivial to explicitly express D x k y S k . Here, we adopt the property of exponential map (which is locally linear) in the ambient space [1], i.e., Exp x ( u ) = x + u + O ( ∥ u ∥ 2 x ) . This requires the use of tangent space projection of ξ in the ambient space as P x ( ξ ) , which is solved for the v such that ⟨ v, ξ ⟩ x = ⟨ u, ξ ⟩ for any ξ ∈ T x M .

For notation simplicity, we denote κ l := L µ and κ ρ := ρ µ . For analysis, we consider κ ρ = Θ( κ l ) .

Lemma 1 (Hypergradient approximation error bound) . Under Assumptions 1, 2, 3, we can bound the error for approximated hypergradient as

HINV: ∥ ̂ G hinv F ( x k ) -G F ( x k ) ∥ x k ≤ ( L + κ ρ M + κ l L + κ l κ ρ M ) d ( y ∗ ( x k ) , y k +1 ) .

CG: ∥ ̂ G cg F ( x k ) - G F ( x k ) ∥ x k ≤ ( L + κ ρ M + L ( 1 + 2 √ κ l )( κ l + Mκ ρ µ )) d ( y ∗ ( x k ) , y k +1 ) + 2 L √ κ l ( √ κ l -1 √ κ l +1 ) T ∥ ˆ v 0 k -Γ y k +1 y ∗ ( x k ) v ∗ k ∥ y k +1 , where v ∗ k = H -1 y g ( x k , y ∗ ( x k ))[ G y f ( x k , y ∗ ( x k ))] .

NS: ∥ ̂ G ns F ( x k ) -G F ( x k ) ∥ x k ≤ ( L + κ l L + κ ρ M + κ l κ ρ M ) d ( y ∗ ( x k ) , y k +1 ) + κ l M (1 -γµ ) T .

AD: Suppose further there exist C 1 , C 2 , C 3 > 0 such that ∥ D x k y s k ∥ y s k ≤ C 1 , ∥ Γ y x P x v -v ∥ y ≤ C 2 d ( x, y ) ∥ v ∥ y and D x Exp x ( u ) = P Exp x ( u ) ( id + D x u ) + E where ∥E∥ Exp x ( u ) ≤ C 3 ∥ D x u ∥ x ∥ u ∥ x for any x, y ∈ U and v ∈ T y M y , u ∈ T x M x . Then,

From Lemma 1, it is evident that the exact Hessian inverse exhibits the tightest bound, which is followed by conjugate gradient (CG) and truncated Neumann series (NS). Automatic differentiation (AD) presents the worst upper bound on the error due to the introduction of curvature constant ζ , resulting in (1 -Θ( µ 2 L 2 ζ )) S = (1 -Θ( 1 κ 2 l ζ )) S for the trailing term, which could be much larger than (1 -γµ ) T = (1 -Θ( 1 κ l )) T for NS and ( √ κ l -1 √ κ l +1 ) T = (1 -Θ( 1 √ κ l )) T for CG. Further, the error critically relies on the number of inner iterations S compared with T for CG and NS, and the constants C 1 , C 2 , C 3 can be large for manifolds with high curvature. We now present the main convergence result with the four proposed hypergradient estimation strategies.

Theorem 1. Denote ∆ 0 := F ( x 0 ) + d 2 ( y 0 , y ∗ ( x 0 )) and L F := ( L µ +1 )( L + τM µ + ρLM µ 2 + L 2 µ ) = O ( κ 3 l ) . Under Assumptions 1, 2, 3, we have the following bounds on the hypergradient norm obtained by Algorithm 1.

HINV: Let η x = 1 20 L F and S ≥ ˜ Θ( κ 2 l ζ ) . We have min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 80 L F ∆ 0 /K .

CG: Let Λ := C 2 v + κ 2 l ( 5 M 2 C 2 0 D 2 µ + 1) , where C v := Mκ ρ µ + Mκ ρ κ l µ + κ 2 l + κ l . Choosing η x = 1 24Λ , S ≥ ˜ Θ( κ 2 l ζ ) , and T cg ≥ ˜ Θ( √ κ l ) , we have min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 96Λ K ( ∆ 0 + ∥ v ∗ 0 ∥ 2 y ∗ ( x 0 ) ) .

NS: Choosing η x = 1 20 L F , S ≥ ˜ Θ( κ 2 l ζ ) , and T ns ≥ ˜ Θ( κ log( 1 ϵ )) for an arbitrary ϵ > 0 , we have min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 80 L F K ∆ 0 + ϵ 2 .

AD: Choosing η x = 1 20 L F and S ≥ ˜ Θ( κ 2 l ζ log( 1 ϵ )) for an arbitrary ϵ > 0 , we have min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 80 L F K ∆ 0 + ϵ 2 .

Complexity analysis. Based on the convergence guarantees in Theorem 1, we have analyzed (in Corollary 1), the computational complexity of the proposed algorithm with four different hypergradient estimation strategies in reaching the ϵ -stationary point. The results are summarized in Table 1. For reference, we also provide the computational cost of Euclidean algorithms which solve bilevel Euclidean optimization problem [40]. We notice that except for CG, the gradient complexity for f (i.e., G f ) matches the Euclidean version. For conjugate gradient, the complexity is higher by O ( κ l ) , which is due to the additional distortion from the use of vector transport when tracking the error of conjugate gradient at each epoch. In terms of gradient complexity for g (i.e., G g ), all deterministic methods require a higher complexity by at least ˜ O ( κ l ζ ) compared to the Euclidean baselines. This is because of the curvature distortion when analyzing the convergence for geodesic strongly convex functions. Similar comparisons can be also made with respect to the computations of cross-derivatives and Hessian vector products.

3.3 Extension to stochastic bilevel optimization

Algorithm 2 Riemannian stochastic bilevel optimization with Hessian inverse. 1: Initialize x 0 ∈ M x , y 0 ∈ M y . 2: for k = 0 , ..., K -1 do 3: y 0 k = y k . 4: for s = 0 , ..., S -1 do 5: Sample a batch B 1 . 6: y s +1 k = Exp y s k ( -η y G y g B 1 ( x k , y s k )) . 7: end for 8: Set y k +1 = y S k . 9: Sample batches B 2 , B 3 , B 4 . 10: Compute ̂ G F ( x k ) . 11: Update x k +1 = Exp x k ( -η x ̂ G F ( x k )) . 12: end for

In this section, we consider the bilevel optimization problem (1) in the stochastic setting, where f ( x, y ∗ ( x )) := 1 n ∑ n i =1 f i ( x, y ∗ ( x )) and g ( x, y ) := 1 m ∑ m i =1 g i ( x, y ) . The algorithm for solving the stochastic bilevel optimization problem is in Algorithm 2, where we sample B 1 , B 2 , B 3 , B 4 afresh every iteration. The batch index is omitted for clarity. The batches are sampled uniformly at random with replacement such that the mini-batch gradient is an unbiased estimate of the full gradient. Here, we denote f B ( x, y ) := 1 |B| ∑ i ∈B f i ( x, y ) and similarly for g . We let [ n ] := { 1 , . . . , n } .

In Step 10 of Algorithm 2, we can employ any hypergradient estimator proposed in Section 3.1. In this work, we only show convergence un- der the Hessian inverse approximation of hypergradient, i.e., ̂ G F ( x k ) = G x f B 2 ( x k , y k +1 ) -G 2 xy g B 3 ( x k , y k +1 )[ H -1 y g B 4 ( x k , y k +1 )[ G y f B 2 ( x k , y k +1 )]] . Similar analysis can be followed for other approximation strategies. The theoretical guarantees are in Theorem 2, where we require Assumption 4, which is common in existing works for analyzing stochastic algorithms on Riemannian manifolds [42, 23, 22].

Assumption 4. Under stochastic setting, Assumption 1 holds and Assumptions 2, 3 are satisfied for component functions f i ( x, y ) , g j ( x, y ) , for all i ∈ [ n ] , j ∈ [ m ] . Further, stochastic gradient, Hessian, and cross derivatives are unbiased estimates.

Theorem 2. Under Assumption 4, consider Algorithm 2. Suppose we choose η x = 1 20 L F , S ≥ ˜ Θ( κ 2 l ζ ) , and |B 1 | , |B 2 | , |B 3 | , |B 4 | ≥ Θ( κ 2 l ϵ -1 ) for an arbitrary ϵ > 0 . Then we have min k =0 ,...,K -1 E ∥G F ( x k ) ∥ 2 x k ≤ 80 L F ∆ 0 K + ϵ 2 and the gradient complexity to reach ϵ -stationary solution is G f = O ( κ 5 l ϵ -2 ) , G g = ˜ O ( κ 9 l ζϵ -2 ) , JV g = O ( κ 5 l ϵ -2 ) .

In Table 1, we compare our attained complexities with that of stocBiO [40], which makes use of a truncated Neumann series. With exact Hessian inverse, we can match the G f and JV g complexities

with stocBio. For the G g complexity, the additional curvature constant is inevitable from the convergence analysis for geodesic strongly convex functions. Nevertheless we observe the same order dependency on κ l . This is mainly due to the analysis where we choose a smaller stepsize η y = Θ( µ L 2 ) compared to Θ( 2 L + µ ) in [40]. The larger stepsize, despite increasing the convergence rate, also increases the variance under stochastic setting. We believe an order of Θ( κ l ) lower can be established for stocBio, following our analysis.

3.4 Extension to retraction

Our analysis till now has been limited to the use of the exponential map. However, the retraction mapping is often preferred over the exponential map due to its lower computational cost. Here, we show that use of retraction in our algorithms also leads to similar convergence guarantees.

Assumption 5. There exist constants c ≥ 1 , c R ≥ 0 such that d 2 ( x, y ) ≤ c ∥ u ∥ 2 x and ∥ Exp -1 x ( y ) -u ∥ x ≤ c R ∥ u ∥ 2 , for any x, y = Retr x ( u ) ∈ U .

Assumption 5 is standard (e.g. in [42, 23]) in bounding the error between exponential map and retraction given that retraction is a first-order approximation to the exponential map.

Theorem 3. Suppose Assumptions 1, 2, 3 and 5 hold and let ˜ L F = 4 κ l c R M +5¯ cL F . Then consider Algorithm 1 with exponential map replaced with general retraction. We can obtain the following bounds.

HINV : Let η x = Θ(1 / ˜ L F ) , S ≥ ˜ Θ( κ 2 l ζ ) . Then min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 16 ˜ L F ∆ 0 /K .

CG : Let η x = Θ(1 / ˜ Λ) , S ≥ ˜ Θ( κ 2 l ζ ) , T cg ≥ ˜ Θ( √ κ l ) , where ˜ Λ = C 2 v ¯ c + κ 2 l ( 5 M 2 C 2 0 ¯ D 2 µ +¯ c ) . Then min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 96 ˜ Λ K (∆ 0 + ∥ v ∗ 0 ∥ 2 y ∗ ( x 0 ) ) .

NS : Let η x = Θ(1 / ˜ L F ) , S ≥ ˜ Θ( κ 2 l ζ ) . Then for an arbitrary ϵ > 0 , T ns ≥ ˜ Θ( κ l log(1 /ϵ )) , we have min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 16 ˜ L F K ∆ 0 + ϵ 2 .

AD : Let η x = Θ(1 / ˜ L F ) , S ≥ ˜ Θ( κ 2 l ζ log(1 /ϵ )) . Then for an arbitrary ϵ > 0 , we have min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 16 ˜ L F K ∆ 0 + ϵ 2 .

Theorem 3 demonstrates that employing a general retraction preserves the same order of convergence and complexity as the exponential map in Theorem 1. This is due to the fact that ˜ L F = Θ( L F ) and ˜ Λ = Θ(Λ) , where L F and Λ are as defined in Theorem 1. In addition, when exponential map is used, Theorem 3 recovers the results in Theorem 1 as c R = 0 and ¯ c = 1 .

4 Experiments

This section explores various applications of bilevel optimization problems over manifolds. All the experiments are implemented based on Geoopt [44] and the codes are available at https: //github.com/andyjm3/rhgd .

4.1 Synthetic problem

We consider the following bilevel optimization problem on the Stiefel manifold St( d, r ) = { W ∈ R d × r : W ⊤ W = I r } and SPD manifold S d ++ = { M ∈ R d × d : M ≻ 0 } (in Appendix A):

where X ∈ R n × d , Y ∈ R n × r , with n ≥ d ≥ r , are given matrices and ν > 0 is the regularization parameter. The above is a synthetically constructed problem that aims to maximize the similarity between X and Y in different feature dimensions. We align X and Y to the same dimension via W ∈ St( d, r ) and also learn an appropriate geometric metric M ∈ S d ++ in the lower-level problem [69]. The geodesic convexity of the lower-level problem and the Hessian inverse expression are discussed in Appendix H.1.

0

25

50

75

100

125

150

175

200

Epochs

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Upper Objective

cg-20

cg-50

ns-20

ns-50

ad-20

ad-50

hinv-20

hinv-50

(a)

0

5

10

15

20

25

30

35

Time

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Upper Objective

(b)

0

25

50

75

100

125

150

175

200

Epochs

10

3

10

2

10

1

Hypergrad error

(c)

0

25

50

75

100

125

150

175

200

Epochs

10

5

10

4

10

3

10

2

Hypergrad error

30-1.0

50-1.0

100-1.0

200-1.0

200-0.5

200-2.0

(d)

Results. We generate random data matrices X , Y with n = 100 , d = 50 , and r = 20 . We set ν = 0 . 01 and fix η x = η y = 0 . 5 . We compare the three proposed strategies for approximating the hypergradient where we select γ = 1 . 0 and T ns = 50 for Neumann series (NS) and set maximum iterations T cg for conjugate gradient (CG) to be 50 and break once the residual reaches a tolerance of 10 -10 . We set the number of outer iterations (epochs) K to be 200 . Figure 1 compares RHGD with different approximation strategies implemented with S = 20 or 50 number of inner iterations.

4.2 Hyper-representation over SPD manifolds

Hyper-representation [54, 61] aims to solve a regression/classification task while searching for the best representation of the data. It can be formulated as a bilevel optimization problem, where the lower-level optimizes the regression/classification parameters while the upper-level searches for the optimal embedding of the inputs. Suppose we are given a set of SPD matrices, D = { A i } n i =1 where A i ∈ S d ++ and the task is to learn a low-dimensional embedding of A i while remaining close to their semantics labels. In particular, we partition the set into a training set D tr and validation set D val .

Shallow hyper-representation for regression. We consider a shallow learning paradigm over D through the regression task. The representation is parameterized with W ⊤ A i W for W ∈ St( d, r ) . The requirement of orthogonality on W follows [38, 30, 33] that ensures the learned representations are SPD. The learned representation is then transformed to a Euclidean space for performing regression, namely through a matrix logarithm (that acts as a bijective map between the space of SPD matrices and symmetric matrices) and a vectorization operation vec( · ) that extract the upper-triangular part of the symmetric matrix. The bilevel optimization problem is

The regularization λ > 0 ensures the lower-level problem is strongly convex. The upper-level problem is on the validation set while the lower-level problem is on the training set. We generate random W , A i and β and construct y with y i = vec(logm( W ⊤ A i W )) β + ϵ i , where ϵ i ∼ N (0 , 1) . We generate 200 A i with |D val | = 100 and |D tr | = 100 . In Figure 2a, we show the loss on validation set (the upper loss) in terms of number of outer iterations. We compare both the deterministic (RHGD) and stochastic (RSHGD) versions of Riemannian hypergradient descent. We again observe that the best performance is attained by either the ground-truth Hessian inverse or the conjugate gradient. NS requires carefully selecting the hyperparameters γ, T , which pose difficulties in real applications. For the stochastic versions, all the methods perform similarly.

Deep hyper-representation for classification. We now explore a 2-layer SPD network [38] for classifying ETH-80 image set [46]. The dataset consists of 8 classes, each with 10 objects. Each object is represented by an image set consisting of images taken from different viewing angles. Here, we represent each image set by taking the covariance matrix of the images in the same set after resizing them into 10 × 10 . This results in 80 SPD matrices A i of size 100 × 100 for classification. Let Φ( A i ) = vec(logm( W ⊤ 2 ReEig( W ⊤ 1 A i W 1 ) W 2 )) be the output of the 2 layer network where ReEig( A ) = U max { ϵ I , Σ } U ⊤ is the eigenvalue rectifying activation with the eigenvectors U and

0

20

40

60

80

100

Epochs

10

1

10 0

Upper loss

RHGD-cg

RSHGD-cg

RHGD-ns

RSHGD-ns

RHGD-ad

RSHGD-ad

RHGD-hinv

RSHGD-hinv

(a) Shallow hyperrep

0

5

10

15

20

25

30

Epochs

0.2

0.4

0.6

0.8

1.0

Val Accuracy

cg

ns

ad

(b) Deep hyperrep

0

20

40

60

80

Time

0.2

0.4

0.6

0.8

1.0

Val Accuracy

cg

ns

ad

(c) Deep hyperrep

0

500

1000

1500

2000

2500

Runtime

0.350

0.375

0.400

0.425

0.450

0.475

0.500

0.525

0.550

Test Accuracy

RHGD-cg

PHGD-cg

RHGD-ns

PHGD-ns

RHGD-ad

PHGD-ad

(d) MiniImageNet

eigenvalues Σ of A . We consider the same bilevel optimization as above except the least-squares loss function becomes the cross-entropy loss. Here we sample 5 samples from each class to form the training set and the rest as the validation set. We set d 1 = 20 , d 2 = 5 , and fix learning rate to be 0 . 1 for both lower and upper problems. Figures 2b and 2c show the good performance on the validation accuracy (upper-level loss).

4.3 Riemannian meta learning

Meta learning [16, 34] allows adaptation of models to new tasks with minimal amount of additional data and training, by distilling past learning experiences. A recent work [64] considers meta learning with orthogonality constraint. In particular, the upper-level optimization searches for the base parameters shared by all tasks while the lower level optimizes over the task-specific parameters to ensure generalization ability. Let P T denote the distribution of meta tasks and for each training epoch, we sample m tasks D ℓ ∼ P T , ℓ = 1 , ..., m . Each task is composed of a support and query set denoted by D ℓ s , D ℓ q , and the task is to learn a set of base parameters Θ such that the model can quickly adapt to the query set from the support set by adjusting only a few parameters w . For each task, the task-specific parameter w ∗ ℓ is learned from the support set, which is used to update the base parameters by minimizing the loss over the query set. In standard settings, w ℓ corresponds to the final linear layer of a neural network [39, 40]. Here, we adopt the setup with w ℓ to be the last layer parameters in the Euclidean space while enforcing Θ on the Stiefel manifold. The problem of Riemannian metalearning is min Θ ∈ St 1 m ∑ m ℓ =1 L (Θ , w ∗ ℓ ; D ℓ q ) s.t. w ∗ ℓ = arg min w ℓ 1 m ∑ m ℓ =1 L (Θ , w ℓ ; D ℓ s ) + R ( w ℓ ) , where D ℓ s , D ℓ q are the support and query sets for task ℓ and R ( · ) is a regularizer that ensures strong convexity of the lower-level problem.

Results. We consider 5-ways 5-shots meta learning over the MiniImageNet dataset [59] where the backbone network is a 4-block CNN with the kernel of the first 2 layers constrained to be orthogonal in terms of the output channel (following [48]). The kernel size is 3 × 3 and we consider 16 output channels with a padding of 1 . Each convolutional block consists of a convolutional layer, followed by a ReLU activation, a max-pooling and a batch normalization layer. Θ , thus, has the dimension (16 ∗ 3 ∗ 3) × 16 = 144 × 16 , which is constrained to the Stiefel manifold.

In Figure 2d, we plot the test accuracy averaged for over 200 tasks. We compare RHGD with an extrinsic update baseline PHGD, which projects the update from the Euclidean space to the Stiefel manifold at every iteration. We observe the RHGD converges faster compared to the extrinsic update PHGD, thereby showing the benefit of the Riemannian modeling.

4.4 Unsupervised domain adaptation

Given two marginals µ ∈ R n , ν ∈ R m with equal total mass, i.e., µ ⊤ 1 n = ν ⊤ 1 m = 1 where we assume unit mass without loss of generality. Let Π( µ , ν ) := { Γ ∈ R n × m : Γ > 0 , Γ1 m =

µ , Γ ⊤ 1 n = ν } be the set of doubly stochastic matrices with strictly positive entries. From [15, 57], it is known that the set forms a Riemannian manifold with the Fisher metric.

Given a supervised source dataset X ∈ R n × d and an unsupervised target dataset Y ∈ R m × d ( n, m ≥ d ), we consider the unsupervised domain adaptation problem to classify target domain instances. Using the optimal transport framework [58, 12], we pose this as a bilevel problem:

where H ( Γ ) = -⟨ Γ , log Γ ⟩ is the entropy regularization, C ( X , Y ) = diag( XX ⊤ ) 1 ⊤ m + 1 n diag( YY ⊤ ) -2 XY ⊤ is the pairwise squared distance matrix, and Π denotes the doubly stochastic manifold [15]. Here, dist is the geodesic distance between SPD matrices and α ∈ [0 , 1] .

The lower-level problem in (3) finds M ∗ which is the weighted geometric mean between X ⊤ X and Y ⊤ Γ ⊤ ΓY [6]. Conceptually, learning of M allows to align the features of the source and target instances. The upper-level problem, on the other hand, minimizes the Mahalanobis distance between the source and target domain instances parameterized by M ∗-1 . An interpretation is that the matrix M ∗ leads to whitening of the data (i.e., XM ∗-1 / 2 is the whitened data) [5].

After the transport plan Γ ∗ is learned, we employ barycentric projection using Γ ∗ to transport the source points to the target domain and and employ the nearest neighbour (1-NN) classifier for the target dataset classification . For barycentric projection, we project the source samples X to the target Y by solving x i = arg min x i ∈ R d ∑ m j =1 Γ ∗ i,j ∥ M ∗-1 / 2 x i -M ∗-1 / 2 y j ∥ 2 = µ -1 i ( ∑ m j =1 Γ ∗ i,j y j ) . Then, a nearest-neighbour (NN) classifier is used to classify the samples in the target given the source labels based on the distance computed with M ∗ , i.e., C ( XM ∗-1 / 2 , YM ∗-1 / 2 ) .

Results. We consider the Caltech-Office dataset [20], which is commonly used for domain adaptation. The dataset contains images from four domains in ten classes, i.e., Amazon (A), the Caltech image dataset (C), DSLR (D), and Webcam (W), each with containing 958, 1123, 157, and 295 samples respectively. Hence, there are 12 domain adaptation tasks, e.g., A → D implies A is the source and D is the target. Each domain has the same ten classes. The goal is to classify images from target domain given source domain. For preprocessing, we normalize the samples to have unit norm and reduce the dimensionality to 128 by mean pooling every 64 columns.

We compare our proposed bilevel approach (3) with single-level optimal transport baselines, i.e., solving min Γ ∈ Π( µ , ν ) ⟨ Γ , C ( X , Y ) ⟩ -λH ( Γ ) , followed by the same barycentric projection. Specifically, the baselines are: (1) optimal transport where λ = 0 (labelled as OT-EMD) and (2) optimal transport with the Sinkhorn algorithm (labelled as OT-SKH). OT-EMD employs the earth mover distance while OT-SKH employs the Sinkhorn distance [13]. We implement the two OT baselines with the POT Python library [17]. λ is tuned in OT-SKH for each source-target pair. The best validation results are obtained by setting λ = 5 × 10 -3 for all the problem pairs except the W → D pair for which λ = 10 -3 gives the best result. For our proposed bilevel approach, we set λ = 0 and α = 0 . 5 .

In Table 2, we observe that the proposed bilevel approach obtains better generalization performance than the baselines across all the tasks. This showcases the utility of learning the whitening metric M -1 in a bilevel setting.

5 Conclusion

In this work, we have proposed a framework for tackling bilevel optimization over Riemannian manifolds. We discussed various hypergradient approximation strategies (conjugate gradients, truncated Neumann series, and automatic differentiation) and provide error bounds. Our proposed algorithms rely only on gradient updates and make use of retraction which scale well across problems. We illustrate the efficacy of the proposal approach in several machine learning applications.

Although in this work, we focus on geodesic strongly convex lower-level problems, our framework can be extended to relax such assumption to (geodesic) convexity with an extra strongly convex regularizer [2], or to (Riemannian) PL condition where a global minimizer exists [9]. Furthermore, we believe there is potential to improve the current results in stochastic bilevel optimization by reducing the strict requirements on batch size. Additionally, the dependency on the curvature constant could also be further optimized.

References

P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds . Princeton University Press, 2008.

Jan Harold Alcantara and Akiko Takeda. Theoretical smoothing frameworks for general nonsmooth bilevel problems. Technical report, arXiv preprint arXiv:2401.17852, 2024.

Foivos Alimisis, Antonio Orvieto, Gary Bécigneul, and Aurelien Lucchi. A continuous-time perspective for modeling acceleration in Riemannian optimization. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2020.

Vincent Arsigny, Pierre Fillard, Xavier Pennec, and Nicholas Ayache. Geometric means in a novel vector space structure on symmetric positive-definite matrices. SIAM Journal on Matrix Analysis and Applications , 29(1):328-347, 2007.

Mikel Artetxe, Gorka Labaka, and Eneko Agirre. Generalizing and improving bilingual word embedding mappings with a multi-step framework of linear transformations. In the AAAI Conference on Artificial Intelligence (AAAI) , 2018.

Rajendra Bhatia. Positive definite matrices. In Positive Definite Matrices . Princeton university press, 2009.

Henri Bonnel, Léonard Todjihoundé, and Constantin Udri¸ ste. Semivectorial bilevel optimization on Riemannian manifolds. Journal of Optimization Theory and Applications , 167:464-486, 2015.

Nicolas Boumal. An introduction to optimization on smooth manifolds . Cambridge University Press, 2023.

Lesi Chen, Jing Xu, and Jingzhao Zhang. On finding small hyper-gradients in bilevel optimization: Hardness results and improved analysis. In Annual Conference on Learning Theory (COLT) , 2024.

Tianyi Chen, Yuejiao Sun, Quan Xiao, and Wotao Yin. A single-timescale method for stochastic bilevel optimization. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2022.

Tianyi Chen, Yuejiao Sun, and Wotao Yin. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems. In Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence , 39(9):1853-1865, 2017.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems (NeurIPS) , 2013.

Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Ahmed Douik and Babak Hassibi. Manifold optimization over the set of doubly stochastic matrices: A second-order geometry. IEEE Transactions on Signal Processing , 67(22):57615774, 2019.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning (ICML) , 2017.

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. POT: Python optimal transport. Journal of Machine Learning Research , 22(1):3571-3578, 2021.

Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. In International Conference on Machine Learning (ICML) , 2018.

Saeed Ghadimi and Mengdi Wang. Approximation methods for bilevel programming. Technical report, arXiv preprint arXiv:1802.02246, 2018.

Boqing Gong, Yuan Shi, Fei Sha, and Kristen Grauman. Geodesic flow kernel for unsupervised domain adaptation. In Conference on Computer Vision and Pattern Recognition (CVPR) , 2012.

Riccardo Grazzi, Luca Franceschi, Massimiliano Pontil, and Saverio Salzo. On the iteration complexity of hypergradient computation. In International Conference on Machine Learning (ICML) , 2020.

Andi Han and Junbin Gao. Riemannian stochastic recursive momentum method for non-convex optimization. In International Joint Conference on Artificial Intelligence (IJCAI) , 2021.

Andi Han and Junbin Gao. Improved variance reduction methods for Riemannian non-convex optimization. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(11):76107623, 2022.

Andi Han, Bamdev Mishra, Pratik Jawanpuria, and Junbin Gao. Learning with symmetric positive definite matrices via generalized Bures-Wasserstein geometry. In International Conference on Geometric Science of Information (GSI) , 2023.

Andi Han, Bamdev Mishra, Pratik Jawanpuria, and Junbin Gao. Nonconvex-nonconcave minmax optimization on Riemannian manifolds. Transactions on Machine Learning Research , 2023.

Andi Han, Bamdev Mishra, Pratik Jawanpuria, and Junbin Gao. Riemannian accelerated gradient methods via extrapolation. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2023.

Andi Han, Bamdev Mishra, Pratik Jawanpuria, Pawan Kumar, and Junbin Gao. Riemannian Hamiltonian methods for min-max optimization on manifolds. SIAM Journal on Optimization , 33(3):1797-1827, 2023.

Andi Han, Bamdev Mishra, Pratik Kumar Jawanpuria, and Junbin Gao. On Riemannian optimization over positive definite matrices with the Bures-Wasserstein geometry. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Pierre Hansen, Brigitte Jaumard, and Gilles Savard. New branch-and-bound rules for linear bilevel programming. SIAM Journal on Scientific and Statistical Computing , 13(5):1194-1217, 1992.

Mehrtash Harandi, Mathieu Salzmann, and Richard Hartley. Dimensionality reduction on SPD manifolds: The emergence of geometry-aware methods. IEEE Transactions on Pattern Analysis and Machine Intelligence , 40(1):48-62, 2017.

Ryuichiro Hataya and Makoto Yamada. Nyström method for accurate and scalable implicit differentiation. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2023.

Mingyi Hong, Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang. A two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application to actorcritic. SIAM Journal on Optimization , 33(1):147-180, 2023.

Inbal Horev, Florian Yger, and Masashi Sugiyama. Geometry-aware principal component analysis for symmetric positive definite matrices. In Asian Conference on Machine Learning (ACML) , 2016.

Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. Meta-learning in neural networks: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(9):5149-5169, 2021.

Zihao Hu, Guanghui Wang, Xi Wang, Andre Wibisono, Jacob D Abernethy, and Molei Tao. Extragradient type methods for Riemannian variational inequality problems. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2024.

Feihu Huang and Shangqian Gao. Riemannian gradient methods for stochastic composition problems. Neural Networks , 153:224-234, 2022.

Feihu Huang and Shangqian Gao. Gradient descent ascent for minimax problems on Riemannian manifolds. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2023.

Zhiwu Huang and Luc Van Gool. A Riemannian network for spd matrix learning. In AAAI Conference on Artificial Intelligence (AAAI) , 2017.

Kaiyi Ji, Jason D Lee, Yingbin Liang, and H Vincent Poor. Convergence of meta-learning with task-specific adaptation over partial parameters. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. In International Conference on Machine Learning (ICML) , 2021.

Michael Jordan, Tianyi Lin, and Emmanouil-Vasileios Vlatakis-Gkaragkounis. First-order algorithms for min-max optimization in geodesic metric spaces. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Hiroyuki Kasai, Hiroyuki Sato, and Bamdev Mishra. Riemannian stochastic recursive gradient algorithm. In International Conference on Machine Learning (ICML) , 2018.

Philip A Knight. The Sinkhorn-Knopp algorithm: convergence and applications. SIAM Journal on Matrix Analysis and Applications , 30(1):261-275, 2008.

Max Kochurov, Rasul Karimov, and Serge Kozlukov. Geoopt: Riemannian optimization in PyTorch. Technical report, arXiv preprint arXiv:2005.02819, 2020.

Jeongyeol Kwon, Dohyun Kwon, Stephen Wright, and Robert D Nowak. A fully first-order method for stochastic bilevel optimization. In International Conference on Machine Learning (ICML) , 2023.

Bastian Leibe and Bernt Schiele. Analyzing appearance and contour based methods for object categorization. In Conference on Computer Vision and Pattern Recognition (CVPR) , 2003.

Jiaxiang Li and Shiqian Ma. Riemannian bilevel optimization. Technical report, arXiv preprint arXiv:2402.02019, 2024.

Jun Li, Fuxin Li, and Sinisa Todorovic. Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform. In International Conference on Learning Representations (ICLR) , 2019.

Jiagen Liao and Zhongping Wan. Inexact restoration methods for semivectorial bilevel programming problem on Riemannian manifolds. Axioms , 11(12):696, 2022.

Jiagen Liao and Zhongping Wan. On the Karush-Kuhn-Tucker reformulation of the bilevel optimization problems on Riemannian manifolds. Filomat , 36(11):3609-3624, 2022.

Zhenhua Lin. Riemannian geometry of symmetric positive definite matrices via Cholesky decomposition. SIAM Journal on Matrix Analysis and Applications , 40(4):1353-1370, 2019.

Bo Liu, Mao Ye, Stephen Wright, Peter Stone, and Qiang Liu. Bome! bilevel optimization made easy: A simple first-order approach. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In International Conference on Learning Representations (ICLR) , 2018.

Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2020.

Luigi Malagò, Luigi Montrucchio, and Giovanni Pistone. Wasserstein Riemannian geometry of Gaussian densities. Information Geometry , 1:137-179, 2018.

David Martínez-Rubio, Christophe Roux, Christopher Criscitiello, and Sebastian Pokutta. Accelerated methods for Riemannian min-max optimization ensuring bounded geometric penalties. Technical report, arXiv preprint arXiv:2305.16186, 2023.

Bamdev Mishra, NTV Satyadev, Hiroyuki Kasai, and Pratik Jawanpuria. Manifold optimization for non-linear optimal transport problems. Technical report, arXiv preprint arXiv:2103.00902, 2021.

Gabriel Peyre and Marco Cuturi. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning , 11(5-6):355-607, 2019.

Mengye Ren, Sachin Ravi, Eleni Triantafillou, Jake Snell, Kevin Swersky, Josh B. Tenenbaum, Hugo Larochelle, and Richard S. Zemel. Meta-learning for semi-supervised few-shot classification. In International Conference on Learning Representations (ICLR) , 2018.

Chenggen Shi, Jie Lu, and Guangquan Zhang. An extended Kuhn-Tucker approach for linear bilevel programming. Applied Mathematics and Computation , 162(1):51-63, 2005.

Daouda Sow, Kaiyi Ji, and Yingbin Liang. On the convergence theory for Hessian-free bilevel algorithms. Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Rhea Sukthanker, Zhiwu Huang, Suryansh Kumar, Erik Endsjo Goron, Yan Wu, and Luc Van Gool. Neural architecture search of spd manifold networks. In International Joint Conference on Artificial Intelligence (IJCAI) , 2021.

Yue Sun, Nicolas Flammarion, and Maryam Fazel. Escaping from saddle points on Riemannian manifolds. Advances in Neural Information Processing Systems (NeurIPS) , 2019.

Hadi Tabealhojeh, Peyman Adibi, Hossein Karshenas, Soumava Kumar Roy, and Mehrtash Harandi. RMAML: Riemannian meta-learning with orthogonality constraints. Pattern Recognition , 140:109563, 2023.

Ioannis Tsaknakis, Prashant Khanduri, and Mingyi Hong. An implicit gradient-type method for linearly constrained bilevel problems. In International Conference on Acoustics, Speech and Signal Processing (ICASSP) , 2022.

Mengdi Wang, Ethan X Fang, and Han Liu. Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. Mathematical Programming , 161:419-449, 2017.

Xi Wang, Deming Yuan, Yiguang Hong, Zihao Hu, Lei Wang, and Guodong Shi. Riemannian optimistic algorithms. Technical report, arXiv preprint arXiv:2308.16004, 2023.

Quan Xiao, Han Shen, Wotao Yin, and Tianyi Chen. Alternating projected sgd for equalityconstrained bilevel optimization. In International Conference on Artificial Intelligence and Statistics (AISTATS) , 2023.

Pourya Zadeh, Reshad Hosseini, and Suvrit Sra. Geometric mean metric learning. In International Conference on Machine Learning (ICML) , 2016.

Dewei Zhang and Sam Davanloo Tajbakhsh. Riemannian stochastic gradient method for nested composition optimization. Technical report, arXiv preprint arXiv:2207.09350, 2022.

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.

Hongyi Zhang and Suvrit Sra. First-order methods for geodesically convex optimization. In Conference on Learning Theory (COLT) , 2016.

Peiyuan Zhang, Jingzhao Zhang, and Suvrit Sra. Sion's minimax theorem in geodesic metric spaces and a Riemannian extragradient algorithm. SIAM Journal on Optimization , 33(4):28852908, 2023.

Appendix

A Riemannian geometries of considered manifolds

Symmetric positive definite (SPD) manifold. The SPD manifold of size d is denoted as S d ++ := { X ∈ R d × d : X ⊤ = X , X ≻ 0 } and the commonly considered Riemannian metric is the affineinvariant metric ⟨ U , V ⟩ X = tr( X -1 UX -1 V ) [6], for U , V ∈ T X S d ++ . Other Riemannian metrics, such as (generalized) Bures-Wasserstein [55, 24], log-Euclidean [4] and log-Cholesky [51] metrics can also be considered. The exponential map is given by Exp X ( U ) = X expm( X -1 U ) where expm( · ) denotes the principal matrix exponential. The corresponding logarithm map is given by log X ( Y ) = X logm( X -1 Y ) . Its Riemannian gradient of a real-valued function f is derived as grad f ( X ) = X ∇ f ( X ) X and the Riemannian Hessian is Hess f ( X )[ U ] = Dgrad f ( X )[ U ] -{ UX -1 grad f ( X ) } S = X ∇ 2 f ( X )[ U ] X + { U ∇ f ( X ) X } S where we use { A } S := ( A + A ⊤ ) / 2 .

Stiefel manifold. The Stiefel manifold is the set of orthonormal matrices, i.e., St( d, r ) := { X ∈ R d × r : X ⊤ X = I } . A common Riemannian metric is the Euclidean inner product. We consider the QR-based retraction in the experiment, which is Retr X ( U ) = qf( X + U ) where qf( · ) extracts the Q-factor from the QR decomposition. Let the orthogonal projection to the tangent space be denoted as P X ( U ) = U -X { X ⊤ U } S . Then, the Riemannian gradient and Riemannian Hessian are given by grad f ( X ) = P X ( ∇ f ( X )) and Hess f ( X )[ U ] = P X ( ∇ 2 f ( X )[ U ] -U { X ⊤ ∇ f ( X ) } S ) .

Doubly stochastic manifold. The doubly stochastic manifold (or coupling manifold) between two discrete probability measures µ, ν with marginals a ∈ R m , b ∈ R n is the set Π( µ, ν ) = { Γ ∈ R m × n : Γ ij > 0 , Γ1 n = a , Γ ⊤ 1 m = b } . It can be equipped with the Fisher information metric, defined as ⟨ U , V ⟩ Γ = ∑ i,j ( U ij V ij ) / Γ ij for any U , V ∈ T Γ Π( µ, ν ) . The retraction is given by Retr Γ ( U ) = Sinkhorn( Γ ⊙ exp( U ⊘ Γ )) where exp , ⊙ , ⊘ are elementwise exponential, product, and division operations. Sinkhorn ( · ) represents the Sinkhorn-Knopp iterations for balancing a matrix [43].

B Important Lemmas

Proposition 2 ([8]) . In a totally normal neighbourhood U ⊆ M , a function f : U → R is µ -geodesic strongly convex, then it satisfies for all x, y ∈ U

If a function f has L -Lipschitz Riemannian gradient, then it satisfies for all x, y ∈ U

Lemma 2 ([63, 26]) . There exists a constant C 0 > 0 such that for any y 1 , y 2 , y 3 ∈ U y , u ∈ T y 1 M y , ∥ Γ y 3 y 2 Γ y 2 y 1 u -Γ y 3 y 1 u ∥ ≤ C 0 d ( y 1 , y 2 ) d ( y 2 , y 3 ) ∥ u ∥ y 1

Lemma 3 (Trigonometric distance bound [72, 71, 28]) . Let x a , x b , x c ∈ U ⊆ M and denote a = d ( x b , x c ) , b = d ( x a , x c ) and c = d ( x a , x b ) as the geodesic side lengths. Then,

where ζ = √ | κ -| ¯ D tanh( √ | κ -| ¯ D ) if κ -< 0 and ζ = 1 if κ -≥ 0 . Here, ¯ D denotes the diameter of U and κ -denotes the lower bound of the sectional curvature of U .

C Proofs for Section 3.1

C.1 Proof of Proposition 1

Proof of Proposition 1. By the first-order optimality condition, y ∗ ( x ) satisfies G y g ( x, y ∗ ( x )) = 0 ∈ T y ∗ ( x ) M y . Based on Theorem 5 in [25], taking the (implicit) derivative of the equality with respect to x yields G 2 yx g ( x, y ∗ ( x ))[ u ] + H y g ( x, y ∗ ( x ))[D y ∗ ( x )[ u ]] = 0 for any u ∈ T x M x . This gives

D y ∗ ( x ) = -H -1 y g ( x, y ∗ ( x )) ◦ G 2 yx g ( x, y ∗ ( x )) . Notice that D y ∗ ( x ) : T x M x → T y ∗ ( x ) M y ( x ) , its adjoint operator (D y ∗ ( x )) † is derived as follows. For any u ∈ T x M x , v ∈ T y ∗ ( x ) M y

where the first equality uses the definition of adjoint operator and the third equality is due to the self-adjointness of Riemannian Hessian (inverse) and the last equality is due to Proposition D.2 in [27] that G 2 xy g and G 2 yx g are adjoint operators. By identification, we have (D y ∗ ( x )) † = -G 2 xy g ( x, y ∗ ( x )) ◦ H -1 y g ( x, y ∗ ( x )) .

Finally by the chain rule, we obtain (from the definition of Riemannian gradient), for any u ∈ T x M x

By identification the proof is complete.

C.2 On Lipschitzness of gradients

Proposition 3. If a bifunction f ( x, y ) has L -Lipschitz Riemannian gradient, then it satisfies ∥G f ( z 1 ) -Γ z 1 z 2 G f ( z 2 ) ∥ z 1 ≤ 2 Ld ( z 1 , z 2 ) , where we let z = ( x, y ) . If an operator G ( x, y ) : T y M y → T x M x is ρ -Lipschitz, then it satisfies ∥G ( z 1 ) -Γ x 1 x 2 G ( z 2 )Γ y 2 y 1 ∥ x 1 ≤ ρd ( z 1 , z 2 ) . If an operator H ( x, y ) : T y M y → T x M x is ρ -Lipschitz, then it satisfies ∥H ( z 1 ) -Γ y 1 y 2 H ( z 2 )Γ y 2 y 1 ∥ y 1 ≤ ρd ( z 1 , z 2 ) .

Proof of Proposition 3. From the definition of Riemannian gradient of product manifold we have

where we use triangle inequality of Riemannian norm.

Similarly, for the other two claims, we verify

The same arguments also hold for H ( x, y ) and hence the proof is omitted.

C.3 Boundedness of ingredients

Lemma 4. Under Assumptions 1, 2, 3, we can show

Proof of Lemma 4. (4.1) First we have for any v ∈ T y M y

where we use the fact that d (Exp y ( ξ ) , y ) = ∥ ξ ∥ y . The operator norm is the same between G 2 xy g ( x, y ) and G 2 yx g ( x, y ) is due to the adjointness. This proves the first claim.

(4.2) We first verify D y ∗ ( x ) can be bounded as

and D y ∗ ( x ) is also Lipschitz as

(4.3) Now suppose we let c : [0 , 1] →M y , defined as c ( t ) := y ∗ ( γ ( t )) where γ : [0 , 1] →M x is a geodesic that connects x 1 , x 2 , i.e., γ (0) = x 1 , γ (1) = x 2 . Then

where we use the fact that the manifold is complete.

(4.4) For the second claim, we first notice for any (invertible) linear operators A,B , A -1 -B -1 = A -1 ( B -A ) B -1 and thus ∥ A -1 -B -1 ∥ ≤ ∥ A -1 ∥∥ A -B ∥∥ B -1 ∥ for some well-defined norm ∥ · ∥ . Here substituting A = Γ y 2 y 1 H y g ( x, y 1 )Γ y 1 y 2 , B = H y g ( x, y 2 ) , we have

where we notice (Γ y 2 y 1 H y ( x, y 1 )Γ y 1 y 2 ) -1 = Γ y 2 y 1 H y ( x, y 1 ) -1 Γ y 1 y 2 and use the isometry property of parallel transport. The same argument applies for ∥H -1 g ( x 1 , y ) -H -1 g ( x 2 , y ) ∥ y . (4.5) we have

From Assumption 2, 3 and Lemma 4, we can obtain

Similarly, we have

and

Combining all the results together, we can show

which completes the proof.

C.4 On strong convexity of the lower-level problem

Lemma 5 (Convergence under strong convexity) . Under Assumptions 1, 2, 3, suppose η y < µ L 2 ζ , where ζ ≥ 1 is a curvature constant defined in Lemma 3, then we have d 2 ( y s +1 k , y ∗ ( x k )) ≤ (1 + η 2 y ζL 2 -η y µ ) d 2 ( y s k , y ∗ ( x k )) .

Proof of Lemma 5. We apply the trigonometric distance bound from Lemma 3 to obtain

where the second inequality is due to geodesic strong convexity and the third inequality is due to ∥G y g ( x k , y s k ) ∥ 2 y s k = ∥G y g ( x k , y s k ) -Γ y s k y ∗ ( x k ) G y g ( x k , y ∗ ( x k )) ∥ 2 y s k ≤ L 2 d 2 ( y s k , y ∗ ( x k )) and the fact that y ∗ ( x k ) is optimal. Here, we require η y < µ L 2 ζ in order for 1 + η 2 y ζL 2 -η y µ < 1 .

C.5 Proof of Lemma 1

Proof of Lemma 1. Hessian inverse : for the Hessian inverse approximation, we let

It can be seen that G F ( x k ) = G f ( x k , y ∗ ( x k )) and ̂ G hinv F ( x k ) = G f ( x k , y k +1 ) . Then for any x ∈ U x , y 1 , y 2 ∈ U y , we have

where we use Assumption 2, 3 and Lemma 4.

Conjugate gradient : we let v ∗ k = H -1 y g ( x k , y ∗ ( x k ))[ G y f ( x k , y ∗ ( x k ))] ∈ T y ∗ ( x k ) M y and let ˆ v ∗ k = H -1 y g ( x k , y k +1 )[ G y f ( x k , y k +1 )] ∈ T y k +1 M y . We first bound

where ∥ v ∗ k ∥ y ∗ ( x k ) ≤ M/µ . From standard convergence result eq. 6.19 in [8], we have

This leads to

where in the last inequality, we use the definition of v ∗ k and ˆ v ∗ k and the Lipschitzness assumptions. Combining the results yield the desired result.

Neumann series : let ̂ H k ( y ) := γ ∑ T -1 i =0 (id -γ H y g ( x k , y )) i . Then we can bound

We now bound

where we use the lower bound on H y g ( x k , y k +1 ) . Substituting the results back the bound yields the desired result.

Automatic differentiation : Given y s +1 k = Exp y s k ( -η y G y g ( x k , y s k )) , we can show its differential is

where

where the second inequality uses Lemma 4.2. Then we bound

where the first equality uses the expression of D x k y ∗ ( x k ) (Proposition 1). The second last inequality follows from Lemma 5 and the last inequality is due to the smoothness of Riemannian gradient and Lemma 5, i.e., ∥G y f ( x k , y s k ) ∥ y s k = ∥ Γ y s k y ∗ ( x k ) G y f ( x k , y ∗ ( x k )) - G y f ( x k , y s k ) ∥ y s k ≤ Ld ( y s k , y ∗ ( x k )) ≤ L (1 + η 2 y ζL 2 -η y µ ) s 2 d ( y 0 k , y ∗ ( x k )) .

Finally, applying the bound recursively, we obtain

where we let ˜ C := ( κ l +1) ρ +( C 2 + η y C 3 ) L ( (1 -η y µ ) C 1 + η y L ) and we note that D x k y 0 k = 0 . Combining the above result with (6) gives

C.6 Proof of Theorem 1

Proof of Theorem 1. By smoothness of F ( x ) (Lemma 4.5), we have

Now we consider the different hypergradient estimator separately. 1. Hessian inverse : Let C hinv := L + κ ρ M + κ l L + κ l κ ρ M . ∥G F ( x k ) -̂ G hinv F ( x k ) ∥ 2 x k ≤ C 2 hinv d 2 ( y ∗ ( x k ) , y k +1 ) ≤ C 2 hinv (1 + η 2 y ζL 2 -η y µ ) S d 2 ( y ∗ ( x k ) , y k ) , (8)

where we notice y k +1 = y S k and apply Lemma 5. Furthermore,

where we apply Lemma 5 and 4.3 in the second inequality.

2 ∗

where we combine (7) and (9) in the first inequality and use κ 2 l ≤ L F in the third inequality. Now setting η x = 1 20 L F , we can simplify the inequality as

where we choose S ≥ log( 80 L F 160 L F +3 C 2 hinv ) / log(1 + η 2 y ζL 2 -η y µ ) = ˜ Θ( κ 2 l ζ ) for the last inequality. Summing over k = 0 , ...K -1 yields

which suggests min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 80 L F ∆ 0 K .

2. Conjugate gradient : Let C cg := L + κ ρ M + L ( 1 + 2 √ κ l )( κ l + Mκ ρ µ ) . Then we can show

where it follows from Lemma 1 and Lemma 5. Then following similar analysis as in Hessian inverse case

Further, noticing ˆ v 0 k = Γ y k +1 y k ˆ v T k -1 , we bound ∥ ˆ v 0 k -Γ y k +1 y ∗ ( x k ) v ∗ k ∥ y k +1 = ∥ ˆ v T k -1 -Γ y k y k +1 Γ y k +1 y ∗ ( x k ) v ∗ k ∥ y k as

where we use Lemma 2 in the first and third inequalities. The second last inequality follows from (5) and d ( y k +1 , y ∗ ( x k )) ≤ d ( y k , y ∗ ( x k )) . The last inequality follows from Lemma 5 and κ l ≥ 1 . Now we bound

where we let C v := Mκ ρ µ + Mκ ρ κ l µ + κ 2 l + κ l . Combining (13) and (12), we obtain

Now we define a Lyapunov function R k := F ( x k ) + d 2 ( y k , y ∗ ( x k )) + ∥ ˆ v 0 k -Γ y k +1 y ∗ ( x k ) v ∗ k ∥ 2 y k +1 . Then

where we use (7), (14) in the first inequality and (11) in the second inequality. In the third inequality, we let Λ := C 2 v + κ 2 l ( 5 M 2 C 2 0 ¯ D 2 µ +1) and because L F = Θ( κ 3 l ) and Λ = Θ( κ 4 l ) , we can without loss of generality have L F ≤ Λ . We also choose η x = 1 24Λ and use (10) for the fourth inequality. The last inequality follows by choosing

Finally, telescoping the inequality, we obtain

where we use the fact that ˆ v 0 k = 0 and the isometry property of parallel transport.

3. Truncated Neumann series : Let C ns := L + κ l L + κ ρ M + κ l κ ρ M . Here we notice that C ns = C hinv . Then by Lemma 1, we see

∥

̂

G

ns

F

(

x

k

)

-G

F

(

x

k

)

∥

2

x

k

2

C

Similar in the previous analysis,

Let the Lyapunov function be R k := F ( x k ) + d 2 ( y k , y ∗ ( x k )) . Then

where we set η x = 1 20 L F and apply (15) in the second inequality.

Now setting S ≥ log( 40 L F 80 L F +3 C 2 ns ) / log(1 + η 2 y ζL 2 -η y µ ) = ˜ Θ( κ 2 l ζ ) and telescoping the results yields

where we set T ≥ -1 2 log( 12 κ 2 l M 2 ϵ ) / log(1 -γµ ) = ˜ Θ( κ log( 1 ϵ )) .

4. Automatic differentiation : Let C ad := 2 M ˜ C µ -η y ζL 2 + L (1 + κ l ) . Then

∥G F ( x k ) -̂ G ad F ( x k ) ∥ 2 x k ≤ 2 C 2 ad (1 + η 2 y ζL 2 -η y µ ) S -1 d 2 ( y k , y ∗ ( x k )) + 2 M 2 κ 2 l (1 -η y µ ) 2 S and similarly

Let the Lyapunov function be R k := F ( x k ) + d 2 ( y k , y ∗ ( x k )) . Then

2

ns

(1 +

η

2

y

2

ζL

-

η

y

µ

)

S

d

2

(

y

∗

(

x

k

)

, y

k

) + 2

κ

2

l

M

2

(1

-

γµ

)

2

T

.

(15)

where we set η x = 1 20 L F and choose S ≥ log 40 L F 80 L F +3 C 2 ad / log(1 + η 2 y ζL 2 -η y µ ) + 1 = ˜ Θ( κ 2 l ζ ) . Telescoping the result gives

by choosing S ≥ 1 2 log( ϵ 12 M 2 κ 2 l ) log(1 -η y µ ) = ˜ Θ( κ 2 l ζ log( 1 ϵ )) . Hence we set S ≥ ˜ Θ( κ 2 l ζ log( 1 ϵ )) for both conditions to hold.

C.7 Proof of Corollary 1

The computational cost of gradient and Hessian for each method for approximating the hypergradient are as follows.

Corollary 1. The complexities of reaching an ϵ -stationary solution are

Hessian inverse: G f = O ( κ 3 l ϵ -1 ) , G g = ˜ O ( κ 5 l ζϵ -1 ) , JV g = O ( κ 3 l ϵ -1 ) , HV g = NA .

Conjugate gradient: G f = O ( κ 4 l ϵ -1 ) , G g = ˜ O ( κ 6 l ζϵ -1 ) , JV g = O ( κ 4 l ϵ -1 ) , HV g = ˜ O ( κ 4 . 5 l ϵ -1 ) .

Truncated Neumann series: G f = O ( κ 3 l ϵ -1 ) , G g = ˜ O ( κ 5 l ζϵ -1 ) , JV g = O ( κ 3 l ϵ -1 ) , HV g = ˜ O ( κ 4 l ϵ -1 log( ϵ -1 )) .

Automatic differentiation: G f = O ( κ 3 l ϵ -1 ) , G g = ˜ O ( κ 5 l ζϵ -1 log( ϵ -1 )) , JV g = ˜ O ( κ 5 l ζϵ -1 log( ϵ -1 )) , HV g = ˜ O ( κ 5 l ζϵ -1 log( ϵ -1 )) .

Proof of Corollary 1. From the convergence established in Theorem 1, we see the iterations in order to reach ϵ -stationary solution are given by

(Hessian inverse) K = O ( L F ϵ -1 ) = O ( κ 3 l ϵ -1 ) , S = ˜ O ( κ 2 l ζ ) .

(Conjugate gradient) K = O (Λ ϵ -1 ) = O ( κ 4 l ϵ -1 ) , S = ˜ O ( κ 2 l ζ ) , T = ˜ O ( √ κ l ) .

(Truncated Neumann series) K = O ( κ 3 l ϵ -1 ) , S = ˜ O ( κ 2 l ζ ) , T = ˜ O ( κ l log( ϵ -1 )) .

(Automatic differentiation) K = O ( κ 3 l ϵ -1 ) , S = ˜ O ( κ 2 l ζ log( ϵ -1 ) ) .

Then based on Algorithm 1, the gradient complexities are G f = 2 K and G g = KS and crossderivative and Hessian product complexities are JV g = K,HV g = KT for CG and NS and JV g = KS , HV g = KS for AD (which we approximate based on the analysis in Lemma 1). We notice here for the Hessian inverse, because we do not compute Hessian vector product, we write NA for Hessian vector product based on the Neumann series. This completes the proof.

D Proofs for Section 3.3

We first show Lemma 4 holds for each f i ( x, y ) , g i ( x, y ) . Further, the variance of the estimate can be bounded as follows. We here use [ · ] to denote all possible derivatives, including x, y, xy, yx .

Lemma 6. Under Assumption 4, we have for any x, y ∈ U , (1) E ∥G [ · ] f i ( x, y ) -G [ · ] f ( x, y ) ∥ 2 [ · ] ≤ M 2 . (2) E ∥G 2 [ · ] g i ( x, y ) -G 2 [ · ] g ( x, y ) ∥ 2 [ · ] ≤ L 2 . (3) E ∥H -1 y g i ( x, y ) -H -1 y g ( x, y ) ∥ 2 y ≤ µ -2 .

For notation, denote the filtration F k := { y 0 , x 0 , y 1 , x 1 , ..., x k , y k +1 } and here we let E k := E [ ·|F k ] . With a slight abuse of notation, we further consider F s k := { y 0 , x 0 , y 1 , x 1 , ..., y k , y 1 k , ..., y s k } and correspondingly let E s k := E [ ·|F s k ] .

Lemma 7 (Convergence under strong convexity and stochastic setting) . Under stochastic setting and under the Assumption that g is geodesic strongly convex, we can show E s k d 2 ( y s +1 k , y ∗ ( x k )) ≤ (1 + η 2 y ζL 2 -η y µ ) d 2 ( y s k , y ∗ ( x k )) + η 2 y ζM 2 |B 1 | and E k -1 d 2 ( y k +1 , y ∗ ( x k )) ≤ (1 + η 2 y ζL 2 -η y µ ) S d 2 ( y k , y ∗ ( x k )) + η y ζM 2 µ -η y ζL 2 1 |B 1 | .

D.1 Proofs for the lemmas

Proof of Lemma 6. Here we only prove one and the rest follows exactly. Due to the unbiasedness of the stochastic estimate, we have

where we use Assumption 4.

Proof of Lemma 7. Similarly from the proof of Lemma 5, we take expectation over F s k

where we use the strong convexity and the fact that E ∥G y g B 1 ( x, y ) - G y g ( x, y ) ∥ 2 y = 1 |B 1 | 2 E ∥ ∑ i ∈B 1 ( G y g i ( x, y ) -G y g ( x, y )) ∥ 2 y = 1 |B 1 | 2 ∑ i ∈B 1 E ∥G y g i ( x, y ) -G y g ( x, y ) ∥ 2 y in the third inequality and Lemma 6 in the last inequality. Further, we telescope the inequality and taking the expectation E k -1 gets

where we use the fact that ∑ S -1 s =0 θ s ≤ 1 1 -θ for 0 < θ < 1 .

where the second inequality uses Lemma 1. Now we bound the first term E k ∥ ̂ G F ( x k ) -̂ G hinv F ( x k ) ∥ 2 x k as follows.

First we bound

where we notice that ∥G y f B 2 ( x k , y k +1 ) ∥ y k +1 ≤ 1 |B 2 | ∑ i ∈B 2 ∥G y f i ( x k , y k +1 ) ∥ y k +1 ≤ M . Hence, we can bound

where we use Lemma 6 in the last inequality. Combining (18) with (17) yields the desired result.

D.2 Proof of Theorem 2

Proof of Theorem 2. From the smoothness of F ( x ) (i.e., (7)) and taking full expectation we obtain, E [ F ( x k +1 ) -F ( x k )] ≤ -( η x 2 -η 2 x L F ) E ∥G F ( x k ) ∥ 2 x k + ( η x 2 + η 2 x L F ) E ∥G F ( x k ) -̂ G F ( x k ) ∥ 2 x k

Further, we can bound

where we use Lemma 7 and 8 in the second inequality.

Next, we construct a Lyapunov function as R k := F ( x k ) + d 2 ( y k , y ∗ ( x k )) . Then

where we choose η x = 1 20 L F in the first equality and S ≥ log( 40 L F 80 L F +3 C 2 hinv ) / log(1+ η 2 y ζL 2 -η y µ ) = ˜ Θ( κ 2 l ζ ) for the last inequality. Telescoping the result gives

where the last inequality follows from the choice that |B 1 | ≥ (24 C 2 hinv + 640 L F ) 8 η y ζM 2 µ -η y ζL 2 /ϵ = Θ( κ 4 l /ϵ ) , |B 2 | ≥ 144 M 2 +576 M 2 κ 2 l ϵ = Θ( κ 2 l /ϵ ) , |B 3 | ≥ 288 M 2 κ 2 l ϵ = Θ( κ 2 l /ϵ ) , |B 4 | ≥ 576 M 2 κ 2 l ϵ = Θ( κ 2 l /ϵ ) in the last inequality.

In order to reach ϵ -stationary solution, we require K = O ( κ 3 l ϵ -1 ) and thus the (stochastic) gradient complexity for f is G f = 2 K |B 2 | = O ( κ 5 l ϵ -2 ) and for g is G g = KS |B 1 | = ˜ O ( κ 9 l ζϵ -2 ) . The complexity for cross-derivative is K |B 3 | = O ( κ 5 l ϵ -2 ) .

E Proofs for Section 3.4

Proof of Theorem 3. We first give a complete proof for the Hessian inverse estimator as follows. For the other estimators, we only provide a proof sketch.

Proof for HINV. (1) First, we derive the convergence under strong convexity using retraction. By the trigonometric distance bound

where we use Assumption 5 in the second inequality and fourth inequality. We require η y < µ ( ζc +2 ¯ Dc R ) L 2 in order to achieve linear convergence. For simplicity, we let τ = µη y -η 2 y ( ζc +2 ¯ Dc R ) . This leads to d 2 ( y s +1 k , y ∗ ( x k )) ≤ (1 -τ ) d 2 ( y s k , y ∗ ( x k )) .

(2) Next, we notice the bound on hypergradient approximation error still holds as ∥ ̂ G hinv F ( x k ) -G F ( x k ) ∥ x k ≤ C hinv d ( y ∗ ( x k ) , y k +1 ) , where C hinv = L + κ ρ M + κ l L + κ l κ ρ M . Further, by L -smoothness,

where in the third inequality, we bound ∥G F ( x k ) ∥ x k ≤ M + L µ M ≤ 2 LM µ .

(3) Then we can bound

where the last inequality follows similarly as (9).

Let a Lyapunov function be R k := F ( x k ) + d 2 ( y k , y ∗ ( x k )) . Then

where we use κ 2 l ¯ c ≤ L F ¯ c and let ˜ L F := 4 κ l c R M + 5¯ cL F in the second last inequality, and we choose η x = 1 4 ˜ L F , S ≥ log ( 16 ˜ L F 32 ˜ L F +3 C 2 hinv ) / log(1 -τ ) = ˜ Θ( κ 2 l ζ ) , in the last inequality. Then telescoping the results yields Finally, we sum over k = 0 , ..., K -1 , which leads to

Thus in order to achieve ϵ -stationary solution, we require K = O ( ˜ L F ϵ -1 ) = O ( κ 3 l ϵ -1 ) and hence the order of gradient and second-order complexities remain unchanged.

Extensions to other estimators. To extend the proof to other hypergradient estimators, we first notice that the convergence of inner iterations for solving lower-level problems is agnostic to the choice of hypergradient estimators, i.e.,

(1) Hypergradient approximation error. For hypergradient estimator based on Hessian inverse, conjugate gradient, truncated Neumann series, the hypegradient approximation error remains the same as in Lemma 1, given no retraction is involved in the computation. That is,

CG : ∥ ̂ G cg F ( x k ) -G F ( x k ) ∥ x k ≤ C cg d ( y ∗ ( x k ) , y k +1 ) +2 L √ κ l ( √ κ l -1 √ κ l +1 ) T ∥ ˆ v 0 k -Γ y k +1 y ∗ ( x k ) v ∗ k ∥ y k +1 , where C cg = L + κ ρ M + L (1 + 2 √ κ l )( κ l + Mκ ρ µ ) .

For hypergradient based on automatic differentiation ( AD ), we first show that there exists a constant C 4 (that depends on C 3 , ¯ c, c R ) such that D x Retr x ( u ) = P Retr x ( u ) (id + D x u ) + E . with ∥E∥ Retr x ( u ) ≤ C 4 ∥ D x u ∥ x ∥ u ∥ x . Such a result can be derived by bounding the difference between retraction and exponential map. Then we follow the analysis for Lemma 1 as follows. Given y s +1 k = Retr y s k ( -η y G y g ( x k , y s k )) , we have

where ∥E s k ∥ y s +1 k ≤ η 2 y C 4 ( (1 -η y µ ) C 1 + η y L ) ∥G y f ( x k , y s k ) ∥ y s k . The rest of the proof follows exactly from Lemma 1, where we replace the convergence of inner iteration with the updated rate. This gives

where ˜ C ′ := ( κ l +1) ρ +( C 2 + η y C 4 ) L ( (1 -η y µ ) C 1 + η y L ) In summary

(2) Objective decrement. This part is also the same across all hypergradient estimators, i.e.,

(3) Lyapunov function decrement. The definition of Lyapunov function depends on the choice of hypergradient estimator.

For CG , we define R k := F ( x k ) + d 2 ( y k , y ∗ ( x k )) + ∥ ˆ v 0 k -Γ y k +1 y ∗ ( x k ) v ∗ k ∥ 2 y k +1 . Then following similar analysis, we first bound

and

Then similarly,

where in the second inequality, we use the bound between retraction and exponential map as well as triangle inequality. Then combining the above two results, gives

Then we can show

where we let ˜ Λ := C 2 v ¯ c + κ 2 l ( 5 M 2 C 2 0 D 2 µ + ¯ c ) and without loss of generality ˜ L F ≤ ˜ Λ . The second inequality is by (19). The last inequality is by appropriately choosing S, T , which is on the same order as the exponential map case. Then telescoping the result yields

For NS , we define R k = F ( x k ) + d 2 ( y k , y ∗ ( x k )) and derive

Then we can follow exactly the same proof as HINV that

by appropriately choosing S, T and η x .

For AD , we define R k := F ( x k ) + d 2 ( y k , y ∗ ( x k )) . Then following the same analysis except for the choice of S , we can show,

Thus the proof is now complete.

F Tangent space conjugate gradient

In Algorithm 3, we show the tangent space conjugate gradient algorithm for solving the linear system H [ v ] = G . Similar to [40], we set the initialization to be the transported output of ˆ v T k -1 from last iteration, where ˆ v T -1 = 0 , which is beneficial for convergence analysis. For practical purposes, we notice setting v 0 = 0 provides sufficient accurate solution without the expensive parallel transport operation.

Algorithm 3 Tangent space conjugate gradient TSCG( H , G , v 0 , T )

1: Set r 0 = G ∈ T x M , p 0 = r 0 .

2: for t = 0 , ..., T -1 do

3: Compute ¯ r t +1 = H [ v t ] .

4: α t +1 = ∥ r t ∥ 2 x ⟨ p t , H [ p t ] ⟩ x .

5: v t +1 = v t + α t +1 p t .

6: r t +1 = r t -α t +1 H [ p t ] .

7: β t +1 = ∥ r t +1 ∥ 2 x ∥ r t ∥ 2 x .

8: p t +1 = r t +1 + β t +1 p t .

9: end for

10: Output : v T

Algorithm 4 Riemannian bilevel solver for min-max optimization

G Extensions: on Riemannian mix-max and compositional optimization

The bilevel optimization considered in the paper (1) generalizes the two other widely studied problems, namely the min-max optimization and compositional optimization.

G.1 Min-max optimization on Riemannian manifolds

Riemannian min-max problems have gained increasing interest over the recent years [37, 41, 27, 73, 67, 25, 56, 35], which takes the form of

and can be seen as a special case of bilevel optimization problem (1) where g ( x, y ) = -f ( x, y ) . Because the problem is nonconvex in x , the order of minimization and maximization matters [73, 25]. Nevertheless, under the assumption where f is geodesic strongly convex in y , the optimal solution x ∗ satisfies G F ( x ∗ ) = 0 , where G F ( x ) = G x f ( x, y ∗ ( x )) due to G y f ( x, y ∗ ( x )) = G y g ( x, y ∗ ( x )) = 0 . Thus Algorithm 1 reduces to alternating gradient descent ascent over Riemannian manifolds, as outlined in Algorithm 4.

Here we adapt the convergence analysis to the min-max optimization setting. Given we no longer require second-order derivatives, we restate assumptions for functions f , g below.

Assumption 6. (1) Assumption 1 holds. (2) Function f ( x, y ) , g ( x, y ) have L -Lipschitz Riemannian gradients. (3) Further, g ( x, y ) is µ -geodesic strongly convex in y .

Under the min-max setup and Assumption 6, we see G F ( x ) = G x f ( x, y ∗ ( x )) and thus the Lipschitz constant can be derived as L F = ( κ l +1) L = Θ( κ l ) . Further we can directly apply Theorem 1 for the Hessian inverse with C hinv = L , which leads to the following convergence result.

Theorem 4. Under Assumption 6, choosing S ≥ ˜ Θ( κ 2 l ζ ) , η x = 1 20 L F , we have min k =0 ,...,K -1 ∥G F ( x k ) ∥ 2 x k ≤ 80( κ l +1) L ∆ 0 K and to reach ϵ -stationary solution, we require gradient complexities as G f = O ( κ l ϵ -1 ) and G g = ˜ O ( κ 3 l ζϵ -1 ) .

G.2 Compositional optimization on Riemannian manifolds

Compositional problems on Riemannian manifolds have been considered in [36, 70], which requires to solve

where ψ : M y → R and ϕ : M x →M y . It is worth noting that in both works [36, 70], the inner function ϕ : M x → R d is vector-valued. In contrast, we consider a general manifold-valued function ϕ . Because the function ϕ can be potentially complex and may be stochastic, we follow [11] to reformulate (20) into a bilevel optimization problem by letting

As long as the squared Riemannian distance is geodesic strongly convex, the reformulation is equivalent to the original problem (20). As formally stated in Lemma 9, this is satisfied for nonpositively curved space, like Euclidean space, hyperbolic manifold, SPD manifold with affine invariant metric. For positively curved space, the strong convexity is guaranteed when restricting the domain relative to the curvature.

Lemma 9. Let U ⊆ M has sectional curvature lower and upper bounded by κ -and κ + respectively. Further U has diameter upper bounded by ¯ D , which satisfies ¯ D < π √ κ + if κ + > 0 . Then let δ = 1 when κ + ≤ 0 and δ = √ κ + ¯ D tan( √ κ + ¯ D ) when κ + > 0 and consider ζ be the same curvature constant as in Lemma 3. Then function H y g ( x, y ) has Riemannian Hessian bounded within [ δ, ζ ] in spectrum.

Proof of Lemma 9. The proof follows from Lemma 2 in [3]. Consider an arbitrary curve γ : [0 , 1] → M , and let f ( x ) = 1 2 d 2 ( x, p ) , for some p ∈ M . From [3], we know that H f ( γ ( t ))[ γ ′ ( t )] = -∇ γ ′ ( t ) Exp -1 γ ( t ) ( p ) and under the conditions, δ ∥ γ ′ ( t ) ∥ 2 γ ( t ) ≤ ⟨ ∇ γ ′ ( t ) Exp -1 γ ( t ) ( p ) , -γ ′ ( t ) ⟩ γ ( t ) ≤ ζ ∥ γ ′ ( t ) ∥ 2 γ ( t ) , where we denote ∇ as the covariant derivative. This immediately leads to

which completes the proof.

Thus, for positively curved manifold, if ¯ D < π 2 √ κ + , we have δ > 0 , which ensures geodesic strong convexity of the inner problem. As shown in Lemma 12 in [3], G y d 2 ( ϕ ( x ) , y ) = 2Exp -1 y ϕ ( x ) and the Riemannian gradient descent on y lead to

which suggests y s +1 k lies on a geodesic that connects y s k and ϕ ( x k ) . When S = 1 and when the lower-level function g is vector-valued, the algorithm recovers the deterministic version of SCGD [66].

̸

However, unlike in the Euclidean space, the Riemannian Hessian does not simplify to the identity operator, but rather the covariant derivative of inverse exponential map and the cross derivatives G 2 xy g ( x, y ) = -(D ϕ ( x )) † .

Assumption 7. (1) Assumption 1 holds and further ¯ D < π 2 √ κ + if κ + > 0 . (2) Function f ( x, y ) has Riemannian gradients that are bounded by M and are L -Lipschitz. (3) Function g has ρ -Lipschitz Riemannian Hessian and cross derivatives.

We notice that for function g we only require second-order derivatives to be Lipschitz because the first-order Lipschitzness can be inferred from Lemma 9.

Theorem 5. Under Assumption 7, Theorem 1 holds with L = ζ, µ = δ .

To prove the convergence, we only need to show Lemma 4 holds. It can be readily proved from Lemma 9 and Assumption 7 that Lemma 4 holds with L = ζ, µ = δ . Hence the convergence follows directly.

H Experimental details

H.1 Synthetic problem

We first verify the lower-level problem is geodesic strongly convex.

Proposition 4. For any A , B ≻ 0 , function f ( M ) = ⟨ M , A ⟩ + ⟨ M -1 , B ⟩ is µ -geodesic strongly convex in U ⊂ S d ++ with µ = λ a, -λ -+ λ b, -λ -λ 2 + , where λ a, -, λ b, -are the minimum eigenvalue of A , B and λ ± are the bounds for maximum and minimum eigenvalues for M ∈ U .

The inverse of Riemannian Hessian of function f ( M ) = ⟨ M , A ⟩ + ⟨ M -1 , B ⟩ is derived as, for any symmetric U , H -1 f ( M )[ U ] = M 1 / 2 GM 1 / 2 where G is the solution to the Lyapunov equation G ( M 1 / 2 AM 1 / 2 + M -1 / 2 BM -1 / 2 ) + ( M 1 / 2 AM 1 / 2 + M -1 / 2 BM -1 / 2 ) G = M -1 / 2 UM -1 / 2 .

Proof of Proposition 4. We first derive the Euclidean gradient and Hessian as

for any U = U ⊤ . The Riemannian gradient and Hessian are derived as

where we let { A } S = ( A + A ⊤ ) / 2 . To show the function is geodesic strongly convex, it suffices to show H f ( M ) is positive definite, which is to show ⟨H f ( M )[ U ] , U ⟩ M ≥ µ ∥ U ∥ 2 M > 0 for any U = U ⊤ . To this end, we vectorize the Riemannian Hessian in terms of U as vec(2 H f ( M )[ U ]) = ( MA ⊗ I + I ⊗ MA + BM -1 ⊗ I + I ⊗ BM -1 )vec( U ) , where ⊗ denotes the Kronecker product. Then, we have

where we let λ a, ± be the maximum/minimum eigenvalues of A and similarly for λ b, ± , λ m, ± . The Hessian inverse can be derived subsequently. This completes the proof.

H.2 Computational time for each hypergradient estimator

This section report the average runtime in seconds (over 10 runs) for single evaluation of hypergradient using four different strategies (for Hypergradient estimation) for the synthetic problem (Section 4.1, Figure 1). The hyper-parameters are set to be the same as the main experiment. We see in general automatic differentiation (AD) is the most efficient strategy. Nevertheless, according to Figure 1(a), it is less accurate compared to other strategies.

H.3 Hyperparameter selection

The selection of hyper-parameters is performed to reflect the best performance. The stepsize is selected from the range [1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1] and for Neumann series is selected from [0.1, 0.5, 1.0, 1.5, 2.0] and number of inner iterations is selected from [5, 10, 30, 50, 100]. Figure 1(d) shows the sensitivity of hypergradient error as we vary and the number of inner iterations.

H.4 Computational Complexity

This section lists out the computational complexity for each task considered in the experiment section. In Table 4, we present an estimate of the per-iteration complexity of computing the gradient, Hessian/Jacobian-vector products. We highlight that we only provide estimates of the complexities given that there may not exist closed form expressions for the gradient and second-order derivatives.

Here, n v , n t denote the size of validation set and training set respectively. For meta learning, m denotes the number tasks and n denotes the number of samples for each task. For domain adaptation, m,n denote the number of samples for two domains, s denotes the number of Sinkhorn iterations.

I Experiment Configurations

All the experiments are conducted on a single NVIDIA RTX 4060 GPU. All datasets used in the paper are publicly available, which are properly cited in the main paper. We include detailed setups for the experiments in the main paper as well as documented in code (provided as supplementary material).

J Broader Impact

This paper proposes new algorithms and are of theoretical in nature. We do not foresee any immediate negative societal impact that we feel obliged to report.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Please refer to Section 3 for theoretical developments and Section 4 for empirical studies.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We include discussions of limitations in the conclusion section, i.e., Section 5. Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We have included all the assumptions and complete proof in the paper. Please refer to Section 3 and Section C, D, E.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We have provided the implementation details for reproducing the experiments, along with the code in the supplementary material.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have provided the code and data as supplementary material to ensure reproducibility.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have included the details in the experiment section.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We have reported standard deviation in Figure 2(b) and (c). For Figure 1, the results barely change as we vary the random seed. For Riemannian meta learning, it becomes time costly to report standard deviation.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have included the detailed configuration in Appendix I. We report the time of execution in Appendix H.2.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The paper conform, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We have included a broader impact section, i.e., Section J.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not release data or models that have high risks for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: In the experiment section, this paper has cited the papers when using exisiting assets, inluding code and data.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not introduce new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
This paper studies Riemannian bilevel optimization, where variables of both lower and upper level problems are constrained on Riemannian manifolds. The authors propose several hypergradient based algorithms via Neumann series and automatic differentiation. 
Convergence analysis is provided for the proposed approaches. Experiments on synthetic problems, hyper-representation learning, meta-learning, unsupervised domain adaptation are provided to demonstrate the effectiveness of the proposed methods.

STRENGTHS:
1.	The studied topic is of interest to bilevel optimization research. As far as I know, this should be the first work to study this type of research.
2.	The analysis is comprehensive. The algorithms cover AID and ITD based hypergradient approximations. 
3.	Quite a few experiments are provided to support the theory.

WEAKNESSES:
1.	More motivation examples should be provided to validate the importance of bilevel optimization on manifolds. In the experiments, for example, in manifold meta-learning, it would be good to compare the method with meta-learning methods without the orthogonality constraint. 
2.	Assumption 1 may be quite strong. It requires all iterates z_1,…, are bounded in a compact space. The analysis in previous standard bilevel optimization does not require this assumption. 
3.	Main results contain complex constants and are a little bit hard to parse. Some simplifications may be helpful here.
4.	The extension in 3.3 will be more interesting if Hessian-vector rather than Hessian-inverse is considered.

QUESTIONS:
See weakness for information.

LIMITATIONS:
See weakness for information.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

comments: SUMMARY:
The paper proposes an RHGD algorithm (Algo 1) to solve the bilevel optimization problems (line 11). 
- Thm 1 proves the convergence to a stationary point of $F$ when using different approximations of the hypergradient.
- Thm 2 shows the convergence under stochastic setting.
- Thm 3 proves the convergence with cheaper retraction.
Both synthetic numerical tests and applications to real problem are provided.

STRENGTHS:
The paper is well-written. The theoretical results are solid and the contributions are new to me. The algorithm is fully implementable. What I like is the convergence is not only established for the Hessian inverse but also its cheaper approximation. The synthetic experiment is carefully designed.

WEAKNESSES:
1. Assumption 1 is too strong: it requires the trajectory of $z$ to be in the unique geodesic neighborhood of $z^*$.
1. Assumption 2 is slightly too strong, does this mean many simple functions like quadratic functions: $f(x, y)=|x|^2+|y|^2$ do not satisfy this assumption?

QUESTIONS:
1. Could you please clarify the synthetic problem in Sec 4.1 satisfies all the assumptions? I only find that Assumption 3 is verified. Assumption 2 is automatically satisfied due to the compactness of the Stiefel manifold. However, could you clarify how to verify Assumption 1?
1. Maybe it is a dumb question, but in Sec. 3.1, why HINV and AD are different?

LIMITATIONS:
Small typo: In line 112 there are 2 negative signs.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a framework for addressing bilevel optimization problems in which variables of both lower- and upper-level problems are constrained on Riemannian manifolds. It introduces multiple hypergradient estimation strategies on manifolds and investigates their estimation error. The paper includes convergence and complexity analyses for the proposed hypergradient descent algorithm on manifolds. Furthermore, it extends these advancements to stochastic bilevel optimization and the utilization of general retraction on the manifolds. The paper also demonstrates the practicality of the proposed framework through several applications.\n\nSTRENGTHS:\n1. The paper is well-written and is easy to read.\n\n2. This paper derives the intrinsic Riemannian hypergradient using the implicit function theorem, proposes four strategies for estimating the hypergradient, and provides estimation error bounds for all the proposed strategies.\n\n3. This paper presents the Riemannian hypergradient descent algorithm for tackling bilevel optimization problems on manifolds, providing convergence guarantees. Furthermore, it extends the framework to the stochastic setting and incorporates retraction mapping on the manifolds.\n\nWEAKNESSES:\n1. Firstly, my primary concern pertains to the novelty of this paper. The algorithm proposed in this paper appears to mirror the framework introduced in a prior work [42]. Therefore, it is crucial to delineate the similarities and disparities between this paper and the referenced study [42].\n\n2. The experimental results should include a comparison with the methods presented in [42].\n\n3. The analysis of retraction on manifolds lacks depth and should include the convergence results of three additional hypergradient estimators.\n\n4. Assumption 1 seems to be excessively strict, which is unusual in common bilevel optimization. Moreover, the notation $D$ in $d(x_1,x_2)\\le D$ may be confusing. Furthermore, Assumptions 2 and 3 could be inferred from Assumption 1, indicating that this paper may not necessitate Assumptions 2 and 3.\n\nQUESTIONS:\n1. The notations $C_{hinv}$, $C_{cg}$, $c_{ns}$, and $C_{ad}$ in Theorem 1 are not utilized in the main body of the paper.\n\n2. Can the stochastic bilevel problem studied in this paper be extended to the general stochastic setting, i.e.,  $\\underset{{x\\in \\mathcal{M}_x}}{\\min} F(x) = \\mathbb{E}[f(x,y)]$ and $y^*(x) = \\underset{y\\in\\mathcal{M}_y}{\\text{argmin}}~ \\mathbb{E}[g(x,y)]$?\n\n3. For other questions, please see the weaknesses above.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper introduces a novel approach for solving bilevel optimization problems where both upper and lower-level variables are constrained on Riemannian manifolds. The authors propose four hypergradient estimation strategies (HINV, CG, NC, AD), analyze their estimation errors with convergence and complexity analysis, and generalize to stochastic bilevel optimization and general retraction. The paper also shows many applications of the framework including hyper-representation over SPD matrices, Riemannian meta-learning, and unsupervised domain adaptation.

STRENGTHS:
1. The proposed framework extends bilevel optimization to Riemannian manifolds, which fills a gap in this field.
2. The paper introduces multiple hypergradient estimation strategies and provides a thorough theoretical analysis.
3. Practical Relevance: Demonstrates the applicability of the framework through several machine learning applications, including hyper-representation and meta-learning.

WEAKNESSES:
1. Assumptions: The framework relies on assumptions such as geodesic strong convexity and Lipschitz continuity, which may limit its applicability in some practical scenarios.
2. The practical performance evaluation is limited to synthetic and small-scale problems, which may not fully represent its utility in large-scale real-world applications.

QUESTIONS:
1. What are the computational trade-offs when choosing between the different hypergradient estimation strategies?

LIMITATIONS:
The authors acknowledge that the framework relies on some strong mathematical properties that limit its broader applicability and may be improved in the future.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses bilevel optimization problems where variables of both the lower and upper-level problems are constrained on Riemannian manifolds. To solve such problems, the authors propose a hypergradient descent algorithm under the key assumption that the lower function is geodesically strongly convex and establish its iteration complexity. To efficiently estimate the hypergradient, several methods with error bounds are also proposed. Additionally, the paper provides several interesting application examples of bilevel optimization on Riemannian manifolds and presents numerical results for these examples.\n\nSTRENGTHS:\n- This paper generalizes the widely used hypergradient descent algorithm to the Riemannian manifold setting and provides a comprehensive analysis of the proposed algorithms.\n- The paper includes four approaches to estimating the inverse of the Riemannian Hessian and a stochastic version of the proposed algorithm.\n- The theoretical analysis is solid, with estimation error bounds and the iteration complexity of the proposed algorithm presented.\n- The numerical results are extensive, showcasing results on four interesting Riemannian bilevel optimization application examples.\n\nWEAKNESSES:\n- A closely related work is [42], and although the authors claimed key differences in Section 1, some points should be clarified. \n\n     - The various hypergradient estimators beyond conjugate gradient and Neumann might be standard. Do the other estimators show better performance? \n\n     - While the paper presents interesting applications (see instances given in Section 4.2), these problems can also be solved by the methods in [42]. The authors should provide a numerical comparison of the methods. \n\n     - Shall the authors include the complexities of the methods in [42] in Table 1?\n\n- The estimations of the inverse Hessian are also widely used in Euclidean space.\n\n- The assumption of geodesic strong convexity for the lower function is too strong for the Riemannian setting, making the application examples in the paper less typical. Additionally, it should be noted that the lower-level problem in Section 4.2 is constrained in Euclidean space.\n\nQUESTIONS:\n1. Why are the performances of “hinv-20” and “hin-50” similar in Figures 1(a) and 2(a)?\n\n2. Why is “hinv” more efficient, as shown in Figure 1(b)?\n\nLIMITATIONS:\nYes, the limitations are acknowledged.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}