Riemannian Proximal Sampler for High-accuracy Sampling on Manifolds

We thank the reviewer for taking the time to engage with our work and for their comments. While we respectfully disagree with some aspects of the assessment (main one being ``total lack of experiments''), we appreciate the opportunity to clarify the points of misunderstanding.

Response to weaknesses:

"However, I remain dubious about the novelty and the broader significance of this work. First, the results seem to be a relatively easy and straightforward extension of proximal sampling techniques from the Euclidean setting. The adaption is perhaps expected, and might limit the originality of the results."

• Actually, the extension is not as easy as it looks. The proximal sampling techniques was interpreted as alternatively run heat flow and its time reversal [1], where a key property is that certain "convexity" (e.g., LSI) is preserved along the heat flow. This is easy to prove in the Euclidean space setting, as one can interpret heat flow as convolution with Gaussian, but on a Riemannian manifold, such strategy doesn't work. See our Proposition 31 for a theoretical justification, where heat flow was interpreted as a Markov operator, which is quite different from the Gaussian convolution treatment.

"A more serious concern arises from a practical perspective. The implementation of the Riemannian proximal sampler relies crucially on rejection sampling to correct the additional bias introduced by the inexact oracles. But, it is folklore that rejection sampling can be computationally intractable in real-world and high-dimensional applications as it might require an exponential number of proposals in dimension. "

• In general, a rejection sampling algorithm would result in an exponential number of proposals in dimension. However, this can be avoided. In Euclidean space, the proximal sampler also requires rejection sampling to implement the RGO step, see [1]. But by carefully choosing a step size $\eta$ and constructing a suitable proposal, the number of proposals is $O(1)$ in dimension. We did theoretically justify this, see Proposition 49 and Proposition 44. Roughly speaking, we showed that for $M = \mathcal{S}^{d}$, the cost for rejection sampling is of constant order, considering the dimension.

"This brings serious issues and may cause the "practical implementations" to be totally infeasible in practice. It might be possible that the inexact method can be made more practical by considering other common techniques like MH used in high-accucary sampling, but the authors do not consider this or offer an explanation. Moreover, it is okay to be without an empirical section for a purely theoretical paper. However, the total lack of experiments is problematic when the practicality of the algorithm is not fully justified by theoretical results alone. Therefore, I would recommend rejection unless these concerns are addressed."

• Please note that we had provided simulation experiments (for both implementations) in Appendix B, see pages 15 and 16.

• Furthermore, based on your questions, we tested the practical applicability by sampling from $\mathcal{S}^{100}$, and the experiment result shows that the (average) number of proposals is not exploded by dimension. We will add this new experimental result to the revised draft as well.

Thank you again for your feedback and questions. We sincerely hope that our response has addressed your concerns and may lead you to reconsider your evaluation. Should our clarifications increase your confidence in the work, we would be truly grateful if you felt it appropriate to adjust the score accordingly.

[1] Yongxin Chen, Sinho Chewi, Adil Salim, and Andre Wibisono. Improved analysis for a proximal algorithm for sampling. In Conference on Learning Theory, pages 2984–3014. PMLR, 2022.

We thank the reviewer for their thoughtful questions and positive evaluation.

Response to weaknesses:

"Some notations should be unified to better understand the results. Due to the lack of space, it is a pity that all the numerical results are given in Appendix."

• Thanks for the suggestion. We will reorganize and try to move some experimental results to the main draft as much as the space allows.

Response to questions:

"Why in Theorem 8, there is no distinction between positive and negative curvature situations, as in Theorem 6? "

• Thanks for pointing out this. Yes there should be a slight distinction, where in case of negative curvature we still need an curvature dependent upper bound for the step size $\eta$. But the iteration complexity will remains to be the same. The revised theorem look like this:
• [Theorem 8] Similar to Theorem 6, let $M$ be a Riemannian manifold without boundary. Assume Assumption 1 holds. For any initial distribution $\rho_{0}^{X}$, to reach $\tilde{\mathcal{O}}(\varepsilon)$ total variation distance with oracle accuracy $\zeta = \zeta_{\mathsf{RHK}} = \zeta_{\mathsf{MBI}} = \frac{\varepsilon}{\log^{2} \frac{1}{\varepsilon}}$ and step size $\frac{1}{\eta} = \tilde{\mathcal{O}}(\log \frac{1}{\varepsilon})$ (for negative curvature, we additionally require $\eta \le 1/|\kappa|$),
  • if $\pi^{X}$ satisfies $\alpha$-$\mathsf{LSI}$, we need $k = \tilde{\mathcal{O}}(\log^{2} \frac{1}{\varepsilon})$ iterations.
  • if $\pi^{X}$ satisfies $\alpha$-$\mathsf{PI}$, we need $k = \tilde{\mathcal{O}}\left(\log \frac{1}{\varepsilon} \log \frac{\chi_{\pi^{X}}^{2} (\rho_{0}^{X})}{\varepsilon^{2}} \right) = \tilde{\mathcal{O}}\left(\log^{2} \frac{1}{\varepsilon}\right)$ iterations.

