PaperlayerView exact revision on arXiv

Source-linked AI summary

diff-soup: GPU-native Differentiable Simulation and Inference for Large-scale Stiff Dynamical Systems

Samuel Degnan-Morgenstern, Martin Z. Bazant

arXiv:2610.37333v1cs.MSphysics.comp-ph

TL;DR

Large, stiff PDE-derived systems demand repeated sparse linear solves and differentiable simulation, while established approaches leave large GPU-native solves difficult to achieve. diff-soup combines sparse-oriented JAX simulation and inference with GPU cuDSS solving. It reports up to 50× faster simulations and 80× faster gradients than established CPU libraries.

  • Problem

    Large stiff PDE-derived systems require repeated sparse linear solves, while established frameworks are largely CPU-based or rely on dense linear algebra for JAX implicit integration.

    Source: PROGRAM SUMMARY · p. 2

  • Method

    diff-soup combines sparse-oriented differentiable JAX simulation and inference with GPU cuDSS solving and multi-GPU parallelism.

    Source: §1 · pp. 4–5

  • Results

    Up to 50× faster simulations and 80× faster gradients than established CPU libraries.

    Source: §6 · p. 48

  • Takeaways & Limitations

    Its end-to-end gradients support PDE-constrained optimization across classical inference and modern learning workflows.

    Source: §6 · p. 48

    In spinodal decomposition, measurements rarely fall inside the spinodal gap, so learned functions there rely primarily on the prior and have considerably larger uncertainty.

    Source: §5.1 · p. 46

Abstract

from arXiv · show

Learning physical models from spatiotemporal data requires simulating systems of partial differential equations (PDEs) and differentiating through their solution. Upon spatial discretization, these models give rise to stiff dynamical systems that require repeatedly solving large, sparse linear systems. Traditional scientific computing frameworks are largely CPU-based and require problem-specific tuning to accelerate, scale, and differentiate. In this work, we introduce diff-soup, a differentiable, sparse-oriented, unified PDE suite for GPU-native simulation and inference of large, stiff dynamical systems. Built in JAX, diff-soup provides 1) a fully differentiable design for problems with sparse structure, 2) scaling to large problems through the NVIDIA cuDSS sparse direct solver with model and data multi-GPU parallelism, and 3) end-to-end PDE-constrained optimization, from classical statistical inference to embedded deep learning. On stiff PDE benchmarks, diff-soup outperforms established CPU- and GPU-based frameworks by one to two orders of magnitude in both forward solves and gradients. We demonstrate both a classical Bayesian inference workflow, identifying constitutive relationships, and a modern deep learning workflow, training a neural field for the catalyst activity of a three-dimensional packed bed reactor. We provide diff-soup as an open-source repository to facilitate research in computational science and engineering.

PROGRAM SUMMARY

diff-soup targets large, stiff PDE-derived systems that require repeated sparse linear solves, a workload not well served by CPU-centered or dense-linear-algebra frameworks. It provides a differentiable JAX suite for modeling, GPU simulation, and inference.

PROGRAM SUMMARY

  • Spatially discretized time-dependent PDEs produce large stiff ODE and DAE systems whose repeated sparse solves are a major computational burden.Established stiff libraries run these operations on CPUs, while general-purpose JAX implicit integrators rely on dense linear algebra.Source: PROGRAM SUMMARY · p. 2
  • diff-soup combines a Recipe API for structured and unstructured domains with finite difference, finite volume, and finite element methods.It implements sparse Jacobian assembly by automatic differentiation with graph coloring, JAX-native VSVO integrators, and a GPU cuDSS sparse direct solver.Source: PROGRAM SUMMARY · p. 2
  • Second-order differentiable simulations support PDE-constrained least-squares inference, including frequentist and Bayesian methods and deep learning workflows.Simulation and inference use batched solves, data parallelism, and model parallelism to scale.Source: PROGRAM SUMMARY · p. 2
  • The described release is diff-soup 1.0.0, distributed under BSD 3-Clause and designed for Linux x86_64 with Python 3.12.GPU execution requires an NVIDIA GPU, CUDA 12 or 13, and cuDSS 1.0 or newer.Source: PROGRAM SUMMARY · p. 2
  • DAE support is limited to index one, while solves can span multiple GPUs and cluster nodes.Without a GPU, solvers use CPU KLU and dense back ends.Source: PROGRAM SUMMARY · p. 3

