SNAP-FM: Sparse GPU Optimization Accelerates Physics-Constrained Generative Models
New method enforces physical laws in generative models faster using sparse GPU solvers.
Generative models are increasingly used as surrogates for physical simulations, but they often fail to respect conservation laws and boundary conditions. Standard constrained sampling methods enforce these at inference time but are computationally expensive due to dense tensor algebra and limited sparse solver support in ML frameworks. SNAP-FM, developed by researchers at MIT, addresses this by leveraging the inherent block-sparse structure in Jacobian and KKT systems arising from sample-wise batching and local PDE couplings. Using Julia's SciML ecosystem for differentiable programming and NVIDIA GPU sparse factorization, SNAP-FM achieves efficient batched nonlinear optimization.
The method is demonstrated on Physics-Constrained Flow Matching (PCFM), a state-of-the-art generative approach for PDEs. On benchmarks with linear, nonlinear, 1D, and 2D constraints, SNAP-FM significantly accelerates nonlinear constraint projection while maintaining exact constraint satisfaction. This sparse GPU nonlinear optimization framework opens the door to practical, physics-compliant generative sampling for scientific machine learning, potentially revolutionizing computational fluid dynamics, structural mechanics, and other engineering domains where physical fidelity is critical.
- Exploits block-sparse Jacobian and KKT system structures from local PDE couplings for efficient GPU optimization.
- Uses Julia's SciML ecosystem and NVIDIA GPU sparse factorization to accelerate batched nonlinear constraint projection.
- Applied to Physics-Constrained Flow Matching, showing speedups on linear, nonlinear, 1D, and 2D benchmark PDE constraints.
Why It Matters
Enables fast, physically accurate generative simulations for scientific and engineering applications.