# Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups

arXiv: https://arxiv.org/abs/2405.16381 · HTML: https://mlsanity.com/paper/2405.16381/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 31 of its 57 references that are in the map.
- First posted on arXiv 25 May 2024.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 31)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Flow Matching on General Geometries](https://mlsanity.com/paper/2302.03660/) (2023)
- [Score-Based Generative Modeling with Critically-Damped Langevin Diffusion](https://mlsanity.com/paper/2112.07068/) (2021)
- [Riemannian Score-Based Generative Modelling](https://mlsanity.com/paper/2202.02763/) (2022)
- [SE(3) diffusion model with application to protein backbone generation](https://mlsanity.com/paper/2302.02277/) (2023)
- [SE(3)-Stochastic Flow Matching for Protein Backbone Generation](https://mlsanity.com/paper/2310.02391/) (2023)
- [Riemannian Continuous Normalizing Flows](https://mlsanity.com/paper/2006.10605/) (2020)
- [A micro Lie theory for state estimation in robotics](https://mlsanity.com/paper/1812.01537/) (2018)
- [Riemannian Diffusion Models](https://mlsanity.com/paper/2208.07949/) (2022)
- [Flows for simultaneous manifold learning and density estimation](https://mlsanity.com/paper/2003.13913/) (2020)
- [Neural Manifold Ordinary Differential Equations](https://mlsanity.com/paper/2006.10254/) (2020)
- [Moser Flow: Divergence-based Generative Modeling on Manifolds](https://mlsanity.com/paper/2108.08052/) (2021)
- [Matching Normalizing Flows and Probability Paths on Manifolds](https://mlsanity.com/paper/2207.04711/) (2022)
- [Improved motif-scaffolding with SE(3) flow matching](https://mlsanity.com/paper/2401.04082/) (2024)
- [Deep Momentum Multi-Marginal Schr"odinger Bridge](https://mlsanity.com/paper/2303.01751/) (2023)
- [Neural Ordinary Differential Equations on Manifolds](https://mlsanity.com/paper/2006.06663/) (2020)
- [Generative Modeling with Phase Stochastic Bridges](https://mlsanity.com/paper/2310.07805/) (2023)
- [Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes](https://mlsanity.com/paper/2310.07216/) (2023)
- [Scaling Riemannian Diffusion Models](https://mlsanity.com/paper/2310.20030/) (2023)
- [A Complete Recipe for Diffusion Generative Models](https://mlsanity.com/paper/2303.01748/) (2023)
- [Variational Optimization on Lie Groups, with Examples of Leading (Generalized) Eigenvalue Problems](https://mlsanity.com/paper/2001.10006/) (2020)
- [Convergence of Kinetic Langevin Monte Carlo on Lie groups](https://mlsanity.com/paper/2403.12012/) (2024)

## Built on by (most cited first, 19 of 19)

- [Diffusion Generative Modeling for Spatially Resolved Gene Expression Inference from Histology Images](https://mlsanity.com/paper/2501.15598/) (2025)
- [Diffuse Everything: Multimodal Diffusion Models on Arbitrary State Spaces](https://mlsanity.com/paper/2506.07903/) (2025)
- [MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control](https://mlsanity.com/paper/2508.10684/) (2025)
- [Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization](https://mlsanity.com/paper/2510.08233/) (2025)
- [Improving the Euclidean Diffusion Generation of Manifold Data by Mitigating Score Function Singularity](https://mlsanity.com/paper/2505.09922/) (2025)
- [TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion](https://mlsanity.com/paper/2509.25171/) (2025)
- [Kinetic Langevin Diffusion for Crystalline Materials Generation](https://mlsanity.com/paper/2507.03602/) (2025)
- [Diffusion Generative Modeling on Lie Group Representations](https://mlsanity.com/paper/2502.02513/) (2025)
- [Explicit Discovery of Nonlinear Symmetries from Dynamic Data](https://mlsanity.com/paper/2510.01855/) (2025)
- [Stiefel Flow Matching for Moment-Constrained Structure Elucidation](https://mlsanity.com/paper/2412.12540/) (2024)
- [Riemannian Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2505.04338/) (2025)
- [Critically-Damped Higher-Order Langevin Dynamics for Generative Modeling](https://mlsanity.com/paper/2506.21741/) (2025)
- [Discovering Symmetry Groups with Flow Matching](https://mlsanity.com/paper/2512.20043/) (2025)
- [Efficient Diffusion Models for Symmetric Manifolds](https://mlsanity.com/paper/2505.21640/) (2025)
- [Inverting Data Transformations via Diffusion Sampling](https://mlsanity.com/paper/2602.08267/) (2026)
- [Explicit and Effectively Symmetric Schemes for Neural SDEs on Lie Groups](https://mlsanity.com/paper/2509.20599/) (2025)
- [Efficient Diffusion Models under Nonconvex Equality and Inequality constraints via Landing](https://mlsanity.com/paper/2604.17838/) (2026)
- [Quotient-Space Diffusion Models](https://mlsanity.com/paper/2604.21809/) (2026)
- [FLARE: Flow Matching with Local Axis-Angle Representations for Stochastic Micromagnetic Evolution](https://mlsanity.com/paper/2609.34070/) (2026)

## Abstract (arXiv)

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.