Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups
Authors: Yuchen Zhu
Tianrong Chen
Lingkai Kong
Evangelos Theodorou
Molei Tao
Venue: ICLR 2025 Poster (2025)
Keywords: non-Euclidean generative modeling
denoising diffusion
Lie group
Primary area: generative models

Abstract: The generative modeling of data on manifolds is an important task, for which diffusion models in flat spaces typically need nontrivial adaptations. This article demonstrates how a technique called `trivialization' can transfer the effectiveness of diffusion models in Euclidean spaces to Lie groups. In particular, an auxiliary momentum variable was algorithmically introduced to help transport the position variable between data distribution and a fixed, easy-to-sample distribution. Normally, this would incur further difficulty for manifold data because momentum lives in a space that changes with the position. However, our trivialization technique creates a new momentum variable that stays in a simple fixed vector space. This design, together with a manifold preserving integrator, simplifies implementation and avoids inaccuracies created by approximations such as projections to tangent space and manifold, which were typically used in prior work, hence facilitating generation with high-fidelity and efficiency. The resulting method achieves state-of-the-art performance on protein and RNA torsion angle generation and sophisticated torus datasets. We also, arguably for the first time, tackle the generation of data on high-dimensional Special Orthogonal and Unitary groups, the latter essential for quantum problems. Code is available at https://github.com/yuchen-zhu-zyc/TDM.

Published as a conference paper at ICLR 2025
TRIVIALIZED MOMENTUM FACILITATES DIFFUSION GENERATIVE MODELING ON LIE GROUPS
Yuchen Zhu'; Tianrong Chen', Lingkai Kong, Evangelos A. Theodorou, Molei Tao Georgia Institute of Technology (yzhu738, tchen 29, 1kong75, evangelos htt.sssthhhhooo mtao) htt.) @gatech.
ABSTRACT
The generative modeling of data on manifolds is an important task, for which diffusion models in flat spaces typically need nontrivial adaptations. This article demonstrates how a technique called 'trivialization' can transfer the effectiveness of diffusion models in Euclidean spaces to Lie groups. In particular, an auxiliary momentum variable was algorithmically introduced to help transport the position variable between data distribution and fixed, casy-to-sample distribution. Nor- mally, this would incur further difficulty for manifold data because momentum lives in space that changes with the position. However, our trivialization technique creates a new momentum variable that stays in simple fixed vector space. This design. together with manifold preserving integrator, simplifies implementation and avoids inaccuracies created by approximations such as projections to tangent space and manifold, which generation with high-fidelity and efficiency. The resulting method achieves state-of- the-art performance on protein and RNA torsion angle generation and sophisticated torus datasets. We also, arguably for the first time, tackle the generation of data on high- dimeesiioii Special Orthogonal and Unitary groups, the latter essential for quantum problems.
typically used in prior work, hence facilitating
Code is available at
zhu- zyc/ TDM
INTRODUCTION
Diffusion-basse (e.g.. Song al. 2020b Ho et al. 2020: Dhariwal & Nichol 2021 and flow-based (e.g., Li.mma et al. 2022: Liu et al. (2023 Albergo & Vanden- (202)) generative models have significantly impacted the landscape of various fields such computer vision, largely due their remarkable ability in modeling data that follow complicated and/or high-ddmenniionl probability distributions. However, many application domains, data explicitly reside on manifolds. Note this is different from the popular data manifold assumption which is implicit; here the manifold is priori fixed due to, e.g.. physics. Such cases occur, for example, in protein modeling (Shapovalov Dunbrack 2011 Yim et al. [223b) 2024 Bose et al.) 2024). cell development (Klimovskaia et al. 2020), geographical sciences (Thornton al. | 2022). robotics (Sola et al. 2018), and high-energ. physics (Weinberg. 1995). The naive application of standard generative models to these cases via embedding data in ambient Euclidean spaces often results suboptimal performance (De Bortoli ct al. 2022). This is partly due to the lack of appropriate geometric inductive biases and potential encounters with singularities (Brehmer & Cranmer 2020).
Pioneering works suggest generalizing (continuous) neural ODE (Chen al. 2018) to manifolds (Mathieu & Nickel 2020 Lou et al. 2020 Falorsi & Forré 2020) with maximum maximum-likelihod training. Rozen al. (2021) Ben-Hamu al. (2022) develop simulation-free algorithms but their objective unscalable or biased (Lou et al. 2023). Recent milestones, such as Riemannian Score-based Model (RSGM) De Bortoli al. (2022). Riemannian Diffusion Model (RDM) (Huang al., (022). and Riemannian Diffusion Mixture (Jo & Hwang 2023) have successfully demonstrated the potential extend diffusion models onto Riemannian manifolds. RSGM explores the effectiveness and complexity of various variants of score matching loss on general manifold and their applicable
"Equal contribution. Corresponning author.

---

Published as a conference paper at ICLR 2025
RFM/RSGM
TDM (ours)
RFM/RSGN RSOM
Ours
5, space
Score
Dff changes
Trivallzed Momantiam
speco
875
Momneniaa
Figure I: Visualization algorithmic intuition TDM. Existing approaches such as RFM and RSGM often model an object that lies on changing langent spaces as the position gt. moves, resulting in inaccuracies when handling complicated manifolds during trajectory simulations. In contrast, TDM only needs to learn the score in simple Euclidean space. Thanks to the special structure of trivialization, TDM guarantees the induced momentum will strictly lie on the tangent space, which improves generation quality and reduces sampling error.
scenarios, and RDM discusses techniques such as variance reduction for the training objective via importance sampling and likelihood estimation. Building upon RSGM and RDM, Riemannian Diffusion Mixture further leverages mixture of bridge processes to achieve significant improvement in training efficiency. These models learn to reverse the diffusion process on a manifold. This is achieved through the employment of Riemannian score- acore-matching methods, which serve simulation- based objectives for the optimization of the model. However, due to the inherent geometric complexity of the data, the training and sampling processes of such models necessitate multiple approximations. In particular, they require the projection of the vector field (i.c. score) to the tangent space which subsequently serves the training label for the neural network during the training phase. Furthermore, to mitigate numerical integration errors during the sampling process, there is a requirement for the projection of samples to the original data manifold. Moreover, among most scenarios, as these models are simulation-base based algorithms, an additional approximation is introduced during the collection of training data from the simulation during training. This process also necessitates the projection of data to the manifold, which is analogous to the sampling phase. The combination of all these three approximations can compromise the quality of generation.
Even more recent advancements, such as Riemannian Flow Matching (RFM) (Chen & Lipman 2024 and RSGMe (Lou et al. 2023). aim to alleviate training complexities and enhance model scalability through the introduction of simulation-fra focusing on Riemannian symmetric spaces, while RFM constructs conditional flows using premetrics. However, it is important to note that both of the approaches still require some of the aforementioned approximations during the training and sampling phases (such as projections onto tangent spaces), which may potentially introduce inaccuracies in the generated results. Furthermore, whether RFM is expressive enough for intricate data distribution was discussed in (Lou et al. 2023).
objectives. Scaled-RSGM achieves this by
In this work, we build upon recent progress in momentum- based optimization (Tao & Ohsawa 2020
and sampling (Kong & Tao 2024) on Lie groups to develop a highly scalable and effective generative model for data on these manifolds, named Trivialized Diffusion Model (TDM). Our approach departs from prior momentum-based generative models (Dockhorn et al. 2021 Chen et al.] [2023] Pandey & Mandi 2023 Chen et al. 2024) due to additional technique called trivialization, which utilizes the additional group structure and enables us to learn score in fixed flat space, while still encapsulating the curved geometry without any approximation.
It's also worth noting that several works have already achieved success in the generative modeling of data distribution on special Lie groups such SE(3) and SO(3), resulting in remarkable performance on applications like protein backbone generation (e.g., Yim et al. 2023b a Bose et al. 2024]. The seminal work of Diff (Yim et al. (023b) extends RSGM SE(3) and SO(3) and takes advantage of the pleasant properties of heat kernel TGSO(33 for (3) to perform denoising score matching. Due to the tractable computation of the heat kernel, Frame also enjoys the benefits of efficient learning in a fixed space well as projection-free simulation on the manifold, much similar to the advantages enjoyed by TDM. However, we need to clarify that the success of FrfmeDiff, while sharing similar spirit with TDM, comes from distinct sources. FifmeDiff leverages the special structures of 53(3) and its algorithm can't 't be easily generalized to other cases. In contrast, TDM holds the advantage of efficient learning and projection-free sampling for general Lie groups (see, e.g., Sec 3 for quantum applications of other Lie groups)
To sum up, our contributions be summarized into the following three bullet points:
2

---

Published as a conference paper at ICLR 2025
1) We introduce TDM that enables manifold data generative modeling through learning trivialized score function a fixed flat space, which dramatically improves the generative performance.
2) We leverage a nontrivial Operator Splitiing Integrator to stay exactly on the manifold in an accurate and efficient way. The reduction of approximations further improves the generation.
3) We outperform baselines by a large margin protein/RNA torsion angle datasets. We achieve much higher quality generation on a newly-introduced challenging problem called Pacman. We present the first results on generating ((z) data corresponding to quantum evolutions, and high dim (0(n) data too; these results are also appealing.
METHOD
In this section, will discuss how to perform generative modeling of data distribution on class of smooth Riemannian manifolds, namely Lie groups, by learning a score function similar to Euclidean one. Our goal is recover the scenario of Euclidean generative modeling to the maximum by leveraging the group structure of the Lie group apart from its Riemannian manifold structure. To achieve this, we explore specific manifold extension Kinetic Langevin dynamics JynamiissNNlson (1967) which contains an additional variable known momentum. Importantly, a direct introduction of the momentum would not simplify the situation, since the momentum lives in changing tangent space as the position moves. Fortunately, the group structure of the Lie group enables the design of trivialized momentum that stays in a Lie algebra for the whole time, which is a simple fixed Euclidean space that suits our needs. In the following, completely avoid challenges posed by the curved geometry in an exact, analytical fashion, without resorting to complicated differential geometry notions such parallel transport, and certainly need for approximations, projections, and retractions.
will discuss how the technique of trivialization can help
In the following, we first introduce forward process that converges with such trivialized momentum. We derive the time reversal of such a process, which can serve as backward generative process. We discuss methods to efficiently learn the drift of the backward process. Finally, we introduce a numerical integrator that achieves high accuracy and preserves the manifold structure of the Lie group.
easy-to-sample distribution
We also provide a brief review of the Euclidean Diffusion Models. kinetic Langevin dynamics, and Lie group. For details, please see Appendix A for more information.
Algorithm TDM (Trivialized Diffusion Model)
Require: Iteration Nieer Total time horizon T', Simulation steps N, time step h = T/N, parameter initialization 0o, Lie group data {g" // TRAINING
M m==1, friction constant ? 0, early-stopping time
for 0, Nisor do
2: Sample g ~ M Σm= 5gm
D Sample initial g from data
3: Sample & by i.i.d. generate Ex ~~ N(0, for s dimg
Sample arbitrary initial Use DSM if possible
4:
if JDSM is tractable then
5:
Sample ~ Uniform(e, (e T), gt. Pejoo E.oogg.Egg.. Elg.
6:
f(0,)) = (Ds {c))
P Compute Denoising score matching objective > Use ISM instead
7:
else
8:
(gt, } = FSOI (.E, 7,7, h, N) f(O,.) (Is (9t, E
> Simuate forward dynamic with Algorithm
9:
> Compute Implicit score matching objective
10:
end
11: Oa+1 optimizer (0,, ((O.))
AdamW optimizer step
12: end for
Set optimal 0- NN
// SAMPLING
Sample (go. Eo) π.
Sample initial condition from stationary measure
15: (9N, EN. BSOII So. N,,,N,
& Simulate backward dynamic with Algorithm
16: return 0*, (gN,En) En)
3

---

Published as a conference paper at ICLR 2025
2.1 TRIVIALIZED KINETIC LANGEVIN DYNAMICS ON LIE GROUP AS NOISING PROCESS
A Lie group is a manifold with a group structure, which gives us an important tool called left- trivialization handle momentum. The left group multiplication Lg G - G is defined as Lg 9 - g9. whose tangent map (also known as differential) TgL. TTGG + Tog is a one-to-one map. As result, for any 9 E G, we can represent the vectors in Tg by TeLy for any & E T.G. where e is the group identity and == T.G is the Lie algebra.
Utilizing such property, Kong & Tao (2024) appropriately added noise to variational Lie group optimization dynamics (Tao & Ohsawa 2020) and constructed the following kinetic Langevin sampling dynamics on Lie groups:
Jt = Le
(I)
de (t) g. (VUl))tt + 22y(t)dWF
where (gt. Et) E G algebra, d the Brownian motion on Lie algebra g. VÙ is the Riemannian gradient of U, and U: G - R is a potential function. & is the left-trivialized momentum at time and TeLg, S is the true momentum.
x g. Vt 2 0, here G denotes Lie group and g denotes its associated Lic
They also proved (Kong& & Tao 2024) that for connected compact Lie groups, which will be our setup, ①T converges, under Lipschitzness of VU, exponentially fast to its invariant distribution, which is
1 (g. 5) = exp ((()) - (E, ) dgdL,
(2)
where dg denotes the Haar measure, de denotes the Lebesgue measure on g. and Z is the normalizing constant. Dynamic T is generalization of the Euclidean kinetic Langevin equation on R. to general Lie groups (R* is a Lie group with vector addition being the group operation).
By Peter- Peter-Weyl Theorem, connected compact Lie group can be represented closed subgroup of Gn, C) (Knapp, 2002). i.e. the group of n x n invertible matrices with entries in C. Such representation
be computed explicitly by, e.g., adjoint representation (Hall 2013). When ge and & are both represented as matrices, the abstractly defined TeLg.Er & in (T) can be calculated explicitly by a matrix multiplication between the matrix representation of 9t and & where we write it as gr In the rest of this work, to avoid confusion and simplify notations, we use Ir and & to refer to their corresponding matrix representation and use grE, to denote Ce
We want to construct forward noising process based on I by choosing a potential U that corre- sponds an easy-to-samee sample distribution. In the case of connected compact Lie groups, we pick the natural choice, which is Ugg) = 0, Vgg corresponding to an invariant distribution whose g marginal the uniform distribution on G, or more precisely, the Haar measure. this case, the following dynamics would be the forward noising process,
5t grEt, dg. (t) 22/WWaW
Important examples of connected compact Lie Groups include but not limited to the Special Orthogonal group S0(n), ) the Unitary group U(n), the Special Unitary group SU(n), etc. Note 1-sphere S1, torus T, and 50(2) are essentially the same thing (isomorphic). Note also the direct product of any two connected compact Lie groups is still a connected compact Lie group, so in general we can consider the Lie group G of form, for GI...., G& connected compact Lie groups
G = G1 x G2 X ... X Gk.
2.2 TIME REVERSAL OF TRIVIALIZED KINETIC LANGEVIN
The following result allows us to revert the time of the forward noising process. Thanks to the introduction of momentum and the fact that it is trivialized, the time reversal will be very similar to the Euclidean version Dockhorn et al. (202T) despite that lives on a manifold. This pleasant feature is because the forward dynamics (1) has no (direct) noise g dynamics and therefore no score-based correction is needed for its reversal, and & on the other hand is simply a Euclidean variable. More precisely, one important implication of the momentum trivialization is that the only score present
4

---

Published as a conference paper at ICLR 2025
in the dynamic is Vg Pr- Et), which now stays in the Lie algebra g (a fixed space and also is isomorphic to Euclidean space). This implies that we manage to get rid of V, log pr- from the dynamic, which is a much more complicated subject than Vg log pr-+ due to being a Riemannian gradient and has complicated geometric dependency. This trivialization technique leads to benefits on numerical accuracy and score representation learning. which is not enjoyed by previous works such as RFM (Chen & Lipman (024), RDM (Huang et al. (2022) and RSGM (De Bortoli et al. 2022). More details of the advantages of the trivialized dynamic will be provided Section
Theorem 1 (Time Reversal of Trivialized Kinetic Langevin on Lie Group). Let T 2 0, W be Brownian motion on the Lie algebra g. Let X, (9t.Cr Et) be the trajectory of the forward dynamics 33. with X, admitting smooth density Pe(Oe, Ee) with respect to the Haar measure on G and Lebesgue measure on g. Then, the solution to the following SDE
{ lt = - -gtEt, dg. ((T - 1)E,dl + 2y(T - 1)Ve
(5)
+ V 27(T l)dWP.
d satisfy Y = (XT- under the notation y == (ge, Ex) and initialization Yo = XT.
Note although a similar time reversal formula has been given for the Euclidean case in (Dochhorn) al., 2021), their results are not applicable due to the presence of the manifold structure. In fact, need& a non-trivial adaptation of the arguments and the proof relies on the Fokker-Planck equation on the manifold G x g. For details of the proof of Theorem see Appendix B
In addition, similar to standard Fuclidean diffusion model, dynamic also has a corresponding probabilistic ODE counterpart, given by the following:
Remark (Probability Flow ODE). The following dynamic has the same marginal as 5
{
gr = - -ieh:
(6)
( 7(T - t) + +T Ve logpr-i(ot. dt
as long initial conditions are consistent.
2.3 LIKELIHOOD TRAINING AND SCORE-MATCHIN FOR TRIVALIZED KINETIC LANGEVIN
To perform generative modeling of data distribution, we would like to simulate and sample from the stochastic dynamic in (5). However, the score Ve log p E.) is intractable and we want to approximate it with a neural network score model sg(gt: Ee, l). We denote the sequence of probability distribution of as the density of L(Y?) with respect to the reference measure, where Y is the trajectory of the following dynamic,
gt = - (dC, = 7T + 2y(T - & + /2 /22y(T-t)dW?
9o. n.. (7)
In order to generate new data with dynamic 77.. we need of Po. which would require learning a score that is close to the true score Vg log Pt {). A natural starting point for learning the score is through Score Matching (SM) between Sg and V log Pe (g. E), but that alone is intractable because pt is not known priori. Instead, we can directly extrapolate some classical tractable variants of SM, such as Denoising Score Matching (DSM) or Implicit Score Matching (ISM), to the Lie group case.
Denoising Score Matching
Note the score matching objective can be rewritten as |Vincent 2011
(s) = Er,po Ep,1
Vg log ()) - solg.E.T -t)l² + Cr
t)]13]
JosM(()
5

---

Published as a conference paper at ICLR 2025
where C1 is constant independent of A. Hence, argming JsM argming JDSM, but evaluating JDSM only requires knowledge of the conditional transition probability Pajo- The question boils down to finding out such condition transition probability induced by the forward dynamic B
Note that the Lie algebra g is a tangent space of G at the identity, so it's a vector space that is isomorphic to Euclidean space IRd where d dim(g). For example, the so(2) is the Lie algebra of the Special Orthogonal group 52(.). 50(2) consists of all the 2 x skew-symmetric matrices. This implies that, for any ç E s2
{
0
so(2)
0
Here, the Brownian motion dW? g should be understood as dW? = E dW; c₁ where (dW?)i=i...d M=i...d independent standard Brownian motions on R and {e:).......dd . is orthogonal basis for g. Therefore, the forward dynamic (3) with initial condition g(0) = Go. (0)) = Eo is equivalent to the following,
ÿt g de: = --jde + 22dd? VI wiiidd.
s.t g(0) = 9o. (0()) c 4A visis d. (8)
Here, without loss of generality, follows is OU process with an explicit solution. This reduces problem 8 to a matrix-valued initial value problem (IVP) for gt, since & can be treated as known function of time. Then the IVP
choose (t) to be constant Y >
0. We notice that each ç
Gt SrSr, (()) = is just a linear system.
Unfortunately, note that even though the linearity ensures linear structure in the solution, namely g(t) = ()) where & is known as fundamental matrix, & in general may not be analytically available closed-form because the linear system has a time-dependent coefficient matrix. This differs from the scalar case where (() would just be exp( {())ds) or the constant coefficient matrix case where (t) would just be expm(gt). Instead, we can represent the solution using geometric tools, resulting in Magnus expansion Magnus 1954) in the following form
gt)) = go (X(m) 2(t) Σ Su(t).
Here { k=......oo is called the Magnus series, which is written in terms of integrals of iterated Lie algebra between ((l) at different times. The first three terms of the Magnus series are given below to illustrate the idea,
1 (-' /. ((ta)),(tt))dtddtt (1( f. (tt)dd1. l2(t)= - -/2/2//
Os(1) /
1
(') /- (1E(6))16(6)) + [g(ta). (E(t)). ((t)tt))tdtd)) diadt2dti 0
In general, the solution given in (9) may not be tractable due to the fact that ((t) is an infinite series with increasing intricacy for each term. However, we want to discuss special yet important case, where the infinite series is reduced to only the first term. In fact, when G is an Abeliean Lie group,
for any ç E g. the Lie bracket [S, E = 0 vanishes identically, and the solution to IVP in (9) reduces g(t = Du exp(So E(s)ds)
Theorem 2 (Conditional transition probability for Abelian Lie Group). Let G be an Abelian Lie group which is isomorphic to SO(2). In this case, the conditional transition probability can be written explicitly
Pt a9 & go. {o) = WN(logm(g0 9:): Mg. 0g) (Ex::
(10)

---

Published as a conference paper at ICLR 2025
where WN(r; H. o²) is the density value of the Wrapped Normal distribution with mean p and variance o², and evaluated at z. For explicit expressions of Hg. og and pę, of and the multivariate case formula. please see Appendix
An important example of Abelian Lie group is the 1D torus T or special orthogonal group (2),),
and any of their direct product. In this case, we can compute the conditional transition probability Pejo(ee: {x) exactly due to the capability of solving the IVP exactly. We summarize the results in Theorem and leave the proof in Appendix[[] Notice that while Theorem[ooly gives the conditional transition probability for G = 50(2), it can be extended to multivariate case where G = (0( since the conditional transition of forward dynamic factorizes over each dimension when conditioned on the initial value. For detailed discussion of the multivariate case, please Appendix C
Implicit Score Matching When G is not Abelian, the conditional transition probability in (T) might not be available. In this case, we resort to another computationally tractable variant of the score- matching loss derived by performing integration by parts, also known as the implicit score-matching objective JISM (Hyvärinen 22055. In fact, we can connect JISM and JSM by.
(O) = Et..e [|00(...00)] 2dvv 2div + C2
(o)
where C2 is a constant independent of 0. Hence, argming JSM = argmin, JDSM- To evaluate JIMM samples approximated distributed as pt are generated through simulation of forward dynamic. Computing also requires evaluating the divergence with respect to &. which the trace the Jaco- bian. For high dimensional problems. stochastic approximations with Hutchinson's 's trace estimator (Hutchinson 1989; Song et al., 2020a) are often employed to improve computational efficiency.
2.4 NUMERICAL INTEGRATION AND SCORE PARAMETERIZATION
To either simulate the forward dynamic for generating trajectories used for evaluating implicit score matching objective JISM or sampling from the backward dynamic for generating new samples, need to integrate the dynamic. To exploit the Euclidean structure of to achieve higher numerical accuracy, we introduce the Operator Splitting Integrator (OSI). Apart from enjoying a better prefactor in terms of numerical errors, such an integrator is also manifold-preservve free. Details of the integrator can be found in Appendix D along with a convergence analysis of OSI (Appendix D.3) which is essentially reproduction of the proof in Kong & Tao (2024).
and projection-
Integrating forward dynamic In order to numerically integrate the forward dynamic we note that the dynamic can be split into the sum of two much simpler dynamics depicted in [1] This is the approach considered by Kong & Tao (202)).
1
{ 9: g dp, d
+
AZ
Jt =0 de. = - + /2ydW?
(11)
While the original forward dynamic does not in general have a simple, closed-form solution for non-Abelian groups, the two smaller systems A and A are essentially linear and both allow exact integration with flosdd-fomm solutions. Therefore, instead of directly integrating the forward dynamic, we can integrate Ag" and A alternatively for each timestep. Another notable property of such integration is that the trajectory of this numerical integration scheme will stay exactly on the manifold G x g. This avoids the use of projection operators the end of each timestep to ensure the iterates stay on the manifold. By performing such a manifold-pressrvvv get rid of the inaccuracy caused by projections but also greatly reduce the implementation difficulties since such projections in general do not admit closed-form formulas.
integration technique, we not only
Integrating backward dynamic dynamic, either directly work with the stochastic backward dynamic in (7) or its corresponding narginally-eeuivaabee probability flow ODE. We discuss mainly the integrators for the stochastic dynamic and defer the discussion of probability flow ODE to Appendix [ Employing similar
To perform generative modeling and sample from the backward
7

---

Published as a conference paper at ICLR 2025
operator splitting scheme, dynamic 7 can be split into the following two simpler dynamics, 9: - - de. = Off
1g {{
+ A Jt de = df + 27 So t)dt + /2ndw?
(12)
While A still allows exact integration and helps preserve the trajectory the Lie group, A no longer has a closed form solution due to the nonlinearity in sg. In this case, we still use exponential integrators conduct the exact integration of the linear component and discretize the nonlinear component by using a left-point rule, i.e. pretending that and { do not change over a short time h.
This treatment is beyond the consideration by Kong & Tao (0024) but it is rather natural extension.
Score parameterization Previous works manifold generative modeling like RFM (Chen & Lipman 2024). RDM (Huang et al., 2022), and RSGM (De Bortoli et al. 2022) often require learning score that belongs to the tangent space at the input, Le.. so(g.t) E T G. This means that the score network at each input needs to adapt individually to the geometric structure at that point. One thus needs to either write explicitly the g depeeddednt isomorphism between G and RRd for each g. or embed T, G in the Euclidean space IR" with n > d and apply projections onto T TGG to obtain valid score. Either way, one needs to handle the geometry of G and/or deal with additional approximation errors and computational costs (e.g., incurred by projections), and learn a hard object in a changing space with structural constraints.
On the other hand, since our approach only needs to approximate the score Ve log Pt. which is an element in the Lie algebra g. we can use a standard Euclidean-value neural network to universally approximate 8g. Thanks to the technique of trivialization, we can enjoy the already demonstrated success of score learning in a fixed Fuclidean space, where the non-Fuclidean effects stemming from the Riemannian geometry The need to parameterize the score function in geometry soometry-dependen space by-passed, without any approximation in this step. The hardwiring of the geometric structural constraints into the dynamics greatly reduces the implementation difficulty, improves the efficiency of score representation learning, releases the flexibility to choose score parameterization users, and potentially makes the generative model more data efficient as there is no more need to learn the geometry from data.
extracted and represented through the left- ceft-multiplidd g position variable.
3 EXPERIMENTAL RESULTS
We will demonstrate accurate generative modeling of Lie group data corresponding to 1) complicated and/or high-dim distribution on torus, 2) protein and RNA structures, 3) sophisticated synthetic datasets on possibly high-dim Special Orthogonal Group, and 4) ensemble of quantum systems, such as quantum oscillator with a random potential or Random Transverse Field Ising Model (RTFIM), characterized by their time-evolution operators. Details of the dataset and training set-up are discussed in Appendix G
Evaluation Methodology: We adhere to the standard evaluation criterion in Riemannian generative modeling, which is Negative Log Likelihood (NLL). A consistent number of function evaluations is maintained as per prior studies. All datasets were meticulously partitioned into training and testing sets using a 9:1 ratio. Details of NI.L estimation procedure are in Appendix [F] note that result not new and only for completeness, but our proof particularly adapted Lie group manifolds, intrinsic and independent of the choice of charts and coordinates.
Complicated and/or High Dimensional Torus Data: We start by
Log-Ilkellhos |N
0.25
comparing our model performance with RFM Chen & Lipman (2024)
0.20
on intricate datasets such as the checkerboard and Pacman on T which
0.15
are discontinuous and multi-modal. Here, Pacman is a dataset newly curated by us to test generation on torus in challenging situations. was noted (Lou et al.) 2023 Fig.3) that RFM produces less satisfactory results when generating complicated patterns on torus, such as the checkerboard with a size larger than x 4. We observed that RFM, although very strong method, ran into a similar issue when generating Pacman, which is arguably more sophisticated. Figure and Figure 4 show that our model consistently exhibited proficiency in generating intricate patterns within the torus manifold. A scalability study shown in Figure confirmed method's good scalability to
0.30
0.05
Yowkkoed D Biliston Mode Remamian Flo Mauching
0.20
semmmmnn Scare beaed
19a N (Dimension)
10
Figure 2: Log likelihood
(t) v.s. Dimensions.
8

---

Published as a conference paper at ICLR 2025
(a) Paeman Maze
(b) TDM
(c) REM
Figure 4: Visualization of Pacman dataset and generated data by TDM on T Pacman maze corresponds to random variable T with complicated distribution corresponding to locations where there is wall.
high-dimensional cases with minimal degradation in performance (NLL). For the scalability study, we adopted the same setting considered in RFM and compared with its results.
Protein/RNA Torison Angles on Torus: We also test on the popular protein Lovell et al. (2003) and RNA MMrray al. (2003) datasets compiled by Huang et al. (2022). These datasets correspond configurations of macro-molecules represented by torsion angles (hence non non-Euclidean), which are 2D or 7D. Results. including generated data of the protein datasets, are presented Table T and Figure 7 The results of RFM were taken from (Chen & Lipman 2024), where RDM was compared to and results of RSGM were not provided. Notably, our model outperforms the baselines by a significant margin, as evidenced by the visualizations of RNA illustrating the alignment of generated data with ground truth via density plots. The empirical results demonstrating our model's 's substantial performance gains are possibly rooted in the proposed simulation-fres training, high-accuracy sampling, and reduced number of approximations.
Figure 3: Visualization of Gen- erated data by TDM on 4 4 and 6 checkerboard.
Table 1: Test NLL (4) over Protein/RNA datasets
Model
General (2D)
Glycine (2D)
Proline (2D)
Pre-Pro (2D)
RNA (7D)
Dataset size
138208
13283
7634
6910
9478
RDM
1.04 + 0.012 1.01 + 0.025
1.97 + 0.012
0.12 0.011
1.24 + 0.004
-3.70 0.592
RFM
1.90 + 0.055
0.15 + 0.027
1.18 4 0.055
55.00 4 0.067
TDM
0.69 + 0.14
1.04 :: 0.27
-0.60 0.15
0.52 + 0.10
-6.86 + 0.46
Generatan
Model
Log likelihood
1.0
0.50
Gmund
RSGM
0.20 + 0.03
6.5
0.25
TDM
0.292 + 0.07
Figure 5: Log likelihood and visualization of generated data for SO(3) with 32 mixture components.
Special Orthogonal Group in High Dimensions: We now evaluate our model's performance on (O(n) data. Notably, our model the first reported one to successfully generate beyond 3. For 5O(3), we generate a difficult mixture distribution in the same way as in (De Bortoli et al. 2022). We also generate data for (O(n) with n > in a similar fashion. With the trivialization technique, bypass the need to compute the Riemannian logarithm map used in RFM training or eigenfunctions of the heat kernel on SO(n), which is needed by RSGM (De Bortoli al. 2022 but in general does not admit a tractable form. The accuracy of our approach can be seen from both the visualization and NLL metric in Figure 5 and Figure 6
Lastly, we experiment with complex-valued Lie group, the unitary group U(n). U(n) holds critical importance in, e.g., high energy physics (Weinberg) 1995) and quantum sciences (Nielsen & Chuang 2010). Our approach, arguably for the first time, tackles the generative modeling of U(z) data and manages to scale to nontrivial dimensions. Specifically, how quantum system evolves is encoded by unitary operator, i.e. an element in U(n), and we consider training data corresponding to the time-evolution operators of an ensemble of quantum systems, and aim generating more quantum systems that similar to the training data.
Learning an Ensemble of Quantum Processes:
9

---

Published as a conference paper at ICLR 2025
True dist.
Generated dist
True dist.
Generated dist.
True dia..
Generated dis..
05
-10
-10
Figure 6: Visualization of Generated (n) data. We scatter plot the marginals of randomly selected dimensions.
Left: 50(4). Middle: 50(6). Right: 50(8)
Glycine
100
100 100
-200
1.00
-200
100
100
1.00
Generated dist.
True
Figure 7: Visualization Generated Protein Torsion Angle. & and & are torsion angles on the protein backbone. We scatter plot the ground truth distribution and overlay it with generated torsion angles. The high overlap suggests good match between the generated distribution and true distribution.
Two examples are tested, respectively quantum oscillators in random potentials, and Transverse Field Ising Model with random couplings and field strength. (Spatial discretization, if needed, of) the time evolution operator of Schrödinger equation for each system gives one U(n) data point in the training set. Fig provides marginals' scatter plots showcase the fidelity of our generated distributions, for Quantum Oscillators. Fig 9 is for Transverse Field Ising Model.
635
633
0050
Figure 8: Visualization of Generated Time-evolution Operator Quantum Oscillator on U(n). Tvolutvol
operators are of the form entR where H. A. Vv, An is the discretized Laplacian and V (x) fw'l | - is random potential function. We scatter plot the marginals of randomly selected dimensions. Left: U(4). Middle: U(6). Right: U(8)
6:15
635
40
-000
08
Figure 9: Visualization Generated Time-evolution Operator of Transverse Field Ising Model (TFIM) on U(r). Time-evolution operators are of the form e d# where H. is the Hamiltonian of the TFIM with random coupling parameter and field strength. We scatter plot the marginals of randomly selected dimensions. Left: 2-qubit TFIM, U(4). Middle: 3-qubit TFIM, U(8). Right: 3-qubit TFIM, U(8)
CONCLUSION, LIMITATION, AND FUTURE POSSIBILITIES
We propose TDM, an approximation free diffusion model on Lie groups, by algorithmically introduc- ing trivialized momentum, which turns a manifold problem to Euclidean-like situation. Compared to existing milestones such as RFM, RSGM and RDM, TDM achieves superior performance on various benchmark datasets thanks to having fewer sources of approximations. However, trivialization strongly leverages group structure and does not directly generalize to general manifolds, although it is possible to extend to homogeneous spaces. Meanwhile, there is potential for improvement in the score parameterization and training, such as by adopting techniques such as preconditioning and exponential moving average tuning (Karras et al. 2022). which will be investigated in the future.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
We thank Michael Brenner, Valentin De Bortoli, and Yuehaw Khoo for encouragements and inspira- tions. YZ, LK, and MT are grateful for partial supports by NSF Grant DMS-1847802, Cullen-Peck Scholarship, Emory-GT Hummanity Award, and Math Works Microgrant.

---

Published as a conference paper at ICLR 2025
REFERENCES
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. ICLR, 2023.
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-326, 1982.
Heli Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Aditya Grover, Maximilian Nickel, Ricky TQ Chen, and Yaron Lipman. Matching normalizing flows and probability paths on manifolds. arXiv preprint TrXiv:2207.04711,, 2022.
Avishek Joey Bose, Tara Akhound Akhound-Sadegh. Kilian Fatras, Guillaume Huguet, Jarrid Rector- Rector-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Stoo)sstochastic flow matching for protein backbone generation. ICLR, 2024.
Johann Brehmer and Kyle Cranmer. Flows for simultancous manifold learning and density estimation. Advances in Neural Information Processing Systems, 33:442 453, 2020.
Ricky TQ Chen and Yaron Lipman. Flow matching on general geometries. ICIR, 2024.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianrong Chen, Jiatao Gu, Laurent Dinh, Evangelos A Theodorou, Josh Susskind, and Shuangfei Zhai. Generative modeling with phase stochastic bridges. arXiv preprint arXiv:2310.07805 2023.
Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos Theodorou. Deep momentum multi- marginal schrödinger bridge. Advances in Neural Information Processing Systems, 36, 2024.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Tch, and Arnaud Doucet. Riemannian score-based generative modelling. Advances Neural Information Processing Systems, 14:2406-2422 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8799 8794, 2021.
Manfredo Perdigao Do Carmo and Flaherty Francis. Riemannian geometry, volume 2. Springer, 1992.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically- damped langevin diffusion. arXiv preprint arXiv:2112.07068 2021.
Albert Einstein. Über die von der molckularkinettschen theorie der wärme geforderte bewegung von in ruhenden flüssigkeiten suspendierten teilchen. Annalen der physik, 4, 1905.
Luca Falorsi and Patrick Forré. Neural ordinary differential equations on manifolds. arXiv preprint arXiv:2006.0666 06663, 2020.
Luca Falorsi, Pim de Haan, Tim R Davidson, and Patrick Forré. Reparameterizing distributions on lie groups. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 3244-3253. PMLR, 2019.
Brian Hall. Lie groups, Lie algebras, and representations. Springer, 2013.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, (166800-65111 2020.
Chin-Wes Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Riemannian diffusion models. Advances Neural Information Processing Systems, 35:2750-2761, 2022.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications Pntiistics-Simuaatiionn and Computation, 1853):1059-1076 1989.
12

---

Published as a conference paper at ICLR 2025
Aapo Hyvärinen. Estimation of non-normllized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Jaehyeong Jo and Sung Ju Hwang. Generative modeling on manifolds through mixture of riemannian diffusion processes. arXiv preprint arXiv:2310.07216, 2023.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 35:26565-26577,, 2022.
Alexander A Kirillov. An introduction to Lie groups and Lie algebras, volume 113. Cambridge University Press, 2008.
Anna Klimovskaia, David Lopez-Paz, Léon Bottou, and Maximilian Nickel. Poincaré maps for analyzing complex hicrarchies in single-cell data. Nature communications, 11(1):2966, 2020.
Anthony W Knapp. Lie groups beyond an introduction. Springer, 2002.
Lingkai Kong and Molei Tao. Convergence of kinetic langevin monte carlo on lie groups. Conference Learning Theory, 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. ICLR, 2023.
Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser Nam Lim, and Christopher M De Sa. Neural manifold ordinary differential equations. Advances in Neural Information Processing Systems, 15117548-17566 2020.
Aaron Lou, Minkai Xu, Adam Farris, and Stefano Ermon. Scaling riemannian diffusion models. Advances in Neural Information Processing Systems, 36:80291-8000 -80305, 2023.
Simon C Lovell, Ian W Davis, W Bryan Arendall III, Paul IW De Bakker, J Michael Word, Michael G Prisant, Jane S Richardson, and David C Richardson. Structure validation by ca geometry: o, v and cB deviation. Proteins: Structure, Function, and Bioinformatics, 50(3):437-450, 2003.
Wilhelm Magnus. On the exponential solution of differential equations for linear operator. Commu- nications on pure and applied mathematics, 7(4):649-677 1954.
Kanti V Mardia and Peter E Jupp. Directional statistics. John Wiley & Sons, 2009.
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in
Neural Information Processing Systems, 33:2503-2515, 2020.
Laura JW Murray, W Bryan Arendall III, David Richardson, and Jane Richardson. Rna backbone is rotameric. Proceedings of the National Academy of Sciences, 100124):13044113000 2003,
Edward Nelson. Dynamical theories of Brownian motion. Princeton University Press, 1967.
Michael Nielsen and Isaac Chuang. Quantum computation and quantum information. Cambridge
university press, 2010.
Kushagra Pandey and Stephan Mandt. A complete recipe for diffusion generative models. In Proceedings of the IEEF/CVF International Conference on Computer Vision pp. 4261 42722, 2023.
Brendan Leigh Ross, Gabriel Loaiza-Ganem, Anthony I. Caterini, and Jesse C Cresswell. Neural implicit manifold learning for topology-aware density estimation. Transactions on Machine Learning Research, 2023.
Noam Rozen, Aditya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence- -bssed generative modeling on manifolds. Advances in Neural Information Processing Systems, 34: 17669-17680, 2021.
13

---

Published as a conference paper at ICLR 2025
Maxim V Shapovalov and Roland L Dunbrack. A smoothed buckbone-dependeer rotamer library for proteins derived from adaptive kernel density estimates and regressions. Structure, 19(6):844-85) 858, 2011.
Joan Sola, Jeremie Deray, and Dinesh Atchuthan. A micro lie theory for state estimation in robotics. arXiv preprini arXiv: 1812001577 2018.
Yang Song. Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. Uncertainty Artificial Intelligence, pp. 574-584. PMLR, 2020a.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13453 13446, 2020b.
Molei Tao and Tomoki Ohsawa. Variational optimization on lie groups, with examples of lead- ing (generalized) eigenvalue problems. International Conference on Artificial Intelligence and Statistics, 2020.
James Thornton, Michael Hutchinson, Emile Mathieu, Valentin De Bortoli, Yee Whye Teh, and Arnaud Doucet. Riemannian diffusion schr)." odinger bridge. arXiv preprint arXiv:2207.03024, 2022.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compula- fion, 23(7):1661-1674, 2011.
Steven Weinberg. The quantum theory of fields, volume 2. Cambridge university press, 1995.
Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Lumm Sarah Lewis, Victor Garcia Satorras, Bastiaan S Veeling, Regina Barzilay, Tommi Jaakkola, et al. Fast protein backbone generation with SE(3) flow matching. arXiv preprint 1rXiv:2310.0529,, 2023a.
Jason Yim, Brian I. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. SE(3) diffusion model with application to protein backbone generation. ICML, 2023b.
Jason Yim, Andrew Campbell, Emile Mathicu, Andrew YK Foong, Michael Gastegger, José Jiménez- Luna, Sarah Lewis, Victor Gareia Satorras, Bastiaan S Veeling, Frank Noé, Regina Barzilay, and Tommi S Jaakkola. Improved motif soaif--caffolding with SE(3) flow matching. TMLR, 2024.
14

---

Published as a conference paper at ICLR 2025
A BACKGROUNDS
Diffusion Generative Model in Euclidean spaces: We first review Diffusion Generative Models (sometimes also referred to as Score-based Generative Model, denoising diffusion model, etc.) (Ho et al. 2020: Song et al. 2020b). Here, we adopt the Stochastic Differential Equations (SDE) description (Song et al. 2020b). Given samples of R random variable Xo that follows the data distribution Po which are interested in, denoising diffusion adopts forward noising process followed by a backward denoising generation process to generate more samples of po- The forward process transports the data distribution to known, easy-to-sample distribution by evolving the initial condition via an SDE,
dx, (Xx,, tjdt V27())W,, Xo ~P-
(13)
In this case, oo will be standard Gaussian AO, 1) with appropriate choice of 7()). The backward process then utilizes the time-reversal of the SDE 13 (Anderson 1982). More precisely, if one considers
dY, - =
logr + (T: Yo~ pr.
(14)
Then we have Y₁ ~ Pr- i.e. Y; Xx in distribution. In particular, the T-time evolution of (14. Yr, will follow the data distribution po- In practice, one considers evolving the forward dynamics for finite but large time T, so that pr A3 (O, I), and then initialize the backward dynamics using Yo (N I) and simulate it numerically till t T obtain approximate samples of the data distribution. Critically, the score function s needs to be estimated in the forward process.
To do so, the score V log Pt is often approximated using neural network sg. For linear forward SDE, it is typically trained by minimizing an objective based on denoising score matching (Vincent 2011. namely
Xx E. (X,,, () V log Xx)
(15)
where V px (X, is the conditional score derived from the solution of (3) with a given initial condition.
Kinetic Langevin dynamics in Fuclidean spaces, and CLD: When Einstein first proposed 'Brown- ian motion', he actually thought of mechanical system under additional perturbations from noise and friction (Einstein 1905). This is now (generalized, formalized, and) known as the kinetic Langevin
dynamics (e.g.. Nelson 1967 i.c.
dQ = =M-1 dP
(16)
adW,
which converges, as l. co, to limiting probability distribution 7-1 eexp((((" V(Q))/T)dPáQ under mild conditions, where M is mass matrix that can be assumed to be without loss of generality, and 7'= = o² /2/(2)) is known as the temperature. If I' is fixed and y oc. one recovers (in distribution and after time rescaling) overdamped Langevin dynamics
dQ vVV(dd + V2ddW.
Just like how overdamped Langevin (often with V being quadratic) can be used as the forward process for diffusion generative model, kinetic Langevin can also be used as the forward process. In fact, in a seminal paper, Dockhorn et al. (Dockhorn et al. 2021) used to smartly bypass the singularity of score function at = 0 when overdamped Langevin is employed as the forward process and data is supported on low dimensional manifolds and called the resulting method CLD. Similar to Equation 14 one can construct the reverse process of Equation 16 follow,
[QQ = --M-iPdt dP {dp
(17)
= = (Q))) + o2V loq P,t)dt + adW,
By endowing P with Gaussian initial condition, p is fully supported in P space, and since the score function only takes the gradient with respect to P. it no longer has the aforementioned singularity issue when / tends to zero. This benefits score parameterization and learning.
15

---

Published as a conference paper at ICLR 2025
Finally, let us contrast the trivialized Lie group kinetic Langevin dynamics (used forward dynamics in this work) with the classical Euclidean kinetic Langevin dynamics by setting o = v and M 1:
[ g..
{Q= P. ( - VU()dd /2,dW,
Lie:
Euclidean:
[g ~ - T-g--
22
Note the main difference the line, i.e. the position dynamics; the 2nd line identical except for conservative forcing, but that has to be different in the manifold case.
Lie group:
group is differentiable manifold that also has group structure, denoted by G. A Lie algebra is a vector space with bilincar, alternating binary operation that satisfies the Jacobi identity, known as the Lie bracket. The tangent space of Lie group at e (the identity element the group) is Lie algebra, denoted as g T.G.
B TIME REVERSAL FORMULA
In this section, we will prove the time reversal formula stated in Theorem We first introduce an important lemma and calculate the adjoint of the infinitesimal generator of general diffusion process. We then apply the lemma to derive the Bokker-Plancee equation for our process interest and finish the proof.
Lemma Given a, > g, let L denote the infinitesimal generator of the following dynamic
و9 T. Lg dt
(18)
de (gg, E) (2)
The adjoint of C is given by
L'p - div div, (pt, dive(p8) + A())apP
Proof of Lemma We first write down the infinitesimal generator L for SDE 18 For any f (G x g). L defined as
[f (g), Eo) = (g+ -(6,)))-((993 {(
Lf(g. == lim
6-0
=
+ (Vef.)) + +tt)Aeff
By definition, L- C +
C2 (the adjoint operator of L) satisties Jaxs ffddg
Saxe ff'ppdddt for any f.pe (G x g). By the divergence theorem, we have
|Cn pCf dgdg= = -f Saxe f( (IV.((IT dive + (t) dgdg
Here T. T stands for the left-invariant vector filed on G generated by & E g. As result, we have
L'p div (VTeLga) dive (pB) + 7())
We are now ready to show that the backward dynamic 5 is the time reversal process the forward dynamic . Fokker-Planck characterizes the evolution of the density of stochastic process: denote the density at time as Pt. we have P1 satisfies 2 ät Ps = Ltp.. Thm. is proved by comparing the
Fokker- Planck equation for Eq. and (5)
Proof of Theorem 7 By denoting the density for SDE following the forward dynamic |I as pt and the density for SDE following backward dynamic (5) as Pt, only need to prove pt Pr-t-
16

---

Published as a conference paper at ICLR 2025
Using Lemma T the Fokker-Plann equation of the forward dynamic Eq. ① is given by,
a P: = - div, (0
and the Fokker-Planck equation for Eq.
(5 is given by
① tr div, divg(- - 27 dive (ee Ve log 7(T didvgg( + 1)
div divg((:IT.L.5) - 7(T div (PrC) 7(T +)
where the last equation holds due divg (f. V. log dive (Veūt) spp- We can calculate the partial derivative the reversed distribution pr-t with respect to i, which gives
0 at PT = - div, div + (() dive ~())epp
Note that this exactly matches the expression for Po = Pr, we deduce that pt = Pr- for all t.
Ø Pr. Together with the same initial condition
C DENOISING SCORE MATCHING FOR ABELIAN LIE GROUP
We first state a detailed version of Theorem 2 in the following,
Corollary 1 (Conditional transition probability for Abelian Lie Group). Let G be an Abelian Lie group which is isomorphic to T SO(2). In this case, the conditional transition probability can be written explicitly as,
Pr. oog,, & go. Eo) = WN([ogm(go 9:): Mg. og) (Eet pç. (g:
(19)
where WN(:; 0²) is the density of the Wrapped Normal distribution with mean a and variance o² evaluate at x, logm is the matrix logarithm with principal root, and pgs Mę, og. of are given by.
1 1- 1
Ng lte
(E: + Eu), Mg E cc
8
0g 2t
+1+1 4, o = 1 21
In this section, we will prove Theorem by proving its detailed version in Corollary [ under the condition that is an Abelian Lie group which is isomorphic to T. or SO Note that this allows us to compute the conditional transition probability for G that is also direct product of these Lie groups. The reason that, for Lie group G that is direct product
and SO(2). we can represent an element in G as (g) ...... gff) where g¹ T or S0(2). The corresponding Lie Algebra can be represented as (..,,.5 For each I < k. (g',e') they follow the following dynamic,
ot = g:
dgt = -tidt + /27dWi (0( gà. (²))) = Ei
(20)
Note that this will not create any confusion since ç² E IR and there's no need for another superscript indicate other elements in E.. Moreover, important consequence of the factorization of the dynamic of (g.E) as k independent smaller dynamic is that we can also factorize the conditional transition probability for g. E as product of & conditional transition probability, each computed from g", E". This means the following.
Peplae &el 90. Eu) = II peolol & (6)
(21)
17

---

Published as a conference paper at ICLR 2025
Based (11) manage to compute the conditional transition probability of any general connected compact Abelian Lie group, since they are necessarily isomorphic to a power of or (2 X 2008). Therefore, just need to compute the conditional transition probability for such base case, which is stated Theorem
Kirillov
From now on, will consider y 1 for simplicity. Generalization of our results time-dependent straightforward. Let's stick to the notation that 90 E S0()), So E R = 5o(2). We slightly abuse the notation in the sense that, when considering the dynamic for &, we considering valid SDE R, while when we are considering the dynamic for g. & should be understood as its matrix representation
in s2(,), which is a 2 skew-symmetric matrix. Let also denote Y, = Jo E,ds for notational simplicity.
Since G is Abelian, 1, = for any č, E g. the Magnus series (t) = in for t: 2 and the solution to the IVP can be written explicitly as,
(Y;) E e e v2 22 (t
(22)
Notice that (gt: {t) is push forward of (&,, Y.). Therefore, to find the condition transition Pejo & 9o: Eo). we first compute the joint distribution of (E.. Yi) conditioned on the (go, Eo), and derive the desired conditional transition probability by computing the probability change of variable.
Since Y1 is the time integral of Et, (Es, Yi) is Gaussian process, with mean and covariance stated in the following Lemma,
Lemma 2. For given t, (Ex, Yi) is distributed according to a bivariate Gaussian,
v. ~ N
(11==")
e e
)
1 -e---- (e
E (22 1)2 (c 1)2 4e- 2 + 2t - ))
(23)
Proof of Lemma 2 To show that the joint distribution (Ec, Yi) as the desired expression, we just need to compute the mean and variance of & and Y, respectively, as well as their covariance. For &
EiGd] B V2 (' e (1(+0)) ccu
Var(E:) Var (V2 (" e-t--) = I' ds -2t
For Y,, since it's the integration of Er, it has the following expression,
Y, / " V2 Jo
/' S" (p--)) c-(o-))
(1 " + v2 f I.e
, p
c - dp
= (1 - - D + v2 / (1-
where we use Stochastic Fubini's theorem exchange the integration order dW, and dp. Therefore. we can compute the mean and variance of Y,.
B|y 1E(1 e e + V2 ( (1 (1(--e- -e- -a-- =(1 - - e))
v2f.
Var(Y:) Var (v2/ (1 (1-e- -(--=)))d = f/ (1-e -(e-s) 4e -t_ 2e + -6-
18

---

Published as a conference paper at ICLR 2025
Finally, need to compute Cov(EL, Yi) to complete the proof, where use Ito's isometry,
Cov(Er, Y,) E o
e-tt=
= 2 / 6 c-((-) (1 - e dl- = e 2 1)2
As corollary of Lemma 2 we can compute the conditional distribution Yils,, here we omit the dependence on Lu. So for simplicity since all probability considered in this section is conditioned on these two value. The conditional distribution between bivariate Gaussian equivalent orthogonal projections,
Corollary 2. For a give I, YelE, has distribution
e- Yik. N ( ( (Er 5o),
8
-4)
(24)
et +11
Proof. Let E and f denotes the variance matrix and the mean vector of V.). The Yelee has conditional mean and variance given by,
B[Y/Ed My Eyg E (5: me)
Var Var(V: Eyy Σy Dg Egy
Plug in the expression for E and and the expressions simplify to the desired ones.
We need the distribution of Yrle, due to the following factorization of pejo & (9005o).
(ejo & go. to) Pa E., 90 Eo) Pejo(E: go. Eo)
Here (Ee 9o. Eo) is known due & being Gaussian, which is hard object since it's a distribution on the Lie group G. However, we can derive its expression by computing the push-forward of Yele,, 9o, Eo by the exponential map. The following theorem characterizes such change of measure given by the exponential map of Lie group as the push-forward,
need to compute (r Et, Jo. Eo),
Theorem (Theorem 3.1 in Falosi et al. (Falorsi et al., 2019)). Let G denotes Lie group and its Lie algebra. Consider distribution m g with density r(E) with respect the Lebesgue measure on the forh-forward of m to G, denoted as cxp.(m) absolutely continuous with respect to the Haar measure G, with density p(g) given by.
p(g) =
Σ
((E)) gEG,
expm
where J(E) det E
(-1)4 (-1) (ade)+).
k=0 (K+TT]
Moreover, when G is 50(2) or T, the scenario simplifies to,
J(E) = 1, Vg E {g g expm(f) g} 15 E g: logm(g) + 2km 1: Z)
Since the exponential map is not injective, computing the density the push- foss-forwdrded measure at g G requires summing density all the pre-images of in g and weighted by the inverse of the Jacobian [. Fortunately, for our considered case, both the pre-images and the Jacobian can be explicitly characterized and computed.
19

---

Published as a conference paper at ICLR 2025
Applying Theorem our case, where is the density of Yilke, Su. 9u. expression for Pt (99 Ee 9o. Eo),
derive the following
Pajo(ge Ee, 90. Eo)
Σ k=-00
Pr.ce (I 9t) + 2km))
(25)
where py,le denotes the density of Yelee, which the density of Gaussian with mean and variance defined in (24. This is also known as the Wrapped Normal distribution (Mardia & Jupp 2009). its name comes from the fact that the density is generated by "wrapping a distribution" on circle. Therefore, denote such a distribution with notation WN(x; μ. o²) denotes such as density evalu- ated at point I, where μr, o² the mean and variance of the normal distribution being wrapped. This finishes the proof of Theorem 2 and Corollary
D PROABILITY FLOW ODE AND OPERATOR SPLITTING INTEGRATOR
In this section, we will first discuss in detail the Operator Splitting Integrator (OSI) and how they help integrate the forward and backward trivialized kinetic dynamics accurately in manifold- preserving, projection-free manner. We then introduce the probability flow ODE of the backward dynamic, which is ODE that is marginally equivalent to dynamic OSI for the probability flow ODE.
We will also introduce the
D.1 OPERATOR SPLITTING INTEGRATOR
In this section, will demonstrate how OSIs are constructed from the dynamies and the benefits they enjoy. We restrict our attention to first-order numerical integrators in the following discussion. However, such an approach can be generalized and we can indeed craft an OSI with arbitrary order of accuracy by following the approach in Tao and Ohsawa (Tao & Ohsawa 2020).
Forward Integrator:
Recall that the forward dynamic can be split into two smaller dynamies,
Ag
{
9: = 9.C d5. - d
A
=
dp - 22
While the original dynamic does not admit a simple closed-form solution, A and Ag can be solved explicitly is shown in the following equations,
A (() g(0) (xpm(t((0)), {(t) = E(O)
A g(t) g(0), exp(->t){(0) + '''
(0ll->/tts)ddW- s))dW?
Therefore, we can integrate A and A alternatively for each timestep h in order to integrate the original forward dynamic. The detailed algorithm can be found in Algorithm 2
Backward Integrator: Recall that the backward dynamic
122 can be split into two smaller dynamics,
A
ÿr - + A dg, = dd
9e
de, yE + 2 (gr, & t)dt W
Again, write out explicit solutions to A and Al in the following equations,
45 s() .(0) ())(0)))))))))). {(() E(()
-)) s))dW? + 27 Ag: g(t) = -x.p((t))xxp())(((0)) + I'' 2xp -
J "xp()(tt - s))so((a Ex, s)ds
Note that, the solution to A though presented in an explicit form, can't be implemented in practice since we can not integrate exactly the neural network sg. We employ an approximation here and discretize the nonlinear component sæ by using left-point rule, i.c. pretending that and do not change over a short time h, and still use the exponential integration technique to conduct the exact integration of the rest of the linear dynamic. The detailed algorithm be found in Algorithm 3
20

---

Published as a conference paper at ICLR 2025
Algorithm 2 Forward Operator Splitting Integration (FOSI)
Require: step size h, total steps N, friction constant ? 0, initial condition 9 E G. E
1: Set 90 = 9. So =
2: for = 1, .,..., N do
3: Sample ii.d -1 N0
1100,-eexp((21) for S i S dimg
4: Si. exp(- ( + G-
Entrywise exponential integration for Lie group preserving update for
5: 9n = 9n-1 expm(hEn)
6: end for
7: return k
b Return whole trajectory
Algorithm 3 Backward Operator Splitting Integration (BSOI)
Require: step size h, total steps N, friction constant 7 0, score network sg, initial condition I G, E g
1: Set 9u = g. So =
2: for .,... N...,Nor do
3: Set - so so([n-1, En-1, (n 10:0-11p 1)h)
4: Sample i.i.d. 6 N( (2)) 1) for Si dim g
5: ei exp(yh) ((exp(11)) - 1)s + - Entrywise exponential integration for
6: 9n = 9n expmm
>> Lie group preserving update for
7: end for
8: return EN
Return final iterate
Advantages of OSI: The benefits of using OSI for integration are threefold.
(1) The first benefit of the OSI is high numerical accuracy. In both the forward and backward dynamic, the linear component of the dynamic is integrated exactly due to the use of the exponential integration technique. This implies that, while OSI is still first-order method in terms of the order of errors, it enjoys smaller prefactor thanks to the reduction in error source compared with the Buler-Maruyann method (EM).
(2) The second benefit of OSI is that the trajectories generated stay on the manifold G x g for the whole time, ((kk)).((k)) E G x for any & 2 0. If we use the EM scheme to integrate the Riemannian component of the dynamic, which is the g dynamic, we would arrive iterates g((k + 1)A) g(kh) + h ((k))g(kh). Note that since ((kh){(kh) is in the tangent space of Lie group G at point (kh), moving arbitrary short time h along such direction would result in g((k 1)h) & G. Therefore, achieve valid trajectory on G, need to employ projection Th onto the manifold, which causes additional numerical errors apart from the time discretization error. If employing OSI, we will be free from such concern of leaving the manifold and also the projection errors.
(3) The third benefit of OSI is that the numerical scheme is projection-free. As we have discussed in point (2), EM method does not respect the Riemannian geometry structure of the Lie group and constantly requires the application of projections to achieve valid iterates. Apart from the numerical error, computing such a projection could be problematic. nonconvex set, which naturally raises concerns about the existence of the closed-form formula for such projections and more generally, how implement them in a fast algorithm. For example, 50(n) is the set of matrices satisfies (X Rax= T X = X by nonlinear constraints. Therefore, finding out the projection onto these Lie groups requires heavy work and needs to be investigated projection functions could be difficult to implement and require heavy computational resources, which certainly not scalable for large-scale applications and nigh-dimensional tasks. On the contrary, by employing OSI, we can enjoy a projection-fis free numerical algorithm and reduce the complexity of both training and generation.
general. Lie groups live in
I.). which is characterized
case-by-case basis. Not mention the possibility that these
21

---

1)h) Published as a conference paper at ICLR 2025
D.2 PROBABILITY FLOW ODE
In this section, we will introduce the OSI for probability flow ODE. We recall that the probability flow ODE is given by,
gr -thet dG. = =>(T- - 1) + +VV Logpr-1(9.. Et
Similar to the SDE setting, this can be split into two smaller dynamics,
17
dt = -.Ct de, O
+A
9: =0 dp, rE,dt ysp( Et, t)dt
Similar to the Backward Operator Splitting Integrator (BSOI). we employ an approximation of the dynamic, discretize the nonlinear component Sø by using a left-point rule, i.e. pretending that g and do not change over short time h, and use the exponential integration technique to conduct the exact integration of the rest of the linear dynamie. The details can be found in Algorithm g
Algorithm 4 Probability Flow ODE
Require: step size h, total steps N, friction constant y score network sø, initial condition
I E G,5 E
1: Set 9l 9,
2: for = 1, 1,.... N do
Set - s6 (g(gn-111 1, Sn-1:
4: Et = exp(?1))n
(exp(7h) - 1)=--1 5: In In-1 expm hh
D Entrywise exponential integration for Lie group preserving update for
6: end for
7:
return JN,EN
Return final iterate
D.3 ERROR ANALYSIS OF OPERATOR SPLITTING INTEGRATOR ON LIE GROUPS
In this section, we reproduce the error analysis for OSI on Lie groups provided in the work of
Kong & Tao 2024). Such analysis justifies the convergence of OSI, which we provide just to be self-contained. For technical details the proof, please kindly refer the original paper of Kong
Tao (024).
In this section, we consider the following general form of OSI update,
0z+1 ÿr exp(1)x+1)
Ex+1 = =exp(-7h){r 1-exp(-7he
p(k+1)h Ig ( + V2r Sth exp(->(h-s)ddWW)
(26)
where (āk, ( denotes the points generated by discrete numerical scheme. We note that by choosing
carefully the potential function U, this OSI update can correspond both the forward operator splitting integration in Algorithm @ and the backward operator splitting integration in Algorithm Therefore, providing a general error and convergence analysis of equation (26) covers situation of
using OSI on the simulation of both the forward noising process and the backward sampling process.
Following the same setting as in Kong & Tao (2024), we introduce the following assumption on the regularity of potential function U,
Assumption 1. For the Lie group G. we assume it is fiiiee-dimesiooall connected, and compact. For the potential U, We assume I--smooth under the Riemannian metric Lemma in Kong & Tao 2024), i.e., there exist constants L (0, oo), s.t.
-ToLy--VUG
1TLg--VU()) Ld(g.0) Vg.g G
We further recall that To is the initial distribution and ग is the stationary distribution. Under these assumptions, we have the following convergence analysis for the OSI:
Theorem (Nonasymptotic error bound for OSI). If the initial condition (go- En) ~ Tg satisfies W,(To, n.) and To is absolute w.r.t. dgdg, then Vk- 1, 2,,,... the density of scheme Eq. (26)
22

---

Published as a conference paper at ICLR 2025
starting from has the following W2 distance from the target distribution:
W2( TT) S Co ( (c-chh (, T.) +
E(h)
exp( -c)
where E(h) O(h)) Note this holds vh O, but E(h) can grow exponentially W. is the Wasserstein distance induced by semi-distance P (given explicitly in Eq. in Kong & Tao (024)
and Cp is constant.
Theorem quantifies the distance between the resulted sampled distribution using OSI and the true distribution in terms of W2 distance. We notice that the upper bound on this Wasserstein shrinks to 0 as timestep tends to suggesting the convergence of OSI under the general setting.
EVOLUTION OF KL DIVERGENCE
In this section, we will discuss the effectiveness of likelihood training in terms of learning the correct data distribution. We will show that the KL divergence between the true data distribution Po and the learned data distribution can be bounded by accumulated ccore-macching errors up to an additional discrepancy error caused by a mismatch in the initial condition.
Recall that X, (gt, &e) is the trajectory of the forward dynamics (3) with X, admitting smooth density Pe(ge: Ee) with respect to the product of Haar measure on and Lebesgue measure on g. Let's denote (0))) (Pr-c)<E(o.rr Note that by construction, do = pr as x² when 7 is large, where n² is the invariant distribution of the forward dynamic 33. Also, qT = Po is the initial condition for the forward dynamic, which in practice is the (partially unknown) joint data distribution Recall that we have denoted the sequence probability distribution of as the density of ((Y?) with respect the reference measure, where Y is the trajectory of the learned backward dynamic in 7
and
We have Theorem 5 regarding the KL divergence between the learnt data distribution of and the true data distribution Po:
Theorem 5. Let Y be the trajectory of the learni backward dynamic 7 under initial con- dition Yi Wo. Y has density (qe )ee When the score is given by solg, E, == Vg log (g, E,t, where (G is the density of the forward dynamic (3) under initial condition qu and satisfies Gr = na. We then have
DKL (po ll dl) axs
pr- () log se(g,E,t) dgdEdt + dgdde DKL(PT TT)
In order to prove Theorem 5 we need the following Lemma that characterizes the time derivative of the KL divergence between sequences of probability distributions that corresponds the time marginal of the same SDE with different initial conditions.
Lemma 3. Given Po, d0 two distributions G x g. We denote the sequence of probability distributions (pe) J:z0 and (2:) the marginals of the forward dynumic (3) with initial conditions po and go respectively. Then, Pt and qe satisfies
DRL. (pr ll 4t) - = 7()prl/Ve log Ve log ( dgdg Dt
Gxx
where Vg ş is the gradient w.r.t. E.
Lemma rellates the time derivative of KL divergence between two distributions to the difference in their score integrated over the manifold x g. To prove this lemma, we need to derive the Bokker-Plancee equation for pt and q, respectively.
23

---

Published as a conference paper at ICLR 2025
Proof of Lemma described in equations,
We prove general case and consider the general form of forward dynamic The evolution pt and ge can be characterized by following Fokker-Plaa
13
0
01 Pt Pr = div, (pr + 7() (dive(PiE) + Aςpt)
ô 0t 91 = L - div + (t) (dive(a:E) + Açqe)
where div, y is the divergence of vector fields on G under the left-invarian metric choose, dive is divergence in g and Ag is the Laplace operator on g. They are well-defined since g linear space. Consequently, we can evaluate the time derivative of KL divergence between pr and qt, where the integration is performed over G x unless specifically noted,
(pr ll qr) Ōt
i
Øt (ʃ Pt. log
0
Pt dgdE) qt
= /
Op:
Pe
dgde ʃ
dqe Pe dgdg 0t at
log
01
/
log
Pt
L prelde - 1 P 2 qndg
Using the divergence theorem, we have for any smooth function ʃ G x + R, we have
J f = f (Vol.P:TeLaa)) (Vee++((Vee))))e) + Ve loggt
Similar results holds for L As a consequence, applying the previous equation with f log Pt. and ʃ = Et respectively, we can write,
0 Ot
Pt. (pe (b) = / V. log Pr
Pt V log ,pr(tt (5 + Ve log po) dgdk q:
+
Pt
Ve Pt ())(t) (5+ Ve log ge) dgdg
4 ToLgE)
9 qi
Pr TeTeL Log
Pt
V
(t) (E + Ve loggpt) dgdg
V 9 qi
Pt
9: TcLgL
V Pt (t) (5 - Ve log q() dgdg
= / 7(t)
VE Pt
2t (Ve log
V log q dgdg
(t)pp. Ve logp: V. log q: dgdE
This finishes the proof of the desired Lemma.
We now ready to present proof for Theorem 5 Note that under the conditions so, in fact
of = -r- We just need to apply Lemma 3 between Pt and d to conclude.
Proof of Thm. S
Lemma gives
DKL (po o
DKL (po go)
DKL (pr ll gr)++ + f"'
T a
DKL (pr ll dt
at
= (pr l1 (r) + L '-' [
T
Jaxa ((t)p:/VV logat Vglog 2 dfddf
24

---

d Published as a conference paper at ICLR 2025
Using the condition se C, 1ogg and @r x², with the choice r(t) arrived at the following equation.
DKL (po ll ( = !'" Jong pr (r e log s.(g, + (pT || T+)
F NLL ESTIMATION WITH INTRINSIC PROOF
In this section, we provide intrinsic proof of the instantaneous change of variables on general manifold, which does not depend on local charts in the proof or the formula. While we are now aware
that the results are not new and has been discussed in (Chen et al. 2018 Lou et al., 2020 Falorsi &
Forré 2020
Chen & Lipman 2024), we still provide proof for self-consistency
Let be random variable whose range is .M and denote its density as po C(M). Given smooth time-dependent vector field (l,,)) -,, i.e.,X (1.-) E r (TM) for any / E [0, T
We consider the
push forward map generated by the flow of X, i.., ft M + M satisfies
d - (s,
Vo E M.
ds
with initial condition fa is the identity map for any We define pt as the density of f6(z). Then have the following theorem,
Theorem 6. [Instaneous Change of Variables on Manifold] Consider p: RxM - R. such that
P: = p(t,,) -) is the density of z(t), where z(t) is the random variable defined by = pushforward along X for time Then we have
d plg
(lt fa(z)) div X(t,
A M
dt
We first review standard result (see e.g., Ross et al. (2023) Lemma and then provide proof for Theorem 6 The following describes the relationship between the density of the push- forward as well the determinant of the differential and corresponding points. We will use this result heavily in our proof.
describes
Lemma 4. For any diffeomorphism f M + M, we have
f+P(f(x)) = p()) (det (f))), VE M
Ty the forh-forward density. On the right-hand side, df(x) T.M + M, denoting the differential of J. is linear map. Consequently, det df is well-defined and independent of the choice of coordinate.
With Lemma we are ready to present our proof for Theorem
Proof of Thm. 6 In this proof, the shorthand notation 2r (xr,, which also induces
2oo = r. Lemma gives
d
pl,, fxxx))
log [16)) (x)))
dt
d
log [po(z) det df)(
dt
log (det df)()))
di
Teee (TM) denotes the of all smooth vector fields on M
25

---

Published as a conference paper at ICLR 2025
Since f is the pushforward map. it has the semi-group strueture, i.e., for fiz flr for any
S t2 s ta, which gives det dft, det dft det dft, and immediately
d dl
log det d/ol = lim (log [det d/d d/)] log [det (x))
c-00+
lim log det (x=)) c-
Consequently,
d
a log det (fz)+(()))
log p(, fú(x))
dt
de
0r det det dstt
det dff (x))
i det dft (xx))
de
c==0
where we use de to denote the right derivative, and the last equation is because of fi is identity.
Jacobi's formula gives & det(A) tr (A) at A = I, which tells
d fx(x)) tr (2e ds. ( c=U de
log p(t, pot dt
Before we proceed, we define two sets of vector fields, {Ei and {YiJ {E is defined as a set of smooth coordinate frame, defined on the whole manifold M. (Y:) is set of vector fields along H, generated by the push forward map ff+e i.e. df is map from Tu, to M, and Y, satisfies
dft (Yi(a++))) = Y((x++)), Vi I,... d, A <
with constraint Yx,,)) = Ex,)) Note that Y, is defined only along the curve x, for E [0,
Since we are considering push forward map / along a time-dependent vector field X (1,,)) rx (TM),, we consider to make it time-independent by considering the problem in a new space M = R x M, and new time ime-independent vector field X (TM)) defined by
(1, (1, X)) LE [0,T), M
Since is the integral curve generated by X, the integral curve of X with initial condition (0, is given by T. (t, i.e., H. satisfies é, = (ā.).
Both (E 1 and {r:)/=1 can be extended to M. For (E.) this is defined by Eg (1, 0) and
È; = (0, E.). Similarly, for (Y:) this defined by (T.)) (0, (xx)) and Y, = (1,
By choosing an arbitrary linear connection V on M, we can also extend V to M, and we denote such induced linear connection as & (see e.g., Do Carmo & Flaherty Francis 1992 Ex. 1 in Chap. 6). Since f is the pushforward generated by X. Je - X(I,)). and we have
d dt log)(t,f(()))
0
= Σ
(ast (F(+)(E()))),, (o(f(*((x=)))(l
Σ dc
(V:()++((x:))) F.ULL)(a))))
Σ V Y;, Ei)
V (Y,, Esy..
1=1
26

---

Published as a conference paper at ICLR 2025
We choose V to be the Levi-civita connection. Because it is compatible with the metric, have
d (ogp(t.
d
VxXY, Fg) + (Y, Vx)
dt
By the definition Lie derivative, we have Lx Y = O. Together with the constraint Yi(x,) = Ex( have VxYi Vv, X VE, X zt. which gives
dt (eg(tt = Σ
Now we show (E,, V E,) 0 using local coordinates:
Vx.VE. = Xxr -
Σ
In Vlgl
de,
where I's are Christoffel symbols corresponding to the local coordinates E,, defined by r V E, Fly, Fx)) Due to our choice of F; orthonormal frames, we have [g| = and (Fi Vx Pa.) = 0
for all i.
By simplifying the expression, arrive at the following desired results, which finishes our proof.
d (ogp(t,f(x))) VE, div
We can choose E, as the left-invariant vector fields generated by e,, i.c. Bi(g) = Te Te where e, E g is a set of orthonormal basis. As corollary, we can compute dynamic for the log probability as the following,
Corollary 3 (NLL estimation on Lie group with left-invariant metric). For SDE in Eq. ( the time-dependent vector field is given by
X(g. {) (TeLgS, - gtt,g,
Using the fact that div, (T 0. we have
θ DE,
log (((x))
Σ
-
9
G TRAINING SET-UP, DATASET
G.1 HYPERPARAMETERS
Hardware: All the experiments are running on one RTX TITAN, one RTX 3090 and one 4090.
Architectural Framework: We employed the score function (g & 0) parameterized by the same network architecture as outlined in the CLD paper Dockhorn et al. 2021), albeit with varying parameter counts for each task. Specifically, and H, as real vectors and embed them into latent vectors of hidden dimension D respectively with trainable MLPs, where the latent vectors denoted as ghid, Énid. We embed time into the sinusoid embedding of dimension D as well, denoted as lxid. The score network output is given by
consider the following architectures. We represent gi
MLP-SKIP CLAP-SKIP(GN((p)) (GN(ghdd + Ehid + tria))
where G.N is the group norm operation, and MLP-SKIP is a k-layer MLP with skip connections. We use SiLU the activation function for all the ML.ss used in the neural network.
27

---

Published as a conference paper at ICLR 2025
For low dimensional experiments such as Torus, SO(3), set D 256. For other experiments, set D = 512. We choose varied & based on problem difficulties, ranging from I: = 3 to h = 5.
Training Hy syperaarameeer:: Throughout our experiments, we maintained the diffusion coefficient 7(t) constant at I, while the total time horizon T varied depending on the task, with a good choice ranging from T 5 to T 15. We use AdamW optimizer to train the neural networks with initial learning rate of 5 x 10 with a cosine annealing learning rate scheduler. We train for at most 200 iterations with batch size of 1024 for each task, and we observe that the model usually converges within 100k iterations.
G.2 DATASET PREPARATION
Protein and RNA Torsion Angles: We access the dataset prepared by Huang et al. (Huang et al.
2022) from the repository of (Chen & Lipman 2024). We further post-processed it and transformed the data into valid elements of T".
Pacman: We take the maze of the classic video game Pacman and extract all the pixel coordinates from the image that corresponds to the maze. We post-processed it and transformed the data into valid elements of T
Special Orthogonal group 5O For SO, we followed the same procedure as the one described in (De Bortoli et al.] 2022) and generate Gaussian Mixture with 32 components, uniformly random mean and variance. For n > 3, we follow a similar procedure with a reduced number of mixture components.
Unitary group U(n): We considered the unitary group data the form expm(-it7H which is the time evolution operator of the following Schrödinger's 's equation for general quantum system, i0 = He, Here v denotes the quantum state vector and H denotes the Hamiltonian operator of the system. We considered the following two types of Hamiltonians,
For quantum oscillator, the Hamiltonian is given by 7 A + V, where A is the Laplacian operator, and Vxx) = f zoll ( is random potential function, where and an are random variables. Note that these are infinite dimensional objects. To obtain a valid element in U(n) for a finite we perform spectral discretization on the Laplacian operator as well as the random potential to get finit-ddimesiioal Hamiltonian operator Hr, with which the time-evolution operator is computed with. We choose & = 1 in this case.
* For Transverse field Ising Model, the Hamiltonian is given by
H =
Σ Jujafa - Σ g.o," (6.)
where of the Pauli matrices, Jij is the coupling parameter and gi is the field strength. Here Jiy and 9i are random variables, which corresponds to the situation of RTFIM. The time-evolution operator is generated with such Hamiltonian at
H ADDITIONAL NUMERICAL RESULTS
In this section, we present additional numerical results on learning time-evolution operators for ensembles of quantum systems, which are data distributions on complex-valued Unitary groups. We randomly select entries to scatter plot the data generated by TDM and compare it against the ground truth distribution. Additionally, we plot against the data distribution generated by RFM to further demonstrate our approach's 's advantage. As shown Figure 10 and Figure [T 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 one produced by TDM. This potentially originates from the approximations adopted by RFM, suggesting again the advantage of our proposed approximation free TDM.
28

---

Published as a conference paper at ICLR 2025
True dist.
RFM generated dist.
TDM generated dist.
1.00
0.75
0.50
0.25
0.00
0.25
-0.50
-1.0
-0.5
0.5
-1.0
0.0
-1.0
1.0
True dist.
RFM generated dist.
TDM generated dist.
-0.2
-0.5
0.0
0.5
-0.5
0.0
0.5
-05
0.0
0.5
True dist.
RFM generated dist.
TDM generated dist.
0.8
0.4
-022
-0.4
0.6
-0.5
0.0
0.5
-0.5
0.0
0.5
-0.
0.0
0.5
Figure 10: Visualization of Generated Time-evolutie Operator of Quantum Oscillator on U(n) by TDM. compared against true distribution and RFM. Plotted entries randomly chosen. Top row: U(1). Mid row: U(6). Bottom row: U(8)
29

---

Published as a conference paper at ICLR 2025
True dist.
RFM generated dist.
TDM generated dist.
1.0
-0.5
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.
0.0
0.5
1.0
True dist.
RFM generated dist.
TDM generated dist.
1.00
0.75
0.50
0.25
0.00
-.25
-0.50
-0.75
-0.5
0.0
0.5
005
0.0
0.5
-05
0.0
0.5
True dist.
RFM generated dist.
TDM generated dist.
0.4
0.2
-022
-0.6
0.88
0.75-0.50-0025
-0775-0.50-0.22 0.50-0.25
0.25 0.50
00.7 0.75-0.50-0.25
0.25
Figure 11: Visualization of Generated Time-evolution Operator of Transverse Field Ising Model (TFIM)
(nn) by TDM, compared against true distribution RFM. Plotted entries are randomly chosen. Top row: U(4). Mid row: U(8). Bottom row: U(S)
30


Reviews:
[{"review_id": "IqvbnTdLNb", "reviewer": "Reviewer_25XD", "summary": "The authors present a method for diffusion on Lie groups. The method is based on the recently published idea of trivialized momentum, that is, diffusion (i.e. noising) is performed on the Lie algebra of the group, and not on the Lie group itself. The advantage of this approach is that the Lie algebra is a vector space isomorphic to $\\mathbb{R}^n$, and therefore it simplifies the problem since all the curvature terms pretty much drop out. This diffusion is brought back into the actual Lie Group G by using the map connecting two tangent spaces of G at different point, which is the tangent of the left multiplication map. \n\nThe authors show that this noising dynamics has an inverse dynamics as well, in very familiar fashion as the usual diffusion dynamics, in which a score function appear, which it is possible to learn with neural networks. \n\nThe authors further argue that it is possible to have a dynamics that is exaclty on the curved manifold at all steps without the need of a projection (I do not agree with this statement, please refer below in the \"Questions\" section). This is done by separating each step into two contributions, i.e., the update of the Lie algebra element $d\\xi$ and the update of the lie group element $dg = g \\xi$. Note that also this procedure is already known in the literature. \n\nThey provide an explicit formula for the Abelian case, while they point out that for the general case, it is necessary to learn the score with implicit score matching, which unfortunately involves the divergence term.", "questions": "- 280-283: The sentence “Since the Brownian motion…” does not seem to be grammatically correct.\n- 292: “Unfortunately, note [that] even though”. Missing “that”.\n- 382: $A_g^B$ used twice, probably copy-paste error.\n- The authors make several times the claim that the formalism on which they based their work does not required projection on the manifold. I think this statement, as is, is false. Group equation is $dg = g_0 \\xi$, which is an equation on the tangent space of the group. Its solution is (neglecting the variation of $\\xi$), $g=\\exp(g_0 \\xi)$. But the exponential map is precisely the projection function on Lie groups, $\\exp: TG —> G$. So to stay on the manifold the authors are performing a projection exactly as all other methods too.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The topic is very relevant, as diffusion-based models achieved quite successful results in many applications, but very little has been explored in terms how inductive bias can be incorporate into them, and diffusion on manifold is surely an important field. \n\nThe notation is fairly consistent and the math is, to the extend that I checked, correct (I did not carefully checked the proofs in the appendix).", "weaknesses": "My two main concerns regard the originality of the work and its practical relevance:\n\n- **Originality**: As far as I can see, the idea and the theory behind trivialized momentum has been developed in previous works, as well as the  splitting discretization technique to exactly integrate the SDEs. 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.\n- **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. \n\nI think the paper would be significantly more relevant if it provided an analogue of Theorem 2 for non-Abelian groups.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Zi7biPjDpF", "reviewer": "Reviewer_AGbj", "summary": "The paper introduces a new framework for constructing diffusion models on manifold-structured data in the particular case where the manifold is a Lie group. This is achieved through the use of a forward noising process defined by kinetic Langevin dynamics on the Lie group and its tangent space jointly. The authors demonstrate that the group structure of this class of manifolds allows for the corresponding reverse process to rely solely on a Euclidean score function, as well as allowing for the manifold-preserving and projection-free sampling of the forward and reverse processes, circumventing the complexities and multiple approximations required in previous work. The authors then apply this to various tasks presenting convincing results.", "questions": "Some possible typos in the text:\n- line 155, missing brackets for g_t, \\xi_t?\n- line 325, should be \"\\sigma_t evaluated at\"?\n- line 1022, the A_\\xi^\\F should be A_\\xi^\\B?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The authors present a novel and clever way to address the computational limitations of previous approaches to manifold-valued diffusion models, by exploiting the group structure of Lie groups through the use of Kinetic Langevin dynamics. This allows for learning a score function in a fixed Euclidean space which reduces the numerical errors associated with the approximations arising from scores defined directly on curved manifolds. This approach is well-fleshed out by the authors who provide a detailed methodology for the application of this for generative modelling, for example, by providing an integration scheme which is manifold-preserving and projection-free, and training objectives for both Abelian Lie groups and non-Abelian Lie groups. This framework is further validated by experiments demonstrating the novel scaling of manifold-valued diffusion models to SO(N), N>3.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5B0COboNJx", "reviewer": "Reviewer_N41G", "summary": "This paper provides a new framework for performing generative modelling on Lie groups which doesn’t require any geometry-specific projections or approximations, utilizing the notion of “trivialized momentum.” The authors provide the denoising score matching loss which can be used for SO(2)^n, as the conditional transition probabilities are tractable and an implicit score matching loss for other connected compact Lie groups. They also provide an integration scheme for the forward and backward dynamics.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is overall well-written and easy to follow. The problem the authors tackle is interesting and important.\n- The proposed technique is interesting and avoids some of the problems mentioned with other existing approaches, such as projections, requiring specific architectures, etc.\n- The mathematical details are explained cohesively and clearly. \n- The results on the torus datasets, as well as the toy tasks for the SO(n) and U(n) groups seem promising.", "weaknesses": "- The main weakness of the paper is the lack of rigorous experiments to demonstrate the effectiveness of the model in comparison with existing approaches. \n- 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 are clearly provided. \n- Secondly, it wasn’t clear me to what approximations one has to make when doing RFM in the torus settings, especially since going from the Lie algebra to the Lie group element is essentially trivial in that case. Could the authors clarify what approximations are required in RFM for SO(2)^n experiments?\n- Lastly, 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.\n- It would be very helpful if the authors provided a more detailed discussion on the approximations made in RFM and RDM which the authors attribute to the worse performance of these models in the torus settings. \n- The paper also provides experiments for the SO(3) setting, however it doesn’t provide a comparison with SE(3) generative models that rely on IGSO(3) distribution. As mentioned in the paper, the special structure of SO(3) can be leveraged, which makes applying TDM in that setting less motivated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bBZNDzECP6", "reviewer": "Reviewer_dYfk", "summary": "- The paper is concerned with tackling the challenge of incorporating diffusion models into a manifold setting\n- Their key takeaway is that on Lie groups for both training and sampling diffusion can take place in the Lie algebra\n- They show that their method performs well on various Lie groups of various dimensions that arise in different applications", "questions": "Questions for clarification \n- 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?\n    - 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\n- 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? \n    - 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.\n- 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? \n    - If not, I would recommend adding this explicitly to let the reader know that there is possibility for error here.\n- 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?\n- Section 5: \n    - 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?\n\nAdditional feedback\n- Section 1: please do a grammar check. A couple of small things that I noticed:\n    - p.1 l.43: change ''manifold’’ in … Neural ODE to manifold with maximum … to ''manifolds’’. This reads better\n    - p.1 l.44: change ''algorithm’’ in Rozen; Ben-Hamu develop simulation free algorithm but… to ''algorithms for the same reason as above\n    - p.1 l.94: add ''a'' right before ''trivialized’’ in …through learning trivialized score function…\n- Section 2: It looks a bit odd that section 2 is just one line. Consider integrating this line in section 3 instead.\n- Section 3:\n    - p.3 l.133: replace ''sequal'' with ''following''. It sounds like this will be discussed in your next paper rather than further down in this one.\n    - p.3 l.147: the authors consider the tangent map. I believe that this is the differential? For the reader who is used to other nomenclature and/or notation, I would suggest to also mention that this is often referred to as the differential.\n    - In equation 8, shouldn’t there be a superscript i for the xi in the dt term?\n    - p.7 l.373: the authors write ``do not admit a closed form.’’ A closed form what?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The authors provide a very nice motivation for the problem and highlight the key challenges that current manifold-valued models are dealing with.\n- The mathematics is displayed in a very clear way and adequately addresses the challenge discussed in the motivation, i.e., how to set up manifold diffusion models on the Lie algebra. They also address different types of Lie groups well. In particular, it is nice that they highlight the special case of Abelian Lie groups learning, where training becomes easier than on non-Abelian ones\n- The numerical experiments nicely display that the method performs well on high-dimensional Lie groups (which is where one would expect issues due to numerical inaccuracies from projections and manifold mappings), but also showcase that for low-dimensional Lie groups performance is superior to baselines.", "weaknesses": "- I feel that for the introduction, there are some mistakes in terms of English grammar \n- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)