跳转至

SplineFlow: Flow Matching for Dynamical Systems with B-Spline Interpolants

会议: NeurIPS 2026
arXiv: 2601.23072
代码: https://github.com/santanurathod/SplineFlow
领域: 优化/理论
关键词: 连续归一化流, 流匹配, 动力系统, B样条插值, 薛定谔桥

一句话总结

针对传统流匹配采用分段线性插值监督信号导致动力系统曲率失真与震荡的问题,本文提出基于 Cox-de Boor 递归 B 样条插值构建条件概率路径的 SplineFlow,兼具确定性 ODE 动力学渐进导数逼近保证与随机 SDE 漂移-得分回归能力。

研究背景与动机

流匹配(Flow Matching, FM)作为近年来兴起的免模拟连续归一化流训练范式,因其无需像神经常微分方程(Neural ODE)那样进行计算开销巨大的正向数值积分反向伴随计算,在生成建模与动力系统拟合领域得到了广泛应用。然而,现有流匹配方法的核心设计初衷集中于端点生成保真度(即仅关注终点 \(t=1\) 分布的采样质量),其所依赖的监督信号普遍由连续观测点之间的线性插值路径构成。对于需要严格刻画全时域演化过程的物理与生物动力系统而言,分段线性条件路径使得各时间段的条件速度场退化为常数标量,其时间一阶导数在观测时刻完全不连续,彻底抹杀了系统的中间动力学曲率,在非线性、振荡或高阶演化轨迹中引入了严重的动力学逼偏。

