Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space

We deeply grateful for your careful reviewing and valuable comments. We address your concerns as below.

Q1: Strong technical conditions (Lipschitz-gradient or log-Sobolev constants).

A1: For the Lipschitz-gradient Assumption 1 and 2, they are only applied to bridge the gap between discrete and continuous methods (Proposition 2, 4, 5), while our convergence results in Theorem 1, 2, 3 do not require them. Besides, the Lipschitz assumption is standard to bridge the discrete and continuous differential equation (as we did) in numerical analysis (see section of convergence in [4]).

For log-Sobolev inequality, we also characterize the convergence rates of our methods without it in Theorems 1, 2, 3, evaluated by the norm of Riemannian gradient as in standard non-convex optimization. Since log-Sobolev inequality roughly implies "strongly-convexity" as clarified in our paper.

Q2: The absence of any empirical validation.

A2: Thanks for your valuable advice. We have made an example (Example 1) of sampling on Gaussian to illustrate the tightness of our theoretical results.

However, we recognize that numerical results will make our conclusions more convincing. We evaluate different sampling algorithms on two experiments. Since the NeurIPS rebuttal do not allow us to convey these results by image, we present them in the following tables (results are averaged by 5 independent runs).

Concretely, we set the each target distribution $\mu_{\xi}$ as Gaussian or mixture-Gaussian, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and SVRG has less gradient computations than Riemannian GD). The results are summarized in Tables 1-4, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1].

As can be seen, under the same update steps, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while under the same gradient computations, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. We promise to add these numerical results in the revised version, and we sincerely hope you may reconsider the concern on numerical results.

Table1. KL and Fisher divergences on Gaussian measured by update steps

Table 2. KL and Fisher divergences on Gaussian measured by gradient computation

Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps

Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations

Q3: The language and grammar used throughout the paper is very challenging to read.

A3: We really grateful for your carefully reviewing all these materials with heavy mathematical notations. We promise we will revise our paper as follows to make our paper more readable.

We will carefully check the paper word by word to fix all unproper phrasing, gramma issues, and convoluted sentence structures.

We will rewrite the part of theoretical background, add more intuitive exposition to make these materials friendly to all audiences.

Q4: Yet, the paper stops short of connecting these continuous flows to concrete algorithmic implementations or demonstrating improvements over existing discrete methods.

A4: By combining our proved convergence rates of continuous methods and the proved gap between discrete and continuous methods can address your concerns. Roughly speaking, by Proposition 6, we know the Wasserstein distance is upper bounded by KL divergence (proved to be converge) $\frac{\gamma}{2}W_{2}^{2}(\pi, \mu) \leq D_{KL}(\pi\parallel \mu)$. Since Propositions 2, 4, 5 imply $W_{2}^{2}(\hat{\pi}_{n}, \pi_{n\eta}) \leq E[||x_{n} - \hat{x}_{n\eta}||^{2}]\leq \mathcal{O}(\eta)$ ($\hat{\pi}_{n}$ and $\pi_{n\eta}$ are respectively obtained by proposed discrete and continuous algorithms), we get the convergence rates of the discrete methods, which further demonstrates the improvements over existing discrete methods. We promise to add these analysis to in the revised version.

Q5: The resulting convergence rates mirror those already known in the Euclidean setting.

A5: Though the convergence rates match the ones in Euclidean space, but the proof techniques are more complicated. For example, it takes us a lot efforts to construct the Lyapunov function in (114), which is critical to prove the convergence rate. The construction relies techniques from both optimization and stochastic analysis, which are more sophisticated compared with the scheme in Euclidean space.

Q6: Extensions to other objective functions more than KL divergence.

A6: Our framework has the potential to be extended to other objectives. As can be seen, our core idea in Figure 2 does not depend on the optimization objective, which guarantees the extensions of flows to the other objectives. However, it worth noting that a practical algorithm may only be obtained under KL divergence.

