A Framework for Bilevel Optimization on Riemannian Manifolds

We thank the reviewer for appreciating the contributions of our work in terms of comprehensive analysis and supportive experiments, as well as providing the constructive comments.

1. (W1) More examples to validate the importance of bilevel optimization on manifolds.

Thank you for the suggestion. We believe we have provided sufficient motivating applications, showcasing the importance of bilevel optimization on manifolds. One such example is in Section 4.4 that motivates the use of bilevel optimization formulation for unsupervised domain adaptation (under the optimal transport framework). As far as we know, this is a novel application of bilevel optimization to domain adaptation, where we learn a whitening matrix $**M**$ in the lower-level problem. Our numerical experiments show that employing Riemannian bilevel optimization yields more suitable adaptation by accounting for the metric structure (through whitening) in unsupervised domain adaptation.

2. (W2) Assumption 1 is strong and is non-standard in previous bilevel optimization.

Standard bilevel optimization in the Euclidean space does not require this assumption. However, for optimization on manifolds, especially when dealing with (geodesic) strongly convex objectives, such an assumption is standard and often unavoidable. See for example [21, 37, 62, 63, 64]. This assumption is to ensure bounded curvature, thus allowing to show fast convergence under strong convexity.

3. (W3) Simplification regarding complex constants.

Thank you for the suggestion. We will simplify the notations in our revision.

4. (W4) Extend Hessian-inverse to Hessian-vector in Section 3.3.

We agree that a full-spectrum analysis covering other hypergradient estimators (in particular the ones using Hessian-vector products) would be interesting. The analysis techniques would be similar to what we already did for Hessian inverse. In order to make the contents concise, we had decided to omit the analysis for other hypergradient estimation strategies. In the revised version, however, we would include a discussion in this regard.

On Assumption 1:

I agreed with the concerns raised by Reviewer 4rfq regarding Assumption 1, which is typically a strong assumption in bilevel optimization. As I am not very familiar with the literature on manifold optimization, I am not very confident in validating the correctness of this assumption. Hope the authors can provide more evidence here. I found several related works on bilevel optimization on manifolds, e.g., [1][2]. I wonder whether similar assumptions are also made there.

[1] Riemannian Bilevel Optimization, Li and Ma, 2024. [2] Riemannian Bilevel Optimization, Sanchayan Dutta, Xiang Cheng, Suvrit Sra, 2024.

We thank the reviewer for the positive feedback on our work, particularly acknowledging that our paper is well-written with solid theoretical results. We also greatly appreciate the constructive comments.

1. (W1) Assumption 1 is strong.

We would like to highlight that Assumption 1 is often unavoidable for analyzing Riemannian optimization algorithms, especially when dealing with geodesic strongly convex functions, see for example [21, 37, 62, 63, 64] that also makes use of such an assumption. This assumption provides bound for the curvature, which is crucial for deploying a trigonometry distance bound (Lemma 3) and achieving linear convergence. This assumption can be satisfied for compact manifolds and by restricting the domain of interest.

2. (W2) Assumption 2 is strong.

Assumption 2 is not strong and can be inferred from Assumption 1 where we consider a compact domain. In this regard, the quadratic function you provided satisfies such an assumption when the domain is bounded.

3. (Q1) Verify the synthetic problem in Sec 4.1 satisfies Assumption 1.

Thank you for the question. Assumption 1 on the bounded domain can be satisfied directly for the Stiefel manifold. For the SPD manifold, we can take the maximal domain encompassing all the iterates considered.

4. (Q2) Why HINV and AD are different.

HINV computes the analytic expression of hypergradient using the Hessian inverse, which does not depend on the inner iterations. However, AD computes an estimate of hypergradient by differentiating through the inner iterations.

We thank the reviewer for the general appreciation of our work as well as constructive comments that we address below.

1. (W1) Similarities and disparities compared to [42].

We would like to emphasize that [42] is a concurrent work and not a prior work. It became available (publicly) after we had completed the draft. Moreover, we would like to highlight several notable differences compared to [42], some of which have already been discussed in lines 48-61.

(1) Our proposed framework is more general and subsumes the developments of [42]. In particular, Algorithms 1 and 2 in [42] are special cases of our proposed RHGD-CG and RSHGD-NS algorithms, respectively. While [42] use exponential map in their Algorithms 1 and 2, we study more general retraction, of which exponential map is one specific instance. It should be noted that our contributions is not limited to CG variants of proposed RHGD algorithm as we study other hypergradient estimation strategies as well. Overall, we provide hypergradient estimation error bound for a variety of estimation strategies and numerically compare them in terms of convergence, estimation error (Figure 1) and runtime (Table 3).

(2) We present many applications that require to solve a Riemannian bilevel optimization problem. Many of them are of independent interest.

(3) In Appendix G, we have shown how the proposed framework also leads to a convergence analysis for Riemannian min-max optimization and Riemannian compositional optimization problems, where they are special instances.

