Paper: Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes (ICLR-2024)

Reviewer comment (Writing):
the presentation is a places not clear. For example, I have a hard time extracting meaning from a sentences such as "In addition, previous diffusion models on the manifold are geometrically not intuitive as their generative processes are derived from the time reversal of the noising process that accompanies score function in the drift which does not explicitly take into account the underlying geometry." I believe the score of a Brownian motion is indeed quite related to the underlying geometry?

Author response:
The purpose of the statement is to explain that the **score function does not provide an explicit geometrical interpretation** for the diffusion of previous methods. In contrast, for our approach, the drift of the Logarithm bridge yields an intuitive geometrical explanation that corresponds to a tangent vector minimizing the geodesic distance toward an endpoint.

We have clarified the statement as “... noising process that accompanies score function in the drift which does not provide explicit geometric interpretation.” in the updated revision.

Reviewer comment (Novelty):
as far as I can read, the class of processes in equation (3) are not new (they are claimed to be new in the paper). In the stochastic analysis literature, they are denoted 'Fermi bridges' see https://arxiv.org/abs/1604.05182 They have been used for bridge simulation in e.g. https://arxiv.org/abs/2105.13190 and https://arxiv.org/abs/2104.03193 . I haven't seen them used in context of generative modelling though.

Author response:
Thank you for the related works on diffusion bridges. However, Fermi bridge [1] describes a bridge process to a submanifold $N$ instead of a single point, and [1] or the related works do not study or present the explicit form of our Logarithm bridge. By constraining the submanifold $N$ to a single point, one can derive our Logarithm bridge from the Fermi bridge, but the derivation is not trivial and was not explored before, to the best of our knowledge. 

Furthermore, in our work, we provide the geometric interpretation for the drift of the Logarithm bridge, i.e., the tangent vector minimizing the geodesic distance toward its endpoint, which is not studied in [1]. Especially, based on this we propose a novel generative framework (Logarithm bridge mixture) which is a novel contribution in the context of generative modeling. 

We have added the comparison with the Fermi bridge in the updated revision (Appendix A.2). Note that we could not find a process similar to our Logarithm bridge or Fermi bridge in [3].

Reviewer comment (Theory):
I am unsure what is the main benefit of the approach: The score matching in denoising diffusion models is replaced by a loss depending on the distance which is also in general expensive to compute (the logarithm map (inverse exponential) on general manifolds is expensive to compute). I am not sure of the general validity of claims such as 'Our method shows superior performance on diverse manifolds that can scale to higher dimensions' because logarithm map computations can be quite non-trivial in high dimensions.

Author response:
Many real-world problems we are interested in consider manifolds that satisfy either (1) have known geodesics, for example, robotics motion planning (SO(3)), protein modeling (SE(3)), neuroscience (SPD matrices), astrophysics (hypersphere),  and computational chemistry (product of tori), or (2) the eigenfunctions of the Laplace-Beltrami operator can be computed, for instance, manifolds represented as meshes. **When dealing with these real-world problems, our work provides practical benefits compared to previous works** as follows:

- Our method achieves superior performance on diverse manifolds.
- Our approach can scale to higher dimensions compared to previous methods that require computation of Riemannian divergence.
- Our method has a significantly faster training time compared to previous works since ours is divergence-free and requires greatly reduced in-training simulation steps.
- Our framework is readily adaptable to general manifolds and arbitrary prior distributions, which is not the case for previous methods.
- Additionally, our work provides an intuitive geometrical interpretation of the diffusion process and further yields predictions of the final result (Eq. (7), Figure 5).


---

Even though computing logarithm maps can be non-trivial and expensive for general high-dimensional manifolds, **many real-world problems (e.g., robotics, biology, chemistry) consider manifolds with known geodesics, and the logarithm maps are easy to compute even for high dimensions**. When dealing with these real-world problems, our framework has significant advantages in faster training and reduced computational cost, compared to previous works that require Riemannian divergence. Also, our method can scale to higher dimensions compared to previous methods as shown from the experiments on high-dimensional tori. Additionally, when the computation of the logarithm map is expensive and non-trivial, we can leverage the spectral bridge process (Eq. (5)) that bypasses computing the logarithm map as the reviewer thankfully acknowledges.