Taking Riemannian gradient as an example, it takes the formulation of $\frac{\partial{\pi_{t}}}{\partial{t}} = \nabla\cdot(\pi_{t}\mathrm{grad} F(\pi_{t}))$. The flow can be implemented in practice only if $\mathrm{grad} F(\pi_{t})$ is efficiently computed. However, this is not true for many objectives. For example, then $F(\pi) = W_{2}^{2}(\pi, \mu)$, $\mathrm{grad} F(\pi) = -2(T_{\pi\rightarrow \mu} - \mathrm{id})$ (see Theorem 1.4.11 in [1]), the $T_{\pi\rightarrow \mu}$ is the optimal transportation map between $\pi$ and $\mu$, which is $\mathcal{O}(e^{d})$ complex in general.

However, for KL divergence, $\frac{\partial{\pi_{t}}}{\partial{t}} = -\nabla\cdot(\pi_{t}\nabla\log{\mu}) + \Delta\pi_{t}$, corresponds to SDE $dx_{t} = \nabla\log{\mu(x_{t})}dt + \sqrt{2}dW_{t}$, which makes it practical. This explains why the existing gradient flow in Wasserstein space is mainly discussed under KL divergence [1, 2, 3].

Q7: Furthermore, there has been a lot of recent work in this area, and it is not entirely clear how much of this work is original in the context of other recent papers.

A7: As clarified in line 72, the existing literature focus on discrete Langevin dynamics and its convergence rates under different criteria [1, 2, 3]. The main contributions in this paper the continuous Riemannian SGD and SVRG flows are not explored yet.

Q8-10: The numerical verification; Grammar issues; The extensions to other divergence.

A8-10: The answers refer to A2, A3, and A6 in above.

Ref:

[1] Kinoshita et al., Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization

[2] Cheng et al., Convergence of langevin mcmc in kl-divergence

[3] Chewi et al., Log-concave sampling

[4] https://en.wikipedia.org/wiki/Numerical_methods_for_ordinary_differential_equations

We thank for your valuable comment and address your concerns as follows.

Q1: I believe some technical assumptions are lacking (e.g. the analysis is restricted to absolutely continuous distributions, but it is never mentioned...)

A1: Thanks for pointing it. In this paper we assume the target distribution $\mu\propto \exp(-V(x))$ in line 102. Then the KL $D_{KL}(\pi\parallel \mu)\propto \int V(x)d\pi(x) - \int \log{\pi(x)}d\pi(x)$, so that the KL divergence exists whenever $\pi(x)$ is absolutely continuous to Lebesgue measure. This is true when the initialization distribution $\pi_{0}$ satisfies this condition, since the discussed $\pi_{t}$ satisfies the F-P equation (2). We will add this discussion in the revised version to make our conclusion more precises.

Q2: There are no numerical results demonstrating the theoretical results

A2: Thanks for your valuable advice. We have made an example (Example 1 in our paper) of sampling on Gaussian distributions to illustrate the tightness of our theoretical results.

However, we recognize that numerical results will make our conclusions more convincing. We evaluate different sampling algorithms on two experiments. Since the NeurIPS rebuttal do not allow us to convey these results by image, we present them in the following tables (results are averaged by 5 independent runs).

Concretely, we set the each target distribution $\mu_{\xi}$ as Gaussian or mixture-Gaussian, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and Riemannian SVRG has less gradient computations than Riemannian GD). The results are summarized in the following tables, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1].

As can be seen, under the same update steps, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while under the same gradient computations, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. We promise to add these numerical results in the revised version, and we sincerely hope you may reconsider the concern on numerical results.

Table1. KL and Fisher divergences on Gaussian measured by update steps

Table 2. KL and Fisher divergences on Gaussian measured by gradient computation

Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps

Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations

Q3: In Section 3 and 4 giving preliminaries on the Wasserstein space and on Wasserstein gradient of the KL divergence, there are no assumptions of regularity on densities.

A3: Yes, you are correct, the Wasserstein space $\mathcal{P}$ in this paper should be $\mathcal{P_{2,ac}}$ which means the set of probability absolutely continuous measures with finite $L_{2}$-norm. These conditions can be satisfied under an absolutely continuous $\pi_{0}$ as mentioned in A1. We will revise it as you suggested. On the other, the definition of Riemannian gradient is similar to the definition of Wasserstein gradient in Section 1.4.1 of [1] and Riemannian gradient in Page 27 of [2]. Mathematically, it can be viewed as sub-gradient. But it is usually named as Riemannian gradient in optimization on manifold community. We will clarify this in the revised version.