1. Introduction

Large-scale PDE modeling increasingly needs simulation and differentiation at the scale of spatially resolved data, but discretization creates stiff systems with demanding sparse linear algebra. diff-soup is designed around that sparse structure and combines GPU-native solving with differentiable simulation and inference.

1. Introduction

  • Spatially resolved data and data-driven physical modeling motivate software that can simulate high-dimensional PDE models at experimental scale and differentiate through them.The introduction cites learned components in physical solvers, neural networks trained through differentiable solvers, and inference of governing equations from data.Source: §1 · p. 3
  • Spatial discretization makes PDE models stiff, with competing time scales that can force integrators to take small steps.Stiff implicit methods require Jacobians and repeated linear solves; differentiating them also entails parameter-sensitivity rules that do not readily compose with stiff solvers.Source: §1 · pp. 3–4
  • Existing GPU approaches include CPU-hosted hybrid solvers, batches of small problems, and physics- or discretization-specific JAX tools, leaving large single stiff solves underserved.The introduction identifies missing native GPU composition of sparse direct solves, VSVO integration, and sparse automatic differentiation for large dynamical systems in Python/JAX.Source: §1 · p. 4
  • diff-soup builds around cuDSS, which performs sparse factorization and triangular solves on the GPU to address a major bottleneck in robust large-scale PDE integration.Its design centers components on the sparse structure of discretized PDEs, rather than extending small-model software to large multiphysics systems.Source: §1 · pp. 4–5
  • The suite combines GPU-native differentiable sparse computation with reduced precision, batched solves, and multi-GPU parallelism for larger problems.It also composes PDE-constrained optimization, second-order optimizers, multiple shooting, and statistical inference with JAX learning and uncertainty-quantification libraries.Source: §1 · p. 5
  • The paper benchmarks forward solves and gradients, then demonstrates Bayesian inference of material constitutive relationships and training a neural field in a packed-bed reactor model.The introduction identifies these as the two end-to-end applications in Section 5.Source: §1 · p. 5

2. Mathematical Background

The mathematical background frames diff-soup around stiff, sparse PDE-derived systems and the numerical methods needed to simulate and differentiate them for PDE-constrained optimization, including state partitioning, gradient accuracy, event-time sensitivities, and discrete second-order derivatives.

2. Mathematical Background

  • The formulation separates local differential and algebraic states from a small set of global algebraic variables, enabling those dense global couplings to be handled separately.The section motivates this split as a way to address the three computational demands of stiff-system integration, sparse computation, and sensitivities for PDE-constrained optimization.Source: §2 · p. 6

2.1. Variable Step, Variable Order Algorithms

  • Three VSVO implicit multistep solvers—JIDA, JBDF, and JNDF—adapt step size and order while caching Jacobians and factors to limit costly recomputation.Their Jacobian-reuse gates refresh after age or convergence failure; QSS NDF and FLC BDF differ in how step-history changes affect coefficients and factor reuse.Source: §2.1 · pp. 6, 9, 11

2.2. Exploiting Sparsity

  • Sparse automatic differentiation reduces Jacobian construction from O(n^2) to O(n), while cuDSS reuses sparse LU factors across Newton iterations and time steps.diff-soup builds sparsity patterns from declared operators and uses greedy CPR coloring; a Schur complement keeps dense global-variable couplings out of the sparse factorization.Source: §2.2 · pp. 12, 14, 15

