Scaling Riemannian Diffusion Models
This work addresses a bottleneck in applying diffusion models to complex geometric spaces, with incremental improvements that enhance scalability for specific domains like physics and machine learning embeddings.
The paper tackled the problem of scaling Riemannian diffusion models to high-dimensional manifolds by improving approximations for the diffusion transition term, resulting in noticeable performance gains on low-dimensional datasets and enabling applications to high-dimensional tasks like modeling QCD densities on SU(n) lattices and embeddings on hyperspheres.
Riemannian diffusion models draw inspiration from standard Euclidean space diffusion models to learn distributions on general manifolds. Unfortunately, the additional geometric complexity renders the diffusion transition term inexpressible in closed form, so prior methods resort to imprecise approximations of the score matching training objective that degrade performance and preclude applications in high dimensions. In this work, we reexamine these approximations and propose several practical improvements. Our key observation is that most relevant manifolds are symmetric spaces, which are much more amenable to computation. By leveraging and combining various ansätze, we can quickly compute relevant quantities to high precision. On low dimensional datasets, our correction produces a noticeable improvement, allowing diffusion to compete with other methods. Additionally, we show that our method enables us to scale to high dimensional tasks on nontrivial manifolds. In particular, we model QCD densities on $SU(n)$ lattices and contrastively learned embeddings on high dimensional hyperspheres.