Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups

TL;DR

This paper introduces a trivialization technique that simplifies diffusion modeling on Lie groups by maintaining a fixed momentum space, enabling high-fidelity, efficient data generation on complex manifolds like proteins, RNA, and quantum groups.

Contribution

The paper presents a novel trivialization approach that transfers Euclidean diffusion model effectiveness to Lie groups, avoiding complex approximations and improving generative performance.

Findings

Achieved state-of-the-art results on protein and RNA torsion angles.

Successfully generated data on high-dimensional SO and U groups.

Simplified implementation with high accuracy and efficiency.

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.