SNAP-FM: 物理约束生成建模中的稀疏非线性加速投影

# SNAP-FM：物理约束生成建模的稀疏非线性加速投影
来源：https://arxiv.org/html/2607.00095
Theodoros Xenakis, Utkarsh Utkarsh, Pengfei Cai, Rafael Gómez-Bombarelli, Alan Edelman, Christopher V. Rackauckas

###### 摘要

生成模型已成为物理仿真的可扩展替代方法，但它们无法保证输出结果遵循支配底层物理学的守恒定律、边界条件和非线性不变量。约束采样填补了这一空白，在推理时无需重新训练即可精确强制此类约束，但计算代价高昂：在采样过程中重复进行投影、校正和轨迹优化步骤，对于非线性约束，这些步骤变得昂贵。标准的机器学习框架加剧了这一问题：它们密集的张量代数和有限的稀疏求解器组合性掩盖了物理约束自然诱导的结构，使得高效的批量非线性优化在实践中难以实现。我们通过利用样本级批处理和局部偏微分方程耦合在投影子问题中诱导的结构（即块稀疏雅可比矩阵和KKT系统）来解决这一瓶颈，使用ExaModels.jl暴露这一结构，并使用MadNLP.jl和GPU稀疏分解求解生成的稀疏非线性规划。将其应用于物理约束流匹配（PCFM），在具有线性、非线性、一维和二维约束的PDE基准测试中，该方法在保持约束满足的同时加速了非线性约束投影。这些结果表明，稀疏GPU非线性优化是科学机器学习中约束生成采样的实用基础。

机器学习, ICML

## 1 引言

