Strong Stochastic Flow Maps¶
会议: NeurIPS 2026
arXiv: 2606.01086
OpenReview: https://openreview.net/forum?id=ppj0Eq4Y6y
论文: https://arxiv.org/abs/2606.01086
代码: 待开源
领域: 其他
关键词: 随机流映射, 神经 SDE, Itô 映射, 粗糙路径理论, 自蒸馏
一句话总结¶
提出强随机流映射（Strong Stochastic Flow Maps, SSFM），通过移位勒让德多项式高效参数化布朗运动并结合陈氏关系与半群自蒸馏目标，首次实现了加性噪声随机微分方程（SDE）强路径解的少步数快速生成。
研究背景与动机¶
基于神经微分方程的深度生成模型（如扩散模型和流匹配模型）在图像生成、分子结构采样及物理系统模拟等领域取得了极佳的生成质量。然而，这类模型在推断时通常需要沿着学到的向量场进行数十乃至数百次数值积分求解（NFEs），高昂的采样计算开销成为实际落地的一大瓶颈。为了打破这一限制，一致性模型、均值流和流映射等确定性解映射方法相继被提出，它们通过自蒸馏直接学习源噪声到目标数据的常微分方程（ODE）积分映射，将采样步数压缩到 1 至 4 步。
然而，现有少步数生成模型绝大多数局限于确定性概率流 ODE。近期一些工作（如 GLASS、Diamond Map 等弱随机流映射）尝试将流映射拓展至随机过程，但它们本质上仅是通过内层流模型近似 SDE 的边际转移核 \(p_{t|s}(X_t \mid X_s)\)。这种近似属于弱解范畴（仅保证边缘分布收敛），其生成过程与底层的布朗运动完全解耦，无法保持样本沿轨迹的路径一致性，亦无法计算依赖于具体随机驱动路径的路径级可观测量（Pathwise Observables）。这导致在需要严格随机动力学一致性的科学仿真（如分子构象平衡态采样）和路径级奖励对齐场景中存在根本性缺陷。
这一缺陷的核心症结在于：布朗运动样本轨道具有无限维且处处不可微的奇异性质，使得直接构建并训练一个以整条布朗路径为输入的神经网络强解算子（Itô 映射）极具挑战。本文从粗糙路径理论切入，将布朗运动投影到正交多项式基底上，将无限维随机路径压缩为少量服从独立正态分布且满足代数拼接法则的系数向量。核心 idea：利用移位勒让德多项式正交展开布朗运动轨道，构建满足切条件和半群性质的欧拉-马山参数化强随机流映射，并通过无模拟自蒸馏目标直接拟合加性噪声 SDE 的强路径解。
方法详解¶
整体框架¶
强随机流映射（SSFM）旨在直接拟合加性噪声 Itô 随机微分方程 \(\mathrm{d}X_t = f(t, X_t)\mathrm{d}t + g(t)\mathrm{d}W_t\) 的强解映射（Itô 映射）\(\Psi_{s,t}: (X_s, W_{[s,t]}) \mapsto X_t\)。该映射以任意起始时间 \(s\) 的状态 \(X_s\) 以及时间区间 \([s, t]\) 内的一条布朗运动实现 \(W_{[s,t]}\) 为输入，直接预测终点时间 \(t\) 的状态 \(X_t\)。
整个学习与推断流程包含三个核心阶段：首先，利用移位勒让德多项式将连续布朗运动增量投影为低维高斯随机变量序列；其次，设计显式满足 SDE 瞬时切条件的欧拉-马山类神经网络参数化架构，由漂移网络和扩散网络分别预测积分项；最后，基于局部小步长单步拟合和宏观跨步半群自蒸馏构建双任务损失函数，实现无需全轨迹数值积分的无模拟训练。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入：初始状态 Xs 与时间区间 [s, t]"] --> B["布朗运动多项式展开<br/>计算正交勒让德系数 I^(N)"]
B --> C["欧拉-马山参数化 Itô 映射<br/>网络预测漂移与扩散积分项"]
C --> D{"时间步长判断：t - s"}
D -->|小步长 t in [s, s+Δt]| E["局部匹配损失<br/>对齐单步 SDE 真实漂移与扩散"]
D -->|大步长 t in [s+Δt, 1]| F["半群自蒸馏损失<br/>陈氏关系拼接二分步与单步目标"]
E --> G["输出：终点状态 Xt 强路径解"]
F --> G关键设计¶
1. 欧拉-马山参数化 Itô 映射：满足 SDE 瞬时切条件
传统神经网络难以直接学习强解映射，因为若任意设定函数形式，模型很难满足当步长趋近于零时导数与 SDE 漂移项、扩散项严格相等的切条件约束。为保证对底层随机过程的强逼近，本文提出了欧拉-马山风格的参数化形式： $\(\Psi_{s,t}^\theta(X_s, W_{[s,t]}) = X_s + f_{s,t}^\theta(X_s, W_{[s,t]})(t-s) + g_{s,t}^\theta(W_{[s,t]})(W_t - W_s)\)$ 其中 \(f_{s,t}^\theta\) 和 \(g_{s,t}^\theta\) 分别作为归一化的漂移积分函数和扩散积分函数。作者在理论上证明了切条件引理（Lemma 3.1 & Proposition 3.2）：当且仅当对角线极限满足 \(f_{t,t}^\theta(X_t) = f(t, X_t)\) 且 \(g_{t,t}^\theta = g(t)\) 时，该解映射在一阶与二阶变差上退化为真实的 SDE。这确保了网络结构先验与随机微分方程微观物理规律的完全契合。
2. 移位勒让德多项式展开与闭式陈氏关系：连续随机路径的紧凑编码
由于连续布朗轨道 \(W_{[s,t]}\) 具有无限自由度，若采用分段折线采样，需要极密采样点，无法作为网络有效输入。本文基于粗糙路径理论，采用移位勒让德多项式 \(\widetilde{P}_n\) 对布朗运动进行正交级数展开，提取前 \(N\) 阶随机积分系数： $\(I_{s,t}^{(n)} = \int_s^t \widetilde{P}_n\left(\frac{u-s}{t-s}\right)\mathrm{d}W_u\)$ 该表示具有三项卓越数学性质：首先，在均方与 \(\alpha\)-Hölder 距离（\(\alpha \in [0, 1/2)\)）下严格强收敛至布朗运动；其次，由 Itô 等距性质可知，各个系数相互独立且精确服从高斯分布 \(I_{s,t}^{(n)} \sim \mathcal{N}(0, \frac{t-s}{2n+1}I)\)，使得训练时能够解析采样，无需任何离散路径模拟；最关键的是，多项式系数满足闭式陈氏关系（Chen Relations）：两个子区间 \([s, u]\) 与 \([u, t]\) 的勒让德系数可通过确定性矩阵变换严格合成为宏观区间 \([s, t]\) 的系数，且反向可精确二分，从而为多步自蒸馏的路径代数拼接提供了坚实基础。
3. 半群自蒸馏优化：无数值模拟的端到端强解训练
在连续动力学中，解映射必然满足半群性质 \(\Psi_{s,t} = \Psi_{u,t} \circ \Psi_{s,u}\)。定理 3.3 证明了只要参数化映射同时满足瞬时切条件与半群条件，它就唯一确定为底层 SDE 的真实 Itô 映射。基于此，本文构建了自蒸馏目标函数： $\(\mathcal{L}_{\text{SD}}(\theta) = \mathcal{L}_{f,g}(\theta) + \mathcal{L}_{\text{D}}(\theta)\)$ 其中，匹配损失 \(\mathcal{L}_{f,g}\) 在小时间步 \([s, s+\Delta t]\) 上将模型预测与局部解析欧拉-马山步进行加权均方对齐；蒸馏损失 \(\mathcal{L}_{\text{D}}\) 则在大时间步 \([s+\Delta t, 1]\) 上抽取中点 \(u=(s+t)/2\)，利用陈氏关系严格构造对应的子区间布朗系数，约束单步大跨步映射预测值逼近以 stop-gradient 固定的两步连续微步输出。对于扩散模型，真实漂移项可通过条件得分函数 \(-\epsilon/\sigma_t\) 显式构建，整个训练过程完全无需运行任何 SDE 数值积分器。
损失函数 / 训练策略¶
算法训练按批次将样本分为两路（比例为 \(\eta\) 与 \(1-\eta\)）： 1. 单步匹配损失（Batch 占比 \(\eta\)）：采样 \(s \sim \mathcal{U}[0, 1]\)，\(t \sim \mathcal{U}[s, s+\Delta t]\)，计算局部目标 \(\hat{X}_t = X_s + f(s, X_s)(t-s) + g(s)(W_t - W_s)\)，最小化加权误差： $\(\mathcal{L}_{f,g} = (t-s)^{-1} \|\hat{X}_t - \Psi_{s,t}^\theta(X_s, I_{s,t}^{(N)})\|_2^2\)$ 2. 多步自蒸馏损失（Batch 占比 \(1-\eta\)）：采样 \(t \sim \mathcal{U}[s+\Delta t, 1]\) 并取中点 \(u = (s+t)/2\)。采样独立的高斯变量 \(I_{s,u}^{(N)}\) 与 \(I_{u,t}^{(N)}\)，由陈氏关系合成 \(I_{s,t}^{(N)}\)。构建目标与预测项： $\(X_{\text{tgt}} = \text{stopgrad}\left(\Psi_{s,t}^\theta(X_s, I_{s,t}^{(N)})\right), \quad X_{\text{pred}} = \Psi_{u,t}^\theta\left(\Psi_{s,u}^\theta(X_s, I_{s,u}^{(N)}), I_{u,t}^{(N)}\right)\)$ 最小化半群自洽误差 \(\mathcal{L}_{\text{D}} = (t-s)^{-1} \|X_{\text{tgt}} - X_{\text{pred}}\|_2^2\)。
模型采用 Adam 优化器配合预热余弦退火学习率，并维护模型权重的指数移动平均（EMA，衰减率 0.999）。在图像生成任务中，漂移与扩散网络均采用 EDM2 架构；布朗展开阶数默认设为 \(N=3\)。
实验关键数据¶
主实验¶
论文在 CIFAR-10 与 CelebA-64 图像生成基准以及 Alanine-Dipeptide（丙氨酸二肽，ALDP）分子构象采样任务上进行了全面评测。在图像生成中，SSFM 与确定性流映射、弱随机流映射进行了少步数 FID 比较；在分子体系中，评估了平均力势能面均方误差（PMF Error）和自由能表面 Jensen-Shannon 散度（JS）。
| 数据集 | 模型类别 | 方法 | NFE=2 | NFE=4 | NFE=8 | NFE=16 |
|---|---|---|---|---|---|---|
| CIFAR-10 (FID \(\downarrow\)) | 确定性流 | Consistency Training | 5.83 | — | — | — |
| CIFAR-10 (FID \(\downarrow\)) | 确定性流 | Flow map (LSD) | 4.37 | 3.34 | 3.33 | 3.57 |
| CIFAR-10 (FID \(\downarrow\)) | 弱随机流 | GLASS | 157.55 | 39.47 | 11.60 | — |
| CIFAR-10 (FID \(\downarrow\)) | 弱随机流 | Diamond Map | 5.80 | 5.80 | 6.73 | — |
| CIFAR-10 (FID \(\downarrow\)) | 强随机流（本文） | SSFM (Ours) | 4.93 | 3.49 | 3.29 | 3.35 |
| CelebA-64 (FID \(\downarrow\)) | 确定性流 | Flow map (LSD) | 5.74 | 3.18 | 2.18 | 1.96 |
| CelebA-64 (FID \(\downarrow\)) | 弱随机流 | GLASS | 95.25 | 51.80 | 26.33 | — |
| CelebA-64 (FID \(\downarrow\)) | 弱随机流 | Diamond Map | 9.16 | 6.74 | 5.78 | — |
| CelebA-64 (FID \(\downarrow\)) | 强随机流（本文） | SSFM (Ours) | 5.65 | 3.89 | 3.60 | 3.79 |
在分子物理平衡态采样（Alanine-Dipeptide）中，评估少步数下的物理势能面恢复精度：
| 方法 | 评估指标 | NFE=2 | NFE=4 | NFE=10 | NFE=20 | NFE=100 | NFE=1000 |
|---|---|---|---|---|---|---|---|
| 标准 Diffusion | PMF Error (\(\downarrow\)) | 16.728 | 15.667 | 11.347 | 3.400 | 0.084 | 0.066 |
| Fokker-Planck 扩散 | PMF Error (\(\downarrow\)) | 15.475 | 15.709 | 11.345 | 3.498 | 0.092 | 0.069 |
| SSFM (Ours) | PMF Error (\(\downarrow\)) | 0.235 | 0.168 | 0.101 | 0.089 | 0.067 | 0.062 |
| 标准 Diffusion | JS 散度 (\(\times 10^{-2}, \downarrow\)) | 48.735 | 48.377 | 38.309 | 16.293 | 0.787 | 0.618 |
| SSFM (Ours) | JS 散度 (\(\times 10^{-2}, \downarrow\)) | 1.990 | 1.480 | 0.950 | 0.860 | 0.640 | 0.590 |
消融实验¶
针对布朗运动展开阶数 \(N\) 以及多步数下强路径收敛特性的消融实验表明，多项式逼近阶数直接决定了强解的精度上限：
| 多项式阶数 / 配置 | 强解逼近误差（非线性 SDE） | 图像生成性能（CIFAR-10 NFE=4 FID） | 说明 |
|---|---|---|---|
| \(N=1\)（仅标准布朗增量） | 相对误差高，大步长下发散明显 | 4.82 | 仅利用 \(W_t - W_s\)，缺乏子区间随机曲率信息 |
| \(N=2\) | 强误差显著下降约 45% | 3.91 | 引入一阶移位勒让德积分，捕捉基础路径起伏 |
| \(N=3\)（默认模型） | 强误差进一步降低约 25% | 3.49 | 平衡精度与计算开销的最佳拐点，系数方差按 \(1/(2n+1)\) 衰减 |
| \(N=4\) | 强误差继续微幅收敛 | 3.46 | 精度提升趋于饱和，网络输入维度增加 |
此外，路径一致性验证实验（固定初值 \(X_0\) 与固定布朗轨道 \(W_t\)）证明，当步数从 2 步增加至 16 步时，SSFM 生成的样本像素几乎完全重合，严格满足强收敛；而当固定 \(X_0\) 并切换不同随机路径 \(W_t\) 时，模型能够生成各具特色但结构合理的多样化样本，证实了真正的随机流行为。
关键发现¶
- 显著超越弱随机流模型：在 CIFAR-10 与 CelebA-64 上，SSFM 在极少步数（2-4 步）下的 FID 彻底碾压以 GLASS 和 Diamond Map 为代表的弱随机流模型（例如 CelebA-64 2-NFE 下 SSFM 为 5.65，而 Diamond Map 为 9.16，GLASS 高达 95.25）。弱解方法由于缺乏真实路径对齐，在少步推断时极易产生累积漂移。
- 分子物理采样效率飞跃：在 Alanine-Dipeptide 构象生成中，SSFM 仅需 2 步即可将 PMF 误差压至 0.235，而传统扩散模型在 20 步时误差仍高达 3.400；SSFM 在 100 步时的表现（0.067）甚至超越了扩散模型积分 1000 步的精度（0.066），实现了高达 10 倍至 50 倍的采样加速比。
- 高阶勒让德系数的边际贡献明确：理论与实证均表明，勒让德积分的方差正比于 \((t-s)/(2n+1)\)，高阶分量能量衰减迅速，因此选取 \(N=3\) 即可充分逼近布朗轨道的一阶与二阶粗糙几何特性。
亮点与洞察¶
- 粗糙路径与深度生成的优雅交汇：巧借粗糙路径理论中的移位勒让德多项式展开，不仅把无限维连续随机扰动解析压缩为独立标准高斯变量，更通过代数陈氏关系打通了时间区间拼接的阻碍，使深度网络能够精确学习处处不可微系统的强积分算子。
- 无模拟自蒸馏破解 SDE 训练瓶颈：摆脱了以往神经 SDE 训练中高开销的伴随方程或伴随状态反向传播，将瞬时物理切线对齐与宏观半群自洽分离，使得强随机流模型的训练复杂度直接与确定性流映射持平。
- 向科学计算与路径奖励对齐的泛化潜力：强解性质意味着生成结果与驱动噪声轨道具备严格的双射因果对应，这使得在分子动力学自由能计算、受控随机控制以及基于路径的强化学习奖励对齐（Pathwise RLVR）中，具有弱解模型无可比拟的梯度稳定性和保真度。
局限与展望¶
- 局限于加性噪声假定：当前理论严格依赖于加性噪声 SDE（\(g(t)\) 不依赖于状态 \(X_t\)），此时 Itô 映射是良态连续函数；若扩展至更通用的乘性噪声流形 SDE，必须引入更复杂的 Lyons 粗糙路径代数，网络参数化与积分计算会急剧复杂化。
- 双网络参数开销：架构中需要独立配置漂移网络 \(f^\theta\) 与扩散网络 \(g^\theta\)（尽管实验中扩散网络通道较小），在显存占用与参数量上较单一向量场网络略有增加。
- 未来方向：探索乘性噪声下的神经粗糙流映射，以及将强随机流图的确定性路径导数直接用于逆问题求解与免梯度分子构象优化。
相关工作与启发¶
- vs Flow Map (LSD / Boffi et al., 2025)：确定性流映射仅能学习单值常微分方程解，无法表达自然随机动力学中的涨落特性；SSFM 是其在随机分析框架下的正统强解推广，且在图像生成上取得了相媲美甚至更优的多样性指标。
- vs Diamond Map (Holderrieth et al., 2026a) & GLASS (Holderrieth et al., 2026b)：两类方法均属于弱随机流模型，仅在边缘分布层面上拟合转移核，采样缺乏路径因果一致性；SSFM 引入实际布朗路径展开，首次达成真正的强路径解。
- vs 神经 SDE 高阶求解器 (Jiang et al., 2025 / Sig-deg)：Sig-deg 等方法依赖预训练的扩散教师模型进行耗时的长轨迹数值采样蒸馏，而 SSFM 提出了一致自蒸馏目标，支持从扩散得分中以 simulation-free 的方式直接从头端到端训练。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将粗糙路径勒让德展开与半群自蒸馏结合，理论严谨地解决了加性 SDE 强流映射的学习难题。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖非线性 SDE 理论验证、高维图像生成与真实生物分子构象采样，数据详实，指标信服力强。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导层层递进，从切条件到半群性质及陈氏关系一气呵成，附录数学证明极其扎实。
- 价值: ⭐⭐⭐⭐⭐ 为少步数随机扩散模型和平衡态分子物理模拟开辟了强解新范式，具有深远的理论与工程启发价值。