Lie Group-Induced Dynamics in Score-Based Generative Modeling

First of all, we thank the reviewer for taking the time to read our manuscript and provide valuable feedback. We look forward to a constructive rebuttal period and are confident that the paper will be improved through this exchange.

In the following, we address the points raised in both the "Weaknesses" and "Questions" sections.

Weaknesses

Motivation

TL;DR: We propose a diffusion process in Euclidean space that retains all the characteristics of curved Lie group diffusion, enabling unconstrained dynamics (unlike equivariant neural networks) while preserving the group action on the data.

The key idea is that while data is typically represented in Euclidean space (e.g., 3D coordinates for point clouds or pixel values for images), its underlying "true" signal coordinates may reside in a curved space (e.g., bond angles for molecules or global transformations of images/solids). However, performing diffusion in curved spaces poses significant challenges:

• Data transformation/preprocessing: Extracting curved coordinates requires extensive preprocessing, which is generally labor-intensive and complex.
• Challenging diffusion dynamics: Diffusion on curved manifolds is challenging due to curvature-related terms and operations (covariant derivatives, etc..). While promising, this area is less developed compared to Euclidean diffusion, where mature theoretical and computational tools are readily available.
• Projection requirement: Ensuring diffusion paths remain on a curved manifold $X$ requires projections (e.g., the exponential map for Lie groups), which are computationally expensive.
• Need of sampling: Exact solutions to the SDE are unavailable for general curved spaces, necessitating simulation to obtain $p(x(t)|x(0))$. In contrast, we provide an exact solution for any Lie group $G$ (see Equation 8).

Thus, our approach aims at comprising the "best of both worlds": having a "curved-like"-dynamics (as if it was on a Lie group) but purely in Euclidean space, so that we do not need to modify the data and we can take over any computational method developed for standard diffusion models.

The "curved dynamics" determined by the Lie group enters in the SDE through a first-order term (the Lie algebra elements) and one second-order term (the Casimir element). In general, the projection map is a Taylor expansion with infinite terms.

Moreover, we wish to point out that standard Euclidean diffusion also corresponds to a choice of Lie group, so it can also be seen as part of our formalism. So, our proposal is not an alternative to standard Euclidean diffusion, but rather an extending framework.

Finally, when the Lie group $G$ satisfies the constraints of Section 2.2, the approach generate an uncostrained dynamics, that is, it is able to learn any distributions with support on $X$, thus achieving the maximum expressiveness while maintaining the inductive bias given by the group action on the data space.

A more general (open) question is the relevance of constrained generative models.

We thank the reviewer for giving us a chance to explain this. Our strategy is not constrained:

• By selecting $G$ according to the sufficient condition of section 2.2, we obtain a fully unconstrained dynamics, whose sub-dynamics though follows the orbits of the given Lie subgroups.
• This is shown explicitly in the experiments in section 5.1 (where we are able to reconstruct symmetry-free distributions with any suitable group), as well as in the QM9 example, where we can generate molecules either with $G=T(3N)$ or with $G=(S(3)\times \mathbb{R})^N$.
• Also recall that we proved that traditional diffusion models are a special case of our formalism with $G=T(N)$, so our formalism includes all known unconstrained models as well.
• In Figure 1(c,d) we tried to depict in a simple case the difference between our generalized Lie Group-induced dynamics and a constrained dynamics (equivariant or invariant).

In this sense, we believe that our approach could bring the benefits of inductive bias into such models (for instance, aligning Lie group subfactor to the desired degrees of freedom of the data, as in Figure 4 for bond and torsion angles), without loosing the expressivity of unconstrained generative models.

This is a very interesting and important point, and we will include it in the final discussion.

We thanks the reviewer for taking the time to read and provide comments to our manuscript. We are looking forward to a constructive rebuttal period and we are confident that the paper will improve through the exchange.

Weaknesses

Can X be curved?

The very short answer to this question is as follows:

• The conditions in section 2.2 "SUFFICIENT CONDITIONS FOR LIE GROUP-INDUCED GENERALIZED SCORE MATCHING" are valid for a curved $X$
• For the main results Theorem 3.1, and the rest of the paper $X$ is flat/Euclidean.

Now, this is not a restriction on our part but a truly motivating factor for our work. Let us be more explicit about it.