"Algorithm 2 and 3, when it is possible to find a suitable t, i.e. to ensure the existence of t and $C_{RHK}$ or $C_{MBI}$ such that … it is unclear whether $C_{RHK}$ and $C_{MBI}$ are negative. "

• We analyzed this theoretically in the appendix, please see for example Proposition 39.

"The idea of Varadnan asymptotics makes sense when $\eta \to 0$, according to Fact 5. However, the rates of inexact oracles rely on large eta in Assumption 1, therefore it is conceptually contradictory to use Varadnan asymptotics in a large eta regime. A clearer explanation about this is needed to ensure that these results are correct."

• Yes, intuitively, when $\eta$ is large, Varadhan's approach could lead to a larger error. An error analysis for Varadhan approach is still open. Our theoretical high-accuracy results are for truncation approach. But we are able to experimentally verify the effectiveness of this approach.

"Clarify the space of g in Definition 4. For all g in which space?"

• $g$ can be any smooth function on $M$.

"The notations of $\pi^{Y|X}$ and $\pi^{X|Y}$ in Algorithm 1 are not consistent with Section 4. Are they the same as $\pi_{\eta}^{Y|X}$ and $\pi_{\eta}^{X|Y}$ in Assumption 1?"

• Yes. See line 170: When there is no ambiguity, we omit the step size $\eta$ ...

"Clarify the notation tilde O (difference to O) used in several places (e.g. thereom 8)."

• Yes, we will clarify the tilde O notation in our paper. In Theorem 8, we used the tilde O notation to emphasize that we are only considering the $\varepsilon$ dependency.

We thank you again for your insightful feedback and questions. If our response above satisfies your concerns, we would greatly appreciate an increase of score to reflect your increased confidence in our work.

We thank the reviewer for their thoughtful questions and positive evaluation.

Response to weaknesses:

"Minor comment: In Section 2.1, Line 133, it would be helpful to specify the metric being used. Assigning a different metric (e.g., the Bures–Wasserstein metric) on the same space can result in a positively curved geometry."

• Thanks for the comment, we will specify the metric. We are using the affine invariant metric.

"The analysis relies on the LSI condition of the distribution on the manifold. However, it is unclear how strict or mild this condition is in the context of Riemannian manifolds. It would have been helpful to illustrate how strict or mild this assumption is by providing examples or criteria that can build the intuition (see also the below question section)."

• Thanks for the suggestion, we will add more examples to clarify. In general, for convex potentials, the celebrated Bakry-Emery condition provides general results for a density to satisfy LSI based on the spectrum of the Hessian of the potential $f$ and the Ricci curvature of the manifold. The question is wide open in the case when the potential is not convex.
• For the Euclidean setting, recent works like [2, 3, 4] for some recent works of translating LSI/PI conditions on $\pi$ to conditions on $f$ like the Polyak-Lojasiewicz condition. We believe that similar conditions could be derived in the manifold setting, although there might be more sophisticated technical challenges to handle. We also emphasize that even in the Euclidean setting this question is not fully resolved.

Response to questions:

"I am not sure whether Lines 601–602 are correct..."

• Let $f(x) = -\kappa \mu^{T} x$. Clearly $x = \frac{\mu}{||\mu||}$ is the minimizer, and $x = -\frac{\mu}{||\mu||}$ is the maximizer. Note that they are the only two critical points. To compute the Riemannian Hessian, we construct a smooth extension of $\text{grad} f$. Let $G(x) = - (I - x x^{T}) \kappa \mu = - \kappa \mu + x x^{T} \kappa \mu$ be a vector field defined on $\mathbb{R}^{d+1}$. The Euclidean directional derivative of $G$ is

$

DG(x)[u] = u x^{T} \kappa \mu + x u^{T} \kappa \mu
= \kappa (x^{T} \mu) u + \kappa (u^{T} \mu) x

$