2.3. Differentiable Simulation and PDE-constrained Optimization

  • Diff-soup differentiates the adaptive solver discretely using forward- or reverse-mode AD, applying the implicit function theorem at each converged step rather than differentiating through Newton iterations.Forward mode costs O(nθ T(n)); reverse-mode discrete AD avoids that parameter multiplier, while continuous adjoints scale as O(T(n) + nθ) but require a stable stiff backward solve.Source: §2.3 · pp. 16, 18, 19
  • The JAX-native engines use a checkpointed discrete adjoint, retaining O(√nsteps) checkpoints and recomputing trajectory segments for O(√nsteps n) memory and O(T(n) + n_theta) gradient cost.The method differentiates accepted steps under a frozen-controller assumption.Source: §2.3 · p. 19
  • Continuous sensitivities approximate the continuous-problem gradient to solver and quadrature tolerances, whereas discrete sensitivities exactly differentiate the frozen discretization up to floating-point precision, with adaptive-map error controlled by local tolerance.Table 2 compares the four sensitivity methods for Eq. (9), reporting costs, memory, accuracy, and stability in terms of state dimension, parameter count, and step ceiling.Source: §2.3 · p. 20, Table 2 · p. 20
  • Custom JVP and VJP event rules differentiate switch times by the implicit function theorem and propagate their shifts through restarted trajectories, matching classical event sensitivities at transversal crossings.These rules address event-time dependence that the frozen-time-grid computational graph would otherwise miss.Source: §2.3 · p. 21
  • Three second-order compositions provide a full Hessian at O(n_theta^2 T(n)), scalar-criterion gradients at O(n_theta T(n)), or Hessian-vector products at O(T(n) + n_theta) without assembling the Hessian.The compositions are forward-over-forward, reverse-over-forward, and forward-over-reverse, respectively.Source: §2.3 · p. 21

3. diff-soup Software Overview

diff-soup separates PDE modeling, differentiable simulation, and inference into connected modules. Together, they take a physical system from specification and discretization through parameter sensitivities to learning model parameters from data.

3. diff-soup Software Overview

  • The three modules connect PDE modeling, differentiable simulation, and inference: modeling constructs a sparse discrete problem, simulation integrates it and provides parameter sensitivities, and inference learns parameters from data.Figure 3 overviews this architecture.Source: §3 · p. 21, Figure 3 · p. 22

3.1. Modeling

  • The Recipe API specifies physical systems on structured grids or unstructured meshes, while the modeling module lowers them into solver-ready problems with residuals, sparse Jacobians, parameters, and controls.Users can declare fields, differential operators, boundary conditions, global constraints, and learnable parameters.Source: §3.1 · pp. 22–24

3.2. Simulation

  • The simulation module provides native variable-step, variable-order stiff solvers and differentiable JAX solves, with checkpointed reverse-mode gradients and forward-over-reverse Hessian-vector products.It also supports batching, GPU and multi-GPU execution, and differentiable events.Source: §3.2 · pp. 24–26

3.3. Inference

  • Users can pose frequentist maximum-likelihood or Bayesian maximum-a-posteriori estimation, then assess uncertainty and practical identifiability using posterior approximations, sampling, and information-based diagnostics.The workflow also supports profile-likelihood analysis by fixing one parameter while refitting the others.Source: §3.3 · pp. 27–28
  • Inference tools support data assimilation, multiple shooting, and function approximation with polynomials, neural networks, neural operators, and spatial-field models.Multiple shooting partitions long trajectories into windows joined by continuity constraints, reducing simulation cost and improving optimization robustness and conditioning.Source: §3.3 · p. 28
  • For hierarchical inference, data parallelism distributes experiments or trajectory windows, while an arrowhead Hessian structure enables a block Newton solver using a Schur complement over shared parameters.Model parallelism distributes a single sparse factorization across multiple GPUs when it exceeds one GPU’s memory.Source: §3.3 · p. 28

4. Benchmarks

Across heat, Brusselator, and Cahn–Hilliard benchmarks, diff-soup combines stiff integration with GPU sparse linear algebra to improve solve and gradient performance. Reduced precision and multi-GPU execution further extend the tested scale, with accuracy and memory trade-offs.

4.1. Comparing time integrators with the heat equation

  • For the heat equation, diff-soup’s JIDA solver tracks the spatial error floor at low, nearly fixed cost, while Kvaerno5 error saturates near 2 × 10−7 from N = 2^12.The nearly fixed cost is about 0.1 s across the tested resolutions.Source: Figure 5 · p. 31, §4.1 · p. 32

