Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups

We thank the reviewer for the supportive review. We are glad the reviewer can appreciate our work, and highlight TDM as well-motivated and the paper as well-structured with clearly displayed mathematics and comprehensive numerical experiments. We provide the detailed response below to your comments and questions.

Weakness 1 (Grammar & Explanations)

I feel that for the introduction, there are some mistakes in terms of English grammar Some parts could have been explained in more detail, but given the page limit restrictions it is natural that the authors aimed to be as concise as possible

We thank the reviewer for the efforts to improve the readability of this paper. We have improved the introduction section by doing grammar checks and also addressing the typos you have spotted in the revision. We also thank the reviewer for understanding that we couldn't explain everything in full detail due to page limit constraints. However, we are happy to help clarify any part of the paper here if you find that helpful.

Question 1 (Numerical Integration Errors)

p.2 l.59: the authors bring up numerical integration errors. It would be helpful to the reader to specify this a bit more. I assume this is because mappings like the exponential mapping can be expensive. Is this only the case in high dimensions or even on low-dimensional Lie groups? a bit more nuance would improve the paper a lot as it would make the actual challenge even clearer, which would highlight your contributions even more

Thanks for a great question. Without our trivialized momentum technique, the ODE to integrate is in general of the form $\dot{g} = f(g)$ for some $f$. Even if one uses the exponential map, namely $g_{k+1}=\exp_{g_k}(h f(g_k))$, it is still an approximation of the actual solution. On the other hand, our approach only needs the solution to $\dot{g}=g\xi$, where $\xi$ is a constant, and therefore the solution $g_{k+1}=g_k \text{expm}(h\xi)$ is not only staying on the manifold but also an exact update for the position variable $g$.

Question 2 (Approximation and Curvature)

p.2 l.78: the authors write ``while still encapsulating the curved geometry without any approximation.’’ There is an approximation error if there is curvature right? It would be great if the authors are a bit more nuanced here as well as it should be clear that this method might introduce new errors.

Thanks for the question. In our framework, there is no approximation error even though the Lie group is curved. This is because the curvature is already taken care of by the exponential map, which can be computed exactly by matrix exponential in our situation. We will make sure this is better clarified.

Question 3 (Wrapped Distribution)

Section 3: In the Lie algebra, the learned distribution wraps around. So ideally, the score should be periodic or should become zero if this cannot be enforced. Is there some mechanism that enforces this? If not, I would recommend adding this explicitly to let the reader know that there is possibility for error here.

In our framework, the distribution does not wrap around on Lie algebra since the Lie algebra is essentially an Euclidean space and the $\xi$ marginal of the probability distribution is simply Gaussian for any time $t$. The score function associated with the Lie algebra element $\xi$ is not periodic, and we do not need to enforce it, and therefore there is no extra source of approximation errors.

Question 4 (Sampling in ISM)

p.6 l.337: In the implicit score matching paragraph, it is not very clear to me how one would go about sampling here (does one have to do a forward simulation for this?). Would the authors clarify this?

Thanks for another great question. To evaluate the implicit score matching objective, one indeed needs to simulate the forward dynamics to generate samples approximately distributed as $p_t$. In practice, this is performed with Forward Operator-Splitting Integrator, which is explained in Algorithm 2 in Appendix D.1.

Question 5 (Relation to Flow Matching)

Section 5: It feels natural to also mention something about flow matching in the outlook. Do the authors have any ideas on this they would like to disclose in the conclusions?

Thanks for the good suggestion. We re-emphasized the comparison with RFM in the conclusion section of the revised version. As we have discussed and compared with Riemannian Flow Matching in both the literature review and experiments section, this makes the paper more coherent.

Other Questions (Typos) We have fixed grammar errors and typos and revised the paper according to your suggestions to improve paper readability and clarity. Thank you for spending the effort to help us improve the paper!

---

Hoping we clarified the misunderstandings and addressed the concerns, we thank you again for helping us improve the clarity of our writing. Please do kindly let us know if anything is still not clear.

We thank the reviewer for the comments. We are glad the reviewer can recognize our contribution, highlighting our work as clearly structured, grounded in well-explained, solid mathematics, and our approach as novel and interesting. We provide a detailed response below to your comments. If further clarifications are needed, please feel free to engage in the discussion.

Weakness 1 (Lack of Experiment details):

The main weakness of the paper is the lack of rigorous experiments to demonstrate the effectiveness of the model in comparison with existing approaches. First, the authors didn’t clarify how the experimental setup and training procedures vary across the different benchmarks. The experimental results would be more convincing if details on the setup were clearly provided. Although the results on the SO(n) and U(n) experiments are interesting and show promise, more rigorous experiments are required to show the scalability of the implicit score-matching loss (as it generally doesn’t perform as well). I think the paper could be strengthened if the authors clearly show better performance compared to RFM and RDM on the higher-dimensional SO(n) and U(n) tasks.