To summarize, our work provides a more general and practical framework than [42].

2. (W2) Experiment comparison with the methods presented in [42].

We have now included numerical comparisons in the one-page supplementary material. Specifically, we compare the performance on the synthetic problem (Figure 1) and hyper-representation task (Figure 2). We see the use of retraction in our case improves the efficiency of the solvers, which is reflected in the reduced runtime to achieve convergence.

3. (W3) The analysis of retraction should include the results of three additional hypergradient estimators.

Thank you for the suggestion. Below we include the convergence results for the three additional hypergradient estimators and will make sure to include them in the final revision.

Theorem. (Extension of Theorem 3) Under Assumptions 1,2,3,5 and let $\tilde L_F = 4 \kappa_l c_R M + 5 \bar c L_F$. Consider Algorithm 1 with the following settings.

• HINV: Suppose $\eta_x = \Theta(1/\tilde L_F), S \geq \tilde \theta(\kappa_l^2 \zeta)$. Then $\min_{k=0,..., K-1} \| \mathcal{G} F(x_k) \|^2_{x_k} \leq 16 \tilde L_F \Delta_0/K$.

• CG: Suppose $\eta_x = \Theta(1/\tilde \Lambda), S \geq \tilde \theta(\kappa_l^2 \zeta), T \geq \tilde \Theta(\sqrt{\kappa_l})$, where $\tilde \Lambda = C_v^2 \bar c + \kappa_l^2 ( \frac{5M^2 C_0^2 D^2}{\mu} + \bar c)$, then $\min_{k=0,..., K-1} \| \mathcal{G} F(x_k) \|^2_{x_k} \leq 96 \tilde \Lambda (\Delta_0 + \| v_0^* \|^2_{y^*(x_0)})/K$.

• NS: Suppose $\eta_x = \Theta(1/\tilde L_F), S \geq \tilde \Theta(\kappa_l^2 \zeta)$. Then for any $\epsilon>0$, $T \geq\tilde \Theta(\kappa_l \log(1/\epsilon))$, we have $\min_{k=0,..., K-1} \| \mathcal{G} F(x_k) \|^2_{x_k} \leq 16 \tilde L_F \Delta_0/k + \epsilon/2$.

• AD: Suppose $\eta_x = \Theta(1/\tilde L_F)$, $S \geq \tilde \Theta(\kappa_l^2 \zeta \log(1/\epsilon))$ for any $\epsilon > 0$, we have $\min_{k=0,..., K-1} \| \mathcal{G} F(x_k) \|^2_{x_k} \leq 16 \tilde L_F \Delta_0/k + \epsilon/2$.

The proof basically contains three parts, where we first show the convergence of inner iterations in terms of retraction, which is agnostic to the choice of hypergradient estimator. Then we bound the hypergradient estimation error where only AD based hypergradient needs different treatment when using retraction because it requires to differentiate through the retraction. For other hypergradient estimators, the bounds in Lemma 1 still hold given no retraction is used in their computation. Lastly we show the defined Lyapunov function is monotonically decreasing.

4. (1) Assumption 1 is strict and unusual in bilevel optimization. (2) The notation $D$ in $d(x_1, x_2) \leq D$ is confusing. (3) Assumptions 2 and 3 could be inferred from Assumption 1 and the necessity to include Assumption 2, 3.

(1) Assumption 1 is often unavoidable in the literature of Riemannian optimization, especially when dealing with geodesic (strongly) convex objectives [21, 37, 62, 63, 64]. This is required to ensure bounded curvature in order to apply the trigonometry distance bound (Lemma 3), which is crucial for achieving linear convergence of the inner iterations. (2) For notation clarity on $d(x_1, x_2) \leq D$, we will change the notation in the revised manuscript. (3) It is correct that Assumption 2 and 3 can be inferred from Assumption 1. Nevertheless, we explicitly state such assumptions for clarity and for properly defining the Lipschitz constants.

5. (Q1) The notations $C_{hinv}$, etc in Theorem 1 are not used in the main text.

We will move such notations to appendix.

6. (Q2) Extension to the general stochastic setting?

Yes, the proposed framework can accommodate the general stochastic setting. To this end, we require to construct unbiased gradient estimate and properly control its variance. A more formal treatment of such general case is left for future exploration.

We thank the reviewer for acknowledging that our theoretical analysis is thorough and our work has practical relevance. We also appreciate the constructive feedback.

1. (W1) Assumption such as geodesic strong convexity and Lipschitz continuity.

The assumption of (geodesic) strong convexity is common for analyzing bilevel optimization, which is also the case even in the Euclidean space [35,40,53]. This is because without strong convexity, the lower-level problem can be ill-defined, where there may exist many local solutions. Studying bilevel problem with non-(geodesic)-strongly-convex lower-level problem is even challenging in the Euclidean space. Nevertheless, it is possible to relax such assumption to (geodesic) convexity with an extra strongly convex regularizer, or to (Riemannian) PL condition where a global minimizer exists. On the other hand, the assumption on Lipschitzness is also standard in the bilevel literature such as [35,40,46].

