SNAP-FM: Sparse Nonlinear Accelerated Projection for Physics-Constrained Generative Modeling

2026-06-30

AI Summary

Existing physics-constrained generative models struggle to efficiently enforce physical laws—such as conservation principles and boundary conditions—during inference, primarily due to the computational intensity of constrained sampling and the underutilization of inherent sparsity. This work proposes a highly efficient, training-free projection method that explicitly exploits the block-sparse structure of the Jacobian and KKT systems induced by batched samples and local PDE couplings. By integrating ExaModels.jl for modeling, MadNLP.jl for nonlinear programming, and GPU-accelerated sparse factorization, the approach dramatically accelerates nonlinear constraint projection. Evaluated on diverse one- and two-dimensional linear and nonlinear PDE benchmarks, the method achieves substantial speedups in sampling while rigorously preserving physical constraints, thereby demonstrating the practical utility of sparse GPU-based nonlinear optimization in scientific machine learning.

Abstract

Generative models have emerged as scalable surrogates for physical simulation, yet they offer no guarantee that their outputs respect the conservation laws, boundary conditions, and nonlinear invariants that govern the underlying physics. Constrained sampling closes this gap, enforcing such constraints exactly at inference time without retraining, but at a computational cost: projection, correction, and trajectory-optimization steps are repeated during sampling, with these steps becoming expensive for nonlinear constraints. Standard ML frameworks exacerbate this: their dense tensor algebra and limited sparse solver composability obscure the structure that physical constraints naturally induce, making efficient batched nonlinear optimization difficult to realize in practice. We address this bottleneck by exploiting the structure that sample-wise batching and local PDE couplings induce in the projection subproblems -- namely, block-sparse Jacobian and KKT systems -- exposing this structure using ExaModels.jl and solving the resulting sparse nonlinear programs with MadNLP.jl and GPU sparse factorization. Applied to Physics-Constrained Flow Matching (PCFM), on PDE benchmarks with linear, nonlinear, one-dimensional, and two-dimensional constraints, this approach accelerates nonlinear constraint projection while maintaining constraint satisfaction. These results show that sparse GPU nonlinear optimization is a practical foundation for constrained generative sampling in scientific machine learning.

Problem

• physics-constrained generative modeling
• nonlinear constraints
• sparse optimization
• constrained sampling
• scientific machine learning

Innovation

• sparse nonlinear optimization
• physics-constrained generative modeling
• GPU acceleration
• block-sparse structure
• constrained sampling