4.2. Comparing performance with the Brusselator and Cahn-Hilliard problems

  • Diff-soup achieves up to 50× speedup over the next-fastest solver at scale, and up to 440× over the same JIDA algorithm with CPU KLU.The benchmark comparison shows gains over scientific computing libraries; the 440× comparison is specifically against CPU KLU.Source: Figure 8 · p. 35, §4.2 · pp. 35–37
  • GPU-parallel Jacobian coloring brings diff-soup’s batched autodiff Jacobian to within about 2.4× of the analytical Jacobian, leaving pencil-matrix factorization as the bottleneck.On CPU, sparse autodiff Jacobians can run up to an order of magnitude slower than analytical ones.Source: §4.2 · pp. 36–37, Table 3 · p. 36
  • Newton–Krylov performance depends on the problem: it closely tracks diff-soup on the Brusselator but falls behind by orders of magnitude on Cahn–Hilliard, where sparse direct solvers perform best.The hybrid GPU comparisons show this same contrast; effective problem-specific preconditioners can make Newton–Krylov methods effective, especially for three-dimensional PDEs.Source: Figure 9 · p. 38, §4.2 · p. 39
  • Reverse-mode gradients are one to two orders of magnitude faster than CPU codes at scale, while forward-mode cost grows with parameter count and reverse mode reaches the largest tested states.For Cahn–Hilliard, every library’s gradient error remains above 10−3; checkpointing bounds reverse-mode memory, while forward mode stores one full-state tangent per parameter.Source: Figure 10 · p. 40, §4.2 · p. 41

4.3. Reduced Precision and Multi-GPU Parallelism

  • At scale, Float32 cuts three-dimensional Brusselator solve cost by an order of magnitude and avoids Float64 factor overflow at 1.8×10^6 degrees of freedom.At smaller sizes, the Float32 final state agrees with Float64 within 4 × 10−6 relative L2 norm; multi-GPU placement roughly halves cost on two devices, with modest further improvement on four.Source: §4.3 · pp. 42–43, Figure 11 · p. 42

5. Applications to PDE-constrained Optimization

The section demonstrates diff-soup workflows for Bayesian inference of material constitutive relationships and deep learning of catalyst activity in a packed-bed reactor.

5. Applications to PDE-constrained Optimization

  • The two workflows illustrate a single package supporting classical Bayesian inference and neural-field training through a PDE model.The examples address constitutive relationships in binary phase-separating materials and three-dimensional catalyst activity in a packed-bed reactor.Source: §5 · p. 43

5.1. Inference with Uncertainty Quantification

  • A Bayesian inverse problem inferred composition-dependent tracer diffusivity and homogeneous chemical potential from one noisy spinodal-decomposition movie.The approach parameterizes both functions with orthogonal polynomials and samples the posterior; multiple shooting starts each prediction window from the preceding measured frame, keeping solves short and accelerating optimization.Source: §5.1 · pp. 43–44
  • Within ten trust-region iterations, the Bayesian fit reached the measurement-noise level, and posterior samples quantified uncertainty in laws lying close to ground truth.Figure 12 shows the noisy movie and predictions, fit error, inferred diffusivity and chemical potential, and 95% credible intervals.Source: Figure 12 · p. 45, §5.1 · p. 45

5.2. Embedded Deep Learning Model

  • A wavelet-based neural field with 7.5 × 10^4 parameters represented catalyst activity in a three-dimensional packed-bed model and was trained against synthetic concentration measurements.The synthetic target combines a log-normal random field, poisoning front, and low-activity streak; dispersion smooths fine-scale activity features out of measured concentrations.Source: §5.2 · p. 46
  • About 5% relative L2 error in catalyst activity accompanies recovery of the poisoning front, low-activity streak, and large-scale random-field texture.Training reached the measurement-noise level within a few hundred steps; 1200 steps took a little over four hours on one NVIDIA L40S GPU.Source: Figure 13 · p. 47, §5.2 · pp. 47–48

6. Conclusion

diff-soup accelerates stiff PDE simulation and gradients, and supports inference and learning workflows. The authors identify expanded benchmarking and software capabilities as directions for future work.

6. Conclusion

  • Up to 50× faster simulations and 80× faster gradients than established CPU libraries are reported for diff-soup.The evaluation compares heat, Brusselator, and chemo-mechanical Cahn-Hilliard problems against SUNDIALS, PETSc, DifferentialEquations.jl, MATLAB, and diffrax.Source: §6 · p. 48
  • The demonstrated workflows infer phase-separating material constitutive relationships with uncertainty quantification and train a neural field in a three-dimensional packed bed reactor.The package differentiates end to end through the discretize-then-optimize approach, with gradients used for PDE-constrained optimization.Source: §6 · p. 48
  • Future work targets more rigorous stiff-PDE benchmarking, stronger distributed and multi-GPU support, general accelerator devices, and broader inference applications.The authors name model predictive control, design of experiments, and constrained optimization algorithms such as IPOPT as prospective applications.Source: §6 · pp. 48–49