Hence when $u \in T_{x}\mathcal{S}^{d}$, we have $\text{Hess} f(x)[u] = \kappa (x^{T} \mu) u$. Recall that the only critical point that is not a global minimum is $x = - \frac{\mu}{||\mu||}$. (Actually, it is the maximizer.) We can get (for all unit tangent vector $u$), the Riemannian Hessian satisfies $\text{Hess} f(x)[u, u] = \kappa (x^{T} \mu) u^{T} u = \kappa (x^{T} \mu) = - \kappa ||\mu||$. We see that the Hessian has negative eigenvalue, and hence $\lambda_{*}$ is positive.

"Related to the previous question: as mentioned in the paper, there exists no globally supported non-constant strongly log-concave density on compact manifolds..."

• Indeed, as mentioned before, this is related to the curvature-dimension conditions (e.g., Bakry-Emery), which is discussed for example in [1](Proposition 5.7.1). For convex potentials, the LSI is satisfied baesd on the Ricci curvature of the manifold. More generally, it is open with recent explorations in the Euclidean setting, as mentioned previously.

"In Theorems 28 or 31, if $\kappa < 0$, then $\alpha < 0$..."

• In our notation, $\alpha = \frac{\kappa}{1-e^{-\kappa t}}$. When $\kappa < 0$, we have $e^{\kappa t} > 1$, hence $\alpha > 0$. For example, if $\kappa = -1$, we have $\alpha = \frac{-1}{1-e^{t}} > 0$.

We thank you again for your intriguing questions and careful evaluation.

[1] Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.

[2] Sinho Chewi and Austin J Stromme. The ballistic limit of the log-sobolev constant equals the Polyak–Łojasiewicz constant. arXiv preprint arXiv:2411.11415, 2024.

[3] Yun Gong, Niao He, and Zebang Shen. Poincare inequality for local log-Polyak-Łojasiewicz measures: Non-asymptotic analysis in low-temperature regime. arXiv preprint arXiv:2501.00429,2024.

[4] August Y Chen and Karthik Sridharan. Optimization, isoperimetric inequalities, and sampling via Lyapunov potentials. arXiv preprint arXiv:2410.02979, 2024.

We thank the reviewer for their thoughtful questions and positive evaluation.

Response to weaknesses:

"LSI somehow corresponds the KL condition in optimization..."

• We would like to clarify that the extension is more involved than it might appear on surface. The proximal sampling algorithm was interpreted alternatively as running heat flow and its time reversal in [1], where a key property is that certain "convexity" (e.g., LSI) is preserved along the heat flow. This is easy to prove in the Euclidean space setting, as one can interpret heat flow as convolution with Gaussian. But on a Riemannian manifold, such a strategy doesn’t work. Please refer to our Proposition 31 for a theoretical justification, where heat flow was interpreted as a Markov operator, which is quite different from the Gaussian convolution treatment. It appears that this subtle yet important point may not have been fully considered in your assessment.
• Also, as opposed to the Euclidean setting, due to the lack of closed-form solution for the heat kernel, both oracles are to be approximated to make the method practical. In our work, we provide high-accuracy rates under this simultaneous inexact oracles setup, which is also different (and relatively non-trivial) from the Euclidean setup.

"The paper also considers analysis under the inexact oracle via kernel truncation..."

• Thanks for the question. Actually, we did provide experiments for both practical approaches; the results were presented in Appendix B, page 16. The first two examples are for hyperspheres, where we implemented the truncation approach. The third one is for the manifold of positive definite matrices, and we implemented the Varadhan approach. Experimental results suggest that our proximal sampler (using both approaches) achieves a higher accuracy compared with the Riemannian Langevin Monte Carlo method.

Response to questions:

"What exactly $d$ in Line 292 refers to? "

• $d$ denotes the geodesic distance on Riemannian manifold. The notation was mentioned in line 81-82. Informally speaking, consider $S^{2} \subseteq \mathbb{R}^{3}$, like our earth. The geodesic distance between two cities is the length of a path that people need to travel. Here the path is "curved" (because the earth itself is curved), but not a Euclidean straight line.

"In Proposition 9, the truncation level is given in order for the inexact sample to achieve certain accuracy. I am wondering whether this should also relies on how fast the reject sampling converges?"

• Yes, in Proposition 9 we mainly consider the truncation level needed for a desired accuracy. On the other hand, the cost for rejection sampling (i.e., how fast rejection sampling converges), is an open problem in general. But we are able to give an theoretical analysis for the cost of rejection sampling when $M = \mathcal{S}^{d}$, see Remarks at line 335 and Corollary 10.

"In Line 144, what does the "minimal solution'' exactly mean? "

