🤖 machine learning

Riemannian Diffusion Models on General Manifolds via Physics-Informed Neural Networks

本文提出了一种黎曼扩散模型的通用框架，该框架通过利用物理信息神经网络来求解流形热方程，克服了闭式热核不可用的挑战，从而实现了在诸如 $S^2$ 、 $SO(3)$ 和 $\mathrm{SPD}(n)$ 等多种流形上有效的基于评分的生成建模。

原作者： Gyeonghoon Ko, Juho Lee

发布于 2026-06-01

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Gyeonghoon Ko, Juho Lee

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

大局观：教 AI 在曲面上行走

想象一下，你正在教一个机器人画画。在标准 AI（称为“欧几里得扩散”）的世界里，机器人生活在一张平坦且无限大的纸上。它通过慢慢“去噪”的过程，将随机的涂鸦变成清晰的图像。这对于像猫或人脸这样的平面事物非常有效。

但如果数据不是平面的呢？如果数据存在于一个球面上（比如地球上的地震位置）、一个扭曲的环上（比如机器人手臂的 3D 旋转），或者一个复杂的形状上（比如分子的结构）呢？

如果你试图用平面的规则强迫机器人在球面上画画，它会感到困惑。它可能会拉伸球面、撕裂球面，或者创造出不可能存在的形状。为了正确完成这项工作，机器人需要理解它所行走的表面的曲率几何（curved geometry）。

问题所在：“热图”缺失了

为了教机器人如何在这些曲面上移动，研究人员需要一张特殊的地图，叫做热核（Heat Kernel）。

你可以把热核想象成一个**“烟雾扩散器”**。如果你在曲面的某个特定点滴下一滴墨水（或一团烟雾），热核会告诉你这团烟雾随时间推移是如何扩散开来的。

为什么我们需要它？ AI 使用这种“烟雾扩散”过程来计算如何逆转过程：从一团烟雾变回最初的那一滴墨水（原始数据）。
难点在于： 对于简单的形状（如平面），我们有完美的公式来描述这种烟雾扩散。但对于复杂的、弯曲的形状（如球面或扭曲的甜甜圈），我们没有公式。 数学计算太难了。以往的方法试图使用捷径来猜测这个公式，但这些捷径只适用于非常特定的简单形状。

解决方案：“物理侦探”（PINN）

作者提出了一种聪明的全新解决方法。他们不再试图寻找一个神奇的公式来描述烟雾扩散，而是将其视为一个物理谜题。

他们知道烟雾扩散的规则（一个著名的物理方程，称为热方程）。他们也知道烟雾在最开始时的样子（一个微小的、集中的点）。

因此，他们使用了一个物理信息神经网络（PINN）。你可以把它想象成一个超级聪明的侦探，通过直接解决谜题来进行学习：

设定： 他们告诉侦探：“这是规则手册（热方程）。这是起点（一小点烟雾）。”
训练： 侦探尝试猜测在中间每一时刻烟雾的样子。如果侦探的猜测违反了物理规则，它就会受到“惩罚”并重新尝试。
结果： 最终，侦探学会了如何预测烟雾在任何曲面上的扩散方式，即使之前从未有人为它写下过公式。

实际应用中的运作方式

论文展示了这个“侦探”在几种不同的“世界”中工作的场景：