We agree that when scaling to high-dimensional general manifolds for which computing the logarithm map is expensive, our framework can also have difficulty like previous works. We have toned down our claim as “our method can scale to higher dimensions when the geodesics are known.“ in the updated revision.

Reviewer comment (Theory):
The authors write "the time reversal of the noising process that accompanies score function in the drift which does not explicitly take into account the underlying geometry." I don't think this is true as the score function on $\mathcal{M}$ uses the \textit{Riemannian} gradient which needs information on the metric $g$.

Author response:
The purpose of the statement is to explain that the **score function (Riemannian gradient) does not provide an explicit geometrical interpretation** for the diffusion of previous methods. In contrast, for our approach, the drift of the Logarithm bridge yields an **intuitive geometrical explanation** that corresponds to a tangent vector minimizing the geodesic distance toward an endpoint.

We have clarified this point by revising the statement into “... noising process that accompanies score function in the drift which does not provide explicit geometric interpretation.” in the revised paper.

Reviewer comment (Theory):
Background: "compact manifolds". Actually, I think you consider compact, connected, orientable Riemannian manifolds. Otherwise, you cannot define a volume form needed for the integration of probability distributions. Furthermore, the inverse of the expmap may not be well defined. Also, hyperbolic spaces do not fit this definition.

Author response:
Thank you for the comment, we do consider **complete, orientable, connected, and boundaryless** Riemannian manifolds, following the assumption of RSGM. We have clarified this point in the revised paper. 

It appears there was a misunderstanding, as the purpose of the experiment on the hyperbolic space (Section 5.5) is to **empirically show that our framework can be extended to non-compact manifolds**. We demonstrate in Figure 9 that our framework can learn the distribution accurately.

Reviewer comment (Writing):
What is $U$ in equation 1?

Author response:
$U$ denotes the potential function. For example, $U(x) = d_g(x, \\mu)^2 + \\log \\|D \\exp^{-1}_{\\mu}(x)\\|$ ensures the stationary distribution to be the wrapped Gaussian distribution with an arbitrary mean location $\\mu\\in\\mathcal{M}$, where $d_g$ denotes the geodesic distance. We have added this explanation to the revised version of the paper.

Reviewer comment (Theory):
The authors write: "the inverse of the exponential map on the manifold," Not True! For example, when the exp map fails to be a diffeomorphism — i.e., at the cut locus. As a concrete example, in the Lie group, the image of the exponential map of the connected but non-compact group $SL2(R)$ is not the whole group.

Author response:
You are right, the exponential map is not injective at the cut locus. Since we are assuming a connected and complete Riemannian manifold, we have clarified in the updated version of the paper that we assume the endpoint is not in the cut locus of the current state.

Reviewer comment (Theory):
The authors write the Bridge process is conditioned on the endpoint $z$. But to build a Brownian bridge you need to pin both endpoints at $t=0$ and $t=T$, otherwise, it's not a bridge. Certainly, that's what is being done when you do two-way bridge matching, so this focus and imprecision in the earlier parts of the text are confusing.

Author response:
Although it is true that we need to fix both the starting point ($t=0$) and the endpoint ($t=T$) to define a bridge process, we have omitted to mention the starting point since **the drift of the bridge process solely depends on the endpoint**, describing the bridge process concisely as a “diffusion process conditioned on the endpoint”. We agree that denoting the fixed starting point would make it clearer, and have clarified this in the updated version of the paper.

Reviewer comment (Theory):
One of the main claims of this paper is that the proposed approach is a lot cheaper computationally because 1.) you do not have to simulate both directions of the SDE and 2.) there is no divergence computation. However, I believe the first part of the claim is a bit misleading because you are effectively trading off variance of the objective for simulation by reusing the same sample in forward and backward matching losses. Moreover, I believe this is exactly why the proposed approach has high variance — because to get a low variance estimate you would need to simulate the SDE multiple times as it is noisy, so reusing samples here is counter-productive.