• Roughly speaking, if there is another function that also solves the heat equation, then it must be greater than or equal to the heat kernel. See [2](Theorem 4.1.5, Proposition 4.1.6) for more details.

"Overall, what properties should $f$ satisfy in order to satisfy the LSI and PI properties? "

• Thanks for the intriguing question. For the Euclidean setting, recent works translated LSI/PI conditions on $\pi$ to conditions on $f$ like the Polyak-Lojasiewicz condition. These works considered the LSI/PI constant for $e^{-f(x)/t}$. For example, in the $t \to 0$ limit, i.e., "low temperature" setting, [3] considered LSI constant and [4] considered PI constant. Another work [5] considered a "Optimizability" condition and analyzed LSI/PI constant for (informally) $t \le O(1/d)$. We believe that similar conditions could be derived in the manifold setting, although there might be more sophisticated technical challenges to handle (and we plan to consider this in our future work, thanks to your question). We also emphasize that even in the Euclidean setting this question is not fully resolved.

"In Line 584, should it be the Riemannian Langevin Monte Carlo Algorithm? "

• Yes. Sorry for the typo, we will correct it.

We thank you again for your insightful feedback and questions. If our response above satisfies your concerns, we would greatly appreciate an increase of score to reflect your increased confidence in our work.

[1] Yongxin Chen, Sinho Chewi, Adil Salim, and Andre Wibisono. Improved analysis for a proximal algorithm for sampling. In Conference on Learning Theory, pages 2984–3014. PMLR, 2022.

[2] Elton P Hsu. Stochastic analysis on manifolds. Number 38 in Graduate Studies in Mathematics. American Mathematical Soc., 2002.

[3] Sinho Chewi and Austin J Stromme. The ballistic limit of the log-sobolev constant equals the Polyak–Łojasiewicz constant. arXiv preprint arXiv:2411.11415, 2024.

[4] Yun Gong, Niao He, and Zebang Shen. Poincare inequality for local log-Polyak-Łojasiewicz measures: Non-asymptotic analysis in low-temperature regime. arXiv preprint arXiv:2501.00429,2024.

[5] August Y Chen and Karthik Sridharan. Optimization, isoperimetric inequalities, and sampling via Lyapunov potentials. arXiv preprint arXiv:2410.02979, 2024.

Thanks for your clarification question.

Please note that Algorithm 1 is stated assuming exact availability of RHK and MBI oracles. Theorem 6 and 8 provides high-accuracy guarantees for Algorithm 1 under the availability of exact and inexact oracles respectively. For the inexact case, the oracles are required to satisfy Assumption 1. Proposition 9 concerns Algorithms 2 and 3, which is one way to implement the RHK and MBI oracles. Specifically we show that this implementation satisfies the condition in Assumption 1. To show, it is enough to determine the truncation level.

Rejection sampling is a part of the implementation in Algorithm 2 and 3, which are both based on truncated heat kernels. The cost of rejection sampling comes in the number of steps of the for loop in both the algorithms.

If we have the exact heat kernel, and we are on the Euclidean space (for which the exact heat kernel is known), then the cost of rejection sampling can be proved to be $O(1)$ [1].

Hypothetically, even if we have the exact heat kernel on a Riemannian manifold, the cost for rejection sampling is actually unknown. This is what we mean in our previous response when we say "the cost for rejection sampling (i.e., how fast rejection sampling converges), is an open problem in general."

Intuitively, one can expect that if we have a highly-accurate evaluation of the heat kernel, the cost of rejection sampling should be the same as that of rejection sampling with the exact heat kernel.

For the case of sphere, we provide the end-to-end result (including the cost of rejection sampling) in Corollary 10. In proving this result, we show that the cost of rejection sampling (even with the inexact heat-kernel based on truncation level as stated in Proposition 9), is $O(1)$ similar to the Euclidean case; See Proposition 39 and 44 for details. Furthermore, we note that empirically the cost of rejection sampling for the Positive Semi Definite manifold in Section B.3 is also $O(1)$. Theoretically proving such a result for more general manifolds is an open question.

To summarize, to show Assumption 1 holds, it is enough to set the truncation level of the heat-kernel. The cost for rejection sampling should be independent of the truncation level, as long as the truncation level is providing high-accuracy heat kernel evaluations.

We will clarify this in detail in our revision.

Please let us know if you have any other questions.

[1] Yongxin Chen, Sinho Chewi, Adil Salim, and Andre Wibisono. Improved analysis for a proximal algorithm for sampling. In Conference on Learning Theory, pages 2984–3014. PMLR, 2022.