将条件路径由线性提升至高阶形式并非易事:朴素的高阶多项式插值(如全局拉格朗日多项式)在多时刻观测约束下极易遭遇龙格现象(Runge's phenomenon),在采样区间边缘引发剧烈的数值震荡与不稳定;经典三次样条(Cubic Spline)虽然保证了全局二次连续平滑,但其推导依赖于全部样本点的全局联立求解,降低了局部控制灵活性并对边界条件极度敏感;而 Hermite 插值则高度依赖往往不可观测的真实导数先验。此外,在非等间隔观测或高缺失率的非规则时序采样(Irregular Sampling)场景下,传统线性流匹配模型对动力学外推和中间态推断的性能会急剧崩溃。

本文的切入视角是利用 B 样条(B-splines)所具备的局部紧支集基函数与可调节多项式阶数特性,将其作为流匹配条件概率路径的均值演化机制。核心 idea:利用 Cox-de Boor 递归公式构建局部稳定且任意光滑度可控的 B 样条条件概率路径,在确定性 ODE 场景下推导闭式条件速度场并提供严格的导数逼近误差上界;在随机 SDE 场景下无缝融合免模拟薛定谔桥(SF2M)联合学习漂移场与得分场,实现高阶动力系统的准确、快速拟合。

方法详解

整体框架

SplineFlow 旨在从离散且可能非规则采样的时序观测轨迹集合 \(\{X_i\}_{i=0}^{N-1}\) 中恢复底层的系统连续演化规律。对于确定性系统,目标是学习一个时间依赖的神经网络速度场 \(u_\theta(t, x)\);对于包含布朗扰动的随机系统,则结合流与得分匹配联合学习概率流速度场 \(u_\theta(t, x)\) 与对数概率密度梯度得分场 \(s_\phi(t, x)\)。方法的整体 Pipeline 包含三大核心步骤:首先基于观测时序数据,利用 Cox-de Boor 递归公式拟合各维度的 \(m\) 阶 B 样条均值函数 \(\mu(t)\);随后根据解析导数构建无偏条件速度向量(或对于随机动力学引入方差调度计算扩散补偿项);最终通过简单的最小二乘回归目标训练神经网络,并在推理阶段通过常微分/随机微分数值求解器生成任意时刻的状态。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["离散观测时序轨迹<br/>X = [x(t0), x(t1), ..., x(tn)]"] --> B["B 样条均值条件路径构建<br/>Cox-de Boor 递归展开与局部基函数加权"]
    B --> C["解析条件速度场与概率流映射<br/>依据样条一阶导数闭式推导速度 ut(x|z)"]
    C --> D{"系统动力学类别分支"}
    D -->|确定性 ODE| E["确定性流匹配回归优化<br/>最小化神经网络速度场与样条导数 MSE"]
    D -->|随机 SDE| F["随机扩展联合漂移与得分回归<br/>免模拟薛定谔桥目标协同学习速度与得分"]
    E --> G["连续动力学轨迹推理生成<br/>ODE/SDE 前向数值积分全时域演化"]
    F --> G

关键设计

1. B 样条均值条件路径构建:摆脱线性粗糙假设与全局多项式震荡

针对线性插值曲率丢失以及高阶多项式全局不稳定的痛点,SplineFlow 采用 \(m\) 阶 B 样条基函数的线性组合来构建高维轨迹在时间维度的插值函数 \(\mu(t)\)。设轨迹在时间序列 \(\{t_0, t_1, \dots, t_n\}\) 上的观测状态为 \(x_{t_j} \in \mathbb{R}^d\),B 样条基函数 \(\mathcal{B}_{j,m}(t)\) 遵循 Cox-de Boor 递推关系: $$ \mathcal{B}{j,m}(t) = \frac{t - t_j}{t} - t_j} \mathcal{B{j,m-1}(t) + \frac{t(t), \quad \forall m \ge 1 $$ 其中 0 阶基函数为区间指示函数 } - t}{t_{j+m+1} - t_{j+1}} \mathcal{B}_{j+1,m-1\(\mathcal{B}_{j,0}(t) = \mathbf{1}_{[t_j, t_{j+1})}(t)\)。对于给定的观测节点,插值函数表示为 \(\mu(t) = \sum_{j=0}^n c_{j,m} \mathcal{B}_{j,m}(t)\),其系数 \(c_{j,m}\) 通过强加节点插值约束 \(\mu(t_j) = x_{t_j}\) 求解线性方程组确定。由于 B 样条基函数具备局部紧支集(Compact Support),单点观测的变化仅影响相邻 \(m+1\) 个子区间,既杜绝了全局拉格朗日多项式的龙格震荡,又保留了通过提升阶数 \(m\) 逼近高阶光滑轨迹的能力。理论上,论文证明了线性插值仅是 \(m=1\) 时的特例,且对于 \(m\) 阶连续可微系统,B 样条不仅在函数值上达到 \(\mathcal{O}(n^{-m})\) 逼近误差(线性插值为 \(\mathcal{O}(n^{-2})\)),更在动力学至关重要的一阶导数上达到 \(\mathcal{O}(n^{-m+1})\) 误差(线性插值为 \(\mathcal{O}(n^{-1})\))。

2. 解析条件速度场与概率流映射:确保边缘连续性与理论渐近界

确定了均值路径后,条件概率路径被建模为高斯分布 \(p_t(x|z) = \mathcal{N}(x; \mu_t(z), \sigma^2 I)\),当标准差极小常数 \(\sigma \to 0\) 时,所诱导的边缘分布 \(p_t(x)\) 在每个观测时间节点 \(t_j\) 严格还原真实数据边缘分布 \(q(x_{t_j})\),并满足连续性方程 \(\frac{\partial p_t}{\partial t} + \nabla_x \cdot (u_t p_t) = 0\)。根据流匹配理论,生成该高斯路径的解析条件速度场完全由均值函数对时间的导数决定:\(u_t(x|z) = \mu_t'(z)\)。论文推导出了 B 样条插值导数的精确闭式表达: $$ u_t(x|z) = \sum_{j=1}^n m \cdot c_{j,m} \left( \frac{\mathcal{B}{j,m-1}(t)}{t} - t_j} - \frac{\mathcal{B{j+1,m-1}(t)}{t \right) $$ 基于此解析速度,论文建立了核心理论误差分解界(定理 4.1):神经网络学得的速度场 } - t_{j+1}\(u_\theta\) 与底层真实动力学速度 \(f'(t)\) 的误差分解为经验回归损失与样条导数逼近误差之和: $$ \mathbb{E}{t,x}[|u\theta(t, x) - f'(t)|] \le \sqrt{\text{Training Loss}} + \mathcal{O}(n^{-m+1}) $$ 当训练损失收敛至极小值时,B 样条阶数 \(m \ge 2\) 显著压缩了残余逼近误差,使学习到的速度场在任意非采样时刻都能紧密贴合系统内在物理规律。

3. 确定性流匹配回归优化:免模拟无偏速度场学习

在确定性 ODE 场景下,训练摆脱了传统连续归一化流依赖伴随敏感度法反向微分的重型仿真流程。对于从轨迹集中抽取的微批次样本,SplineFlow 统一在时间域 \(t \sim \mathcal{U}(t_0, t_n)\) 上随机采样查询时刻,计算对应高斯扰动状态 \(x_t \sim \mathcal{N}(\mu(t), \sigma^2 I)\),并将目标对准解析导数 \(u(t) = \mu'(t)\)。速度网络 \(u_\theta(t, x)\) 的回归损失为: $$ \mathcal{L}{\mathrm{ODE}}(\theta) = \mathbb{E} \left[ | u_\theta(t, x) - \mu_t'(z) |_2^2 \right] $$ 这一优化过程在训练前预计算样条基与导数向量,单次预计算开销仅为数秒,随后完全转变为标准前馈反向传播。推理阶段只需给定初始条件 \(x_{t_0}\),利用任意高阶数值积分器(如 Runge-Kutta 45)对 \(u_\theta\) 沿时间轴推进,便可以极快速度完成长程动力学模拟。

4. 随机扩展联合漂移与得分回归:免模拟薛定谔桥赋能 SDE

对于服从朗之万或带有常数扩散系数 \(g(t) = \sigma\) 的随机动力系统 \(dx_t = u_t(x_t)dt + \sigma dw_t\),概率流 ODE 的等价漂移与原始系统的漂移之间存在由得分函数决定的偏移项:\(u_t(x_t) = u_t^o(x_t) + \frac{\sigma^2}{2} \nabla_x \log p_t(x_t)\)。SplineFlow 将 B 样条路径无缝延伸至免模拟薛定谔桥(SF2M)架构中。引入时变方差构造样条桥(Spline-Bridge)\(p_t(x|z) = \mathcal{N}(x; \mu_t(z), \sigma_t^2(z) I)\),当加权因子 \(\lambda(t) = \sigma_t(z)\) 时,双网络目标简化为: $$ \mathcal{L}{\mathrm{SplineFlow}}(\theta, \phi) = \mathbb{E} \left[ | u\theta(t, x) - (\epsilon \sigma_t'(z) + \mu_t'(z)) |^2 + | \lambda(t) s_\phi(t, x) + \epsilon |^2 \right] $$ 其中 \(\epsilon \sim \mathcal{N}(0, I)\)。针对不同物理与生物系统,作者设计了两种方差方案:常数方差方案 \(\sigma_t(z) = \sigma\)(此时 \(\sigma_t' = 0\),适用于标准物理 SDE)以及分段二次方差方案 \(\sigma_t^2(z) = \frac{\sigma^2}{(t_{j+1} - t_j)^2} (t - t_j)(t_{j+1} - t)\)(在每个观测点收敛至 0,特别适用于单细胞非配对轨迹插值)。在生成时,借助由伊藤扩散二次变差解析估算的常数 \(\sigma\),利用 Euler-Maruyama 离散格式即可直接模拟随机动力学演变。

损失函数 / 训练策略

SplineFlow 统一采用标准多层感知机(MLP,深度为 4,宽度对于低维系统设为 256,对于 Lorenz 与高维 HopperPhysics 系统设为 2048)。优化器采用 AdamW,初始学习率为 \(5 \times 10^{-4}\),搭配余弦退火(Cosine Annealing)学习率调度器,训练 10,000 轮。B 样条的阶数 \(m\) 可通过在训练集上留出子集交叉验证预测误差直接确定(与最优阶数的 Spearman 相关系数达到 \(\rho = 0.85\))。

实验关键数据

主实验

论文在 6 种确定性动力系统(指数衰减 Exp-Decay、谐振子 Harmonic、阻尼谐振子 Damped Harmonic、Lotka–Volterra 捕食者-猎物模型、MuJoCo Hopper 刚体物理系统以及混沌 Lorenz 吸引子)和不同时序缺失率 \(p \in \{0, 0.25, 0.5, 0.75\}\) 下,全面对比了 NeuralODE、LatentODE、多边缘流匹配(MMFM)、轨迹流匹配(TFM)与 SplineFlow。

系统 (缺失率 \(p=0\)) 指标 SplineFlow TFM (线性FM) LatentODE (伴随法) 相对 TFM 优势
指数衰减 (Exp-Decay) MSE 8.90e-5 6.60e-4 1.00e-4 误差降低 86.5%
简谐振子 (Harmonic) MSE 3.70e-4 1.81e-1 8.20e-4 提升接近 3 个数量级
阻尼谐振子 (Damped) MSE 9.50e-5 1.50e-2 1.20e-4 误差降低 99.4%
捕食者-猎物 (LV) MSE 8.60e-2 1.84e-1 1.56e-3 优于线性流匹配两倍以上
刚体运动 (HopperPhysics) MSE 1.41 1.55 1.89 保持最优线性特化
混沌系统 (Lorenz) MSE 6.39e-1 2.44 1.16 显著领先伴随与线性基线

在训练耗时方面,伴随法 LatentODE 在上述 6 个系统上的平均单次训练时长高达 \(2.80\text{e}4 \sim 1.21\text{e}5\) 秒,而 SplineFlow 仅需 \(190 \sim 279\) 秒,取得了 147 倍至 597 倍的训练加速比。

在 SDE 随机系统主实验中,以捕食者-猎物(Lotka–Volterra)与混沌 Lorenz 吸引子系统为例(5 次独立运行平均值 \(\pm\) 标准差):

数据集 (\(p=0\)) 模型 轨迹 MSE (SDE) Wasserstein 距离 \(\mathcal{W}_2\) MMD 能量距离 Energy
Lotka–Volterra SF2M (线性桥) \(0.460 \pm 0.097\) \(0.443 \pm 0.050\) \(0.125 \pm 0.024\) \(0.465 \pm 0.094\)
Lotka–Volterra SplineFlow (\(m=2\)) \(0.273 \pm 0.014\) \(0.294 \pm 0.003\) \(0.042 \pm 0.003\) \(0.156 \pm 0.013\)
混沌 Lorenz SF2M (线性桥) \(2.159 \pm 0.008\) \(0.632 \pm 0.008\) \(0.130 \pm 0.003\) \(0.404 \pm 0.011\)
混沌 Lorenz SplineFlow (\(m=4\)) \(1.227 \pm 0.056\) \(0.180 \pm 0.015\) \(0.005 \pm 0.002\) \(0.037 \pm 0.008\)

消融实验

在阻尼谐振子与刚体 HopperPhysics 系统上对 B 样条阶数 \(m \in \{1, 2, 3, 4, 5\}\) 与采样缺失率 \(p\) 进行消融研究:

实验配置 \(p=0\) (完全采样) MSE \(p=0.5\) (50%缺失) MSE \(p=0.75\) (75%缺失) MSE 规律分析
阻尼谐振子 \(m=1\) (线性退化) 1.97e-2 3.19e-2 1.71e-1 线性路径曲率严重失真
阻尼谐振子 \(m=2\) (二次样条) 6.50e-5 5.40e-5 4.15e-4 显著改善周期曲率
阻尼谐振子 \(m=4\) (四次样条) 7.20e-5 4.70e-5 3.40e-5 极度稀疏采样下高阶抗噪性能最优
阻尼谐振子 \(m=5\) (五次样条) 9.50e-5 4.70e-5 4.80e-5 阶数过高边界收益趋于饱和
HopperPhysics \(m=1\) (线性) 1.41 1.83 3.34 近似线性动力学中 \(m=1\) 表现最优
HopperPhysics \(m=3\) (三次) 12.5 6.38e4 1.73e4 在线性动力学中强套高阶样条过拟合

关键发现

  • 非线性与高缺失下高阶样条收益最大:对于包含正弦、指数与捕食相互作用的振荡/非线性系统(如阻尼振子、Lotka–Volterra),当时间点缺失率上升至 \(p=0.75\) 时,线性 TFM 预测完全发散,而 4 阶 B 样条 SplineFlow 凭借局部平滑与高阶连续性将误差抑制在 \(3.4 \times 10^{-5}\),表现出极强的轨迹补全与插值鲁棒性。
  • 避免高阶盲目套用:在近乎刚体运动或一阶衰减系统中,底层动力学本身为一阶线性关系,此时 \(m=1\)(即恢复为线性插值)性能最佳,强行采用 \(m \ge 3\) 的高阶 B 样条反而会因边界条件对噪声的局部过度敏感而导致误差放大。
  • 单细胞生物学轨迹插值优势:在 Axolotl 脑损伤再生 Stereo-seq 空间转录组与胚胎演化时序数据中,结合 entropic 最优传输与 \(m=5\) 阶 B 样条桥的 SplineFlow 在插值任务中将 MMD 与能量距离相较于 SF2M 与 MOTFM 降低了 38% 以上,准确重现了复杂的分化流形。

亮点与洞察

  • 将 B 样条理论与流匹配理论完备自洽地融合:不仅停留在经验替换,而是从数学上证明了连续性方程、边缘恢复性以及闭式速度场的存在性,推导了流匹配回归损失与 B 样条一阶导数逼近误差之间的理论泛化上界。
  • 破解连续归一化流的算力与精度两难悖论:传统观念认为要精准拟合复杂物理动力系统必须依靠 NeuralODE 的伴随积分模拟,而 SplineFlow 证明仅通过优化条件路径质量,纯免模拟回归即可取得比伴随法更优或相当的精度,同时带来两个数量级的训练速度飞跃。
  • 对不规则采样的天然免疫力:B 样条的 Cox-de Boor 算法天然支持非等间隔节点序列(Non-uniform Knots),无需插值前重采样即可直接拟合真实世界中常见的非规则生物与医疗时序。

局限与展望

  • 节点放置优化尚未端到端参数化:当前 B 样条节点序列(Knot Placement)完全固定于观测时间戳上,在动力学梯度变化急剧的时区缺乏自适应节点加密机制,未来可探索自适应节点放置学习。
  • 随机系统局限于常数扩散调度:目前理论与实验扩展仅涵盖加性常数噪声的 SDE 系统(\(g(t) = \sigma\)),对生物学与金融系统中普遍存在的状态依赖型非线性扩散矩阵(State-dependent Diffusion)尚无法直接求解。

相关工作与启发

  • vs Trajectory Flow Matching (TFM) [Zhang et al., 2024]: TFM 首次将流匹配应用于时序多观测,但局限于分段线性条件概率路径,在连续振荡或强非线性系统中速度导数失真;SplineFlow 将其广义化至任意阶 B 样条,理论证明线性仅为 \(m=1\) 特例,显著修正了动力学偏差。
  • vs Multi-Marginal Flow Matching (MMFM) [Rohbeck et al., 2025]: MMFM 尝试引入三次样条耦合,但受制于最小曲率假设且对多项式阶数缺乏可调控性,在近似线性系统(如 Hopper)上发生过拟合;SplineFlow 支持阶数灵活调整并提供解析导数与严谨逼近误差界。
  • vs NeuralODE / LatentODE [Chen et al., 2018a; Rubanova et al., 2019]: 伴随敏感度法反向传播极其耗时,大规模系统扩展困难;SplineFlow 属于免模拟算法,训练耗时缩短 150-600 倍,同时消除了长时间积分带来的梯度弥散与数值震荡。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统性推导并融合 B 样条与条件流匹配,建立了严谨的动力学导数逼近误差理论。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 大物理/生物系统、确定性与随机性场景、规则与不规则采样,附带详尽的消融与 500 倍加速比评测。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备,从问题暴露、理论证明到算法闭环一气呵成。
  • 价值: ⭐⭐⭐⭐⭐ 为科学机器学习(AI for Science)与单细胞动态建模提供了极具工程实用价值与理论坚实度的高速求解范式。