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

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.

• The paper introduces a sparse GPU infrastructure that combines ExaModels.jl, MadNLP.jl, and cuDSS to accelerate the nonlinear projection bottleneck in physics-constrained flow matching.
• Benchmarks across six PDE problems show the approach is fastest overall, delivering up to 53× GPU speedups over CPU execution and remaining tractable for constraint sets where other solvers fail.
• The method preserves hard-constraint feasibility with residuals near solver tolerance and scales approximately linearly with batch size, although GPU benefits depend on problem size and constraint complexity.

Overview

SNAP-FM addresses the principal computational bottleneck of hard-constrained generative sampling: the repeated nonlinear projection of intermediate states onto constraint manifolds during inference. The work builds on Physics-Constrained Flow Matching (PCFM) (Utkarsh et al., 4 Jun 2025), which enforces arbitrary nonlinear constraints zero-shot on pretrained flow matching models by interleaving Gauss-Newton projections with the sampling dynamics. While PCFM guarantees feasibility, each sampling step requires solving a constrained nonlinear program, and for nonlinear conservation laws this optimization dominates end-to-end cost. SNAP-FM's contribution is not a new sampling algorithm but an infrastructure-level one: exploiting the block-sparse Jacobian and KKT structure induced by sample-wise batching and local PDE couplings, and solving the resulting sparse NLPs on GPU with ExaModels.jl, MadNLP.jl, and NVIDIA's cuDSS sparse factorization.

Background: PCFM and the projection bottleneck

PCFM transports samples from a prior $u_0 \sim \pi_0$ to the data distribution via a learned velocity field $v_\theta(u,\tau)$ over flow time $\tau \in [0,1]$. At each of $N$ substeps, the sampler performs three operations: a forward shoot to terminal time, a projection of the candidate onto the manifold $\mathcal{M} = \{u : h(u) = 0\}$ via

$$u_1 = \arg\min_u \tfrac{1}{2}\|u - \hat{u}_1\|^2 \quad \text{s.t.} \quad h(u) = 0,$$

and a linear interpolation back to $\tau'$ along the optimal-transport displacement. The authors deliberately omit PCFM's penalized correction term, finding the three-step variant sufficient in their experiments.

The key structural observation is that the projection is the only constraint-dependent operation; the remaining kernels are fixed arithmetic. For nonlinear constraints (flux conservation, integral invariants, nonlinear boundary conditions), the projection requires iterative nonlinear optimization, repeated across sampling steps and batch elements. Standard ML frameworks, optimized for dense batched tensor algebra, lack an end-to-end stack for batched nonlinear programming with GPU-resident sparse KKT factorization, which is the infrastructure gap SNAP-FM targets.

Jacobian and KKT structure

The projection NLP is solved via Newton iterations on the KKT system

$$\begin{bmatrix} W & J_{\mathcal{H}}^T \ J_{\mathcal{H}} & 0 \end{bmatrix} \begin{bmatrix} \Delta x \ \Delta \lambda \end{bmatrix} = -\begin{bmatrix} \nabla f + J_{\mathcal{H}}^T\lambda \ h(u) \end{bmatrix}.$$

The paper establishes two levels of sparsity. Within a single sample, physically motivated constraints (initial/boundary conditions, mass conservation, local flux residuals discretized by finite differences or Riemann sums) couple only local spatial or temporal degrees of freedom, so each Jacobian row contains $\mathcal{O}(N_x)$ or $\mathcal{O}(N_t)$ nonzeros rather than $v_\theta(u,\tau)$0. Across a batch of $v_\theta(u,\tau)$1 samples, the Jacobian is block diagonal because sample constraints are independent, so batch growth adds independent blocks without cross-sample coupling — a property the scaling experiments confirm empirically. The sparsity pattern is fixed across optimizer iterations, permitting reuse of symbolic factorization structure. An appendix works through a small $v_\theta(u,\tau)$2 example making the block structure explicit.

Implementation stack

SNAP-FM composes three components. ExaModels.jl symbolically compiles the projection NLP into sparsity-preserving, SIMD-parallel GPU kernels, evaluating objectives, constraints, Jacobians, and Hessians in a single pass — avoiding the AD overhead that dominates general-purpose modeling layers such as JuMP.jl. MadNLP.jl provides a primal-dual interior-point solver that exploits the fixed sparsity pattern across iterations. cuDSS performs the sparse $v_\theta(u,\tau)$3/$v_\theta(u,\tau)$4/Cholesky factorizations on GPU. The benchmark suite compares this stack against five baselines designed to isolate each axis: ExaModels+MadNLP on CPU (isolating GPU hardware), JuMP+MadNLP on CPU (isolating the modeling layer), JuMP+Ipopt, Optimization.jl+IPNewton, and Optimization.jl+L-BFGS. The L-BFGS baseline converts the problem to a penalized unconstrained form, so its runtimes are not comparable as solutions to the hard-constrained projection.

Benchmark results

