Constraint-Aware Flow Matching: 面向约束采样的决策对齐端到端训练

# 面向约束采样的决策对齐端到端训练
来源：https://arxiv.org/html/2605.12754
## 约束感知流匹配：面向约束采样的决策对齐端到端训练

Jacob K\. Christopher 弗吉尼亚大学 csk4sr@virginia\.edu &James E\. Warner 美国宇航局兰利研究中心 james\.e\.warner@nasa\.gov &Ferdinando Fioretto 弗吉尼亚大学 fioretto@virginia\.edu

###### 摘要

深度生成模型在广泛应用中展现出最先进的性能，近期研究也表明其在科学与工程领域的适用性日益增强。尽管已有大量文献关注将基于物理的约束整合到生成过程中，现有方法仍难以在保持样本质量的同时强制执行严格的约束满足。特别是，无训练的约束采样方法虽然提供了每个样本的可行性保证，但引入了训练目标与约束采样过程之间的根本性不匹配，常常导致性能下降。本文识别出这种训练-采样错位是当前约束生成建模方法的核心局限，并提出了一种新颖的端到端框架——**约束感知流匹配（Constraint-Aware Flow Matching）**，该框架明确地将约束投影纳入训练目标。通过将模型学习到的动力学与约束采样过程对齐，该方法缓解了由投影校正引起的分布偏移，从而实现了高质量的约束生成。所提出的方法在三个具有挑战性的现实世界基准上进行了评估，展示了该方法的通用性和有效性。

## 1 引言

