JAX-DIPS: Neural bootstrapping of finite discretization methods and application to elliptic problems with discontinuities
Not on the map. Why?
Abstract (arXiv)
We present a scalable strategy for development of mesh-free hybrid neuro-symbolic partial differential equation solvers based on existing mesh-based numerical discretization methods. Particularly, this strategy can be used to efficiently train neural network surrogate models of partial differential equations by (i) leveraging the accuracy and convergence properties of advanced numerical methods, solvers, and preconditioners, as well as (ii) better scalability to higher order PDEs by strictly limiting optimization to first order automatic differentiation. The presented neural bootstrapping method (hereby dubbed NBM) is based on evaluation of the finite discretization residuals of the PDE system obtained on implicit Cartesian cells centered on a set of random collocation points with respect to trainable parameters of the neural network. Importantly, the conservation laws and symmetries present in the bootstrapped finite discretization equations inform the neural network about solution regularities within local neighborhoods of training points. We apply NBM to the important class of elliptic problems with jump conditions across irregular interfaces in three spatial dimensions. We show the method is convergent such that model accuracy improves by increasing number of collocation points in the domain and predonditioning the residuals. We show NBM is competitive in terms of memory and training speed with other PINN-type frameworks. The algorithms presented here are implemented using textttJAX in a software package named textttJAX-DIPS (https://github.com/JAX-DIPS/JAX-DIPS), standing for differentiable interfacial PDE solver. We open sourced textttJAX-DIPS to facilitate research into use of differentiable algorithms for developing hybrid PDE solvers.
Key facts. Cited by 7 papers in the map (data through 14 Oct 2026). Builds on 21 of its 76 references that are in the map. First posted on arXiv 25 Oct 2022.
Builds on
21 of 76 references.
References · 21
- 2014Adam: A Method for Stochastic Optimizationcited by 37,240
- 2020Fourier Neural Operator for Parametric Partial Differential Equationscited by 2,908
- 2020Neural Operator: Graph Kernel Network for Partial Differential Equationscited by 734
- 2017DGM: A deep learning algorithm for solving partial differential equationscited by 678
- 2021Characterizing possible failure modes in physics-informed neural networkscited by 643
- 2017The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problemscited by 583
- 2021Machine learning accelerated computational fluid dynamicscited by 471
- 2020Model Reduction and Neural Networks for Parametric PDEscited by 258
- 2020Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solverscited by 205
- 2021Efficient and Modular Implicit Differentiationcited by 203
- 2018Measuring the Effects of Data Parallelism on Neural Network Trainingcited by 193
- 2020Learning to Control PDEs with Differentiable Physicscited by 132
9 more
- 2022JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flowscited by 36
- 2020Solving inverse-PDE problems with physics-aware neural networkscited by 27
- 2019Extracting Interpretable Physical Parameters from Spatiotemporal Systems using Unsupervised Learningcited by 23
- 2019Data driven approximation of parametrized PDEs by Reduced Basis and Neural Networkscited by 22
- 2021A Discontinuity Capturing Shallow Neural Network for Elliptic Interface Problemscited by 20
- 2022Statistical Properties of the log-cosh Loss Function Used in Machine Learningcited by 16
- 2017Neural network augmented inverse problems for PDEscited by 11
- 2020A three layer neural network can represent any multivariate functioncited by 10
- 2021A non-parametric shape optimization approach for solving inverse problems in directed self-assembly of block copolymerscited by 2
Built on by
Citing papers per year: 2023: 2; 2024: 3; 2025: 2.
7 papers.
- 2024RandONet: Shallow-Networks with Random Projections for learning linear and nonlinear operatorscited by 16
- 2024GoRINNs: Godunov-Riemann Informed Neural Networks for Learning Hyperbolic Conservation Lawscited by 8
- 2025JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modelingcited by 5
- 2023Neural-Integrated Meshfree (NIM) Method: A differentiable programming-based hybrid solver for computational mechanicscited by 4
- 2025A Multiple Transferable Neural Network Method with Domain Decomposition for Elliptic Interface Problemscited by 2
- 2023Zero Coordinate Shift: Whetted Automatic Differentiation for Physics-informed Operator Learningcited by 0
- 2024Invariant deep neural networks under the finite group for solving partial differential equationscited by 0
- 2025JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modeling6 Jul 2025
- 2025A Multiple Transferable Neural Network Method with Domain Decomposition for Elliptic Interface Problems27 Feb 2025
- 2024GoRINNs: Godunov-Riemann Informed Neural Networks for Learning Hyperbolic Conservation Laws29 Oct 2024
- 2024Invariant deep neural networks under the finite group for solving partial differential equations30 Jul 2024
- 2024RandONet: Shallow-Networks with Random Projections for learning linear and nonlinear operators8 Jun 2024
- 2023Neural-Integrated Meshfree (NIM) Method: A differentiable programming-based hybrid solver for computational mechanics21 Nov 2023
- 2023Zero Coordinate Shift: Whetted Automatic Differentiation for Physics-informed Operator Learning1 Nov 2023