生成模型已成为物理仿真的灵活替代方法，学习偏微分方程（PDE）的解分布，并跨不同物理条件摊销推理（Price等，2023 (https://arxiv.org/html/2607.00095#bib.bib22)；Yuan等，2023 (https://arxiv.org/html/2607.00095#bib.bib23)；Huang等，2024 (https://arxiv.org/html/2607.00095#bib.bib21)；Utkarsh等，2025a (https://arxiv.org/html/2607.00095#bib.bib1)）。然而，它们在科学场景中的部署受到一个根本缺口的限制：无约束生成模型本身无法保证物理保真度。质量、动量和能量守恒、非线性边界条件以及与控制方程相关的不变量是经典数值模拟的核心，但除非明确强制，否则通常会被学习到的替代模型违反（Raissi等，2019 (https://arxiv.org/html/2607.00095#bib.bib15)；Li等，2021 (https://arxiv.org/html/2607.00095#bib.bib19)）。在不牺牲生成模型摊销推理优势的情况下弥合这一差距，是物理约束生成建模的核心挑战。

生成采样中的约束强制可大致分为软方法和硬方法。软方法包括训练时的惩罚损失（Baldan等，2025 (https://arxiv.org/html/2607.00095#bib.bib29)；Huang等，2024 (https://arxiv.org/html/2607.00095#bib.bib21)）、PINN风格残差正则化、物理信息神经算子以及架构级归纳偏置（Greydanus等，2019 (https://arxiv.org/html/2607.00095#bib.bib27)；Richter-Powell等，2022 (https://arxiv.org/html/2607.00095#bib.bib28)），它们近似地鼓励约束满足。这些方法在计算上具有吸引力，并且自然适用于现代深度学习流水线，但它们通常不提供精确的可行性保证，并且在分布偏移下可能表现出增加的约束违反。

硬约束方法则旨在精确强制可行性，要么通过推理时的校正和优化（Utkarsh等，2025a (https://arxiv.org/html/2607.00095#bib.bib1)；Christopher等，2024 (https://arxiv.org/html/2607.00095#bib.bib24)；Cheng等，2025 (https://arxiv.org/html/2607.00095#bib.bib17)；Römer等，2024 (https://arxiv.org/html/2607.00095#bib.bib25)；Yuan等，2023 (https://arxiv.org/html/2607.00095#bib.bib23)），要么通过端到端的约束公式（Utkarsh等，2025b (https://arxiv.org/html/2607.00095#bib.bib31)）。在最简单的情况下，可以应用最终的后期投影将生成的样本映射到可行流形上，但这种延迟校正可能会引入显著的分布失真。为了减少这种不匹配，许多方法将校正与采样动力学交错进行，反复引导或投影中间状态走向可行性（Utkarsh等，2025a (https://arxiv.org/html/2607.00095#bib.bib1)；Cheng等，2025 (https://arxiv.org/html/2607.00095#bib.bib17)；Christopher等，2024 (https://arxiv.org/html/2607.00095#bib.bib24)；Ben-Ham等，2024 (https://arxiv.org/html/2607.00095#bib.bib18)）。松弛变体进一步延迟或软化早期校正，反映了早期迭代通常仍类似噪声，可能因过于严格的约束强制而受损。

尽管存在这些算法差异，硬约束测试时方法共享一个共同的计算负担：在采样过程中通过重复优化来强制可行性（Utkarsh等，2025a (https://arxiv.org/html/2607.00095#bib.bib1)；Cheng等，2025 (https://arxiv.org/html/2607.00095#bib.bib17)；Christopher等，2024 (https://arxiv.org/html/2607.00095#bib.bib24)）。在基于投影的方法中，这种负担尤为明显：每次校正通过求解一个约束优化问题将中间状态映射到约束流形上或更接近它。对于线性约束，此投影通常简化为一次矩阵求解。对于支配许多物理系统的非线性守恒律，包括通量约束、积分不变量和非线性边界条件，同样的步骤需要迭代非线性优化。由于此优化必须在采样步骤之间以及生成的样本批次中重复执行，它可能主导硬约束生成的总成本。因此，核心问题不仅在于是否强制硬约束，还在于如何使由此产生的投影和校正步骤足够快，以便在大规模下实用。

大规模优化中已知的算法响应是：利用约束诱导的稀疏性，并使用稀疏线性代数求解所得的KKT系统（Nocedal and Wright，2006 (https://arxiv.org/html/2607.00095#bib.bib8)；Pacaud and Shin，2024 (https://arxiv.org/html/2607.00095#bib.bib12)；Rennich等，2014 (https://arxiv.org/html/2607.00095#bib.bib32)；Lu and Yang，2025 (https://arxiv.org/html/2607.00095#bib.bib33)）。然而，这一策略难以在现代生成采样流水线内部部署。标准ML框架主要针对密集批处理张量代数进行优化，尽管存在稀疏张量支持，但它尚未为批处理非线性规划、内点全局化、GPU驻留稀疏KKT分解以及与采样循环的集成提供成熟的端到端栈（Paszke等，2019 (https://arxiv.org/html/2607.00095#bib.bib34)；Bradbury等，2018 (https://arxiv.org/html/2607.00095#bib.bib35)）。约束优化实践与深度生成建模基础设施之间的这一差距激发了我们的方法。

批处理硬约束采样中出现的投影NLP具有高度结构性。跨批次，约束雅可比矩阵是块对角化的，因为每个样本的约束相互独立。在每个块内，雅可比矩阵是稀疏的，因为物理守恒律和PDE离散化仅耦合局部空间或时间自由度。这种两级结构产生了适合GPU驻留稀疏分解的稀疏KKT系统，避免了否则会使非线性投影变得极其昂贵的密集张量化。

我们通过SNAP-FM实现了这一想法，这是一个基于物理约束流匹配（PCFM）（Utkarsh等，2025a (https://arxiv.org/html/2607.00095#bib.bib1)）的稀疏非线性投影框架。SNAP-FM使用ExaModels.jl符号编译结构化的投影NLP，使用MadNLP.jl作为内点求解器，并使用GPU稀疏分解求解所得的KKT系统。基于PCFM（在预训练流模型中零样本强制任意非线性约束），SNAP-FM针对约束强制的主要计算瓶颈：采样过程中重复的非线性投影。在热传导、反应扩散、Burgers方程和二维Navier-Stokes基准测试中，相对于通用优化基线，SNAP-FM加速了投影步骤，同时保持了约束满足。

## 2 物理约束流匹配

本工作直接建立在Utkarsh等人（2025a (https://arxiv.org/html/2607.00095#bib.bib1)）提出的物理约束流匹配框架之上，目标是将其扩展到当前实现不支持的一个领域：大规模非线性约束的高效强制。在近期约束感知的生成模型中，PCFM是我们工作的自然基础，也显示出扩展到不同科学领域（如原子生成模型）的有前景拓展（Cai等，2026 (https://arxiv.org/html/2607.00095#bib.bib36)）。它在推理时精确确保约束满足，无需重新训练或修改底层流模型的架构，完全以后验方式操作。它在推理时将中间流状态投影到约束流形上，无需训练期间的梯度信息，从而实现对预训练流匹配模型的机器精度级零样本约束强制。与之前局限于线性或非重叠约束的零样本方法不同，它在单一框架内允许任意非线性和耦合约束。这些特性使PCFM成为构建更通用约束采样机制的坚实基础。

### 2.1 采样

PCFM中使用的采样算法是一种约束引导算法，将轻量级约束校正与边际一致的流更新交错进行。该过程包括四个主要步骤：前向射击、高斯-牛顿投影、反向更新和松弛约束校正。

设vθ(u,τ)表示预训练流模型，它定义了一个ODE，将样本u0∼π0从可处理先验传输到类似解的输出u1∼π1，流时间τ∈[0,1]。给定约束函数h(u)，PCFM通过将[0,1]离散化为N个均匀子步骤来强制h(u1)=0。在每个时间步，τ→τ′=τ+Δτ，执行三个操作：使用学习到的速度向前外推一个步骤到终端时间，将得到的候选解投影到约束流形M={u:h(u)=0}上，以及沿最优传输位移线性插值回到τ′（Utkarsh等，2025a (https://arxiv.org/html/2607.00095#bib.bib1)）：

û1 = uτ + (1-τ) vθ(uτ, τ), (1)
u1 = arg min_u ½ ‖u - û1‖² s.t. h(u)=0, (2)
uτ′ = u0 + τ′ (u1 - u0). (3)

Utkarsh等人（2025a (https://arxiv.org/html/2607.00095#bib.bib1)）的完整PCFM算法还允许在每个步骤进行惩罚校正：arg min_u ‖u - ûτ′‖² + λ‖h(u + (1-τ′)vθ(u, τ′))‖²，以补偿粗离散化下的非线性。然而，我们省略了它，因为在我们的实验中，简化的三步变体已足够。

本工作中利用的关键结构特征是上述循环与约束无关。除了步骤2中的投影外，每个操作都是固定的算术内核，问题学习的整个规格都被汇集到投影子问题中。因此，单次采样过程的成本可分为两个主要部分：约束投影和速度评估。前向射击和OT回拉与所强制的约束无关，而投影步骤的规模随施加的约束数量及其复杂性而变化，导致投影步骤成为采样的瓶颈，尤其是在这些更复杂的情况下。本工作的其余部分通过利用雅可比矩阵中的稀疏性来优化投影步骤，因为约束在样本之间是独立的，且时间耦合通常是局部的。

## 3 理论

在本节中，我们讨论PDE系统常见的约束，并推导相关约束雅可比矩阵的结构特性。

我们采用类似于Utkarsh等人（2025a (https://arxiv.org/html/2607.00095#bib.bib1)）的记号。具体来说，我们考虑有界时空域Ω×[0,T]上的偏微分方程，其中Ω⊂R^d。真实的PDE解是一个场u: Ω×[0,T]→R，我们使用离散化张量u∈R^{Nx×Nt}来近似，其中Nx和Nt分别是空间和时间网格点的数量。

除了满足支配PDE外，解还必须满足一系列物理约束。这些约束表示为

(H_k(u))_k = 0，其中H_k是作用于解u的第k个约束。为方便起见，我们令H(u)表示所有施加约束的串联。进一步隐含约束算子也可以作用于离散化解u。

虽然本节讨论集中于一维PDE，但主要思想自然扩展到更高维设置。

### 3.1 约束

为PDE施加的约束通常源于物理原理，如守恒律、初始条件和边界条件（LeVeque, 2004 (https://arxiv.org/html/2607.00095#bib.bib11), 1992 (https://arxiv.org/html/2607.00095#bib.bib10)）。这些约束可以是u的线性或非线性，并且可以在空间和时间上局部或全局作用。常见约束的选择总结在表1 (https://arxiv.org/html/2607.00095#S3.T1)中。

表1：PDE常见约束类型概述。实践中遇到的大多数约束在时间上是局部的。即，它们仅依赖于单个时间步或少量相邻时间步。例如，初值约束仅适用于t=0，而质量守恒约束通常独立适用于每个时间步。即使是涉及时间导数的非线性守恒律，也可以使用仅耦合少量相邻时间步的有限差分离散化，从而保持时间局部性（Quarteroni等，2007 (http