流匹配和扩散生成模型在广泛场景中展现出最先进的性能，代表了图像和视频生成[49 (https://arxiv.org/html/2605.12754#bib.bib34)， 30 (https://arxiv.org/html/2605.12754#bib.bib45)， 31 (https://arxiv.org/html/2605.12754#bib.bib46)]、工程[55 (https://arxiv.org/html/2605.12754#bib.bib29)]、材料科学[18 (https://arxiv.org/html/2605.12754#bib.bib35)， 14 (https://arxiv.org/html/2605.12754#bib.bib47)， 58 (https://arxiv.org/html/2605.12754#bib.bib49)]以及其他科学应用[16 (https://arxiv.org/html/2605.12754#bib.bib9)， 3 (https://arxiv.org/html/2605.12754#bib.bib48)， 28 (https://arxiv.org/html/2605.12754#bib.bib50)]的前沿。尽管这些模型具有巨大潜力，但其随机性常常导致生成结果类似于训练分布，却缺乏真实世界数据所必需的整体要求。特别是在科学领域，样本必须精确满足领域特定的约束和物理定律才能保持任何实际意义，生成模型无法遵守这些要求是阻碍其广泛采用的关键障碍。

针对这些挑战，近期文献开始探索科学领域的约束生成建模。相关工作可分为三类：(1) **约束引导**方法将生成过程条件化于约束集，但这些方法不提供正式的可行性证明，因此严重依赖拒绝采样[12 (https://arxiv.org/html/2605.12754#bib.bib30)， 42 (https://arxiv.org/html/2605.12754#bib.bib31), 10 (https://arxiv.org/html/2605.12754#bib.bib51)]；(2) **物理信息训练**通过物理残差损失增强生成模型训练，使采样分布与约束集对齐，但未能强制执行每个样本的满足[56 (https://arxiv.org/html/2605.12754#bib.bib1)， 4 (https://arxiv.org/html/2605.12754#bib.bib2)]；(3) **约束采样**为约束满足提供每个样本的保证，因此这些方法已被最广泛地采用[15 (https://arxiv.org/html/2605.12754#bib.bib3)， 35 (https://arxiv.org/html/2605.12754#bib.bib5)， 53 (https://arxiv.org/html/2605.12754#bib.bib6)， 59 (https://arxiv.org/html/2605.12754#bib.bib4)， 13 (https://arxiv.org/html/2605.12754#bib.bib26)]。然而，尽管约束采样显示出重大潜力，但与无约束采样方法相比，生成质量往往会下降。

本文认为这种性能权衡源于训练目标与约束采样过程之间的错位。现有的约束采样方法是无训练的，仅修改采样阶段，利用预训练的流匹配或扩散模型。虽然这使得约束集成成本低廉，无需重新训练模型，但也导致模型训练时的采样方式与部署时的使用方式存在根本性不同。一般来说，最先进的方法依赖于在采样过程中集成到可行集的投影，将样本推向可能在训练过程中未学习到的低密度区域[15 (https://arxiv.org/html/2605.12754#bib.bib3)， 35 (https://arxiv.org/html/2605.12754#bib.bib5)， 53 (https://arxiv.org/html/2605.12754#bib.bib6)]。图1(b)提供了一个示例：虽然初始预测样本落在数据分布的高密度区域，但投影在恢复可行性时会将样本推向低密度区域。本文通过引入约束感知训练程序来解决这一差距，确保训练与采样之间的对齐。

参见图注图1：**约束感知流匹配**与标准流匹配的可视化对比。标准流匹配的干净状态预测\( \hat{z}_1 \)落在分布的高密度区域，但投影会降低保真度。相反，我们的约束感知目标优化了下游任务，学习预测\( \hat{z}_1 \)使得投影落在高密度区域。**贡献**。本工作通过以下贡献解决了约束流匹配方法中训练与采样过程之间的脱节问题：(1) 提出了一种新的约束生成任务端到端公式，分析了无训练约束建模的失败模式（如图1所示）。(2) 受可微优化和决策聚焦学习[57 (https://arxiv.org/html/2605.12754#bib.bib12)， 54 (https://arxiv.org/html/2605.12754#bib.bib55)， 23 (https://arxiv.org/html/2605.12754#bib.bib58)， 21 (https://arxiv.org/html/2605.12754#bib.bib56)， 1 (https://arxiv.org/html/2605.12754#bib.bib16)， 5 (https://arxiv.org/html/2605.12754#bib.bib60)， 8 (https://arxiv.org/html/2605.12754#bib.bib61)， 22 (https://arxiv.org/html/2605.12754#bib.bib59)， 9 (https://arxiv.org/html/2605.12754#bib.bib62)， 40 (https://arxiv.org/html/2605.12754#bib.bib57)]领域的最新进展，推导了一种约束感知训练目标，将训练过程与下游生成任务对齐，提供了第一个用于约束采样方法的端到端训练方法。(3) 在三个现实世界科学场景（包括PDE约束生成、微天气预报和微结构生成）上进行了实证评估，所引入的**约束感知流匹配（CAFM）**在这些领域报告了最先进的性能。

## 2 相关工作

**约束引导与物理信息训练**。生成模型在受约束领域的早期应用寻求使用模型条件化方法来控制样本动力学，依赖基于学习的方法——要么训练独立分类器，要么直接利用扩散模型来引导生成[26 (https://arxiv.org/html/2605.12754#bib.bib27)， 27 (https://arxiv.org/html/2605.12754#bib.bib28)]。随着这些方法越来越多地应用于科学任务[55 (https://arxiv.org/html/2605.12754#bib.bib29)， 12 (https://arxiv.org/html/2605.12754#bib.bib30)， 42 (https://arxiv.org/html/2605.12754#bib.bib31)]，文献中开始探索更正式的约束感知训练概念，将物理信息神经网络的原理应用于生成语境[48 (https://arxiv.org/html/2605.12754#bib.bib32)， 47 (https://arxiv.org/html/2605.12754#bib.bib33)]。(2024 (https://arxiv.org/html/2605.12754#bib.bib2)) 提出将物理约束直接纳入扩散模型的训练目标中。Warner 等人 (https://arxiv.org/html/2605.12754#bib.bib1) (2026 (https://arxiv.org/html/2605.12754#bib.bib1)) 证明，物理和统计约束可以作为残差损失用于学习潜在流匹配的潜在空间[19 (https://arxiv.org/html/2605.12754#bib.bib52)]。然而，尽管这些工作为基于条件化的约束执行提供了更强的替代方案，但它们未能提供每个样本的保证，仅在分布层面上对齐采样过程，并且严重依赖神经网络来近似真实约束。

**约束采样**。最近，Christopher 等人 (https://arxiv.org/html/2605.12754#bib.bib3) (2024 (https://arxiv.org/html/2605.12754#bib.bib3)) 引入了扩散模型的约束采样方法，通过在生成过程中投影到可行集来合并约束校正。后续研究表明，约束采样在机器人技术[35 (https://arxiv.org/html/2605.12754#bib.bib5)]、生物学[16 (https://arxiv.org/html/2605.12754#bib.bib9)]和材料科学[18 (https://arxiv.org/html/2605.12754#bib.bib35)]等广泛领域提供了最先进的性能。采样时的方法已扩展到潜在扩散模型[59 (https://arxiv.org/html/2605.12754#bib.bib4)]以及流匹配参数化[53 (https://arxiv.org/html/2605.12754#bib.bib6)， 36 (https://arxiv.org/html/2605.12754#bib.bib36)]。然而，尽管这类方法在约束生成建模任务中贡献了最强的性能，但它们本质上仍是无训练的，训练与采样之间的这种脱节通常会导致按传统指标衡量的样本质量下降。

**可微优化**。我们的方法建立在优化过程可以嵌入可微计算流水线的一般原理之上。通过约束优化问题的微分早在 Gould 等人 (https://arxiv.org/html/2605.12754#bib.bib37) (2016 (https://arxiv.org/html/2605.12754#bib.bib37)) 就被认为对机器学习设置具有重要意义。这一研究路线主要应用于端到端学习设置，使机器学习模型能够将约束规划整合到其训练过程中。Amos 和 Kolter (https://arxiv.org/html/2605.12754#bib.bib15) (2017 (https://arxiv.org/html/2605.12754#bib.bib15)) 提出使用二次规划层来在网络的前向传播中求解约束优化问题。虽然这种方法在后向传播中利用了KKT条件的隐式微分，但可微优化层已被广泛推广[1 (https://arxiv.org/html/2605.12754#bib.bib16)， 33 (https://arxiv.org/html/2605.12754#bib.bib7)， 57 (https://arxiv.org/html/2605.12754#bib.bib12)， 39 (https://arxiv.org/html/2605.12754#bib.bib17)， 24 (https://arxiv.org/html/2605.12754#bib.bib18)， 46 (https://arxiv.org/html/2605.12754#bib.bib19)， 50 (https://arxiv.org/html/2605.12754#bib.bib20)， 43 (https://arxiv.org/html/2605.12754#bib.bib21)]。尽管许多方法仅限于凸约束集，Kotary 等人 (https://arxiv.org/html/2605.12754#bib.bib7) (2023 (https://arxiv.org/html/2605.12754#bib.bib7)) 和 Blondel 等人 (https://arxiv.org/html/2605.12754#bib.bib63) (2022 (https://arxiv.org/html/2605.12754#bib.bib63)) 展示了通过顺序二次规划和不动点微分对非凸约束的适用性，从而实现了更广泛的用例通用性。

## 3 预备知识：流匹配

**流匹配**[37 (https://arxiv.org/html/2605.12754#bib.bib39)]。流匹配是一种生成建模框架，其中神经网络学习一个时间相关的速度场 \( v_\theta: \mathbb{R}^d \times [0,1] \to \mathbb{R}^d \)，该场定义了一个确定性流 \( \psi_t: \mathbb{R}^d \to \mathbb{R}^d \)，将样本从基础分布 \( p_0 \) 传输到目标分布 \( p_1 \)。给定初始条件 \( z_0 \sim p_0 \)，生成轨迹通过求解ODE获得：

\[
\frac{d}{dt} \psi_t(z_0) = v_\theta\!\bigl(\psi_t(z_0), t\bigr), \qquad \psi_0(z_0) = z_0,
\] 
(1)

生成的样本为 \( z_1 = \psi_1(z_0) \)。流匹配中的训练使用从配对端点 \( (\mathbf{z}_0, \mathbf{z}_1) \) 构建的**参考传输路径**，其中 \( \mathbf{z}_0 \sim p_0 \) 和 \( \mathbf{z}_1 \sim p_1 \) 是耦合的（例如，通过最优传输计划）。该参考路径上的时间 \( t \) 点记为 \( \mathbf{z}_t \)（例如，\( \mathbf{z}_t = (1-t)\mathbf{z}_0 + t\mathbf{z}_1 \)，可选地添加少量噪声），对应的参考速度为 \( \mathbf{z}_1 - \mathbf{z}_0 \)。然后通过回归到该参考速度来训练速度网络 \( v_\theta \)：

\[
\mathcal{L}_{FM} = \| \overbrace{v_\theta(\mathbf{z}_t, t)}^{\text{预测}} - \overbrace{(\mathbf{z}_1 - \mathbf{z}_0)}^{\text{真实速度}} \|^2.
\] 
(2)

**符号说明**。粗体变量 \( \mathbf{z}_0, \mathbf{z}_t, \mathbf{z}_1 \) 表示由**训练**耦合/路径（直接从配对端点构建）诱导的状态，而非粗体变量 \( z_0, z_t, z_1 \) 表示由 (1) 中的**学习生成流**诱导的状态。这一区别很重要，因为训练时在从参考路径采样的 \( \mathbf{z}_t \) 上评估 \( v_\theta \)，而生成时模型在积分学习动力学产生的 \( z_t \) 上被查询。

## 4 约束感知采样

约束生成的目标是生成满足一组可行性要求 \( C \) 的样本 \( z_1 \sim p_{\text{data}} \)。虽然真实数据分布 \( p_{\text{data}} \) 通常可能已经通过构造满足相关约束（例如，物理系统通常遵守底层控制定律），但实际挑战在于，学习到的生成模型 \( v_\theta \) 尽管近似 \( p_{\text{data}} \)，但可能将概率质量分配给环境空间的不可行区域。如果 \( \mathbbm{1}(\cdot) \) 表示指示函数，则**理想化的约束目标**为：

\[
p_C(z_1) \propto p_{\text{data}}(z_1) \mathbbm{1}\{z_1 \in C\},
\] 
(3)

实际目标通过将 \( p_{\text{data}} \) 替换为 \( v_\theta \) 获得的学习近似 [16 (https://arxiv.org/html/2605.12754#bib.bib9)]。在不同应用中，约束集 \( C \) 的选择基于