Our goal is to define a dynamics which can be represented by any continuous transformation (Lie group) but acting on flat space. In other words: How can we obtain a diffusion process that behaves like it is curved (as the orbits/trajectories of the generalized score functions are not in general "lines", i.e., geodesics of Euclidean space) while remaining on flat space?

The motivation for the above goal is as follows: while data is typically represented in Euclidean space (e.g., 3D coordinates for point clouds or pixel values for images), its underlying "true" signal coordinates may reside in a curved space (e.g., bond angles for molecules or global transformations of images/solids). However, performing diffusion in curved spaces poses significant challenges:

• Data transformation/preprocessing: Extracting curved coordinates requires extensive preprocessing, which is generally labor-intensive and complex.
• Challenging diffusion dynamics: Diffusion on curved manifolds is challenging due to curvature-related terms and operations (covariant derivatives, etc..). While promising, this area is less developed compared to Euclidean diffusion, where mature theoretical and computational tools are readily available.
• Projection requirement: Ensuring diffusion paths remain on a curved manifold $X$ requires projections (e.g., the exponential map for Lie groups), which are computationally expensive.
• Need of sampling: Exact solutions to the SDE are unavailable for general curved spaces, necessitating simulation to obtain $p(x(t)|x(0))$. In contrast, we provide an exact solution for any Lie group $G$ (see Equation 8).

Our approach, therefore, aims at obtaining an equivalent dynamics to curved space diffusion, while remaining in Euclidean space. Thus we have the property of the former, with all the benefits of the latter. In particular, the data does not need to be transformed, we can use all the algorithm developed for Euclidean diffusion, and there is no need for an expensive projection map.

The "curved dynamics" determined by the Lie group enters in the SDE through a first-order term (the Lie algebra elements) and one second-order term (the Casimir element). In general, the projection map is a Taylor expansion with infinite terms.

Moreover, we wish to point out that standard Euclidean diffusion also corresponds to a choice of Lie group, so it can also be seen as part of our formalism. So, our proposal is not an alternative to standard Euclidean diffusion, but rather an extending framework.

Finally, when the Lie group $G$ satisfies the constraints of Section 2.2, the approach generate an uncostrained dynamics, that is, it is able to learn any distributions with support on $X$, thus achieving the maximum expressiveness while maintaining the inductive bias given by the group action on the data space.

We will make this clearer in the introduction where we are motivating our work.

Counterexample to $X$ only being flat

While the main Theorem is derived for flat $X$ for the reasons above, we derived the sufficient conditions for generalized score matching in Section 2.2 for a general $X$.

A counterexample as the reviewer asks is the following. Consider $X=S^2$ the 2-sphere and $G=SO(3)$. This satisfies Proposition 2.2 (thus condition 1), as the stabilizer $G_x$ at each point $x\in X$ is 1-dimensional (rotations around that line defined by the origin and $x$ fix $x$), and while $\dim G =3$ we have that $\dim G/G_x = 2$ and $n=dim X =2$ thus, $\dim G/G_x < n$ holds nowhere on $S^2$.

Homogeneity holds because any two points on the sphere can be connected by an SO(3) rotation around the axis perpendicular to the plane of the great circle between them.

Additionally, the vector fields on $S^2$ generated by the Lie algebra elements commute (even if the Lie algebra elements themselves do not!), as shown in Appendix A.3. These actions, restricted to $SO(3)$, preserve the radius and descend naturally to $S^2 \subset \mathbb{R}^3$.

The reviewer's proof fails because condition 3 must hold only locally, not globally; the orthogonal tangent vectors spanning $G/G_x$ vary with the point $\mathbf{x}$ on $S^2$, as explicitly shown in Equation (25).

Questions

Elaborate on the specific advantages it offers over traditional algorithms.

An example that exemplifies the utility of our approach is the MNIST example in section 5.2. We wanted to learn a mapping between $p_T=$ rotated MNIST and $p_0=$ standard MNIST. Each digit is represent by a picture $x=\mathbb{R}^{28\times28}$, and with traditional diffusion we need to learn a $28\times28=784$ dimensional bridge problem (since we need to learn a map between two non-trivial distributions), even though we know intuitively it is a 1-dimensional problem.

