Hybridization of Neural Networks and Numerical Solvers in JAX with Differentiable Physics

Value of the Information & Strength of the Argument

The talk provides a comprehensive and insightful comparison between numerical solvers and neural networks, offering a clear framework for hybridization. The argumentation is solid, supported by examples and references to recent research. The speaker’s strict definition of hybridization is well-justified and helps clarify the landscape. The discussion of challenges and costs is valuable, providing a balanced view. The case studies and practical preparation enhance the practical value of the talk.

Scientific Rigor, Source Quality, Title Accuracy

The talk demonstrates high scientific rigor, with references to peer-reviewed papers and benchmarks. The sources are relevant and support the claims made. The title accurately reflects the content, and the talk is well-structured. The speaker’s expertise is evident, and the content is presented in a clear and logical manner.

The title accurately reflects the content, which focuses on the hybridization of neural networks and numerical solvers using differentiable physics in JAX.

9/10

The talk is given by an expert in the field, with references to peer-reviewed papers and benchmarks. The content is well-structured and technically accurate, though it reflects the author's perspective and definitions.

• 00:00Intro & agenda
• 04:41Part 1: Solvers vs. neural nets
• 05:40The compute graph of a solver
• 09:07...and of a ConvNet
• 10:34Where the parameters come from
• 12:28Navier-Stokes & implicit solves
• 14:28Higher-level primitives in NNs
• 15:24Levels of autodiff granularity
• 22:33Spatial processing & APEBench
• 27:18Part 2: Hybridization defined
• 28:11Not PINNs, Neural ODEs or DEQs
• 31:18Full-field correction
• 32:21Spectrum of learned components
• 35:29Unrolling & BPTT
• 40:34Part 3: Why hybridize
• 43:00Combining paradigms
• 46:17Catch 1: non-smooth operations
• 48:13Catch 2: implicit differentiation
• 50:35Catch 3: legacy solvers
• 55:12Part 4: BPTT curse & blessing
• 56:32Gradient cuts
• 58:53O(1) memory & checkpointing
• 1:02:37Progressive refinement
• 1:05:59Part 5: Solver-in-the-loop
• 1:07:25ML-accelerated CFD
• 1:10:23Part 6: Practical prep
• 1:13:20One-step supervised training
• 1:16:04Tooling & objectives
• 1:18:21Conclusion

• APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs — Benchmark comparing neural emulators and numerical solvers.
• Neural Ordinary Differential Equations — Reference for Neural ODEs, discussed as not hybridization.
• Hamiltonian Neural Networks — Reference for Hamiltonian Neural Networks, discussed as not hybridization.
• Lagrangian Neural Networks — Reference for Lagrangian Neural Networks, discussed as not hybridization.
• Learning to Simulate Complex Physics with Graph Networks — Seminal work on neural-hybrid correctors.
• Learning to Correct for Conditional Modeling in Fast and Slow — Influential work on neural-hybrid correctors achieving speedups.
• Implicit Differentiation for Hyperparameter Optimization — Study on implicit differentiation in JAX.
• Training Neural-Hybrid Correctors without a Differentiable Simulator — Approach for building neural-hybrid correctors without differentiable simulators.
• Scheduling the Fidelity of Coarse Simulators for Neural-Hybrid Correctors — Reducing training cost by scheduling fidelity.
• SoftJAX: Softened Operations for Differentiable Physics — Library providing softened modifications for non-differentiable operations.
• Autodiff Table — Collection of primitive rules for explicit autodiff.
• Implicit Autodiff Table — Collection of autodiff primitive rules for implicit autodiff.
• Corrector Configurations — Collection of ways to couple neural networks and coarse solvers.
• Predictor Learning Setups — Collection of SVG schematics for training autoregressive emulators.
• Hybridization in JAX (GitHub) — Repository with all material from the talk.
• IPC: Incremental Potential Contact — Technique for contact mechanics, used as example of non-smooth operations.

• APEBench — Empirical evidence on neural emulators vs. numerical solvers.
• Learning to Correct for Conditional Modeling in Fast and Slow — Supports the benefits of neural-hybrid correctors.

• Physics-Informed Neural Networks — The talk explicitly excludes PINNs from hybridization, which some researchers might consider as a form of hybridization.