Q4: In Remark 1, it is stated "Unfortunately, unlike our method, this method is not applied to the stochastic optimization problem in Section 5, because the aforementioned minimization problem can not involve stochastic gradient as expected." However, stochastic proximal algorithm also exist, and could be used to solve the same problem as Section 5, equation (11). While it may not have been done yet, I don't think the statement is true.

A4: Our expression in Remark may leads to a little bit misunderstanding. We say solving $\pi_{n + 1} = \arg\min_{\pi}D_{KL}(\pi\parallel \mu) + W_{2}^{2}(\pi, \pi_{n}) / 2\eta$, when $\eta\to 0$ "can not generalize to stochastic problem" means the method can not induce a continuous flow in stochastic version. Let us check this in Euclidean space will be more clear.

Consider $x_{n + 1} = \arg\min_{x} f(x) + \frac{||x – x_{n}||^{2}}{2\eta}$, we get $(x_{n + 1} – x_{n}) / \eta = -\nabla f(x_{n})$ so that $\eta\to 0$ implies the gradient flow $\frac{dx_{t}}{dt} = -\nabla f(x_{t})$. However, when the objective is stochastic i.e., $E_{\xi}[f_{\xi_{\eta}}]$, the sub-problem becomes $x_{n + 1} = \arg\min_{x} f_{\xi_{n}}(x) + \frac{||x – x_{n}||^{2}}{2\eta}$, then $(x_{n + 1} – x_{n}) / \eta = -\nabla f_{\xi_{n}}(x_{n})$, which implies $\frac{dx_{t}}{dt} = -\nabla f_{\xi_{t}}(x_{t})$ with $\xi_{t}$ is a sample of $\xi$. As can be seen, the drift term has randomness, which is not expected in Wasserstein space (a deterministic metric space consists of probability measures). Because the direction $v$ of curves $\frac{\partial{\pi_{t}}}{\partial{t}} = -\nabla\cdot(\pi_{t}v)$ in Wasserstein space should be deterministic. We will make this more clear in the revised version.

Q5: In this work, you analyze only the convergence of the continuous dynamic. Would it be possible to study also the convergence of the discrete schemes?

A6: The proof techniques of our continuous methods can be generalized to discrete schemes, by reusing our proposed Lyapunov function. A more straightforward extrapolation is directly combining our convergence rates of continuous methods with the proved gap between discrete and continuous methods in Propositions 2, 4, 5. Though this can be done when log-Sobolev inequality is satisfied i.e., $\frac{\gamma}{2}W_{2}^{2}(\pi, \mu) \leq D_{KL}(\pi\parallel \mu)$ (see our Proposition 6 in Appendix).

Q7: Other objectives are solved using a stochastic Wasserstein gradient descent. Could the analysis done in this work be extended to these objectives?

A7: Yes, the key idea is using our transformation in Figure 2 to link the constructed SDE and curve in Wasserstein space.

Q8: It could be nice to add experiments demonstrating the rates, even in the Gaussian case where everything may be known in closed-form?

A8: Thanks for pointing this, we have added the numerical experiments as clarified in Q2.

Ref:

[1] Kinoshita et al., Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization

Q1: “No numerical performance comparison in specific sampling problems is shown”

A1: Thanks for your valuable advice. We have made an example (Example 1 in our paper) of sampling on Gaussian distributions to illustrate the tightness of our theoretical results.

However, we recognize that numerical results will make our conclusions more convincing. We evaluate different sampling algorithms on two experiments. Since the NeurIPS rebuttal do not allow us to convey these results by image, we present them in the following tables (results are averaged by 5 independent runs).

Concretely, we set the each target distribution $\mu_{\xi}$ as Gaussian or mixture-Gaussian, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and Riemannian SVRG has less gradient computations than Riemannian GD). The results are summarized in the following tables, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1].

As can be seen, under the same update steps, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while under the same gradient computations, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. We promise to add these numerical results in the revised version, and we sincerely hope you may reconsider the concern on numerical results.

Table1. KL and Fisher divergences on Gaussian measured by update steps

Table 2. KL and Fisher divergences on Gaussian measured by gradient computation

Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps

Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations

Q2: “Does the technology for constructing gradient flow (including SGD and SVRG flows) mentioned in this paper largely depends on the special properties of exponential mapping”