2. (W2) Empirical evaluation is limited to synthetic and small-scale problems.

We would like to emphasize that in addition to the synthetic and small-scale setting, we evaluate our approach on real-world scenarios such as unsupervised domain adaptation on Office-Caltech dataset (Section 4.4) and meta-learning where we optimize a convolutional neural network (CNN) with orthogonal constraints on mini-Imagenet dataset (Section 4.3). In particular, the meta-learning example leads to a large-scale problem that involves training a 3-layer CNN, with each layer kernel parameters constrained to Stiefel manifold. This results in a parameter set of $3 \times 144 \times 16$ parameters. In addition, we train over 20,000 tasks where each task contains 50 images of size $84 \times 84$ (5-ways, 5-shots). We will make sure to include the details on experiment setup in the revised version.

3. (Q1) Computational trade-offs when choosing between the different hypergradient estimators.

We have discussed computational trade-offs in Table 1 and Appendix H.4 (together with convergence plots in Figure 1), where we compare the runtime for evaluating hypergradient for different estimation strategies. In summary, CG is the most computational demanding but yields the best approximation while AD requires the least computational effort but gives a poor hypergradient estimation. NS balances the computation and estimation error.

Thanks so much for the responses. I will keep my rating and raise my confidence.

We thank the reviewer for positive comments on recognizing our solid theoretical analysis and extensive numerical results. We also greatly appreciate the constructive feedback and comments.

1. (W1) Clarification regarding [42]: (1) Do other estimators show better performance? (2) Provide numerical comparisons with [42]. (3) Include complexities of methods in [42].

Thank you for raising these points. Below, we address each one individually.

(1) There exists a accuracy-runtime trade-off for these estimators. The trade-off can be observed from Figure 1 where we compared various estimators in terms of convergence, hypergradient error and runtime. We have also additionally compared per-hypergradient runtime for these estimators in Table 3 (Appendix H.4). From the figures and table, we see hypergradient based on automatic differentiation requires the least runtime while suffers from a poor hypergradient approximation accuracy. Neumann series based estimators balance the trade-off while conjugate gradient requires the most runtime but has the lowest approximation error. Choosing a suitable estimator depends on specific applications.

(2-3) We highlight that the methods presented in [42] are covered by the framework that we propose, i.e., Algorithm 1 in [42] coincides with RHGD-CG and Algorithm 2 in [42] is RSHGD-NS that we propose, except that they use exponential map for the update, while our framework works for general retraction. Thus, the methods in [42] can be considered as special cases in our framework where we have already included the complexities.

Please find numerical comparisons with [42] in the one-page supplementary material. In particular, we compare on the synthetic problem (Figure 1) and shallow and deep hyper-representation (Figure 2). We observe the use of retraction in our case results in more efficient update and thus reduces the runtime for convergence.

2. (W2) The estimations of the inverse Hessian are also widely used in Euclidean space.

While it is true that the estimation strategies for inverse Hessian have been explore in the Euclidean space, in more general manifold setting, certain more care is needed when using second-order derivatives. For example, $\mathcal{G}^2_{xy}$ (used in Algorithm 1) is a second-order cross derivative relating two different manifolds (one for x and the other for y) and defining this is tricky. Additionally, Hess is now an operator on the tangent space instead of a matrix in the Euclidean space. Consequently, its inverse needs to be properly characterized and computed.

3. (W3) (1) Assumption of geodesic strong convexity is strong. (2) Lower-level problem in Section 4.2 is in Euclidean space.

(1) Analyzing bi-level problem with non-(geodesic)-strongly-convex lower-level problem is tricky (even in the Euclidean setting) because of non-unique solution of the lower-level problem. Non-unique solution would give rise to several challenging cases such hypergradient becoming ill-defined, invertibility of the Hessian, etc. However, it is possible to relax this assumption to geodesic convex lower objective, by adding a strongly convex regularizer term to the objective function.

(2) Since Euclidean space is a special case of Riemannian manifold, it is not unreasonable to consider problems where lower-level problem is in the Euclidean space. However, we would like to emphasize that Sections 4.1 and 4.4 explore lower-level problems constrained to the SPD manifold.

4. (Q1) Why performances of hinv-20 and hinv-50 similar?

It should be noted that the numbers 20 and 50 refers to the number of inner iterations for solving the lower-level problem. The similar performance of hinv-20 and hinv-50 is because the lower-level problem for such synthetic problem is easy to solve and 20 iterations already provides a good close-to-optimal solutions.

5. (Q2) Why hinv more efficient as in Figure 1(b)?

In this synthetic example, we can derive an analytic form of the Hessian inverse. In addition, the problem dimension is not large (d = 50) and thus computing the Hessian inverse directly is more efficient in this setup.