PaperlayerView exact revision on arXiv

Source-linked AI summary

Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes

Jaehyeong Jo, Sung Ju Hwang

arXiv:2310.07216v2cs.LGstat.ML

TL;DR

Generative modeling for manifold-valued data matters in scientific applications, but existing approaches face computational and scalability constraints. Riemannian Diffusion Mixture builds manifold diffusion from bridge processes with tangent-direction drift, without heat-kernel estimates. It reports strong performance across diverse manifolds and uses only 5% of CNF's in-training simulation steps on general manifolds.

  • Problem

    Existing manifold generative models face expensive divergence computation or heat-kernel approximations, while many approaches are difficult to extend to general geometries and high dimensions.

    Source: §1 · p. 1

  • Method

    Riemannian Diffusion Mixture constructs a manifold diffusion from endpoint-conditioned bridge processes, with drift given by weighted tangent directions and a scalable regression-based training objective.

    Source: §1 · p. 2

  • Results

    On diverse real-world and synthetic manifolds, the method outperforms or matches state-of-the-art baselines and uses only 5% of CNF's in-training simulation steps on general manifolds.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The framework targets non-Euclidean applications including sphere-valued earth and climate data, SE(3) protein structures and robotic movements, and manifold-valued shapes.

    Source: §1 · pp. 1–2

    The Logarithm bridge is constrained on general manifolds where the logarithm map may not be available.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

Learning the distribution of data on Riemannian manifolds is crucial for modeling data from non-Euclidean space, which is required by many applications in diverse scientific fields. Yet, existing generative models on manifolds suffer from expensive divergence computation or rely on approximations of heat kernel. These limitations restrict their applicability to simple geometries and hinder scalability to high dimensions. In this work, we introduce the Riemannian Diffusion Mixture, a principled framework for building a generative diffusion process on manifolds. Instead of following the denoising approach of previous diffusion models, we construct a diffusion process using a mixture of bridge processes derived on general manifolds without requiring heat kernel estimations. We develop a geometric understanding of the mixture process, deriving the drift as a weighted mean of tangent directions to the data points that guides the process toward the data distribution. We further propose a scalable training objective for learning the mixture process that readily applies to general manifolds. Our method achieves superior performance on diverse manifolds with dramatically reduced number of in-training simulation steps for general manifolds.

1. Introduction

Data with intrinsic non-Euclidean geometry are important across scientific applications, but existing manifold generative models face computational and geometric limitations. Riemannian Diffusion Mixture builds a process from bridge mixtures and tangent-direction drift, with reported performance and training-step benefits across manifolds.

1. Introduction

  • Manifold-valued data include sphere-based earth and climate data, SE(3)-represented protein structures and robotic movements, and 3D shapes identified as general closed manifolds.The introduction describes modeling such geometrically structured data as underexplored despite its relevance to applications.Source: §1 · p. 1
  • Previous manifold diffusion models can rely on imprecise heat-kernel approximations or expensive divergence computation, while general-manifold CNFs may require many simulation steps for accurate trajectories.Source: §1 · p. 1
  • The framework constructs manifold diffusion as a mixture of endpoint-conditioned bridge processes, using weighted tangent directions for drift and a regression-based objective that avoids heat-kernel estimation and divergence computation.The drift is described as guiding the process toward the data distribution.Source: §1 · p. 2
  • On diverse real-world and synthetic manifolds, the method outperforms or matches state-of-the-art baselines and uses only 5% of CNF's in-training simulation steps on general manifolds.The paper also reports scalability to higher dimensions with significantly faster training than previous diffusion models.Source: §1 · p. 2
  • On general manifolds, the method outperforms CNF with greatly reduced in-training simulation steps, which the authors present as demonstrating the necessity of stochasticity.Source: §1 · p. 2

2. Background

The background introduces the geometric objects and diffusion processes used in the paper. It defines manifold metrics and tangent spaces, then describes Brownian motion, its stationary distribution, and geodesic random-walk simulation.

2. Background

  • A Riemannian manifold is equipped with a metric that defines inner products of tangent vectors; each point has a tangent space, and the text also introduces gradient, divergence, Laplace-Beltrami, and exponential and logarithm maps.Source: §2 · p. 2
  • Brownian motion on a Riemannian manifold is generated by half the Laplace-Beltrami operator, and its transition distribution is the heat kernel.In Euclidean space, the heat kernel coincides with the Gaussian distribution.Source: §2 · p. 2
  • A manifold diffusion can converge to a stationary distribution proportional to e−U(x), and a geodesic random walk simulates it by taking small tangent-space steps in the drift direction.Source: §2 · p. 2

