[论文解读] Scaling Riemannian Diffusion Models

Aaron Lou, Minkai Xu+1 more|arXiv (Cornell University)|Oct 30, 2023

Machine Learning in Healthcare被引用 4

一句话总结

该论文提出了一种可扩展的黎曼扩散模型，利用黎曼对称空间实现热核量的高精度计算，克服了分数匹配中先前难以处理的近似问题。通过精确计算热核并利用概率流ODE实现精确采样，该方法在高维流形上实现了最先进性能，包括$SU(3)$晶格（128维）和超球面对比嵌入（127维），首次成功将基于微分方程的生成模型应用于数百维的真实世界任务。

ABSTRACT

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.

研究动机与目标

• 解决由于复杂几何结构导致的黎曼扩散模型中去噪分数匹配的不可行性问题。
• 实现在高维流形（尤其是与物理学和表征学习相关的流形）上可扩展且精确的黎曼扩散模型训练。
• 克服先前分数匹配中的近似方法，这些方法会降低性能并限制向高维的可扩展性。
• 证明基于微分方程的生成建模在真实世界高维流形（如$SU(n)$晶格和超球面对嵌入）上的可行性。
• 通过建模真实的超球面分布而非依赖倒数第二层特征，提升对比学习中的模型可解释性和性能。

提出的方法

• 利用黎曼对称空间实现热核的高效且精确计算，这对去噪分数匹配至关重要。
• 将已知的热核技术适配到黎曼分数匹配的具体情境中，实现快速且稳定的训练。
• 通过概率流ODE开发了从热核精确采样的方法，可在流形的最大阿贝尔子群上使用标准ODE求解器求解。
• 利用最大阿贝尔子群降低计算复杂度，并实现在对称空间上的高效无模拟采样。
• 将精确的热核估计器集成到去噪分数匹配目标中，以提升训练稳定性和性能。
• 将该方法应用于高维任务，如在$4\times4$ $SU(3)$晶格上学习威尔逊作用量，以及在$S^{127}$上建模对比学习嵌入。

(a) $S^{2}$

实验结果

研究问题

• RQ1 黎曼扩散模型能否扩展到高维流形（如$SU(n)$晶格和高维超球面）？
• RQ2 在对称空间上精确的热核计算能否消除分数匹配中不准确近似的需要？
• RQ3 通过概率流ODE实现的精确采样是否能提升流形上的模型性能和训练稳定性？
• RQ4 黎曼扩散模型能否从对比学习得到的超球面对嵌入中恢复出忠实的分布，从而提升可解释性和分布外检测性能？
• RQ5 是否可能通过基于微分方程的方法在非平凡流形上实现高维生成建模的最先进性能？

主要发现

• 在$4\times4$ $SU(3)$晶格上学习威尔逊作用量时，所提方法的有效样本量达到0.62，该流形为128维的高维流形。
• 在127维超球面上的对比学习嵌入中，黎曼扩散模型在分布外检测性能上与或优于使用倒数第二层特征的最先进方法。
• 该模型成功学习了平面上的复杂棋盘分布，性能优于RFM和流形匹配方法，后者在复杂度增加时性能下降。
• 若缺乏精确的热核估计器，去噪分数匹配会产生次优结果，证明了精确几何计算的必要性。
• 该方法首次成功将基于微分方程的生成模型应用于真实世界、数百维的高维任务。
• 该模型恢复了对比学习中坍塌的投影头表示，通过建模真实的超球面分布提升了可解释性。

(b) $S^{127}$