diff-soup: GPU-native Differentiable Simulation and Inference for Large-scale Stiff Dynamical Systems
Abstract
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.
Related
- Axio-Dilaton Dark Energy: A Dynamical Systems and Bayesian Inference Analysis
- Inference of Unknown Dynamical Components Using Next Generation Reservoir Computing: From Chaotic Systems to Climate Data
- Fundamental Dynamical Units for Physics-Informed Structural Inference from Perturbation Time-Series in Networked Systems
- Resolving Structural Uncertainty in Stochastic Dynamical Systems via Ensemble-Based Bayesian Inference
- JZ-FMM: GPU-native differentiable N-body simulations with the Fast Multipole Method
BibTeX
@misc{degnanmorgenstern2026diffsoup,
title = {{diff-soup: GPU-native Differentiable Simulation and Inference for Large-scale Stiff Dynamical Systems}},
author = {Samuel Degnan-Morgenstern and Martin Z. Bazant},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.37333},
eprint = {2610.37333},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.37333},
note = {CC BY-NC-SA 4.0},
}
RIS
TY - UNPB
TI - diff-soup: GPU-native Differentiable Simulation and Inference for Large-scale Stiff Dynamical Systems
AU - Samuel Degnan-Morgenstern
AU - Martin Z. Bazant
PY - 2026
DA - 2026/10/27
PB - arXiv
DO - 10.48550/arxiv.2610.37333
UR - https://arxiv.org/abs/2610.37333
AB - 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.
ER -