地球（球面 $S^2$ ）： 他们模拟了火山和地震等自然灾害的扩散。AI 学会了生成现实的地图，展示这些事件可能发生的位置，同时遵循地球的圆形形状。
3D 旋转器 ($SO(3)$）： 他们模拟了物体在 3D 空间中的旋转。这对于机器人技术和 3D 建模至关重要。
交通流 ($SPD$ 矩阵)： 他们模拟了纽约市的交通模式。这里的数据不仅仅是点；它们是不同城市区域之间复杂的相互关系。AI 学会了生成真实的交通流模式。
分子谜题 ( $R^{k \times n}/S_n$ )： 这是最复杂的部分。分子由原子组成，但原子的排列顺序并不重要（无论你先列出原子 A 还是 B，分子都是同一个）。AI 通过将分子视为一种“点云”，且顺序无关紧要，从而学会了生成新的分子，有效地解决了置换谜题。

结果：效果如何？

研究人员将他们的“侦探”与其它方法进行了对比测试：

准确性： 在已知答案的形状（如球面）上，他们的方法与现有的最佳方法一样准确，但不需要预先写好的公式。
多功能性： 它在其他方法因数学难度过大而无法处理的形状上也能正常工作。
质量： 生成的数据（分子、交通图、地震位置）看起来非常真实，并且遵循正确的几何规则。

局限性

论文也诚实地说明了缺点：

计算量大： 训练这个“侦探”需要大量的计算能力和时间。这就像是每次想要生成数据时，都要进行一次大规模的物理模拟。
复杂性高： 为一个新形状设定“规则手册”需要具备数学专业知识。它目前还不是一个点击即可使用的“即插即用”按钮。
权衡取舍： 在分子实验中，AI 生成的分子在化学上是非常正确的，但不如其他方法那样具有“新颖性”或“独特性”。这是一种在保证正确性与保持创造力之间的权衡。

总结

简而言之，这篇论文为 AI 提供了一个在曲面上行走的工具。AI 不再需要一张预先写好的地图（对于大多数形状来说，这类地图并不存在），而是通过求解物理扩散方程，在运行过程中实时学习地图。这使得 AI 能够生成从全球气候事件到 3D 分子的各种现实数据，而不会破坏几何规则。

技术摘要：基于物理信息神经网络（PINN）的通用流形黎曼扩散模型

1. 问题陈述

黎曼扩散模型通过定义基于黎曼布朗运动（Riemannian Brownian motion）的前向扩散过程，将基于评分的生成模型推广到存在于光滑流形上的数据。训练和采样这些模型的一个关键瓶颈是需要访问流形热核（heat kernel，即布朗运动的转移密度）及其梯度。具体而言，训练目标（去噪评分匹配，Denoising Score Matching）需要评估条件评分 $\nabla_x \log p_{t|0}(x|x_0)$ ，而前向采样则需要从热核中抽样。

虽然热核在高度对称的流形（如欧几里得空间、球面、环面）上具有闭式解，但对于一般的黎曼流形，这种闭式解极少存在。现有的近似方法（如短时渐近展开 Varadhan 公式或谱展开）通常受限于特定区间（如小时间 $t$ ）或需要较强的结构性假设（如紧致李群），这限制了它们在多样化几何环境中的实际应用能力。

2. 方法论

作者提出了一种通用的框架，通过直接使用**物理信息神经网络（PINN）**求解流形热方程来近似热核。该方法通过以下四个阶段进行：

A. 坐标系选择与嵌入
该方法并非使用局部坐标图（local coordinate charts）进行拼接，而是将目标流形 $M$ 嵌入到一个更大的环境流形 $\tilde{M}$ （通常为 $\mathbb{R}^N$ ）中，将其作为子流形或商流形。

子流形（Submanifolds）： 由约束条件 $f_i(x)=0$ 定义。度量和微分算子通过将环境空间映射到 $M$ 的切空间的投影矩阵 $P(x)$ 来导出。
商流形（Quotient Manifolds）： 通过在群作用下识别点来定义（例如置换对称性）。该方法将轨道点视为相同的状态，通常使用规范代表元。

B. 热方程的推导
利用环境坐标和投影矩阵，作者推导出了拉普拉斯-贝尔特拉米算子 $\Delta_M$ 以及相应的福克-普朗克（Fokker-Planck，即热）方程的显式形式：
$\partial_t p_{t|0}(x|x_0) = \Delta_M p_{t|0}(x|x_0)$
该方程完全基于环境度量和投影矩阵进行表达，从而实现了无需内在坐标图的基于坐标的计算。

C. 初始条件近似
由于无法在数值上施加 Dirac $\delta$ 初始条件 $p_{0|0} = \delta_{x_0}$ ，作者使用基于原理的短时渐近近似在微小时间 $t_0 > 0$ 时对 PINN 进行初始化。他们利用参数法（Minakshisundaram–Pleijel 递归）来获得热核的高阶展开：
$p_{t_0}(x|x_0) \approx G_t^{(n)}(x, x_0) = G_t(x, x_0) \left( u_0 + u_1 t + \dots + u_{n-1} t^{n-1} \right)$
其中 $G_t$ 是类高斯项，而 $u_k$ 是由几何性质（如体积畸变和曲率）导出的系数。

D. PINN 训练
训练一个神经网络 $\phi_\theta(t, x)$ 来近似对数热核 $\log p_{t|0}$ 。训练目标最小化以下各项的组合：

初始/边界损失： 网络输出与在 $t=t_0$ 处的渐近初始化值之间的差异。
PDE 残差损失： 热方程 $\partial_t \phi - \Delta_M \phi - \|\nabla \phi\|^2 = 0$ （针对对数密度重写后）的残差。

作者采用了自适应损失权重和课程学习（逐步扩大时间间隔）来确保稳定性。

E. 集成至扩散模型
训练好的 PINN 作为热核的代理，实现两个关键操作：

前向采样： 通过 MCMC 方法（如测地线随机游走）利用 PINN 评估的密度生成样本。
评分评估： 通过对 PINN 求导来计算条件评分 $\nabla_x \log p_{t|0}$ ，进而通过去噪评分匹配（DSM）训练逆向时间扩散模型。

3. 核心贡献

通用的基于 PDE 的代理模型： 本文引入了一种通过数值求解热方程来近似任意流形热核的方法，绕过了对闭式解或限制性谱假设的需求。
无坐标依赖的推导： 该框架提供了一种系统的方法，用于推导定义为欧几里得空间子流形或商流形的流形的热方程，使其适用于复杂的几何体，如 $S^2$ 、$SO(3) $、$ SPD(n)$ 和置换商化的点云。
稳定性分析： 作者提供了理论证明（定理 4.1），表明从近似初始条件求解热方程所得到的解在所有 $t \ge t_0$ 时误差是有界的，从而确保了 PINN 代理的可靠性。
置换不变性扩散： 该方法被应用于商流形 $\mathbb{R}^{k \times n}/S_n$ ，实现了能够显式消除原子置换影响的扩散过程，这在标准欧几里得扩散无法尊重结构对称性的场景中至关重要。

4. 实验结果

该方法在多种流形和数据集上进行了评估：

球面数据 ( $S^2$ )： 对地学事件（火山、地震、洪水、野火）进行密度估计。所提方法取得了与依赖 Varadhan 近似和谱方法的黎曼评分生成模型（RSGM）相比，具有竞争优势或更优的测试对数似然值。
旋转数据 ($SO(3)$)： 合成混合模型显示出比 RSGM 基准更高的对数似然值。
正定矩阵 ($SPD(n)$)：
- 交通分析 ($SPD(10)$)： 从纽约出租车数据中进行交通流矩阵的条件生成。与 SPD-DDPM 相比，该方法降低了与真实值的 Frobenius 距离。
- 脑连接性 ($SPD(n)$)： EEG 连接矩阵的条件生成。该方法取得了极高的精确度和 F1 分数，足以与使用欧几里得空间微分同胚映射的 DiffeoCFM 竞争，并优于 LowTriDDPM。
分子生成 ( $\mathbb{R}^{k \times n}/S_n$ )： 在置换商空间上生成分子。该方法生成的分子比方差爆炸型（VE）EDM 基准具有更高的有效性和稳定性，尽管新颖性较低，反映了已知的有效性-新颖性权衡。

验证： 作者通过将学习到的对数密度和评分与参考近似（在可用时）进行比较，并报告归一化的 PDE 残差和边界误差，验证了 PIN-Surrogates 满足热方程。

5. 重要性与主张

本文声称，基于 PDE 的热核近似为黎曼扩散提供了一种实用且通用的替代方案。通过将热核视为已知 PDE 的解，该方法将无漂移黎曼扩散扩展到了不存在解析热核的流形上。

作者将这项工作定位为一种“求解热方程”的视角，拓宽了流形扩散模型的适用范围。然而，他们也谦虚地承认了局限性：

计算成本： 由于高阶自动微分和配置点评估，PINN 训练需要大量资源。
稳定性： 在高维或复杂几何结构下，训练可能不稳定，可能需要更强的归纳偏置。
复杂度： 该流程比完全不需要热核估计的 Flow Matching 等方法更为复杂。

这项工作并不声称要取代 Flow Matching，而是旨在为评分模型提供必要的后端支持——特别是在 Flow Matching 对测地线插值或指数映射的依赖可能不足，或者需要特定热核结构的流形场景下。