7. Methods

The methods evaluate diff-soup on stiff-PDE benchmarks and demonstrate Bayesian inference and neural-field training in two worked examples. The benchmarks and examples use specified solver settings, data, and physical-model conditions.

7.1. Benchmarking

  • The benchmark suite evaluates heat, Brusselator, and elastic Cahn-Hilliard problems, using finite differences for heat and finite volumes for the other two.Problems are posed on Cartesian domains; forward comparators generally use eight CPU cores, while hybrid GPU lanes use one L40S GPU and one CPU core.Source: §7.1 · p. 49
  • Benchmark timings are medians of warm calls after one cold call, with a 1,800-second wall cap and accuracy measured against reference solutions or gradients.Diff-soup uses cuDSS on GPU; comparator versions, configurations, and reference validations are recorded in the companion materials.Source: §7.1 · pp. 50–51

7.2. Worked Example 1

  • The Bayesian example fits regular-solution laws to six noisy movie frames and samples their posterior using four replicates of four chains.The fit took 37 seconds on one NVIDIA A100; all sixteen chains mixed, with split R-hat below 1.01 and no divergent transitions.Source: §7.2 · p. 51

7.3. Worked Example 2

  • The reactor example models a three-dimensional packed bed whose Brinkman-derived superficial velocity profile channels flow along the walls.The reactor uses wall-dependent porosity and specified inlet, outlet, and no-flux wall conditions.Source: §7.3 · pp. 51–52, Figure 14 · p. 52
  • A WIRE neural field is trained for 1,200 steps on noisy reactor data generated from an activity field with a poisoning front and a low-activity pellet streak.Each training step integrates the trajectory with JNDF and cuDSS in Float32 precision, using a discrete checkpointed adjoint on one NVIDIA L40S.Source: §7.3 · p. 52

CRediT authorship contribution statement

Samuel Degnan-Morgenstern is credited with the broad range of research and implementation work, while Martin Z. Bazant is credited with writing, supervision, project administration, and funding acquisition.

CRediT authorship contribution statement

Appendix A. The three diff-soup VSVO solvers

The three diff-soup VSVO solvers show broadly similar performance across the appendix benchmarks. JIDA is selected for the main-text results for robustness and generality.

Appendix A. The three diff-soup VSVO solvers

  • The three solvers generally show similar performance at matched error, with no significant leader.Their work-precision diagrams share a performance-accuracy frontier, and their scaling is similar within a small factor.Source: Appendix A · p. 54, Figure A.15 · p. 55
  • JIDA is used for the main-text results because it is generally the most robust stiff implicit integrator and handles the most general problem form.Source: Appendix A · p. 54
  • Gradient timings are within a small factor across solvers and differentiation modes, with marginal solver advantages varying by problem.JBDF and JNDF are marginally faster on the Brusselator closure, while JIDA is fastest on the Cahn-Hilliard closure.Source: Figure A.15 · p. 55

Appendix B. One-dimensional Brusselator comparison

For the one-dimensional Brusselator, CPU solvers lead diff-soup in work-precision below the GPU crossover, while diff-soup shows fixed overhead across the scaling ladder.

Appendix B. One-dimensional Brusselator comparison

  • The fastest CPU solver beats diff-soup by over an order of magnitude in work-precision, while only multigrid-preconditioned PETSc trails diff-soup.The comparison is for the one-dimensional Brusselator.Source: Figure B.16 · p. 56
  • Diff-soup exhibits fixed overhead across the scaling ladder, whereas CPU-based solvers exhibit power-law growth.The scaling diagram compares diff-soup with SUNDIALS, PETSc, Julia DifferentialEquations.jl, and MATLAB.Source: Figure B.16 · p. 56
  • Below the crossover, GPU fixed overhead dominates for this one-dimensional problem, and CPU-based solvers lead.Data points are omitted when a solve fails to converge or exceeds the wall cap.Source: Figure B.16 · p. 56

Related papers

Loading 2610.37333v1…