[论文解读] Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes

Jaehyeong Jo, Sung Ju Hwang|arXiv (Cornell University)|Oct 11, 2023

Generative Adversarial Networks and Image SynthesisComputer Science被引用 3

一句话总结

本文提出Riemannian Diffusion Mixture（RFM），一种新颖的生成建模框架，该框架在一般Riemann流形上构建随机过程，作为终点条件化的桥接过程的混合，避免了先前扩散模型的时间反演和得分匹配的缺陷。该方法实现了高效、可扩展的训练，显著减少了训练过程中的模拟步数——在包括高维流形和一般闭流形（如双曲面和3D网格）在内的多种流形上实现了最先进性能。

ABSTRACT

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.

研究动机与目标

• 解决现有流形上生成模型的局限性，这些模型依赖于近似热核或计算成本高昂的散度计算。
• 克服基于得分的扩散模型和散度密集型CNF模型在高维和一般流形设置下的可扩展性和几何不一致性问题。
• 开发一种基于原理、几何感知的生成过程，通过Riemann桥接过程的混合，直接建模从先验分布到数据分布的传输。
• 通过引入一种双向桥接匹配目标，避免时间反演，减少模拟步数，实现在一般流形上的高效训练。
• 证明在复杂流形（尤其是高维）上，随机过程相较于确定性流是准确建模所必需的。

提出的方法

• 提出Riemannian扩散混合框架，将生成过程建模为桥接过程的混合——即固定终点条件的扩散过程——而非依赖于时间反演去噪。
• 引入适用于一般流形的可处理Riemann桥接过程族，无需热核近似或难以处理的散度计算。
• 定义一种双向桥接匹配目标，通过匹配前向和后向桥接过程，实现高效、可扩展的训练，适用于任意Riemann几何。
• 通过基于流形几何引导至最可能终点的漂移项，构建从先验分布（例如流形上的均匀分布）到数据分布的随机映射。
• 利用连续归一化流（CNF）与所提框架的特例之间的理论等价性，验证其通用性。
• 使用JAX实现该方法，采用高效的神经网络参数化方式表示漂移向量场，支持高维推理。

实验结果

研究问题

• RQ1 能否在不依赖去噪过程时间反演的前提下，构建流形上的生成模型，从而避免得分匹配和散度计算？
• RQ2 Riemann桥接过程的混合能否为一般流形上的现有扩散模型和CNF模型提供更几何直观且可扩展的替代方案？
• RQ3 所提出的双向桥接匹配目标是否能实现比先前方法更快速、更准确的训练，且训练中模拟步数更少？
• RQ4 Riemannian Diffusion Mixture在包括高维和非紧致几何在内的多种流形上，与最先进基线相比表现如何？
• RQ5 在复杂流形上，准确建模是否必须依赖随机生成过程，还是确定性流已足够？

主要发现

• 所提出的Riemannian Diffusion Mixture（RFM）在包括球面、环面、双曲面和3D网格在内的多种流形上，于合成数据集和真实世界数据集上均优于或匹配最先进基线。
• 在斯坦福野兔和Spot the Cow等一般闭流形上，RFM仅需先前CNF模型所需训练模拟步数的5%，即实现更优性能。
• 该方法在高维场景下具有良好的可扩展性，相比依赖昂贵散度计算的基于得分的扩散模型，训练时间显著缩短。
• 双向桥接匹配目标使Volcano和斯坦福野兔数据集上仅需15次模拟步即可实现准确建模，而单向方法则需显著更多的步数。
• 在双曲面上，RFM成功学习了复杂的多模态分布（例如Wrapped Gaussians的混合），并通过Poincaré圆盘投影可视化了准确的密度估计。
• 消融研究证实，采用所提目标函数的时间加权训练优于均匀时间采样，凸显了目标函数设计选择的重要性。