Improving Function Space Flow Matching with Kernel Optimal Transport¶
会议: NeurIPS 2026
arXiv: 2609.38049
代码: https://github.com/heraclixus/KFFM
领域: 图像生成
关键词: 函数空间流匹配, 核最优传输, Hilbert Sinkhorn Divergence, 神经算子, 路径签名核
一句话总结¶
针对函数空间流匹配中独立端点配对导致速度场方差大、难以捕获无限维几何结构的问题，本文提出基于核诱导代价的端点最优传输配对方法（kFFM），在不改变神经算子架构的前提下大幅提升了时序路径与 PDE 场分布生成的质量与训练稳定性。
研究背景与动机¶
函数值数据（如物理系统中的偏微分方程解快照、连续时间序列轨迹）本质上生活在无限维函数空间中。为了直接对这类连续对象进行生成建模，函数空间流匹配（Functional Flow Matching, FFM）将欧氏空间中的流匹配框架推广至可分 Hilbert 空间，通过在紧支集空间上定义高斯测度路径，并使用傅里叶神经算子（FNO）等网络直接参数化速度场，实现了分辨率不变的连续生成能力。
然而，现有 FFM 继承了标准有限维 Flow Matching 的独立端点配对机制：在训练的每个 mini-batch 中，高斯先验采样与真实数据函数是完全随机且独立组合的。这种独立配对迫使条件概率桥不仅要学习穿越数据集共享的全局结构，还要同时消除单对样本间巨大的个体残差，导致回归目标方差极大。在有限维中，最优传输（OT）耦合已被证明能显著拉直轨迹并改善分布拟合，但在无限维函数空间中直接构建最优传输极难落地——无限维空间缺乏类 Lebesgue 的通用参考测度导致概率密度无定义，经典的 Cameron-Martin 定理与 Feldman-Hájek 测度等价条件在非高斯真实物理场中极为严苛，而简单将离散网格展平采用平坦 \(L^2\) 欧氏距离又会彻底忽视函数空间的固有几何（如 Sobolev 范数中的导数平滑性、时间序列的路径粗糙度与重参数化不变性）。
本文的核心切入点在于引入可微且有界的核诱导希尔伯特空间作为代理几何：利用希尔伯特 Sinkhorn 散度（HSD）底层的熵正则最优传输，将先验函数与真实数据函数在特征空间中进行几何感知的最优端点耦合。核心 idea：保持原有 FFM 概率路径与神经算子架构不变，在 mini-batch 内引入核诱导代价矩阵（针对 PDE 场采用 Sobolev-RBF 核，针对连续路径采用路径签名核）执行熵正则最优传输配对，从理论上证明其有界性与离散不变性，使网络只需回归几何结构紧凑的小残差。
方法详解¶
整体框架¶
kFFM 整体遵循两阶段设计：在训练前向传播之前，先对先验样本批次与真实数据批次计算几何感知核代价矩阵，求解批内离散熵最优传输解算得到软传输方案并完成端点配对；随后以配对后的先验作为条件基底噪声，构造连续线性概率流并计算目标速度场，由傅里叶神经算子（FNO）进行回归拟合。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["先验采样 f0 ~ μ0 与数据采样 f1 ~ μ1"] --> B["核诱导代价矩阵构建<br/>c_κ(f0, f1) 嵌入 Sobolev/签名空间"]
B --> C["熵正则化最优传输耦合<br/>Log-Domain Sinkhorn 求解配对方案 π"]
C --> D["端点耦合概率流构建<br/>g_t = t f1 + σ_t f0 线性插值"]
D --> E["神经算子条件流匹配训练<br/>FNO 拟合速度场目标 u = f1 - (1-σ_min)f0"]
E --> F["连续生成采样<br/>ODE 求解器实现网格无关函数生成"]关键设计¶
1. 核诱导代价矩阵构建：将函数空间几何嵌入传输代价
针对函数空间无法定义全局 Lebesgue 密度且平坦欧氏距离忽略导数和粗糙度信息的痛点，kFFM 将函数值样本映射至再生核希尔伯特空间（RKHS）\(\mathcal{H}_\kappa\)。根据核技巧，无需显式构造无限维特征映射，两函数 \(x, y \in \mathcal{X}\) 之间的传输代价即可表示为 RKHS 中的平方距离： $\(c_\kappa(x, y) = \|\phi(x) - \phi(y)\|_{\mathcal{H}_\kappa}^2 = \kappa(x, x) + \kappa(y, y) - 2\kappa(x, y)\)$ 在具体几何实现上，对于生活在 Sobolev 空间 \(H^k(\Omega)\) 的偏微分方程快照，采用 Sobolev-RBF 核 \(\kappa_{\mathrm{RBF}}^{(k)}(x, y) = \exp(-\|x - y\|_{H^k}^2 / (2\sigma^2))\)，将高阶导数信息注入耦合；对于生活在连续有界 \(p\)-变差路径空间 \(C_p([0,T], \mathbb{R}^d)\) 的时序数据，则采用结合时间增强的路径签名核（Signature Kernel）\(\kappa_{\mathrm{sig}}(x, y) = \langle S(x), S(y) \rangle_{\mathcal{T}}\)，利用迭代 Young 积分捕获高阶通道交互并赋予其重参数化不变性。
2. 熵正则化最优传输耦合：良态有界且稳定的批量配对
针对非有界 \(L^2\) 距离在重尾批次上极易导致 Sinkhorn 算法数值溢出或发散的问题，kFFM 采用对数域（log-domain）Sinkhorn 求解熵正则化 OT 方案。论文在理论上证明了对于任意有界正定核（\(\sup \kappa(f,f) \le B_\kappa < \infty\)），核代价函数在任意实 Banach 空间上一致有界：\(0 \le c_\kappa(f, g) \le 4B_\kappa\)。这保证了 Gibbs 核 \(\exp(-c_\kappa / \epsilon)\) 远离 0 和 1，使得 Sinkhorn 求解器在不同批次间具备一致且快速的收敛速度，且对正则化参数 \(\epsilon\) 横跨三个数量级（\(10^{-3} \sim 10^0\)）表现出极高鲁棒性。在实际实现中，先对批内代价矩阵进行中位数归一化，求解离散传输方案 \(\pi\) 后直接从中采样配对对 \((f_0, f_1) \sim \pi\)，彻底避免了重心投影（Barycentric Projection）导致的样本方差坍缩。
3. 端点耦合概率流构建与离散不变性保证：回归紧凑小残差
传统的独立配对让条件速度目标 \(u \approx f_1 - f_0\) 表现为两个互不相关的独立高斯场与真实场相减，方差居高不下；而核最优传输配对将数据点 \(f_1\) 与几何上最为邻近的先验样本 \(f_0\) 对齐，网络只需学习结构化的微小残差。同时，实际训练输入依赖于有限网格离散化，kFFM 在理论上建立了离散不变性定理（Discretization Invariance）：当使用秩为 \(N\) 的截断傅里叶空间 \(V_N\) 逼近连续紧致 Sobolev 球 \(B_R^{H^\alpha} \subset H^k(\Omega)\) 时，计算出的离散核 OT 目标与连续核 OT 目标之间的误差满足： $\(|\mathsf{KOT}_\varepsilon^{(k)}(\mu, \nu) - \mathsf{KOT}_\varepsilon^{(k)}(\mu_N, \nu_N)| \le \frac{4R^2}{\sigma^2} N^{-2(\alpha-k)/d}\)$ 该上界以 \(O(h^{2(\alpha-k)})\) 的网格步长多项式速率指数衰减，从数学上严格保证了批内核传输配对能够无缝继承 FNO 的网格分辨率无关性。
损失函数 / 训练策略¶
给定采样的最优传输配对 \((f_0, f_1) \sim \pi\) 以及时间采样 \(t \sim \mathcal{U}(0, 1)\)，插值场为 \(g_t = t f_1 + \sigma_t f_0\)，其中 \(\sigma_t = 1 - (1 - \sigma_{\min})t\)。条件速度场定义为： $\(v_t^{f_0, f_1}(g_t) = f_1 - (1 - \sigma_{\min}) f_0\)$ 神经算子参数 \(\theta\) 通过最小化连续条件流匹配损失进行优化： $\(\mathcal{L}_{\mathrm{CFM}}(\theta) = \mathbb{E}_{t \sim \mathcal{U}[0, 1], (f_0, f_1) \sim \pi, g_t}\left[ \| v_t^{f_0, f_1}(g_t) - v_\theta(g_t, t) \|_{\mathcal{F}}^2 \right]\)$ 在整个训练过程中，神经网络的输入与输出均保留在物理离散网格空间，样本无需投影至 RKHS；核与最优传输仅在每次反向传播前作为无梯度的批内预配对操作介入，训练额外开销在 \(128^2\) 高分辨率下仅为 5%–7%。
实验关键数据¶
主实验¶
论文在多组一维序列基准（AEMET 气象温度、Gene Expression 基因表达、World Economy 经济时序、Heston 随机波动率模型）以及一维/二维 PDE 状态快照（KdV、随机 KdV、Navier–Stokes、随机 Navier–Stokes）上进行了系统评测。主要评估指标包括无偏 MMD-RBF（越低越好，由于无偏估计可能出现极小负值）与非核统计量 Sliced Wasserstein 距离。
| 数据集 | 指标 | 本文 kFFM (selected) | 独立配对基线 FFM | 离散向量最优传输 CFM-OT(\(L^2\)) | 扩散算子基线 DDO/NCSN | 提升幅度 (相对最强非 kFFM 基线) |
|---|---|---|---|---|---|---|
| Economy | MMD-RBF | \(2.89 \times 10^{-3}\) (Sig/gp) | \(7.30 \times 10^{-3}\) | \(8.46 \times 10^{-3}\) | \(1.57 \times 10^{-1}\) | -60.4% |
| Gene Expr. | MMD-RBF | \(7.11 \times 10^{-3}\) (Sig/gp) | \(1.78 \times 10^{-2}\) | \(2.80 \times 10^{-2}\) | \(3.42 \times 10^{-2}\) | -60.1% |
| Stoch. KdV | MMD-RBF | \(1.00 \times 10^{-4}\) (Sob-RBF/gp) | \(8.80 \times 10^{-4}\) | \(6.58 \times 10^{-4}\) | \(4.48 \times 10^{-2}\) | -74.0% |
| Stoch. NS | MMD-RBF | \(2.80 \times 10^{-3}\) (RBF/wn) | \(9.68 \times 10^{-2}\) | \(4.22 \times 10^{-3}\) | \(2.95 \times 10^{-1}\) | -33.6% |
| Navier–Stokes | MMD-RBF | \(1.63 \times 10^{-3}\) (RBF/wn) | \(1.31 \times 10^{-1}\) | \(1.78 \times 10^{-3}\) | \(7.35 \times 10^{-1}\) | -8.4% |
| Gene Expr. | Sliced-W | 0.087 ± 0.006 | 0.106 ± 0.006 | - | - | -17.9% |
| Economy | Sliced-W | 0.0233 ± 0.0012 | 0.0291 ± 0.0019 | - | - | -19.9% |
| Navier–Stokes | Sliced-W | 0.081 | 0.315 | - | - | -74.3% |
消融实验¶
为了理清性能提升究竟来源于“有界代价的数值稳定性”还是“函数空间几何感知核设计”，论文在固定骨干网络和先验基底的前提下进行了端点配对代价函数与强湍流物理诊断消融（在 \(\nu=10^{-5}, \mathrm{Re} \approx 2000\) 的 2D 湍流 Navier-Stokes 上进行 5-seed 评测）。
| 配置 / 代价机制 | 涡度分布 Vort-PDF \(W_1\) (\(\times 10^{-2}\)) | 拟能分布 Enstrophy \(W_1\) (\(\times 10^{-2}\)) | 涡度偏度误差 Skew (\(\times 10^{-2}\)) | 对数能谱误差 log-Spec (\(\times 10^{-2}\)) | 说明 |
|---|---|---|---|---|---|
| kFFM (selected: RBF/wn) | 3.8 ± 1 | 0.155 ± 0.02 | 0.27 ± 0.3 | 2.0 ± 0.8 | 有界核代价 + 白噪声基底 |
| CFM-OT(\(L^2\)) (raw \(L^2\)/wn) | 4.1 ± 1 | 0.157 ± 0.01 | 0.40 ± 0.2 | 2.2 ± 0.7 | 平坦欧氏无界代价 + 白噪声基底 |
| kFFM-Sob† (Sobolev-RBF/gp) | 10.4 ± 3 | 0.194 ± 0.05 | 4.6 ± 2 | 6.5 ± 1 | 粗糙场误用平滑高斯过程先验 |
| kFFM-Euc† (raw \(L^2\)/gp) | 10.4 ± 4 | 0.197 ± 0.06 | 5.0 ± 3 | 6.7 ± 1 | 欧氏代价 + 平滑高斯过程先验 |
| FFM† (gp, independent) | 10.7 ± 4 | 0.201 ± 0.06 | 4.9 ± 3 | 6.5 ± 1 | 原始独立配对，能谱与大尺度偏离 |
| DDO/NCSN (Score diffusion) | 122 ± 1 | 1.07 ± 0.02 | 0.83 ± 1 | 93.7 ± 2 | 湍流能谱严重扁平化坍缩 |
关键发现¶
- 数据几何决定最优核选取：路径型时间序列数据在签名核（Signature Kernel）下获益最大，在 Gene Expression 上自相关误差较平坦欧氏 OT 降低了 66 倍；而在光滑场（PDE）中，Sobolev-RBF 与欧氏 RBF 占据主导。
- 有界代价提供优异的数值良态性：相比于无界的平坦 \(L^2\) 距离，有界 RBF 核代价使 Sinkhorn 算法迭代提速约 8 倍，全流程训练提速 15%–25%，且在正则化参数 \(\epsilon\) 横跨 \(10^{-3} \sim 10^0\) 时性能保持平稳，彻底免疫梯度异常与重尾爆发。
- 低正则度物理极端场景的先验基底效应：在强湍流 Navier-Stokes（\(\nu=10^{-5}\)）下，物理场的粗糙度远超光滑高斯过程先验的 Cameron-Martin 空间范畴，此时先验基底从平滑 GP 切换为网格白噪声是性能飞跃的前提，而核 OT 配对则在此基础上进一步逼近真实的双峰重尾涡度分布。
亮点与洞察¶
- 无缝即插即用的模块化解耦：kFFM 巧妙地将函数空间最优传输隔离在 Mini-batch 的端点配对阶段，网络架构无需做任何修改，完全保留了 FNO 算子原生的无限维与超分辨率生成特性。
- 严谨完备的无限维测度理论闭环：通过引入 Banach 空间上的 Hilbert Sinkhorn 散度，严格证明了核代价的有界性以及 \(O(h^{2(\alpha-k)})\) 的网格离散不变性误差上界，为深度生成模型在函数空间的理论分析提供了典范。
- 端点配对改变了优化流形的本质：揭示了 Flow Matching 中独立采样的方差来源，通过核最优传输显式建立结构化残差回归目标，以极低的计算代价（5%–7% 耗时增量）换取了生成收敛速度与精度的跨越式提升。
局限与展望¶
- 签名核的长序列内存瓶颈：路径签名核在时间点长度增加时张量计算开销显著增加（序列长度由 100 增至 512 时显存从 0.63 GB 飙升至 12.5 GB），长序列任务只能退化回 RBF 或欧氏代价。
- 无显式直线流（Straightness）保证：作者明确指出函数空间中的非欧几何与高斯测度并不直接继承有限维 OT-CFM 的欧几里得直线轨迹定理，kFFM 的增益主要来自端点残差方差的缩减而非函数空间轨迹的严格拉直。
- 强湍流物理小尺度统计仍有改进空间：在极度湍流小尺度扰动下，核代价由于距离超出带宽而容易趋于饱和，未来可探索针对多尺度非高斯物理场定制的特种非局部核函数与多尺度传输算法。
相关工作与启发¶
- vs FFM (Kerrigan et al., 2024): FFM 首次实现了 Hilbert 空间流匹配，但采用独立端点耦合，条件残差方差极大；kFFM 引入核最优传输进行端点预对齐，显著降低回归方差并保持神经算子结构不变。
- vs CFM-OT(\(L^2\)) (Tong et al., 2024): CFM-OT 仅适用于有限维向量并使用无界 \(L^2\) 代价；kFFM 将其推广至无限维函数空间，采用有界核代价解决了数值发散，并通过特征核引入 Sobolev 和路径拓扑结构。
- vs Functional Rectified Flow (Zhang & Scott, 2025): FRF 依赖多轮重训练生成轨迹对来拉直函数空间流；kFFM 仅在批内依据输入数据几何直接进行单轮最优传输配对，两者机制互补且可正交结合。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [首次将核最优传输与 Hilbert Sinkhorn 散度引入函数空间流匹配，构思严密优雅]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖经典时序、金融波动率、PDE快照与极端强湍流流体，具备严格的种子配对统计显著性检验]
- 写作质量: ⭐⭐⭐⭐⭐ [理论推导脉络清晰，数学符号与定理体系完备，消融实验直指核心机理]
- 价值: ⭐⭐⭐⭐⭐ [为连续物理系统、复杂时间序列与科学计算领域的函数级生成建模树立了全新标杆]