3. Riemannian Diffusion Mixture

Riemannian Diffusion Mixture builds a generative process from manifold bridge processes without heat-kernel estimation. Its weighted tangent-vector drift guides the process toward data, and two-way bridge matching provides a training objective that avoids divergence computation.

3.1. Bridge Processes on Manifold

  • Brownian bridges require an intractable heat kernel on general manifolds, while logarithm maps may lack a closed form and be costly to compute.The heat kernel is known in very limited cases; even on the sphere it is represented as an infinite sum.Source: §3.1 · p. 3
  • The Logarithm Bridge Process follows the logarithm-map shortest-path direction to its endpoint, and in R^d reduces to the Euclidean Brownian bridge.Its drift grows as t approaches T, forcing convergence to the endpoint.Source: §3.1 · p. 3
  • The Spectral Bridge Process uses the tangent direction that minimizes spectral distance to the endpoint, with eigenvalues and eigenfunctions computed before training.Replacing spectral distance with geodesic distance yields the Logarithm bridge.Source: §3.1 · p. 3

3.2. Generative Process on Riemannian Manifold

  • The mixture drift is a weighted mean of tangent vectors toward data points, and its terminal distribution equals the data distribution regardless of the initial distribution.For Logarithm Bridge Mixtures, projecting the drift onto the manifold gives a prediction of the most probable endpoint given the current state.Source: §3.2 · p. 4

3.3. Two-way Bridge Matching

  • Two-way bridge matching simulates a fixed-endpoint bridge forward before T/2 and backward otherwise, bypassing terminal drift explosion and halving computational cost; drift regression avoids divergence computation.Minimizing the objective is guaranteed to minimize KL divergence between the data distribution and the parameterized process’s terminal distribution.Source: §3.3 · p. 5
  • Up to ×34.9 speed up for training over one-way simulation and ×12.8 speed up in training over RFM are reported; the latter is without sacrificing performance on manifolds with non-trivial curvature.In the complex-distribution setting, RFM completely fails.Source: §3.3 · pp. 5–6

4. Related Work

Prior manifold generative models use projection, diffusion, or continuous normalizing flows, but each approach has limitations in geometry, training cost, or scalability. Diffusion mixtures provide an alternative to denoising by connecting initial and terminal distributions with bridge processes.

  • Earlier diffusion-mixture methods are limited to Euclidean space, whereas this work extends the framework to manifold-valued data.Source: §4 · p. 6
  • Projected Euclidean models can be numerically unstable because projection to a manifold is not bijective; manifold flows instead model transport directly on the manifold.Source: §4 · p. 6
  • Manifold diffusion models rely on score matching that requires heat-kernel approximations or divergence computation, while diffusion mixtures avoid both and allow arbitrary priors.Source: §4 · p. 6
  • Manifold CNFs can require costly simulation-based likelihood training, while simulation-free flows scale poorly to high dimensions and may not extend to general geometries.Source: §4 · p. 6

5. Experiments

Experiments cover real-world and synthetic manifold data, with reported gains or parity against baselines and faster training in several settings. On general closed meshes, the method fits complex distributions with fewer in-training simulation steps than RFM.

5. Experiments

  • The experiments evaluate real-world and synthetic manifold datasets and compare against manifold CNF, score-based, diffusion, and flow-matching baselines.Source: §5 · pp. 6–7

5.1. Earth and Climate Science Datasets

  • On sphere-based earth and climate datasets, the method significantly outperforms baselines on Volcano and Fire, and matches RFM on Earthquake and Flood.Source: §5.1 · p. 7

5.2. Protein Datasets

  • On protein datasets represented on tori, the method outperforms or matches RDM, makes marginal improvements over RFM, and is comparable to RFM on General and Pre-Pro datasets.Source: §5.2 · p. 8

5.3. High-Dimensional Tori

  • On high-dimensional tori, the method outperforms RSGM especially at higher dimensions, maintains performance through dimension 104, and achieves up to 5× training speedup over RSGM with implicit score matching.Source: §5.3 · p. 8

