扩散模型之（七）Manifold 视角下 Flow Matching

0.概述在上一篇中主要介绍了flow matching的数学原理本篇将从流形manifold的角度进一步理解flow matching重点参考论文 Riemannian Flow Matching on General Geometries.在前面的博客中我们主要研究了欧几里得空间的流匹配。然而现实世界的数据例如全球天气数据球面上、蛋白质结构具有旋转和平移对称性的空间或机器人关节角度环面上都具有复杂的几何结构流形无法在欧几里得空间中自然地表示。本文将探讨如何将强大的流匹配框架扩展到这些非欧几里得空间特别是黎曼流形。本文将介绍流形上的概率分布、流和速度场等概念并探讨如何推广欧几里得空间中的流匹配思想。具体而言本文将重点关注如何利用测地线流形上的“直线”构建条件流以及如何利用更一般的预度量来设计流旨在帮助解流匹配的理论及其在具有复杂几何约束的数据中的应用可能性。1. 黎曼流形弯曲空间的几何首先简要地了解一下黎曼流形它是扩展流匹配的基础。流形局部看起来像欧几里得空间的空间。例如地球表面一个球体看起来像我们周围的一个二维平面但从整体上看它是一个弯曲的二维流形。黎曼度量它是对流形上每个点的“切空间”一种与流形在该点相切的欧几里得空间上向量的长度和角度进行“内积”定义的一种光滑方法。黎曼度量使我们能够研究流形上的几何量例如曲线、距离和曲率。黎曼流形配备黎曼度量的流形。流形 (): 图像背景中的灰色曲面代表一个流形这是一个局部看起来像欧几里得空间的拓扑空间。图1. 流形示意图黎曼流形上的流匹配相应地需要扩展原始flow matching中的源分布以及目标分布将其视为流匹配内容的转化形式来学习。2. 流形上的概率场、流场和速度场将欧几里得空间的概念推广到流形上。概率密度对于流形上的概率密度函数为具有非负性并且在流形上的整个积分体积元是一个返回值为 1 的函数.流随时间变化的流 实现的变换也就是各时刻t从流形到流形的微分同胚映射。速度场速度场定义为时刻t为每个点在该点的切空间分配一个属于该空间的一个向量即.流形上的一个常微分方程的解也就是满足2.1 流形上的连续性方程和瞬时变量变换在上一篇博客中看到的连续性方程和瞬时变量变换公式也可以自然地推广到黎曼流形上。黎曼连续性方程MetaAI 式5.1其中div是黎曼度量关于散度的算子.黎曼瞬时变量变换MetaAI 式5.3这些公式在形式上与欧几里得空间中的公式相同但散度算子被替换为考虑黎曼度量的算子。这使得能够从理论上计算流形上流模型的似然性并处理速度场和概率路径之间的关系。2. 2 流形上的概率路径和边缘化技巧流匹配的核心——随机路径构造和边缘化技巧也可以扩展到黎曼流形上。流形上的随机路径MetaAI 式 5.6边缘概率路径是通过将条件概率路径按目标分布取平均值而构建的结果。条件路径即边界条件满足如下要求MetaAI 式 5.8​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​流形上的边缘化技巧定理 10与欧几里得空间的情况类似在适当的正则性假设假设 2下​​​​​​​条件概率路径生成的是条件速度场.如果对应的是边缘概率路径生成的是边缘速度场​​​​​​​ ​​​​​​​ ​​​​​​​这证明了即使在流形上针对条件量的学习也是合理的。3. 黎曼流匹配损失Meta AI 5.5节用于训练的损失函数也考虑了流形的结构。黎曼条件流匹配RCFM的损失函数是的切空间上定义的Bregman散度对应的期望值利用MetaAI 式 5.13其具体定义如下也可以表示为其中与都是切空间上定义的向量。Bregman 散度的向量它可以使用黎曼度量导出的平方范数来定义。具体的定义如下​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​ ​类似地在欧几里得空间中该 RCFM 损失的梯度与针对边缘量的黎曼流匹配 (RFM) 损失的梯度一致MetaAI 定理 11。3.1 流形上的条件流测地线和预度量设计流匹配的关键在于创建一个既易于处理又有效的条件流​​​​​​​这将决定相应的条件速度场。3.1.1 欧式空间距离在欧式空间Euclidean space中3.1.1 测地线流在黎曼流形中与欧几里得空间中的“直线”等价的概念是“测地线”。测地线给出了两点之间局部最短路径。利用这条测地线我们可以构造条件流。首先考虑流形上的两个点, 考虑连接两点的唯一或主测地线如下图表示点和之间的连接即对应的测地线Geodesics。​​​​​​​ ​​​​​​​ ​​​​​​​图2. 测地线的示意图指数映射点​ 的切空间向量,初始速度是处的速度这个映射它将测地线上的一个点映射到从一点开始向前移动一个时间步的点.对数映射表示将其映射到测地线的初始切向量的映射指数映射的局部逆映射)。​ 是表示从到的方向和距离的向量。闭式测地线Closed form geodesics的角度利用这些测地线条件流可以定义如下Meta AI式 5.17​​​​​​​ ​​​​​​​其中是一个单调递增的调度函数满足这意味着沿着从到的测地线上的路径是按照时间参数正向前进的。这是欧几里得空间中的线性插值的自然推广​​​​​​​ ​​​​​​​ ​​​​​​​事实上在欧几里得空间中,这种测地线流允许在简单流形例如球面、双曲空间上以闭合形式计算指数对数映射时进行无模拟流匹配这与流形上的扩散模型形成对比后者通常需要复杂的分数计算和 SDE 模拟实现。3.1.2 带有预度量的流然而对于一般的流形指数对数映射可能难以计算测地线可能不唯一例如穿过割线或者测地线可能不是所需的路径例如数据可能以满足某些约束的方式嵌入流形中但测地线可能违反这些约束。因此为了更好地解决问题有人提出了一种更灵活的方法即引入一个一般的类距离函数(称为预度量)来代替测地距离并设计条件流使其能够保持相对于该预度量的距离(Chenand Lipman,2024。MetaAI 5.6)具体而言满足以下条件MetaAI 式 5.18。​​​​​​​ ​​​​​​​其中是一个单调递减函数满足比如),如下图​​​​​​​ ​​​​​​​ ​​​​​​​ ​​​​​​​满足此条件的流动的最小范数条件速度场由下式给出MetaAi 式 5.19​​​​​​​ ​​​​​​​其中对于所有条件下都有.这种方法的优势在于可以选择除测地距离之外的其他距离函数。例如使用谱距离或神经网络学习到的距离函数可以设计出更适合数据的流。然而在这种情况下速度场取决于预度量的梯度因此预度量的设计就显得至关重要。此外通常无法显式计算这种流的路径可能需要使用速度场进行常微分方程ODE模拟。​​​​​​​ ​​​​​​​3.2 测地线流问题及解决方案测地线流和准度量流可能会出现奇异性问题。例如球面上的测地距离在计算对应的对向点时不可微。此外对于一般的光滑函数即使梯度存在临界点例如最大值点和最小值点在这些临界点上速度变为零上述速度场也无法定义。为解决这些问题可以考虑采取以下措施调整调度函数调整或者不仅依赖于 t还依赖,从而平滑地调整流以避免出现奇点MetaAI P39。非退化条件的放宽实际上如果预度量的梯度为零的点集的测度为零通常不会有问题MetaAI P39 页“非退化放宽”。在MetaAI Code 8 中给出了一个使用球面上的测地线流进行流匹配的训练示例其中涉及将模型输出的速度矢量投影到当前点的切空间上这是在流形上处理速度场时的一个重要实际步骤。在本文中我们实现了如何将流匹配框架从欧几里得空间扩展到黎曼流形。黎曼流形是局部欧几里得空间但整体弯曲的空间其几何结构由黎曼度量定义。即使在流形上概率分布、流动、速度场、连续性方程和瞬时变量变换等概念也自然地被定义并且流动匹配的理论基础得以保持。边缘化技巧也适用于流形证明了针对条件量RCFM 损失的学习的合理性。在流形上构造条件流的主要方法是测地线流沿着测地线流形上的“直线”移动点。如果指数映射和对数映射可计算则无需模拟的学习也是可能的。预度量流通用距离函数并设计速度场使流保持该距离。这种方法更灵活但可能需要速度场计算或常微分方程模拟。流形特有的挑战例如奇点的存在但也正在研究处理这些挑战的实用方法。该扩展使流动匹配成为一种强大而灵活的生成建模工具可以应用于更广泛的数据特别是具有复杂几何结构的数据预计在蛋白质结构生成、分子动力学模拟、机器人和计算机图形学等各个领域都有应用。4.参考资料Flow Matching Guide and Codehttps://zenn.dev/doctorin/articles/flow-matchinghttps://neurips.cc/media/neurips-2024/Slides/99531.pdfhttps://arxiv.org/pdf/2506.02070https://www.emergentmind.com/topics/riemannian-flow-matchingAn Introduction to Flow Matching and Diffusion Modelshttps://www.bilibili.com/video/BV1Zm421u7Nt/?spm_id_from333.1391.0.0vd_source87b4c7d8e5af5521141c2285ecd78702