We thank the reviewer for bringing this point up. To provide more details on the training procedures and make the numerical results more convincing, we expand the paragraph "Architectural Framework" and "Training Hyperparameters" in Appendix G.1 and include a detailed training configuration in the revision. We include the choice of all the important training hyperparameters and a more detailed explanation of the neural architectures employed in our experiments. We will also soon make the code implementation public, which contains full hyperparameter configurations for different experiments.

To better present the advantage of our approach compared with existing approaches, we also include additional numerical results in Appendix H in the revision. We additionary compare TDM with RFM on two types of distribution on Unitary groups ranging from $\mathsf{U}(4)$ to $\mathsf{U}(8)$, namely the time evolution operator of Quantum oscillator and Random Transverse Field Ising Models. As is evident from the plots, TDM captures the complicated patterns of distribution more accurately compared with RFM. While in general RFM also well captured the global shape of the data distribution, the details near the distribution boundary tend to be noisier than the ones produced by TDM. This potentially originates from the approximations adopted by RFM, suggesting again the advantage of our proposed method.

We sincerely thank the reviewer for the supportive review, recognizing our work as novel, convincing, and well-validated with comprehensive numerical experiments. We provide a detailed response to your comments below.

Weakness 1 (Training flexibility)

The framework can only provide simulation-free training for the manifold SO(2) and its product space which is less flexible than Riemannian Flow Matching. However, this setting is still important for applications.

Thank you very much for acknowledging that $\mathsf{SO}(2)$ and its product space are still important for applications. Regarding Riemannian Flow Matching, we would like to additionally comment that while it potentially enjoys simulation-free training in many cases, it comes with the cost of the need to compute the logarithm map for defining the conditional flow. Unlike the exponential map, the logarithm map as its inverse is much more numerically unstable and harder to compute. TDM does not require the computation of the logarithm map during training, so it is better in this regard, but it is indeed not always simulation-free. Therefore, which one is theoretically better is not yet completely clear to us.

Questions (Typos)

We have fixed grammar errors and typos and revised the paper according to your suggestions to improve paper readability and clarity. Thank you for spending this effort in improving the paper!

We thank the reviewer for the comments. We are glad the reviewer can recognize our contribution to manifold diffusion models. We provide the detailed response below. If further clarifications are needed, please feel free to engage in the discussion.

Weakness 1 (Originality):

As far as I can see, the idea and the theory behind trivialized momentum have been developed in previous works, as well as the splitting discretization technique to exactly integrate the SDEs.

Previous works employed trivialized momentum for optimization and sampling, and our work is the first that uses it in generative modeling. We're unsure if the novelty of generative modeling algorithm should be criticized for gaining inspiration from existing sampling/stochastic analysis literature. For example, score-based generative model (SGM) [e.g., Yang et al. 2020] turned Ornstein–Uhlenbeck process into a generative model, Critically-damped Langevin Diffusion (CLD) [Dockhorn et al. 2021] turned kinetic langevin dynamics into a generative model. We think they had significant originality and impact, and our case is the same: TDM turned trivialized Lie-group sampling dynamics into a generative model.

Regarding splitting discretization, similarly, while this technique has been studied in the literature of numerical analysis, we are the first to use it for Lie group generative modeling. We also do not claim any credit over its convergence analysis, and present it in the paper just for self-consistency.

The authors prove a theorem for the inverse SDE of the system, but given that the non-trivial equation (about the variation of the Lie algebra) is essentially in flat space, it reduces pretty much to the flat-space formula.

Regarding Theorem 1 in the paper (which is about the reverse process of TDM), we agree that the result itself is not surprising. That is why our approach is good -- it just gives results like in the Euclidean case, which is the intention of our design. However, the proof of Theorem 1 (see Appendix B) is not straightforward due to the curved geometry. In addition, we also presented many other novel and non-trivial theoretical results such as Theorem 2 and Theorem 5.

Weakness 2 (Practical Relevance):

The main drawback is in the non-Abelian case, which is the most interesting since the Abelian case is essentially SO(2) (or copies of it). For non-Abelian groups, there's no conditional transition probability available, requiring implicit score matching. This involves computing the divergence with respect to the data space (more precisely, the Lie algebra space, but as dim data = dim $G$ = dim $\mathfrak{g}$ its computationally as expensive). This limitation significantly reduces the method's practicality and makes it unlikely to be widely adopted.

We thank the reviewer for pointing it out. We agree with the reviewer that the extension of denoising score matching (DSM) to general non-Abelian Lie groups would be exciting. However, we also note that existing Riemannian Diffusion Models, such as RSGM [De Bortoli et al. 2022] and RDM [Huang et al. 2022], also heavily rely on implicit score matching when dealing with general Lie groups. In fact, the success on simulation-free training of manifold diffusion models has been very limited in the literature, and we're unaware of any denoising score-matching result for general non-Abelian groups.