arXiv preprint ·
diff-soup: GPU-native Differentiable Simulation and Inference for Large-scale Stiff Dynamical Systems
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
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.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Mathematical Software
Beyond atol and rtol: Sound Test Oracles for Low-Precision GPU Kernels
Building Mathematical Software -- Lessons from the Development of GAP
Integer-Quantized Hierarchical Matrices with Adaptive Bit Allocation
ArcTN: Efficient Tensor-Network Contraction Planning for Quantum Circuit Simulation
The Rise of Low Precision Floating Point
Get the few papers worth your time in Mathematical Software every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.