5.4. General Closed Manifolds

  • On triangular meshes from Stanford Bunny and Spot the Cow, the method fits complex distributions with 15 in-training simulation steps, uses only 5% of RFM’s steps, and trains 12.8× faster than RFM.Source: §5.4 · p. 9

5.5. Further Analysis

  • Uniform-time training significantly reduces performance; two-way simulation produces trajectories close to exact trajectories in 15 steps, unlike one-way simulation.With a small noise scale, accurate trajectories require many steps, which the authors connect to RFM’s failures in the mesh experiments.Source: §5.5 · p. 9

6. Conclusion

The Riemannian Diffusion Mixture builds manifold generation from bridge processes whose drift is a weighted mean of tangent directions, without heat-kernel approximation. The conclusion emphasizes scalable training and strong performance across diverse manifolds with fewer in-training simulation steps.

6. Conclusion

  • The method combines bridge processes with a weighted mean of tangent directions, avoiding heat-kernel approximation while guiding generation toward the data distribution.Its training scheme uses simple drift regression and is described as highly scalable on general manifolds.Source: §6 · p. 9
  • On general manifolds, the method uses only 5% of the in-training simulation steps used by CNF models.The paper also reports superior performance across diverse manifolds.Source: §6 · p. 9
  • The authors identify protein design and high-energy physics as scientific fields their work may help advance.Source: §6 · p. 9

A. Derivations

The derivations establish bridge processes and combine them into a manifold diffusion mixture, then derive its probability flow and training objective. They also relate the construction to other bridge processes and manifold generative methods.

A.1. Diffusion Process on Riemannian Manifold

  • The appendix introduces Brownian bridges as Brownian motion conditioned on fixed endpoints and describes Langevin dynamics with a stationary distribution determined by a potential.A specified potential yields a wrapped Gaussian stationary distribution.Source: §A.1 · p. 13

A.2. Logarithm Bridge Process

  • The Logarithm bridge is shown to converge to its fixed endpoint and is derived using a change of time; its derivation also provides a more general form of manifold bridge processes.Source: §A.2 · p. 14

A.3. Other Bridge Processes

  • The Logarithm bridge is a Fermi bridge conditioned on a single point, and the paper identifies its use in generative modeling as a novel contribution.Source: §A.3 · p. 14

A.4. Diffusion Mixture on Riemannian Manifold

  • The Riemannian diffusion mixture combines bridge processes to obtain a process whose marginal density is the mixture of their marginal densities.The derivation uses the Fokker–Planck equation to characterize the mixture process.Source: §A.4 · p. 15

A.5. Probability Flow ODE

  • The probability flow is derived from the mixture process and its time-reversed bridges, and its marginal density follows the mixture process’s Fokker–Planck law.In Euclidean space, the derived flow corresponds to the probability flow for Schrödinger bridges.Source: §A.5 · pp. 15–16

A.6. Bridge Matching on Riemannian Manifold

  • Two-way bridge matching trains forward and backward drift models using one simulated bridge, halving simulation cost; time importance sampling is reported to stabilize training and improve generation quality without extra computation or training time.The proposal distribution is described as simple and easy to sample.Source: §A.6 · p. 17

A.7. Other Relevant Works

  • The related-work discussion contrasts divergence-intensive iterative training and geometry-specific methods with the paper’s approach, and leaves Brownian bridges using improved heat-kernel estimators for future work.It also describes stochastic flow matching on SO(3) as a special case of the Logarithm bridge.Source: §A.7 · p. 17

B.1. Implementation Details

The implementation section specifies evaluation settings, model and likelihood computation choices, and the simulation procedure. Its figures and table captions describe visualizations, convergence and ablation comparisons, and protein-dataset test NLL results.

B.1. Implementation Details

  • Experiments follow prior data-split settings, use five runs with seeds 0–4, and divide data into 80% training, 10% validation, and 10% test sets.Test NLL is computed from the checkpoint with the best validation NLL.Source: §B.1 · p. 17
  • The models use multilayer perceptrons with time concatenated to the input, tangent-space projection, and 15 in-training Geodesic Random Walk simulation steps.Most experiments use 512 hidden units; the number of layers varies from 6 to 13.Source: §B.1 · pp. 17–19
  • Figure 4 encodes test samples as red dots and generated samples as green dots, with darker green indicating higher modeled likelihood; Figure 8 shows protein likelihood contours and test samples.Source: Figure 4 · p. 18, Figure 8 · p. 19
  • Figure 5 compares trajectory-to-final geodesic distance and prediction convergence; Figures 6 and 7 describe test-NLL and bridge-sampling ablations, with the latter reporting no additional variance from same-bridge sampling.Source: Figure 5 · p. 18, Figure 6 · p. 18, Figure 7 · p. 18
  • Likelihoods use a Dormand–Prince solver with 1e−5 absolute and relative tolerance, except mesh likelihoods, which use 1,000 Euler steps with projection after each step.Source: §B.1 · pp. 19–20
  • Figure 9 visualizes generated samples and learned density on mesh datasets, while Table 3 reports mean test NLL across five runs on protein datasets and marks best and statistically comparable results.The captions do not provide the table’s individual NLL values.Source: Table 3 · p. 19, Figure 9 · p. 20