If we wanted to perform diffusion on a curved manifold, we would need to transform the MNIST images $x$ in such a way that they lie in a curved $SO(2)$ manifold, which is very complicated.

With our methods, we can simply apply a $G=SO(2)$ generalized diffusion on the original data, leveraging the data structure we have and the true dimensionality of the problem.

The only other option for a baseline approach is to train a bridge-model in a $784$ dimensional space. Therefore, we trained a Brownian Bridge Diffusion model (BBDM) [1] on the rotated MNIST dataset from the experimental Section 5.2. The BBDM implements a continuous-time stochastic process in which the probability distribution during the diffusion process is conditioned on the start $x(0) \sim q(x_0)$ and end point $x(T) \sim p(x_T)$, where an intermediate point $x(t) \sim p(x_t| x_0, x_T, t) = N(x_t, \mu(t, x_0, x_T), \Sigma(t)$ can be simply drawn since the mean function depends on the pairs $(x_0, x_T)$, see Eq. 3 and 4 in reference [1].

We observe that the BBDM model is, in some cases, not able to reconstruct a rotated MNIST digit, see Figure 6a, digits: (2,4,5,6,7) but removes/adds pixels in the image. In some cases it even transitions from a rotated digit to another digit as shown in the example (from 9 to 4) in the last row of Figure 6b in the updated manuscript. This artifact is most likely due to the way the BBDM is trained by pinning two endpoints and allowing to change the (latent) represenation in the full 784-dimensional space.

[1] BBDM: Image-to-image Translation with Brownian Bridge Diffusion Models. Bo Li, Kaitao Xue, Bin Liu, Yu-Kun Lai, CVPR 2023

Other advantages of our method compared to Riemannian diffusion were highlighted in our summary of the motivation for this work above. For example, compared to GeoDiff [1], our trained model for conformer generation on QM9 is faster, requiring only T=100 diffusion timesteps, whereas GeoDiff was trained with T=5000 (a speed increase of 50 times!). This is significant because a higher number of timesteps is a major drawback of diffusion models and has driven the development of alternative approaches like flow matching.

[1] Xu, Minkai, et al. "Geodiff: A geometric diffusion model for molecular conformation generation."

Homogeneity

First a comment about why we impose homogeneity: homogeneity simply states that any two points in $X$ can be connected by a group transformation, since these determined the dynamics. If the opposite would be the case, it could happen that the endpoint of the dynamics would depend on the starting point, which is something we do not necessarily want to introduce. Other than that, everything would work even without this constraints. We will make this clearer in the manuscript.

To the specific points:

Traditional Euclidean methods also satisfy homogeneity. We prove that traditional diffusion model are a subclass of our formalism where we take $G=T(N)$, and $X=\mathbb{R}^N$, but $X$ is homogeneous with respect to $G$ for this choice: given $x,y\in \mathbb{R}^n$, there is always a translation element $v=y-x \in T^(N)\simeq \mathbb{R}^N$ connecting the two points.

Thus, homogeneity is not an extra requirement with respect to traditional Euclidean methods.

We hope to have addressed the remaining concerns raised by the reviewer, particularly those related to the scope of the paper. Specifically, we emphasize that $X$ being Euclidean is not a weakness but rather the main advantage achieved by of our work. The reviewer has also acknowledged that we successfully resolved their other concerns, including the assumptions required for the validity of the conditions in Section 2.2 (including the counterexample) and the need for homogeneity.

Additionally, we conducted several new experiments to address the reviewer’s request to elaborate on the specific advantages our framework offers over traditional algorithms. Since these points were not revisited in the second round of comments, we assume the reviewer is satisfied with our revision.

We would like to kindly ask the reviewer to reevaluate our work in light of the discussions and the fact that we feel we have thoroughly addressed all the concerns raised.

Should the reviewer have any further questions or require additional clarifications, we would be more than happy to provide them.

Best regards

Dear Authors

I deeply appreciate appreciate for your reply.

Some of my concerns are fully resolved:

Q2 for X in rotateMNIST: The explanation helps a lot and now I fully understand your setting. I would think it would significantly help the readers if you could add those to the final version. I agree the conditions are not violated.

Q3: this fully solves my concerns, thank you! The new experiment is really important and helpful.

Q5:

The theoretical analysis and the experiment are not aligned

The authors are totally correct.

Dimensionality reduction

The authors replies for the MNIST experiment they do not lose expressivity. I totally agree with this. However, if you consider under this setting, then RSGM also has the same benefits of no need of a tradeoff, as we retain the expressivity of unconstrained models with the benefits of group inductive bias.. I think at least in the rotate MNIST experiment, this is because of the experiment setting but not the algorithm. But I am partially convinced.

I still have the following concerns:

Q1: I am so sorry I am still hard to follow. I spent a lot of time reading the interesting paper and I am getting lost during this rebuttal period. The main contribution is still not unclear to me, due to the fact that the authors are updating the paper, e.g., $X_{diff}$ is newly introduced and did not show up anywhere in the original version. The main contribution troubles me:

Rotated MNIST and CrossDocked2020 are low dim problems, since if you use RSGM, they are 1-dim and 4-dim, respectively.

QM9 and 2d, 3d and 4d distributions: both these 2 tasks can be solved by standard DSM in Euclideans and the improvement is unclear to me.

Q4:

This is my main concern: it is unclear to me what G and what group action enables me to use this algorithm.

• $G=U(d)$ and $X=C^d$: the authors claim In this case we have a complex Lie group. I agree this is not a good example. However, the difficulty is not from complex numbers but complicated group structures. Let's focus on the next examples.

• $G=GL$: The author states All the conditions are trivially satisfied since $G'=SO(3)\times R_+$ is a subgroup. No this is definitely not that trivial. Please take a look at condition 3. GL(n) has dim $n^2$. In the case $n=3$, you need to check all the 9 vector fields commutes. Simply saying the vectors in subgroup G' commutes is not enough.

• What's more, the group action can be very different from matrix-vector product, e.g., G=GL(n) and $X=R^{n\times n}$, the group action is $(A, x)\to Ax$ where $A\in G$ and $x\in X$ and $Ax$ means matrix multiplication of a $n\times n$ matrix (A) and another $n\times n$ matrix (x reshaped to $n\times n$ matrix). Note that X is still an Euclidean space and I am only defining the group action. How can you check this case? I can also have more examples, like conjugation: $(A, x)\to AxA^{-1}$. See https://math.stackexchange.com/questions/1970623/showing-that-conjugation-is-a-group-action

Q5:

Higher dim cases Again, in our revision we provided an experiment on SO(4) and we provided the proof that such setup satisfied our conditions above. We provided the proof that the conditions for SO(4) (and for SO(n)) are satisfied in our answer to Q4.

I am so sorry but are you claiming the calculation in Sec A.5 is finished? I only see the parametrization but not checking the vector field commutes.

I am so sorry I cannot improve my rating. I am grateful for your effort but I am still confused in many aspects. I believe some readers may have similar confusion.

Summary to AC:

Dear AC,

I appreciate your service and hard work in holding such a great conference. In case of I cannot edit my comment after the discussion ends, I will summarize it here:

Strengths: I think the experiment shows the work performed better than RSGM in some cases, and the formulation of the problem is novel.

Weaknesses:

• The writing needs to be significantly improved. I spend a lot of time trying to read the paper and the main contribution is unclear to me until now. The paper is still being updated during the rebuttal period.

• It is unclear in what cases the conditions are satisfied and the conditions are hard to check. This fact also stops higher dim Lie groups and other more complicated Lie groups. The paper only checked SO(3) act on $R^3$ for matrix-vector product, but not higher-dim SO group nor other more complicated cases.

• I am still confused about the problem setup and the main contribution. To me, it seems all the problems solved here can be solved by existing tools. The numerical experiment is good enough to show the algorithm works, but not enough to show the performance is better than existing algorithms, e.g., Euclidean DSM in QM9.

• I would say the work provides a good result. however, this is fully based on a strong assumption (Condition 3) that is hard to check.

Regarding your statements:

It is unclear in what cases the conditions are satisfied, and the conditions are hard to check.

[...] this is fully based on a strong assumption (Condition 3) that is hard to check.

We find these comment a bit difficult to reconcile with the rest of your feedback. The conditions are stated explicitly in the form of Theorems and Lemmas, complete with rigorous proofs. Furthermore, you have indicated a confidence level of 5, suggesting that you have carefully reviewed and are satisfied with the correctness of these results. Additionally, you acknowledged in earlier comments that our counterexample resolved your concerns about the validity of these conditions.

We also wish to highlight that the conditions were explicitly verified for all the groups and examples analyzed in the paper. These include SO}(2), SO(3), SO(4) (please refer to our latest comment above), T(N), the dilation group in $N$-dimensions, and SE(3). These groups are the standard settings considered in almost all the works on diffusion on manifolds (SO(4) being a relatively novel inclusion). Given this, we feel that the assertion that the conditions are "hard to check" is somewhat unfair, as we have demonstrated their applicability to all common cases and explicitly verified them in every experiment included in the paper.

