Primal-Dual Flow Matching for Sample-Wise Constrained Generation¶
会议: NeurIPS 2026
arXiv: 2610.12699
OpenReview: https://openreview.net/forum?id=5hNRSOnqjv
代码: https://github.com/ZhengyanHuan/PDFM
领域: 图像生成
关键词: 流匹配、样本级约束生成、原对偶优化、确定性策略梯度、流形约束
一句话总结¶
针对科学与工程中样本必须满足硬性物理或几何约束的生成难题，本文提出原对偶流匹配（PDFM），仅依赖二元归属预言机反馈，通过确定性策略梯度构建可微时空代理判别器，实现对样本级约束满足率的直接概率控制与严格全变差收敛理论保证。
研究背景与动机¶
生成模型（特别是扩散模型与流匹配）在物理科学、分子设计与复杂时空混沌动力学仿真中展现出强大的分布建模能力。然而在这些高可靠性领域中，仅匹配训练集的边际数据分布往往远不足以保证实用价值：生成的动力学轨迹必须满足守恒律与偏微分方程（PDE），生成的生物分子必须符合化合价稳定结构，生成的指纹图像必须保持连通脊线几何拓扑。传统生成模型受限于有限容量与离散逼近误差，即使在全部训练样本均完全满足约束的理想场景下，生成轨迹依然会频繁溢出可行域。
要对单样本施加硬性或概率性约束，现有方法面临极其苛刻的前置假设瓶颈。一类方法依赖几何拓扑的强先验，例如要求约束集合具备严格凸性、闭式双射映射（如 Mirror Diffusion Models, MDM）或已知光滑的反射边界（如 Reflected Flow Matching, RFM）；另一类推理阶段修正方法（如 PCFM、PDM、HardFlow）则要求约束具有全可微的距离函数、解析雅可比矩阵或高昂的在线前向打靶投影求解开销。在实际的科学计算中，物理仿真器（例如流体力学步进器、分子价键推断规则、图像拓扑骨架提取器）大多是高度非线性、黑盒且完全不可微的程序代码，仅能针对输入样本给出“是否满足约束”的二元布尔反馈（Membership Oracle）。
面对仅有零阶二元判定预言机的极简反馈，现有探索方法（如 FM-RE）通常通过向生成轨迹强行注入高斯噪声来构造似然比梯度估计，这不仅破坏了常微分方程（ODE）轨迹的确定性与保真度，而且只能依靠手工调节约束惩罚权重进行间接折中，完全无法精确控制生成样本的达标概率。本文的核心 idea 是：将样本级约束生成严格建模为基于负对数似然的总体级原对偶凸优化问题，并借助确定性策略梯度定理设计沿流轨迹满足贝尔曼一致性的可微时间连续代理判别器，仅凭二元成员预言机反馈即可实现对目标约束满足率的直接闭环调控与严格 TV 范数理论收敛界。
方法详解¶
整体框架¶
PDFM 的目标是从参考分布 \(Q\) 学习一个生成分布 \(P\)，使其在最小化与 \(Q\) 的 KL 散度（等价于最小化负对数似然）的同时，确保样本落入闭约束集 \(\mathcal{C}\) 的概率不低于预设阈值 \(b \in (0, 1)\)。整个框架建立在条件流匹配（Conditional Flow Matching, CFM）的速度场参数化与对偶变量的上升迭代之上，其宏观运行机制包含：(1) 基于当前速度场前向积分生成样本并由黑盒预言机评估二元可行性；(2) 训练时空连续的代理判别器拟合轨迹可行性概率；(3) 利用代理判别器回传的梯度更新速度场参数；(4) 根据蒙特卡洛经验达标率与目标 \(b\) 的差距更新对偶乘子 \(\lambda'\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["初始噪声分布 P0 与训练样本 x1~Q"] --> B["条件流匹配基础场<br/>线性插值与基准速度场初始化"]
B --> C["前向欧拉离散积分<br/>当前速度场生成终端轨迹"]
C --> D["黑盒成员预言机评估<br/>非可微求解器输出二元标签 y"]
D --> E["代理判别器训练<br/>时空分类器拟合贝尔曼一致性"]
E --> F["确定性策略梯度反传<br/>代理梯度联合 CFM 损失更新速度场"]
F --> G["投影超梯度对偶更新<br/>经验约束间隙修正对偶乘子 λ"]
G -->|未收敛继续迭代| C
G -->|对偶平衡收敛| H["输出约束流匹配生成模型"]关键设计¶
1. 原对偶概率约束优化建模与总体级理论最优解：将硬性截断转为可控连续概率重配
直接过滤非可行数据并训练在截断分布 \(Q'=Q(\cdot \mid \mathcal{C})\) 上的传统思路，在噪声较大或可行域较窄时会大幅损失有效样本量，且有限模型容量仍会导致生成样本越界。PDFM 从总体级凸优化切入，建立带约束优化问题： $\(\min_{P} \mathbb{E}_{x \sim Q}[-\log p(x)] \quad \text{s.t.} \quad b - \int_{\mathcal{C}} p(x) dx \le 0\)$ 引入拉格朗日乘子 \(\lambda \ge 0\)，在严格可行性假定（Slater 条件）下，利用 KKT 条件解析求解出对偶变量 \(\lambda\) 下的最优密度函数形式。在 \(0 < Q(\mathcal{C}) < b < 1\) 条件下，最优分布 \(P^{(\lambda^*)}\) 显式表现为训练分布 \(Q\) 在可行域 \(\mathcal{C}\) 与补集 \(\mathcal{C}^c\) 上的精确分段重加权： $\(p^{(\lambda^*)}(x) = \frac{b\,q(x)}{Q(\mathcal{C})}\mathbf{1}_{\mathcal{C}}(x) + \frac{(1-b)q(x)}{1-Q(\mathcal{C})}\mathbf{1}_{\mathcal{C}^c}(x), \quad \lambda^* = \frac{1-Q(\mathcal{C})}{1-b} - \frac{Q(\mathcal{C})}{b}\)$ 该解析形式揭示了约束生成的本质是在可行域与不可行域之间重新分配测度质量，为后续利用参数化流模型逼近真实最优解奠定了理论基石。
2. 基于确定性策略梯度的可微时空代理判别器：攻克非可微成员预言机的梯度回传瓶颈
由于成员指示函数 \(\mathbf{1}_{\mathcal{C}}(x)\) 在可行域内部梯度几乎处处为零、在边界处不可微，直接通过采样路径对参数 \(\theta\) 求导无法进行。传统方法（如 FM-RE）依赖高斯随机扩散引入可微似然比，但方差巨大且破坏了 ODE 确定性轨迹。受确定性策略梯度（DPG）思想启发，PDFM 引入辅助神经网络 \(f_\phi(x, t) \in [0, 1]\)，作为终端指示函数在沿流时空路径上的可微代理。\(f_\phi\) 满足贝尔曼一致性条件 \(f(x, t) = f(x + u_\theta(x, t)\Delta t, t + \Delta t)\)，并在终端满足 \(f(x, 1) = \mathbf{1}_{\mathcal{C}}(x)\)。
在训练中，固定当前速度场 \(u_\theta\) 生成完整的欧拉采样轨迹 \(\{x_{i\Delta t}\}_{i=0}^N\)，黑盒预言机仅在终点打上二元标签 \(y = \mathbf{1}_{\mathcal{C}}(x_1)\)。网络 \(f_\phi\) 通过最小化沿轨二元交叉熵损失进行有监督更新： $\(\text{grad}_\phi = \nabla_\phi \mathbb{E}_{i, x_{i\Delta t}}\left[-y \log f_\phi(x_{i\Delta t}, i\Delta t) - (1-y)\log(1 - f_\phi(x_{i\Delta t}, i\Delta t))\right]\)$ 当 \(f_\phi\) 充分逼近状态转移一致性时，终端约束满足率 \(\mathcal{F}_{u_\theta}(\mathcal{C}) = \int p_0(x_0) f(x_0, 0) dx_0\) 对速度场参数 \(\theta\) 的全微分，可精确拆解为单步局部速度扰动沿时空梯度的链式积分累积。
3. 速度场确定性导引与投影超梯度对偶循环：交替求解原变量与校准对偶乘子
基于代理判别器 \(f_\phi\)，速度场更新不需要进行全轨迹反向传播，而是将传统条件流匹配损失与约束满足率代理梯度有机融合，构造可操作的代理拉格朗日目标函数。速度网络参数 \(\theta\) 接收来自 CFM 回归项与判别器导引项的合成更新： $\(\text{grad}_\theta = \nabla_\theta \mathbb{E}_{t, x_0, x_1}\left[\|u_\theta(x_t, t) - (x_1 - x_0)\|^2\right] - \lambda' N \mathbb{E}_{i, x_{i\Delta t}}\left[\nabla_\theta u_\theta(x_{i\Delta t}, i\Delta t) \left.\nabla_u f_\phi(x_{i\Delta t} + u\Delta t, (i+1)\Delta t)\right|_{u=u_\theta(x_{i\Delta t}, i\Delta t)}\right]\)$ 其中 \(\lambda' = \lambda / w\) 为缩放对偶变量。在速度场参数更新后，模型在当前流场下采样一批终端样本，利用预言机评估经验满足率 \(P_{u_\theta}(\mathcal{C})\)，进而对对偶变量执行超梯度投影上升： $\(\lambda'(i+1) = \left[\lambda'(i) + \eta_i (b - P_{u_\theta}(\mathcal{C}))\right]_+\)$ 当生成样本的达标率低于目标 \(b\) 时，超梯度为正，推动 \(\lambda'\) 增大，迫使速度场在后续步更激进地推向可行区域；一旦达标率超过 \(b\)，\(\lambda'\) 自动衰减，使模型重新聚焦于保真度匹配，从而实现无需人工精调超参的自适应闭环概率控制。
4. 渐进全变差（TV）有界性与概率达标误差控制理论：参数化逼近与对偶失配的严格解耦
论文给出了严格的收敛性理论证明。通过将实际参数化模型与总体最优解之间的全变差距离分解为“流匹配参数逼近误差”与“对偶变量失配误差”两项，证明了在对偶函数 \(g(\lambda)\) 具有 \(k\)-强凹性的保证下，渐进 TV 误差与约束达标界满足： $\(\limsup_{i \to \infty} \|P_{u'_\theta}^{(\lambda(i))} - P^{(\lambda^*)}\|_{\mathrm{TV}} \le \mathcal{E}_{\mathrm{PDFM}}, \quad \liminf_{i \to \infty} P_{u'_\theta}^{(\lambda(i))}(\mathcal{C}) \ge b - \mathcal{E}_{\mathrm{PDFM}}\)$ 误差界 \(\mathcal{E}_{\mathrm{PDFM}} = e_{\max} + \frac{2 e_{\max}}{k ((1-Q(\mathcal{C}))^{1/3} + Q(\mathcal{C})^{1/3})^3}\) 显式依赖于速度场拟合误差 \(\nu\)、局部优化次优度 \(\varepsilon\) 以及强凹常数 \(k\)。这意味着只要参数化网络容量充足且代理判别器充分收敛，生成分布将严格逼近理论最优重配测度，且约束满足率与目标值 \(b\) 之间的偏差由显式上界严格约束。
损失函数 / 训练策略¶
训练分为两个阶段交替进行： 1. 预训练/基准初始化：首先以标准 CFM 目标更新速度场参数 \(\theta\)，使速度场初步具备生成目标拓扑的基本流形结构；随后独立初始化时空代理判别器 \(f_\phi\)，利用欧拉采样轨迹与预言机二元标签预训练至交叉熵收敛； 2. 原对偶交替迭代：设置速度学习率 \(\eta^{(\theta)}\)（如 \(1\times 10^{-5}\)）、判别器学习率 \(\eta^{(\phi)}\)（如 \(1\times 10^{-4}\)）及对偶学习率 \(\eta_i\)（视任务规模设为 2 至 500）。每次主迭代中交替更新 \(\theta\) 与 \(\phi\) 各若干步，随后利用蒙特卡洛评估生成终端样本的经验通过率，以步长 \(\eta_i\) 投影更新对偶变量 \(\lambda'(i+1) = [\lambda'(i) + \eta_i(b - P_{u_\theta}(\mathcal{C}))]_+\)。若约束极难满足导致正样本匮乏，算法要求初始速度场至少能产生非零的可行样本以启动判别器分类边界的学习。
实验关键数据¶
主实验¶
论文在多维合成数据、混沌时空 PDE 仿真、生物分子生成及复杂图像拓扑等多个任务上进行了系统评测。
在 Kuramoto–Sivashinsky (KS) 混沌偏微分方程约束生成中，偏微分方程通过非可微前向数值步进器计算归一化残差 \(\delta_{L_n}\)，当 \(\delta_{L_n} < \tau\) 时判定为物理自洽。Case 1 采用纯净数据（训练集 \(Q(\mathcal{C})=1\)，设定阈值 \(\tau=0.033\)，目标达标率 \(b=98\%\)）；Case 2 采用空间每隔 9 个栅格进行线性插值破坏的退化数据（训练集 \(Q(\mathcal{C})\approx 39.47\%\)，设定阈值 \(\tau=0.05\)，目标达标率 \(b=90\%\)）。比较指标包括经验达标率 \(P(\mathcal{C})\)、平均 PDE 残差 \(\text{Avg } \delta_{L_n}\) 以及统计直方图误差 \(\mathrm{HE} = (\mathrm{HE}^{(1)}, \mathrm{HE}^{(2)})\)（分别对应一阶与二阶时空导数的边际分布匹配度）。
| 方法 | Case 1: \(P(\mathcal{C})\) (↑) | Case 1: Avg \(\delta_{L_n}\) (↓) | Case 1: \(\mathrm{HE}\) (↓) | Case 2: \(P(\mathcal{C})\) (↑) | Case 2: Avg \(\delta_{L_n}\) (↓) | Case 2: \(\mathrm{HE}_{\text{clean}}\) (↓) |
|---|---|---|---|---|---|---|
| Vanilla FM (\(Q\)) | 14.19% \(\pm\) 0.73% | 0.0409 \(\pm\) 0.0002 | (0.0087, 0.0453) | 9.82% \(\pm\) 0.77% | 0.0622 \(\pm\) 0.0005 | (0.0094, 0.0465) |
| Truncated FM (\(Q'\)) | N/A | N/A | N/A | 33.24% \(\pm\) 0.47% | 0.0539 \(\pm\) 0.0001 | (0.0135, 0.0477) |
| PCFM (推理打靶) | 6.03% \(\pm\) 0.38% | 0.0484 \(\pm\) 0.0001 | (0.0156, 0.1547) | 0.06% \(\pm\) 0.05% | 0.1022 \(\pm\) 0.0010 | (0.0173, 0.0729) |
| FM-RE (随机扰动) | 43.17% \(\pm\) 1.50% | 0.0339 \(\pm\) 0.0002 | (0.0117, 0.0329) | 46.33% \(\pm\) 1.01% | 0.0513 \(\pm\) 0.0003 | (0.0118, 0.0558) |
| PDFM (本文) | 97.74% \(\pm\) 0.53% | 0.0235 \(\pm\) 0.0001 | (0.0107, 0.0269) | 90.48% \(\pm\) 0.53% | 0.0406 \(\pm\) 0.0002 | (0.0118, 0.0435) |
消融与高维合成实验¶
在 8 维与 20 维混合高斯分布的 \(\ell_2\) 球约束任务（\(\mathcal{C}=\{x: \|x\|_2 \le 1\}\)，训练集无噪声 \(Q(\mathcal{C})=1\)）中，测试模型抵抗高维维度灾难并在保持切片瓦瑟斯坦距离（SWD）分布保真度的同时满足约束的能力。为公平对比，PDFM 将目标 \(b\) 设为 FM-RE 所达到的经验达标率（8-D 设 \(b=99.923\%\)，20-D 设 \(b=99.748\%\)）。
| 方法 | 约束类型/先验要求 | 8-D: \(P(\mathcal{C})\) (↑) | 8-D: SWD (↓) | 20-D: \(P(\mathcal{C})\) (↑) | 20-D: SWD (↓) |
|---|---|---|---|---|---|
| Vanilla FM | 无先验 / 无反馈 | 97.633% \(\pm\) 0.143% | 0.0193 \(\pm\) 0.0029 | 90.918% \(\pm\) 0.318% | 0.0087 \(\pm\) 0.0008 |
| MDM | 严格凸性 + 闭式双射映射 | 100% \(\pm\) 0.000% | 0.0217 \(\pm\) 0.0030 | 100% \(\pm\) 0.000% | 0.0144 \(\pm\) 0.0012 |
| NAMM | 可微连续距离函数 | 99.986% \(\pm\) 0.008% | 0.0228 \(\pm\) 0.0037 | 99.711% \(\pm\) 0.053% | 0.0205 \(\pm\) 0.0008 |
| PDM | 显式解析投影算子 | 100% \(\pm\) 0.000% | 0.0200 \(\pm\) 0.0031 | 100% \(\pm\) 0.000% | 0.0257 \(\pm\) 0.0018 |
| PCFM | 可微残差与雅可比矩阵 | 100% \(\pm\) 0.000% | 0.0284 \(\pm\) 0.0042 | 100% \(\pm\) 0.000% | 0.0148 \(\pm\) 0.0009 |
| FM-RE | 二元成员预言机 (加噪探索) | 99.923% \(\pm\) 0.025% | 0.0194 \(\pm\) 0.0029 | 99.748% \(\pm\) 0.052% | 0.0132 \(\pm\) 0.0010 |
| PDFM (本文) | 仅二元成员预言机 (确定性策略) | 99.906% \(\pm\) 0.026% | 0.0172 \(\pm\) 0.0028 | 99.755% \(\pm\) 0.005% | 0.0101 \(\pm\) 0.0007 |
关键发现¶
- 局部自洽无法免费泛化至生成分布：在 KS 系统的 Case 1 中，即使训练数据 \(100\%\) 服从物理规律，Vanilla FM 的生成满足率也断崖式下跌至 \(14.19\%\)；而依赖梯度打靶的 PCFM 在高度非线性混沌动态下发散（仅 \(6.03\%\)）。PDFM 在零阶黑盒反馈下将满足率拉升至 \(97.74\%\)，成功闭合了生成模型与物理规律的鸿沟。
- “青出于蓝”的抗噪纯化特性：在退化数据的 Case 2 中，训练样本自身仅有 \(39.47\%\) 合规，截断拒绝采样丢弃了约 \(60\%\) 的数据，样本量缩水导致严重过拟合且生成达标率仅有 \(33.24\%\)。PDFM 借助对偶约束自适应抑制不可行区域流形，实现了 \(90.48\%\) 的高合规率，且动力学直方图误差（\(\mathrm{HE}_{\text{clean}} = 0.0118\)）显著低于截断模型（\(0.0135\)），证明其具备从受损数据中逆向提纯高品质物理样本的能力。
- 训练效率与超参稳定性双重飞跃：由于不需要对整条采样轨迹进行高斯噪声前向重参数化，在单轮迭代耗时上，PDFM 相比 FM-RE 提速 \(2 \sim 3\times\)（8-D 实验由 4 分钟降至 1.5 分钟，KS Case 1 由 70 分钟降至 25 分钟）。更重要的是，PDFM 直接以目标满足率 \(b\) 作为显式输入，彻底摒弃了 FM-RE 需反复试凑惩罚项权重的巨大痛点。
亮点与洞察¶
- 将强化学习 DPG 思想升华入流匹配流形导引：通过为指示函数训练一个沿流路径保一致性的可微时间连续时空分类器 \(f_\phi\)，巧妙规避了黑盒预言机不可微、不可投影的死局，使得速度场更新在任意中间时刻都能获取精准的局部推进梯度。
- 原对偶框架在连续生成模型中的严格闭环控制：不同于启发式的正则化惩罚，PDFM 继承了凸优化中拉格朗日乘子的自适应纠偏特质，乘子 \(\lambda'\) 随经验超梯度实时升降，在“满足约束”与“保留原分布多模态结构”之间动态达成全局帕累托最优。
- 对不可微科学仿真的广阔泛化潜力：现代工程与工业软件中充斥着历史遗留的 Fortran/C++ 模拟器、商业黑盒求解器以及离散规则引擎。PDFM 仅需 True/False 的判定信号即可训练神经流模型，彻底解放了物理约束生成对自动微分工具链的强依赖。
局限与展望¶
- 冷启动极端负反馈依赖：作为分类器架构的代理判别器 \(f_\phi\)，要求在训练轨迹中必须同时存在正例和反例。如果目标约束极端狭窄且初识速度场生成的样本通过率为 \(0\)，判别器将退化为全负预测，失去空间导引梯度。未来可引入定向热启动引导或由粗到细的退火约束集。
- 数据分布与约束流形的拓扑重叠度制约：理论假设支撑集存在非空可行交集。当无约束训练数据分布 \(Q\) 与可行集 \(\mathcal{C}\) 几乎处于完全分离的几何空间时，CFM 似然损失与对偶惩罚将形成剧烈对抗，引发优化误差 \(\varepsilon\) 剧增。
相关工作与启发¶
- vs FM-RE (Huan et al., 2026)：两者均仅需二元归属预言机，但 FM-RE 采用随机扩散打靶与似然比估计，方差大、训练慢且只能间接调权重；PDFM 采用基于 DPG 的确定性时空代理与原对偶动态乘子，训练快 \(2\sim 3\times\) 且支持直接指定达成率 \(b\) 并具备 TV 收敛界。
- vs PCFM (Utkarsh et al., 2025) / NAMM (Feng et al., 2025)：PCFM 与 NAMM 强烈依赖解析可微的残差函数或显式距离场，在非线性混沌偏微分方程中极易陷入病态局部极值或梯度爆炸；PDFM 将非可微求解器视作完全黑盒，鲁棒性与适用面全面超越。
- vs MDM (Liu et al., 2023) / RFM (Xie et al., 2024)：镜像与反射类方法从理论上能达到 \(100\%\) 严格合规，但其应用严格受限于凸集或简单边界几何；PDFM 放宽至概率满足，换取了对任意非凸、多连通甚至离散组合约束集的泛化支撑。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立流匹配下针对二元黑盒预言机样本级约束的原对偶框架与全变差收敛保证。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D/高维合成、混沌 KS-PDE 仿真、SOCOFing 图像拓扑以及 QM9 分子稳定性，定量指标详实。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，从总体解析解到参数化误差分解的理论链条极其清晰。
- 价值: ⭐⭐⭐⭐⭐ 完美攻克了复杂物理工程科学仿真中非可微求解器参与深度生成建模的核心卡点。