Author response:
We would like to clarify that apart from these two reasons, **our two-way approach also contributes to greatly reducing the computational cost** during training, since it requires significantly fewer in-training simulation steps compared to the one-way approach, achieving $\\times$34.9 speed up.

Our method **does not trade the variance using the two-way approach** (i.e., reusing the same sample in forward and backward matching). Reusing the same sample does not result in higher variance since we simulate the bridge processes with different starting and end points sufficiently many times.

We empirically show in the table below that we **obtain almost similar results and variances when using different samples for both forward and backward matching losses**. Further, we would like to clarify that our method does not have high variance compared to the baselines as shown in Tables 1, 2, and Figure 2.

|                                                        | Earthquake                    | Proline                       |
|--------------------------------------------------------|-------------------------------|-------------------------------|
| Ours w/ different sample for forward/backward matching | -0.29 $\\pm$ 0.08 | 0.14 $\\pm$ 0.025 |
| Ours w/ same sample for forward/backward matching      | -0.30 $\\pm$ 0.06 | 0.14 $\\pm$ 0.027 |

Reviewer comment (Novelty):
The authors comment that the time-rescaled proposal distribution $q(t)$ is different than Huang et al. 2022. I disagree; this is the exact same idea that you can resample time instead of uniformly sampling it. Huang et al. 2022 do this using a 1D flow which is quite cheap and can be done as a preprocessing step, while here it is a constant. These are essentially the same, but the proposed approach is a cruder approximation.

Author response:
The **proposal distribution used for the importance sampling in our framework and the one used in RDM are not the same**. We leverage $q(t)\\propto\\sigma^{-2}_t$ as a proposal distribution that is tractable and $t\\sim q$ can be sampled easily with ignorable cost. 

On the other hand, RDM uses an intractable proposal distribution and relies on an approximation via an additional neural network that requires separate training. Especially, the training of the importance sampler happens in between the training of the diffusion model $a$, which is not done as a pre-processing step or not done in parallel, and therefore results in extra computational cost and an increase in the training time of the framework. 

We would like to emphasize that our approach is much simpler compared to that of RDM while effective in stabilizing the training and improving the generation quality, without the need for additional computation or training time. 

We have clarified this point in the revised paper as “While the idea of the importance sampling for the time distribution was also used in Huang et al. (2022), our approach leverages a simple and easy-to-sample proposal distribution $q$, which is effective in stabilizing the training and improving the generation quality, without the need for additional computation or training time.”

Reviewer comment (Experiments):
The Table 1 results are problematic for two reasons. First, these earth science datasets have known issues and lead to wildly different results based on the subset used for training and test sets. This largely means that we cannot take table numbers from other papers. A quick check shows that the baseline numbers are copied from the papers, and looking at the provided code it seems the baselines were not implemented. This leads me to believe that the authors did not try the baselines themselves on the same training and test splits as their own reported models. As a result, this table is not useful and cannot be compared. I actually think the proposed method will likely be even better if all baselines are compared properly, so I encourage the authors to include this in the updated paper.

Author response:
We have **used the same data split with the reproducible works, RSGM and RFM**, by using the same seed values of 0-4 for the data split. Unfortunately, **the results from some previous works (CNFM, RDM) were not reproducible** as they do not provide open-sourced codes and do not provide details of the data split, and we had no choice but to take the numbers from their paper. Note that we obtained similar numbers to the reported results for RSGM and RFM when running their open-source codes. As the reviewer acknowledges, we believe that our method will likely be even better if all baselines are compared properly, and we are happy to update the results if the codes for other baselines are open-sourced.

Reviewer comment (Experiments):
A major claim in the paper is the ability of the method to start from any prior distribution on the manifold to hit any target. This is a strength that vanilla diffusion models do not have. However, I could not find an experiment in the main paper where this is exploited.