B.2. Earth and Climate Science Datasets

The Earth and climate science experiments average results across five data splits and use a uniform sphere prior for baseline comparison. The convergence analysis is conducted on Volcano models and measures geodesic distance over discretized time steps.

B.2. Earth and Climate Science Datasets

  • The experiments average five runs using seeds 0–4 and use a uniform sphere prior for fair baseline comparison.The convergence analysis uses Volcano-trained models and measures geodesic distance between final samples and trajectory states at discretized times, as well as prediction convergence.Source: §B.2 · p. 20

B.3. Protein Datasets

These appendices describe the data representations and experimental setups used for protein, high-dimensional torus, and general closed-manifold evaluations.

B.3. Protein Datasets

  • Protein structures use three torsion angles represented on a 2D torus, while RNA sequences use seven torsion angles represented on a 7D torus.The dataset contains 500 high-resolution proteins and 113 selected RNA sequences; evaluations average five data splits, with a uniform prior on each torus.Source: §B.3 · p. 20

B.4. High-Dimensional Tori

  • High-dimensional torus experiments use wrapped Gaussians with uniformly sampled means and scale 0.2, evaluating dimensions up to 2000.Methods are trained for 50k iterations with batch size 512; the model parameter counts are matched to baselines, and log-likelihood per dimension and training time are measured.Source: §B.4 · pp. 20–21
  • Figure 10 projects learned densities for mixtures of four and six wrapped Gaussians on the hyperboloid onto a Poincare disk.Source: Figure 10 · p. 21

B.5. General Closed Manifolds

  • General closed-manifold experiments use triangular meshes for Spot the Cow and Stanford Bunny, with target densities constructed from thresholded mesh eigenfunctions.The Stanford Bunny mesh has 5000 triangles; the prior is uniform on the mesh, and training time is compared with RFM.Source: §B.5 · p. 21

B.6. Further Analysis

Further analysis examines simulation-step requirements, trajectory comparisons, and a time-scaled training objective, alongside a hyperboloid density visualization.

  • The time-scaled objective is compared with uniform-time training on the Volcano dataset using the earth and climate science experimental setup.Source: §B.6 · p. 21
  • The trajectory analysis compares discretized LogBM samples with 500-step near-exact samples using MMD and Wasserstein distance.Absolute MMD scales across Figure 11 panels (a)–(d) are not directly comparable because they use different reference distributions.Source: §B.6 · p. 21
  • The hyperboloid experiment visualizes learned densities by projecting onto a Poincare disk.It uses synthetic distributions formed from wrapped Gaussian mixtures.Source: §B.6 · p. 21

C. Limitations

The paper identifies limits in theoretical guarantees for non-compact manifolds and potential scaling difficulty on complex manifolds. Its simulation-step analysis also reports that accuracy depends on the approach and noise scale.

  • Theoretical guarantees for constructing the mixture process on non-compact manifolds are insufficient, and complex manifolds may be difficult to scale to when logarithm maps or eigenfunctions are expensive.The paper also reports experimental modeling of non-compact manifolds with negative curvature and scaling to higher-dimensional hyperspheres and tori.Source: C. Limitations · p. 21
  • The simulation analysis compares MMD and Wasserstein distances between discretized LogBM samples and 500-step near-exact samples.Source: Figure 11 · p. 22, C. Limitations · p. 23
  • Fifteen steps nearly match exact simulation for the two-way approach and suffice for Stanford Bunny, but small-noise simulation requires more than 100 steps.The one-way approach needs significantly more steps for accurate trajectories.Source: C. Limitations · p. 22, Figure 12 · p. 23
Loading 2310.07216v2…