A2: Yes, you are correct, our conclusion highly depends on the linear structure of exponential map. Roughly speaking, which makes the corresponded “update moving” of random vectors in Euclidean space is linear. Based on this, we can take $\eta\to \infty$ and derive the SDE and corresponded flows in Wasserstein space.

Q3: “The exponential mappings are different when using different metrics. Hence, if using another Riemannian metric instead of the inner product in the Euclidean space, the resulting exponential may not possess the ``linearity’’. Then, how do we construct the gradient flow in that case?”

A3: This is a good question. Technically, the Riemannian metric induced the distance in Wasserstein space, when the inner product is replaced (the definition above (3)), the distance in probability measure space is no longer Wasserstein space. As in Euclidean space the inner product $\langle a, b\rangle = a^{\top}b$ is replaced by $\langle a, b\rangle_{H} = a^{\top}Hb$ for some positively definite matrix $H$, then the distance between vectors $a$ and $b$ becomes $||a - b||_{H}^{2}$.

Back to your question, when the inner product is changed and specifying the corresponded exponential map, we can similarly construct the Riemannian GD flow by following the idea in Figure 2. However, as you said, the new exponential map may not possess the ``linearity’’. Therefore, it may take some efforts to construct some differential structure in the new update rule. Fortunately, when replacing the inner product $\langle u, v\rangle_{\pi} = \int \langle u, v\rangle d\pi$ with $\langle u, v\rangle_{\pi} = \int \langle u, v\rangle_{H} d\pi$ for some positively definite matrix, the resulting exponential map still poses the linearity.

Q4: How do the results extend to non-smooth or non-log-concave distributions, which are common in practice?

A4: Firstly, we prove the convergence rates under non-log-concave distributions in the first parts of Theorem 1, 2, 3, which address your second concerns.

Secondly, the smoothness Assumptions 1 and 2 are used to bridge the gap between discrete optimization methods and the continuous Riemannian flows, i.e., Propositions 2, 4, 5, which is unavoidable due to the theory of numerical ODE. However, without the smoothness Assumptions does not influence the convergence rates of our proposed flows. The conclusions in Theorems 1, 2, 3 do not rely on the smoothness Assumptions.

Q5: Small size of equation (25)

A5: Thanks for pointing out this, we will properly resize it.

Ref: [1] Kinoshita et al., 2022. Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization.

Dear reviewer,

Thanks for your applying, we would like to clarify the details of our numerical results, and we will add these and more experimental results in our revised version. Firstly, we summarize the hyperparameters of Gaussian and Mixture-Gaussian sampling tasks in Table 1. For Gaussian sampling, we sample from a finite sum $N = 1000$, i.e., $\exp(\frac{1}{N}\sum_{i=1}^{N}\log{\mu_{\xi_{i}}})$ with $u_{\xi_{i}}$ are sampled from $N(u_{i}, I)$.. This is the case in Example 1 of our paper.

For Mixture-Gaussian sampling, we sample from we sample from a finite sum $N = 5$, i.e., the distribution $\exp(\frac{1}{N}\sum_{i=1}^{N}\log{\mu_{\xi_{i}}})$ with $\mu_{\xi_{i}}\sim \sum_{k=1}^{K}N(u_{i, k}, I)$, i.e., each $\mu_{\xi_{i}}$ is a mixture-Gaussian with $K=2$ clusters. For both methods, we sample $M = 1000$ samples to verify our methods. Please notice that the we set $N = 5$ here for mixture-Gaussian, because sampling from such mixture-Gaussian is a highly non-convex optimization problem [1], a large $N$ will make the sampling process really difficult.

Table 1: Hyperparameters of Sampling Tasks.

For our three algorithms, their hyperparameters are summarized as in Table 2 and Table 3. Notably, the results for Gaussian are numerical solution of our derived closed-form solutions in Example 1, i.e., equation (20) and line 322.

Table 2: Hyperparameters for Gaussian in Table 1 of our first rebuttal. (Epochs means the SVRGLD are implemented for 1 epochs)

Table 3: Hyperparameters for Gaussian in Table 2 of our first rebuttal. (counted by gradient computations)

Table 4: Hyperparameters for Mixture-Gaussian in Table 3 of our first rebuttal.

Table 5: Hyperparameters for Mixture-Gaussian in Table 4 of our first rebuttal (counted by gradient computations)