Author response:
By the construction of the Riemannian diffusion mixture (Eq. (6)), it is **theoretically guaranteed that we can use any prior distribution to build a diffusion process to hit the given target distribution**. However, in order to fairly compare with the previous methods that are limited to trivial prior distributions, we have used the same prior distributions with the baseline methods during the experiments. As finding a prior distribution to improve the generation quality is not the main goal of our work, we leave it as future work.

Reviewer comment (Experiments):
Another claim in the paper is that the method can scale to high-dimensional problems, but the largest dimensionality problem is a Torii of $10^3$ dimension. This is hardly a high-dimensional manifold as Torii can be written as a product group $S^1 \times S^1 \times \dots \times S^1$ and each manifold in the product is simple and does not actually utilize the difficulty that comes with high dimensions. I would consider toning down this claim.

Author response:
Our framework can scale to high dimensions (dimension 2000), which is a notable result considering the limited number of model parameters (MLP with 2 hidden layers), while previous methods (Moser Flow, RSGM, and RFM) fail on lower dimensions (~500). We observe that our method produces consistent results (log p) for higher dimensions when using more model parameters, scaling fairly well even to the tori of dimension 10,000.

Although it is true that our framework has advantages on high-dimensional tori since the logarithm map can be easily computed due to the fact that tori can be written as a product group, previous works cannot exploit this advantage as the **difficulty of modeling high-dimensional data with previous methods comes from the computation of Riemannian divergence**. 

We would like to emphasize that many real-world problems we are interested in consider manifolds with known geodesics, for example, robotics motion planning (SO(3)), protein modeling (SE(3)), neuroscience (SPD matrices), astrophysics (hypersphere),  and computational chemistry (product of tori), and for these problems, our work can provide practical benefits of faster training and reduced computation, as seen in the case of high-dimensional tori.

However, we agree that our framework may also have difficulty when scaling to high-dimensional general manifolds for which computing the logarithm map is expensive. Therefore, we have toned down our claim as “our method can scale to higher dimensions when the geodesics are known.” in the revision.

Reviewer comment (Theory):
The noisy paths learned by this approach do not appear to be optimal in the sense of matching the empirical marginal densities. Can we make this better and perhaps more numerically stable by using Optimal transport paths like done in OT-Flow matching (Tong et al. 2023) for Euclidean geometries?

Author response:
We believe one can construct a diffusion process satisfying optimal transport property by leveraging the tools from Schrodinger bridge [1] or stochastic optimal control [2] to the Riemannian setting, which would be a promising direction for future work.

[1] Peluchetti, Diffusion Bridge Mixture Transports, Schrodinger Bridge Problems and Generative Modeling, arXiv 2023

[2] Chen et al., Likelihood Training of Schrodinger Bridge using Forward-Backward SDEs Theory, ICLR 2022

Reviewer comment (Experiments):
The experimental validation of the method may not be comprehensive and may not cover all possible scenarios. The details of the experimental setup and results may not be easy to follow.

Author response:
We respectfully disagree, as we conduct **extensive experiments in diverse settings on both real-world and synthetic datasets**, including various manifolds (sphere, tori, manifold with non-trivial curvature, hyperbolic space) and different dimensions, to validate the effectiveness of our method.

Reviewer comment (Experiments):
The training time of the proposed framework is not extensively discussed or compared with other methods. While it is mentioned that the proposed method allows for faster training compared to previous diffusion models, a more detailed analysis of the computational efficiency and scalability of the framework would be beneficial.

Author response:
We have provided a **comparison of the training time with the baselines in Table 2 (Right)** and discussed in Sections 5.3 and 5.4. Our method achieves up to $\\times$5 speed up for training compared to RSGM based on the Riemannian divergence and further achieves $\\times$12.8 speedup compared to RFM on mesh datasets.

Reviewer comment (Experiments):
Does the author try such method on higher dimensional manifolds?

Author response:
Yes, we show that our method can scale to high dimensions through **experiments on high-dimensional tori (Section 5.3)**.  Figure 2 (Right) demonstrates that our method can scale fairly well even for high dimensions, while previous works fail as the dimension increases showing a significant drop in the performance.