If the formal conditions are too abstract for the reviewer, here is a more descriptive summary of what they imply:

1. Identify the space $X$:
   Consider $X$ as the space where diffusion is performed. This is not necessarily the space where the data resides, but rather the space where the network outputs. We called this $X_{diff}$ in our previous replies while we named $X_{data}$ where the data live ($\dim X_{data}$ is the dimension for the input of the network) For instance:

   • $X$ has dimension 1 for the MNIST experiment,
   • $X$ has dimension 6 for $\mathrm{SE}(3)$ in CrossDocked,
   • $X$ has dimension $3 \times \text{\# nodes}$ for QM9.

2. Dimension matching of the group action:
   Ensure that the dimension of the image of the group action satisfies $\dim \rho_X(G) \geq \dim X$. This step is often straightforward, as it involves computing the kernel of the group action (i.e., the group elements whose action is the identity element on $X$) and ensuring the dimension of its complement matches or exceeds $\dim X$, $\dim \rho_X(G) \geq \dim X$.

3. Group coverage of the space:
   Verify that the group "covers" the entire space $X$. For example, T(N) generates all of $\mathbb{R}^N$, and SO(2) generates all rotations of an image. This property is usually so obvious that it doesn’t require explicit mention.

4. Lie algebra differential operator commutators:
   Compute the action of the Lie algebra differential operators and check that the commutators vanish. While this calculation can sometimes be lengthy, it follows systematically from the representation of the group action.

We hope that the points above, while slightly informal and not capturing the full mathematical rigor of our proofs, help convey the intuition behind our conditions. Most importantly, we aim to show that these conditions are natural and not inherently "strong" or "hard to check."

Your suggestion to include additional examples involving more unusual groups, such as GL(n) or the Unitary group, is certainly interesting. However, this suggestion comes very late in the rebuttal process, at a point when we can no longer update the paper. While we acknowledge the merit of exploring such examples, we also believe this request may go beyond the scope of what is reasonable within the current review cycle, especially considering the extensive effort we have already made to address your concerns and perform additional experiments.

We thanks the reviewer for taking the time to read and provide comments to our manuscript. We are looking forward to a constructive rebuttal period and we are confident that the paper will improve through the exchange.

We have improved the presentation of the paper. The changes can be seen in the newly updated version of the paper (highlighted in blue). Among them:

• We unified the notation for action of a group (representation) on a manifold.
• We defined $\boldsymbol(\Pi)$ more clearly in line 143.
• We defined $\boldsymbol{A}$ in line 178
• We corrected the notation $(x,y) \rightarrow (x_1,x_2)$ in the SO(2) example around equations (2) and (3)
• We unified the notation $x \rightarrow \boldsymbol{x}$ for vector quantities.
• We made clear the assumptions about the curvature of $X$. The conditions of Section 2.2 hold for any differentiable $X$, while the Theorem 3.1 shows how we can obtain a curved Lie group dynamics while remaining in Euclidean space.
• New Experiments
  • Higher dimension Lie group: We added an experiment with G=SO(4) to show that our formalism works with any Lie group beyond the usual SO(2) and SO(3)
  • We added a benchmark to the MNIST experiment. We trained a Brownian Bridge Diffusion model (BBDM) [1] on the rotated MNIST dataset from Section 5.2. The BBDM implements a continuous-time stochastic process in which the probability distribution during the diffusion process is conditioned on the start $x(0) \sim q(x_0)$ and end point $x(T) \sim p(x_T)$, where an intermediate point $x(t) \sim p(x_t| x_0, x_T, t) = N(x_t, \mu(t, x_0, x_T), \Sigma(t)$ can be simply drawn since the mean function depends on the pairs $(x_0, x_T)$, see Eq. 3 and 4 in reference [1]. We observe that the BBDM model is in some cases not able to reconstruct a rotated MNIST digit, see Figure 6a, digits: (2,4,5,6,7) but removes/adds pixels in the image and in some cases transitions from a rotated digit to another digit as shown in the example from 9 to 4 in the last row of Figure 6b in the updated manuscript. This artifact is most likely due to the way the BBDM is trained by pinning two endpoints and allowing to change the (latent) representation in the full 784-dimensional space.

We added these new experiments in the main text of the paper.

[1] BBDM: Image-to-image Translation with Brownian Bridge Diffusion Models. Bo Li, Kaitao Xue, Bin Liu, Yu-Kun Lai, CVPR 2023

First of all, we thanks the reviewer for taking the time to read and provide comments to our manuscript. We are looking forward to a constructive rebuttal period and we are confident that the paper will improve through the exchange.

We will address in the following both the comments raised in the "Weaknesses" and "Questions" sections.

Weaknesses

Motivation

TLDR: We derive a diffusion process in Euclidean space that possesses all the features of (curved) Lie group diffusion. We thus obtain an unconstrained dynamics (as opposed to equivariant neural networks) that maintains the knowledge of the group action on the data.

The reviewer asks about the motivation of our work. Here we provide a more detailed explanation that hopefully makes the motivation behind our work intuitive.

The idea is that while often the representation of the data is presented in Euclidean space (3d coordinates for point clouds, pixel values for images, etc.), the underlying "true" coordinates might live in a curved space (bond and torsion angles for molecules, global transformation of images/solids, etc.). However, performing diffusion in curved space, while possible, has notable challenges:

• Data transformation/preprocessing. The data needs to be preprocessed and transformed to extract the curved coordinates. This is laborious and difficult to do in general
• Diffusion is more complicated. Diffusion on curved manifolds includes additional difficulties that related to the curvature of the manifold itself. This is a very interesting field, but it is at its infancy compared to Euclidean diffusion, where lots of theoretical and computation results are readily available and we can leverage them.
• Need of projection. In order to perform diffusion on a general manifold $X$, we need to impose the condition that the diffusion path is contained within $X$ at all times. This implies the need of a projection (the exponential map in case of Lie groups), which is in general very expensive to compute.
• Need of sampling. For general curved space there is no exact solution to the SDE. That means that the forward solution (in training) must also be simulated in order to obtain $p(x(t)|x(0))$. We instead derived an exact solution for any group $G$ (equation 8).

Our approach, therefore, aims at obtaining an equivalent dynamics to curved space diffusion, while remaining in Euclidean space. Thus we have the property of the former, with all the benefits of the latter. In particular, the data does not need to be transformed, we can use all the algorithm developed for Euclidean diffusion, and there is no need for an expensive projection map.

The "curved dynamics" determined by the Lie group enters in the SDE through a first-order term (the Lie algebra elements) and one second-order term (the Casimir element). In general, instead, the projection map is a Taylor expansion with infinite terms.

Moreover, we wish to point out that standard Euclidean diffusion also corresponds to a choice of Lie group, so it can also be seen as part of our formalism. So, our proposal is not an alternative to standard Euclidean diffusion, but rather an extending framework.

Finally, we point out that when the group satisfies the condition of Section 2.2, we obtain an unconstrained dynamics, that is, not restricted to invariant/equivariant representations of the data. Our approach can therefore bring the benefits of group inductive bias into such models (for instance, aligning Lie group sub-factors to the desired degrees of freedom of the data, as in the CrossDocked experiment for global symmetries), without loosing the expressivity of unconstrained generative models. This is particularly relevant in light of recent work [1], where it has been shown that unconstrained model outperform equivariant ones. Our framework unifies the best of both worlds: we obtain an unconstrained dynamics (the network is unconstrained) but we still keep track of the group action on the data.

[1] Abramson, Josh, et al. "Accurate structure prediction of biomolecular interactions with AlphaFold 3." Nature (2024): 1-3.

How to effectively choose $G$?

We recall that if the conditions of Section 2.2 are satisfied, any choice of $G$ allows the learning of any distribution with support on $X$ (exactly like standard diffusion). So, our formalism is not a constrained learning.

Practically, it makes sense to choose $G$ as a product $G=G_1\times G_2 \times G_3\times \cdots$, where the specific subgroups act on "meaningful underlying coordinates" of the data. For instance $G_1$ could be related to a specific bond angle (see Figure 4). This means that the corresponding component of the score $s_1$ will guide the learning of that specific angle (all while remaining in Euclidean space!).

Another example for a useful choice of $G$ is if we wish to perform conditional learning. This is exemplified in the MNIST experiment. There we chose G=SO(2) to model global rotation of the images (note that these are not symmetries of the distributions). In this case we can learn the conditional distribution $p(MNIST | rotated MNIST)$. This is not possible with standard diffusion, as it necessitates to build a bridge between two non-trivial distributions. To compare our framework with a known benchmark, we trained a Brownian Bridge Diffusion model (BBDM) [1] on the rotated MNIST dataset from Section 5.2. The BBDM implements a continuous-time stochastic process in which the probability distribution during the diffusion process is conditioned on the start $x(0) \sim q(x_0)$ and end point $x(T) \sim p(x_T)$, where an intermediate point $x(t) \sim p(x_t| x_0, x_T, t) = N(x_t, \mu(t, x_0, x_T), \Sigma(t)$ can be simply drawn since the mean function depends on the pairs $(x_0, x_T)$, see Eq. 3 and 4 in reference [1].

We observe that the BBDM model is in some cases not able to reconstruct a rotated MNIST digit, see Figure 6a, digits: (2,4,5,6,7) but removes/adds pixels in the image and in some cases transitions from a rotated digit to another digit as shown in the example from 9 to 4 in the last row of Figure 6b in the updated manuscript. This artifact is most likely due to the way the BBDM is trained by pinning two endpoints and allowing to change the (latent) representation in the full 784-dimensional space.

We added this new experiment in the main text of the paper.

[1] BBDM: Image-to-image Translation with Brownian Bridge Diffusion Models. Bo Li, Kaitao Xue, Bin Liu, Yu-Kun Lai, CVPR 2023

Technical assumptions

The reviewer is correct that this distinction was not consistently emphasized throughout the work. As outlined earlier, our primary goal is to achieve (curved) Lie group dynamics within Euclidean space. However, in the process, we derived sufficient conditions for generalized score matching, which also hold when $X$ is curved. To clarify:

• The conditions presented in Section 2.2, "Sufficient Conditions for Lie Group-Induced Generalized Score Matching", are valid for a curved $X$ and they are a novel and nontrivial result on its own;
• However, the main results, including Theorem 3.1, as well as the focus of the rest of the paper, assume $X$ is flat/Euclidean.

We have updated the introduction and relevant sections to make this distinction explicit, ensuring there is no further confusion.

Numerical comparison to existing methods

We are not exactly sure what the reviewer means by numerical performance, but in the QM9 experiments, we trained the standard Fisher $G_0=T(3)^N$ and proposed diffusion model on $G_1=(SO(3)\times \mathbb{R}^+)^N$ with the same network architecture (but different forward as well as backward SDE dynamics) to model molecules' conformers in $X=\mathbb{R}^{3N}$, where $N$ is the number of atoms in a molecule. Hence, the computational complexity with respect to runtime within the same model class should be around the same, where the standard diffusion model on $G_0$ is likely to have faster runtime since less FLOPs are required compared to $G_1$. The purpose of the QM9 experiment was to show that both approaches are applicable to train a generative model for conformer sampling and that therefore we can achieve an unconstrained dynamics with Lie group different than for standard diffusion.

In regards to comparison with other methods:

• Compared to GeoDiff [1] our trained models are faster since both were trained on $T=100$ diffusion timesteps as opposed to GeoDiff that was trained with $T=5000$;
• Compared to the Torsional Diffusion model [2], their runtime is faster using 20 steps only since they operate on the curved manifold for the torsion angles $X=(\-\pi, \pi)^M$, reducing the dimensionality of the problem, where M is the number of rotatable bonds with $M<3N$ in a molecule. We instead perform diffusion in the whole $\mathbb{R}^{3N}$ space

[1] Xu, Minkai, et al. "Geodiff: A geometric diffusion model for molecular conformation generation."

[2] Jing, Bowen, et al. "Torsional diffusion for molecular conformer generation."

Paper writing clarity.

We thank the reviewer for carefully reading the paper and catching the missing figure caption. We also added in the same figure an experiment with a higher dimension Lie group, namely for $SO(4)$, showing that our methods trivially applies to Lie groups of any dimension.

We also corrected some inconsistent notation, especially in the $SO(2)$ example in Section 2.

Questions

How to choose the drift function.

The drift function can be chosen according to the following criteria:

• Choice of prior distribution: the prior distribution, from which we sample during inference, depends on the drift terms.
• Tractability of score function: in order to be able to have a tractable score function, we can choose the drift term to be affine in the flow coordinates. This is the assumption we make around equation (10) and (11).
• A more general drift function is also possible. However, in the score function is not tractable anymore, and we need to perform implicit score matching, along the lines of [1]

In the experiments in Section 5.1, 5,3 and 5.4 we used the "variance-preserving" drift function which resembles an Ornstein-Uhlenbeck process. The forward SDE is stated in Eq. 53 in the Appendix of the submitted manuscript. The time-varying beta function is chosen from a cosine scheduler proposed in [2]. By the choice of such drift function, the the distribution of the flow coordinates $\boldsymbol{\tau}$ will converge to a standard Gaussian, i.e. $p(\boldsymbol{\tau}_T | \boldsymbol{\tau}(\mathbf{x}_0)) \sim N(0, I)$.

[1] Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems.

[2] Prafulla Dhariwal, Alexander Quinn Nichol. Diffusion Models Beat GANs on Image Synthesis

RSGM is a special case of the proposed framework?.

We think that the answer to this question is already implicit in our description of the motivation above, but we address it again here.

Our framework is not a special case of the diffusion on Riemannian manifolds. If $M$ is a Riemannian manifold, it is more like that when $M=G$, a Lie group, we obtain a (curved) dynamics similar to diffusion on the Lie group $G$, but with the benefit of always staying on Euclidean space.

We see therefore our method as a more practical framework than RSGM, since we do not need to involve the projection map and we can use all the results and techniques developed for Euclidean score matching/diffusion.

We made this point clearer in the text when discussing our relationship with that line of work.

Connection between this work and diffusion models on general Lie groups.

The reviewer mentioned the paper [1] Zhu, Yuchen, et al. ”Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups.”, and we agree that we should include a discussion of how our work relates to it. A shorted version of our discussion here is added in the related work section.

The method in that paper leverages the concept of trivialized momentum, where diffusion (noising) is performed on the Lie algebra of the group rather than the Lie group itself. This approach, as is the case for ours, also allows to reduce diffusion on an Euclidean space, since the Lie algebra, being a vector space, is isomorphic to $\mathbb{R}^n$, which also eliminates curvature terms. The diffusion is mapped back to the Lie group $G$ via the (derivative of the) left multiplication map, which connects tangent spaces at different points on $G$.

They also propose a framework for dynamics that remain entirely on the curved manifold without, according to their claims, requiring projections. This is achieved by separating each step into two parts: updating the Lie algebra element $d\xi$ and updating the Lie group element $dg = g \xi$.

This method has, however, severe drawbacks:

• The main drawback lies in the non-Abelian case, which is the most relevant since the Abelian case essentially reduces to SO(2) or copies of it. For non-Abelian groups, the absence of a conditional transition probability necessitates implicit score matching, requiring divergence computation in the data space which is very expensive. Our method includes both theoretically and practically non-Abelian groups, as demonstrated in the QM9 and CrossDock examples (which include SO(3) factors).
• Despite the authors' claim that their formalism avoids projection onto the manifold, this is however implicitly still the case. The equation $dg = g_0 \xi$ ($g\in G, \xi\in \mathfrak{g}$) operates on the tangent space of the group, and its solution (neglecting the variation of $\xi$ is $g = \exp(g_0 \xi$. Thus, we need to know and apply the exponential map $\exp: TG \to G$, which is precisely the projection onto the Lie group. Thus, their method also needs the projection onto the group. In our framework, this is not the case.