We would like to address any of your further concerns.

Thanks for noting this. The fast convergence rate only holds for Gaussian case. In fact, for sampling on from Gaussian, as clarified in line 319 of paper. SVRG becomes GD but with the same computational cost with SGD. Thus, in Table2 of initial rebuttal, the SVRG is actually GD with much more update steps (5000 in Table 2 v.s. 100 in Table 1). The larger update steps and exponential convergence rate (Theorem 3) together lead to the 1e-16 accuracy. Please also see Table 3 in our second response for more details (more update steps).

To Reviewer mLox:

We thank for your valuable comment and address your concerns as follows.

Q1: “Somehow lack numerical validations on the efficiency of Riemannian SGD and SVRG on sampling”

A1: Thanks for your valuable advice. We have made an example (Example 1 in our paper) of sampling on Gaussian distributions to illustrate the tightness of our theoretical results.

However, we recognize that numerical results will make our conclusions more convincing. We evaluate different sampling algorithms on two experiments. Since the NeurIPS rebuttal do not allow us to convey these results by image, we present them in the following tables (results are averaged by 5 independent runs).

Concretely, we set the each target distribution $\mu_{\xi}$ as Gaussian or mixture-Gaussian, and conduct sampling by the proposed algorithms. Notably, our convergence rates are respectively evaluated by update steps ($T$) and gradient computation (Riemannian SGD and Riemannian SVRG has less gradient computations than Riemannian GD). The results are summarized in the following tables, the KL divergence and Fisher divergence are estimated by combining MCMC and an estimated density as in [1].

As can be seen, under the same update steps, Riemannian SVRG and Riemannian GD have better convergence results than Riemannian SGD (see Theorem 1, 2, 3), while under the same gradient computations, Riemannian SGD and Riemannian SVRG (the two algorithms have less gradient computations in each update step) have better convergence results (see discussion after Theorem 3). These results are consistent with our conclusion in paper. We promise to add these numerical results in the revised version, and we sincerely hope you may reconsider the concern on numerical results.

Table1. KL and Fisher divergences on Gaussian measured by update steps

Table 2. KL and Fisher divergences on Gaussian measured by gradient computation

Table 3. KL divergence and Fisher divergenceson on mixture Gaussian measured by update steps

Table 4. KL divergence and Fisher divergenceson on mixture Gaussian measured by gradient computations

Q2: “The overall technical details differ not too much from the Euclidean case on the convergence of the flows”

A2: Actually, our theoretical techniques are quite different from the ones in Euclidean space. Firstly, the construction of Riemannian SGD, SVRG i.e., Propositions 2, 4, 5 are mainly based on the techniques from stochastic analysis, which are quite new in machine learning community. On the other hand, we admit the convergence analysis borrow the ideas from Euclidean space, while their proof techniques are more complicated. For example, the construction of Lyapunov function (the core idea of proving convergence results) for Riemannian SVRG (114), which takes us a lot of efforts (introducing the distance penalization).

Q3: A discussion on the random noise $\xi$ in objective, can Riemannian reduced into Riemannian GD.

A3: Firstly, as you said, our conclusion is independent of $\xi$ for Riemannian SGD (to apply Riemannian SVRG $\xi$ should be discrete). We use a constant $\xi$ as a special case to address your concern about the connection between Riemannian SGD and GD. As can be seen in Riemannian SGD (12), when $\xi$ is constant, the variance $\Sigma_{SGD}$ becomes zero, which makes the Riemannian SGD reduce to Riemannian GD as you predicted. We will add this discussion in the revised version.

Q4: In Lemma 1, better be ``Lemma 1 ([29,42]).'', similar for Theorem 1

A4: Thanks for your suggestion, we will revise it.

Q5: What are the key differences from the existing proof for the establishment of Theorem 1?

A5: Our proof technique is a little bit difference with the existing result. Concretely, we prove the theorem by borrowing the notations in Riemannian derivate (see (28) and (40)), and prove the theorem under the idea of optimization on manifold. However, the existing literature e.g., [1] prove this theorem by computing the derivate to $t$ directly.

Q6: Eq. (25) too small?

A6: Thanks for pointing out this, we will make (25) larger in the revised version.

Ref: [1] Kinoshita et al., 2022. Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization.