跳转至

Total Variation Rates for Riemannian Flow Matching

会议: NeurIPS 2026
arXiv: 2602.05174
OpenReview: https://openreview.net/forum?id=TFfnvGlQ2C
领域: 图像生成
关键词: 黎曼流匹配、总变差收敛率、非渐近分析、流形常微分方程、欧拉离散化

一句话总结

本文首次建立了流形上流匹配(RFM)在连续时间ODE及离散欧拉采样下的非渐近总变差(TV)收敛率界,明确解耦了数值离散误差与速度场学习误差,并在紧流形(超球面)与非紧Hadamard流形(SPD流形)上证明了多项式迭代复杂度。

研究背景与动机

流匹配(Flow Matching, FM)通过构建连接简单基分布与复杂目标数据分布的测度插值,并训练神经网络拟合驱动该输运过程的时间依赖切向量场,已成为生成建模领域的重要范式。在地球物理、机器人构型空间、生物大分子骨架及脑连接网络等应用中,数据天然分布在黎曼流形上。为适应此类几何先验,黎曼流匹配(RFM)通过测地线插值与内在流形常微分方程(ODE)推进粒子演化,天生保持流形几何结构。更关键的是,与基于黎曼布朗运动的黎曼扩散模型相比,RFM 是纯确定性的动力学系统,无需模拟缺乏解析转移核的流形布朗运动,从根本上绕开了复杂的局部扩散近似瓶颈。

然而,尽管 RFM 在经验应用中取得了显著成功,其在分布度量下的非渐近收敛性理论此前仍处于完全空白状态。在欧氏空间中,扩散概率流与流匹配的 TV 误差收敛性已逐步建立起多项式依赖的理论框架;但在一般黎曼流形上,不同位置的切空间不再天然等同,连续与离散轨迹的切向量无法跨点直接相减,并行输运、曲率张量与测地线共轭点/割迹(Cut Locus)引发的体积畸变与导数奇点交织在一起。同时,条件流匹配向量场在终点时刻伴随明显的奇异因子 \(1/(1-t)\),轨迹偏离会导致向量场急剧刚性化。这一系列几何与动力学障碍,使得传统欧氏误差分析方法完全失效。

为克服上述几何分析障碍,本文建立了基于连续性方程扰动分析的内蕴流形输运微分不等式,将采样误差严格分解为切空间速度场偏差与参考流 score 函数的内积项,以及速度场偏差的流形散度项。核心 idea:通过沿欧拉插值测地线构造平行输运诱导的内蕴连续时间 ODE,并借助 Jacobi 场变分方程与曲率比较定理严密控制逆指数映射的局部反转条件与发散度畸变,首次给出了显式分离步长 \(h\) 与学习精度 \(\varepsilon\) 的全多项式 TV 采样误差界。

方法详解

整体框架

本文针对黎曼流形 \((M, g)\) 上的流匹配生成过程建立端到端的非渐近理论分析框架。系统输入为基分布 \(\pi_0\) 与数据分布 \(\pi_1\),通过测地线桥 \(X_t = \mathrm{Exp}_{X_0}(t \mathrm{Log}_{X_0}(X_1))\) 定义插值过程,其总体极小化向量场为条件期望 \(v(t,x) = \frac{1}{1-t}\mathbb{E}[\mathrm{Log}_x(X_1) \mid X_t=x]\)。实际采样时,算法使用神经网络学习的速度场 \(\hat{v}(t,x)\) 并结合流形欧拉积分步进 \(x_{t_{k+1}} = \mathrm{Exp}_{x_{t_k}}(h \hat{v}(t_k, x_{t_k}))\)。为度量真实分布 \(\pi_t\) 与采样分布 \(\hat{\pi}_t\) 之间的总变差距离,理论分析沿着连续时间插值轨迹推进,通过积分曲率敏感的微分不等式推导最终收敛率。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入:先验分布 π₀ 与数据分布 π₁<br/>设定测地线插值路径"] --> B["测地线插值与总体流匹配向量场<br/>v(t,x) = (1-t)⁻¹ 𝔼[Log_x(X₁) | X_t=x]"]
    B --> C["平行输运连续时间插值 ODE<br/>dy_t = P_{y_k}^{y_t} v̂(t_k, y_k) dt"]
    C --> D["流形输运总变差演化微分不等式<br/>∂_t TV(X_t, Y_t) ≤ 𝔼[|div(ṽ - v)|] + 𝔼[‖grad log p_t‖ · ‖ṽ - v‖]"]
    D --> E["曲率驱动的散度畸变与截断控制<br/>结合 Jacobi 场比较定理与 score 正则性积分"]
    E --> F["输出:非渐近 TV 收敛边界<br/>TV ≤ C_Lip · h + C_eps · ε (紧流形/非紧流形)"]

关键设计

1. 平行输运诱导的连续时间欧拉插值与可逆性条件:解决离散步与连续输运在切空间上的错位

在离散流形欧拉格式中,更新步由指数映射给出:\(y_{t_{k+1}} = \mathrm{Exp}_{y_k}(h \hat{v}(t_k, y_k))\)。若直接冻结切向量,速度属于 \(T_{y_k}M\),而连续时间插值轨迹上的即时速度必须属于当前流形点切空间 \(T_{y_t}M\)。针对这一矛盾,本文设计了沿插值测地线的平行输运连续场: $\(\tilde{v}(t, y_t) := P_{y_k}^{y_t} \hat{v}(t_k, y_k), \quad \text{其中} \quad y_t = \mathrm{Exp}_{y_k}((t-t_k)\hat{v}(t_k, y_k)).\)$ 为了将 \(\tilde{v}\) 转化为全流形上的良定全局时变向量场 \(\tilde{v}(t,x)\),必须能从当前位置 \(x\) 唯一回溯其前序生成足点 \(x_k\)。为此,定义映射 \(F_{t_k, t-t_k}(z) := \mathrm{Exp}_z((t-t_k)\hat{v}(t_k, z))\)。本文利用 Jacobi 场沿测地线的变分动力学方程 \(\ddot{Y}(t) + P_{\gamma(t)}^{\gamma(0)} R(P_{\gamma(0)}^{\gamma(t)}Y(t), \dot{\gamma})\dot{\gamma} = 0\),并结合 Rauch 比较定理证明:当步长 \(h\) 严格受限于流形单射半径 \(\mathrm{inj}(M)\)、速度场协变导数界 \(L_\nabla\) 以及曲率界 \(L_R, K_{\min}\) 时,\(dF_{t_k, h}(x)\) 的最小奇异值严格大于零,保证了 \(F_{t_k, h}\) 的全局微分同胚性与可逆性,从而确立了 \(\tilde{v}(t,x) = P_{F_{t_k, t-t_k}^{-1}(x)}^x \hat{v}(t_k, F_{t_k, t-t_k}^{-1}(x))\) 的严格数学良定性。

2. 黎曼连续性方程下的总变差演化微分不等式:将分布距离直接锚定到速度偏差与几何量

为突破流形上无法直接做欧氏差值相减的拓扑障碍,本文拓展了 Marsden 输运定理与流形连续性方程 \(\partial_t p_t + \mathrm{div}_g(p_t v_t) = 0\)。考虑两个独立过程 \(X_t\)(由真实场 \(v\) 驱动)与 \(Y_t\)(由插值场 \(\tilde{v}\) 驱动),设 \(\Omega_t := \{x \in M : p(t,x) - q(t,x) > 0\}\)。通过对符号差积分沿区域边界与内部的微分展开,边界项利用 \(p=q\) 自然消去,内部项借助黎曼散度乘积法则 \(\mathrm{div}_g(f W) = f \mathrm{div}_g W + \langle \mathrm{grad} f, W \rangle\) 严格转化为: $\(\frac{\partial \mathrm{TV}(X_t, Y_t)}{\partial t} = \int_{\Omega_t} p(t,x) \Big[ \mathrm{div}_g(\tilde{v} - v) + \langle \mathrm{grad}_g \log p_t(x), \tilde{v} - v \rangle \Big] dV_g(x).\)$ 利用 Cauchy-Schwarz 不等式,该微分界进一步收敛为三项核心几何积分:黎曼 Fisher 信息项(score 正则性)\(\mathbb{E}[\|\mathrm{grad}_g \log p_t(X_t)\|^2]\)、速度场切空间偏差二阶矩 \(\mathbb{E}[\|\tilde{v}(t, X_t) - v(t, X_t)\|^2]\) 以及流形散度差期望 \(\mathbb{E}[|\mathrm{div}_g(\tilde{v} - v)|]\)。这为后续将离散误差与统计估计误差彻底解耦奠定了核心工具。

3. 平行输运微分与曲率张量耦合的散度畸变界定:显式刻画内在弯曲对测度演化的几何效应

散度项 \(\mathbb{E}[|\mathrm{div}_g(\tilde{v} - v)|]\) 是流形设置下最具挑战性的几何部分。欧氏空间中插值速度场的散度仅依赖于空间差分,但在流形上,\(\tilde{v}\) 的协变导数不仅包含前序点处 \(\hat{v}\) 的导数,还包含对平行输运算子沿测地线变分所诱导的曲率累积。本文利用正交基场沿变分测地线面展开,推导出平行输运沿基向量的协变导数基本恒等式: $\(D_s P_{c(0)}^{c(h)} V(z) = P_{c(0)}^{c(h)} D_s V(z) - P_{c(0)}^{c(h)} \int_0^h P_{c(\tau)}^{c(0)} R(\partial_s \Lambda(0,\tau), \partial_t \Lambda(0,\tau)) P_{c(0)}^{c(\tau)} V(z) d\tau.\)$ 在迹运算求散度时,曲率张量与测地线切向量的耦合项贡献了 \(L_R (t - t_k) d \|\hat{v}(t_k, z)\|^2\) 形式的内蕴畸变。结合截断点处速度场的 Lipschitz 性与时间变化率界,成功将散度差精确控制在 \(\varepsilon + (t - t_k) C\) 之内,使得曲率界 \(L_R\) 和维数 \(d\) 显式且紧致地进入了最终 TV 收敛率的常数项中。

4. 针对非紧 Hadamard 流形的指数矩补偿机制:克服负曲率发散与无界域积分

紧流形(如超球面 \(\mathcal{S}^d\))由于单射半径下界良好且切向量有界,易于建立一致性导数界;但在非紧非正曲率的 Hadamard 流形(如对偶空间对称正定矩阵流形 \(\mathrm{SPD}(n)\))上,截面曲率下界 \(K_{\min} < 0\) 导致模型函数 \(s_{K_{\min}}(r) = \kappa^{-1} \sinh(\kappa r)\) 随距离 \(r\) 指数级爆炸,且点集测地距离无界。为此,本文设计了基于流形高斯先验 \(p_0(x) \propto \exp(-d \cdot d_g(x, I)^2)\) 的期望正则性分析体系,并要求目标数据分布满足指数矩条件 \(\mathbb{E}[\exp(\lambda_1 d_g(X_1, I))] \le M_{\lambda_1}\)(其中 \(\lambda_1 = 24 \max\{1, \kappa\}\))。通过在切空间利用欧氏极坐标变分分割积分核,证明了即使在负曲率指数发散效应下,高斯先验的超二次衰减项仍能强力压制曲率增长,从而首次在非紧流形上确立了期望意义下的多项式 TV 采样界。

实验关键数据

主实验

论文对超球面 \(\mathcal{S}^d\)(紧流形代表)与对偶仿射不变度量下的对称正定矩阵流形 \(\mathrm{SPD}(n)\)(非紧 Hadamard 流形代表,流形内在维数 \(d = n(n+1)/2\))进行了详尽的理论复杂度推导。设定目标采样精度为 \(\varepsilon_{\mathrm{target}}\),提早停止截断时间为 \(T = 1 - \delta\)(\(\delta > 0\)),向量场估计误差 \(\varepsilon\) 满足充分小精度要求。

目标流形类别 测试流形 步长策略 (Step-size Schedule) TV 采样误差界形式 达到 \(\varepsilon_{\mathrm{target}}\) 精度所需的迭代复杂度 \(N\) 关键几何常数依赖
紧流形 (Compact) 超球面 \(\mathcal{S}^d\) (\(d \ge 4\)) 常步长 \(h \equiv \mathrm{const}\) \(\mathcal{O}\big(h \frac{d^2}{(1-T)^2} + \varepsilon \frac{d}{1-T} \log \frac{1}{1-T}\big)\) \(\mathcal{O}\left(\frac{d^2}{\varepsilon_{\mathrm{target}} (1-T)^2}\right)\) 截面曲率 \(K \equiv 1\),\(L_R = 1\),直径 \(\pi\)
紧流形 (Compact) 超球面 \(\mathcal{S}^d\) (\(d \ge 4\)) 几何递减步长 \(t_i = 1 - (1 + \eta i)^{-2}\) \(\mathcal{O}\big(\eta \frac{d^2}{\sqrt{1-T}} + \varepsilon \frac{d}{1-T} \log \frac{1}{1-T}\big)\) \(\mathcal{O}\left(\frac{d^2}{\varepsilon_{\mathrm{target}} (1-T)}\right)\) 依赖密度界比例 \((M_1 / m_1)^{3/2}\)
非紧流形 (Hadamard) 对称正定矩阵 \(\mathrm{SPD}(n)\) 常步长 \(h \equiv \mathrm{const}\) \(\mathcal{O}\big(h \frac{d^{24} L_R^3 M_{\lambda_1}^{3/2}}{(1-T)^3} + \varepsilon \frac{d^7 L_R M_{\lambda_1}^{1/2}}{1-T}\big)\) \(\mathcal{O}\left(\frac{d^{24} L_R^3 M_{\lambda_1}^{3/2}}{\varepsilon_{\mathrm{target}} (1-T)^3}\right)\) 截面曲率 \(K_{\min} = -1/2\),依赖指数矩 \(M_{\lambda_1}\)
扩散基准 (Prior Baseline) 黎曼得分扩散 (RSBGM) 离散布朗运动模拟 定性多项式 (Qualitative Poly) 缺乏显式参数依赖 / 非多项式收敛 依赖流形热核离散近似

消融实验

下表给出了在超球面 \(\mathcal{S}^d\) 上步长调节策略、曲率比较项及割迹(Cut Locus)平滑性处理对理论收敛阶与复杂度界限的定性与定量分析:

理论机制配置 (Theoretical Config) 迭代步数对 \((1-T)\) 的阶数依赖 散度估计项控制机制 逆映射可逆性步长约束要求 \(h_{\max}\) 理论评价与说明
完整模型 (几何非均匀步长) \(\mathcal{O}\big((1-T)^{-1}\big)\) 平行输运曲率抵消 + 散度梯度控制 \(\mathcal{O}\big((1-T) / (d \cdot M_1/m_1)\big)\) 最优方案:时间调度在终点前加密,精确抵消 \(1/(1-t)^3\) 刚性累积
退化配置 1 (全局常步长采样) \(\mathcal{O}\big((1-T)^{-2}\big)\) 同上 \(\mathcal{O}\big((1-T)^2 / d^2\big)\) 误差界出现 \((1-T)^{-2}\),步数增加一个数量级数量级
退化配置 2 (忽略平行输运曲率项) 无法证明 TV 稳定性 散度导数发散,误差不闭合 无定义 导致连续性方程右端项无法用内蕴量控制,理论崩溃
几何对照 (无割迹边界平滑因子) 出现奇异脉冲测度 \(J_t(x \mid x_1)\) 边界处导数非光滑 需要极小阶数截断 超球面上 \(\sin^{d-1}\) 因子恰好在割迹处抹平了指示函数的不连续性

关键发现

  • 离散时间加密对奇异性的抵消作用:条件流匹配向量场在终点处天然存在 \((1-t)^{-1}\) 的发散因子,导致散度梯度以 \((1-t)^{-3}\) 速度发散。通过设计二次幂时间调度 \(t_i = 1 - (1 + \eta i)^{-2}\),步长在接近终点时以几何级数收缩,使得离散积分离散项从常步长的 \(\mathcal{O}((1-T)^{-2})\) 直接优化至 \(\mathcal{O}((1-T)^{-1})\)。
  • 割迹与奇异因子的自然光滑化:在超球面上,尽管极坐标变换在对跖割迹处引入了非光滑指示函数 \(\mathbf{1}_{\{d(x, x_1) < (1-t)\pi\}}\),但被积测度中的行列式因子 \((\sin(r/(1-t)) / \sin r)^{d-1}\) 在边界 \(r \to (1-t)\pi\) 处具有 \((d-1)\) 阶零点,当维数 \(d \ge 4\) 时保证了被积函数至少为 \(C^2\) 光滑,确保了流的唯一存在性。
  • 曲率驱动的阶数代价:在 \(\mathrm{SPD}(n)\) 流形上,由于负曲率造成的测地线发散,Jacobi 场解呈现双曲正弦 \(\sinh(\kappa r)\) 的指数放大,导致最终迭代复杂度中的维数依赖达到 \(d^{24}\)。这表明在负曲率空间中,虽然无割迹阻碍,但体积爆炸效应构成了主要的计算复杂度瓶颈。

亮点与洞察

  • 流形输运 TV 演化恒等式的建立:将欧氏概率流 ODE 的 TV 导数分析推广至任意黎曼流形,把时间导数严密转化为流形内蕴梯度内积与内蕴散度之和。这一微分工具不仅适用于流匹配,还能直接复用到流形上的任何确定性输运采样分析中。
  • 几何与统计误差的纯粹解耦:定理明确将总误差拆为 \(h \mathsf{C}_{\mathrm{Lip}} + \varepsilon \mathsf{C}_{\mathrm{eps}}\)(紧流形附带 \(\varepsilon^2\) 高阶项),彻底分清了由数值积分步长引入的确定性几何离散误差和由神经网络训练引入的速度场逼近误差。
  • 避免流形热核近似的本质优势:相较于黎曼扩散模型需要估算难以求出闭式解的流形热核,RFM 将问题转化为纯 ODE 动力学与平行输运积分,避开了布朗运动离散化的复杂伊藤修正,证明了流匹配在几何生成任务中的理论优越性。

局限与展望

  • 收敛界止步于提前停止截断时间 \(T < 1\):当前收敛界在 \(T \to 1\) 时发散,需要借助 \(W_1\) 距离过渡并引入 \(1 - T\) 的早停误差截断,无法直接得出端点处 \(T=1\) 的无条件 TV 收敛。
  • 非紧流形维数依赖较粗糙:在 \(\mathrm{SPD}(n)\) 上的维数多项式指数高达 \(d^{24}\),源于对曲率模型函数多层 Cauchy-Schwarz 放大以及松弛的指数矩放缩,未来仍有进一步紧致优化的空间。
  • 未覆盖更复杂的流形拓扑:针对紧流形的具体收敛率目前仅在具有高度对称性的超球面 \(\mathcal{S}^d\) 上完成显式常数闭式推导,对于一般李群(如 \(\mathrm{SO}(3)\))或复杂复流形,跨越割迹的高阶正则性仍待进一步验证。

相关工作与启发

  • vs 黎曼扩散模型 (Riemannian Diffusion Models, De Bortoli et al., 2022; Xu et al., 2026):黎曼扩散模型必须离散化流形布朗运动,往往难以获得显式多项式 TV 收敛率且转移核高度依赖渐近展开;本文证明了纯 ODE 驱动的流匹配能在 TV 距离下实现严格多项式复杂度,计算与理论都更加清洁。
  • vs 欧氏流匹配理论 (Euclidean Flow Matching, Li et al., 2025; Liu et al., 2025):欧氏分析严重依赖全局平移不变性与平凡切空间;本文建立了平行输运插值和曲率张量耦合控制框架,成功攻克了非平凡几何结构中的体积畸变难题。
  • vs 测地流形插值采样 (Wu et al., 2025; Guan et al., 2026):先前流形流匹配文献多停留在算法构建与实证评估,本文提供了首个全非渐近收敛性理论基石,确立了步长选取与学习精度的理论指导准则。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次建立黎曼流匹配的非渐近 TV 收敛理论,填补了流形生成模型领域的关键理论空白。
  • 实验充分度: ⭐⭐⭐⭐☆ 理论推导极为扎实完整,详尽覆盖了紧流形(超球面)与非紧流形(SPD)两大典型体系,缺乏实机数值验证代码。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学表达严密精细,几何物理图像清晰,附录推导层层递进。
  • 价值: ⭐⭐⭐⭐⭐ 为流形上的生成式 AI 提供了坚实的理论依据,对结构生物学、材料科学等几何学习场景具有深远指导意义。