Six PDE problems span linear and nonlinear constraints in 1D and 2D: heat equation (two constraint regimes), reaction–diffusion, Burgers (two regimes), and 2D Navier–Stokes in vorticity form. All runs use 32 generated samples (2 for Navier–Stokes) on an NVIDIA L40S, with the neural operator always on GPU to isolate projection cost.

Problem	Best method	Time (s)	Runner-up	Time (s)
Heat (IC, mass)	Exa+MadNLP GPU	11.73	Exa+MadNLP CPU	14.29
Heat (IC, mass, energy, PDE)	Exa+MadNLP GPU	110.26	Exa+MadNLP CPU	3501.11
Burgers (BC, mass)	Exa+MadNLP CPU	109.74	Exa+MadNLP GPU	153.53
Burgers (IC, mass, flux)	Exa+MadNLP GPU	38.83	Exa+MadNLP CPU	2048.83
Reaction–diffusion	Exa+MadNLP CPU	24.58	Exa+MadNLP GPU	35.68
Navier–Stokes (2D)	Exa+MadNLP GPU	45.69	JuMP+Ipopt CPU	59.38

The ExaModels+MadNLP combination is fastest in all six problems, and the JuMP-vs-ExaModels comparison on identical solvers isolates a consistent speedup from compile-time sparsity exploitation — roughly 12× on the simplest heat problem (137.94 s vs 14.29 s). The GPU-vs-CPU gap widens dramatically with constraint nonlinearity: on the heat equation with the energy-evolution constraint, GPU execution is ~32× faster (110 s vs 3501 s), and on Burgers with Godunov flux constraints, ~53× faster (39 s vs 2049 s). The authors attribute this to the harder nonlinear optimization making per-iteration linear algebra dominant.

The failure modes are equally informative. On Burgers with initial-value and flux constraints, every method except ExaModels+MadNLP failed to converge within a six-hour budget. IPNewton additionally failed on Burgers with boundary constraints and on Navier–Stokes. In the constraint-count scaling study on Burgers, JuMP+MadNLP, Ipopt, IPNewton, and L-BFGS all fail beyond the IC+mass regime, while ExaModels+MadNLP on CPU fails at Flux(10); the GPU backend is the only method tractable across the full range, growing from 6.90 s (IC-only) to 62.95 s (IC, mass, Flux(10)).

Constraint satisfaction and correctness

Infeasibility is measured as the average magnitude of per-constraint violations, normalized within constraint type. Residual trajectories for Burgers mass conservation show that structure-exploiting methods (ExaModels+MadNLP on both devices, JuMP+MadNLP) oscillate around zero at magnitude $v_\theta(u,\tau)$5, whereas IPNewton exhibits systematic undershoot-then-overshoot bias and L-BFGS drifts monotonically negative — a qualitative distinction beyond magnitude. Generated heat-equation samples from ExaModels+MadNLP reproduce the analytic reference to visual accuracy, while JuMP, IPNewton, and L-BFGS produce visibly different samples from identical input noise, indicating that solver choice affects the generated sample beyond runtime.

An important caveat applies throughout: constraints are solved only to the interior-point solver's default feasibility tolerance of $v_\theta(u,\tau)$6, so reported infeasibility reflects post-convergence residual violation, not exact satisfaction. The paper notes explicitly that tighter tolerances might reduce violation at the cost of runtime or robustness, and that tuning this trade-off is left open. Notably, the GPU variant reports higher infeasibility than CPU in some regimes (e.g., 3.19e-6 vs 3.23e-7 on the first heat problem; 0.920 vs 1.594 on Burgers IC), and the paper does not fully explain this variation. Two scaling results round out the empirical picture: runtime grows approximately linearly with batch size for all methods, consistent with the block-diagonal structure, and method ordering is preserved across batch sizes.

Limitations and open questions

Several limitations are conceded directly. The GPU advantage is not universal — the CPU variant is faster on Burgers with boundary constraints and reaction–diffusion, suggesting GPU sparse factorization overheads can dominate for smaller or less nonlinear problems, and the paper does not characterize the crossover. The feasibility-tolerance trade-off is unresolved. The Navier–Stokes experiment uses only 2 samples due to state dimension, leaving large-batch 2D performance untested. The streamlined three-step PCFM variant (omitting the penalized correction) is validated only on these benchmarks; behavior under coarse temporal discretization of highly nonlinear constraints is not examined. Finally, the claim that the approach generalizes beyond 1D rests on a single 2D problem at modest resolution.

Conclusion

SNAP-FM demonstrates that the projection subproblem of hard-constrained flow matching sampling can be accelerated by one to two orders of magnitude by compiling the structured NLP with ExaModels.jl, solving it with MadNLP.jl's interior-point method, and factorizing the sparse KKT systems on GPU. The approach maintains constraint satisfaction comparable to CPU execution, is the only tested backend that remains tractable under the largest nonlinear constraint sets, and exhibits batch-linear scaling consistent with the block-diagonal Jacobian structure. The results establish sparse GPU nonlinear optimization as a practical substrate for constrained generative sampling, while leaving open tolerance tuning, the GPU-vs-CPU crossover regime, and scaling to higher-dimensional states with large batches.