Differentiable Simulation and Scientific Machine Learning
Fast Solving and Automated Model Construction
Abstract:
Scientific machine learning (SciML) methods allow for the automatic discovery of mechanistic models by infusing neural network training into the simulation process. In this talk we will start by showcasing some of the ways that SciML methods, like universal differential equations, are being used. Demonstrations of the automated discovery of relativistic corrections to black hole physics to the construction of earthquake-safe buildings showcase the successes of the techniques throughout scientific domains. From there, we will discuss some of the increasingly advanced computational techniques behind the training process, focusing on the numerical issues involved in handling differentiation of highly stiff and chaotic systems. We will then show how the solutions to these numerical issues are being productionized in industry, explaining how these computational techniques led to improved trajectory planning in track-side computers for Formula 1 races, accelerated the clinical trials of the Covid vaccine, and allowed for crash testing computers to predict forces in a way that saved days of customer time on multi-million dollar machines. The audience will leave with an understanding of how these latest compiler techniques are being infused into the next generation simulation stack to increasingly automate the process of developing mechanistic models and industrial processes.
Bio:
Dr. Chris Rackauckas is the VP of Modeling and Simulation at JuliaHub, the Director of Scientific Research at Pumas-AI, Co-PI of the Julia Lab at MIT, and the lead developer of the SciML Open Source Software Organization. For his work in mechanistic machine learning, his work is credited for the 15,000x acceleration of NASA Launch Services simulations and recently demonstrated a 60x-570x acceleration over Modelica tools in HVAC simulation, earning Chris the US Air Force Artificial Intelligence Accelerator Scientific Excellence Award. See more at https://chrisrackauckas.com/. He is the lead developer of the Pumas project and has received a top presentation award at every ACoP in the last 3 years for improving methods for uncertainty quantification, automated GPU acceleration of nonlinear mixed effects modeling (NLME), and machine learning assisted construction of NLME models with DeepNLME. For these achievements, Chris received the Emerging Scientist award from ISoP.
Summary:
Goal: automatic discovery of models and physical laws
Focus of the talk: the SciML effort (https://sciml.ai/) and automatic differentiation
Scientific Machine Learning
Use the available prior knowledge as part of the modeling process
This will hopefully be much more data efficient than regular machine learning
Key tool: Differentiable Simulation (Hard but very much worth it)
Idea: Universal (Approximator) Differential Equation
Neural networks are a category of functions that can approximate any other functions (in the limit)
Ordinary Differential Equations (ODEs) are a powerful techniques for approximating physics
Can fuse these by replacing portions of known ODEs with neural nets and then training the neural nets to recover unknown physics
This is a form of auto-complete on models
Scientific approximation
Normal goal: identify the parameters of the model
In SciML: identify the shape of the unknown function terms; not the actual functional form, just the first few Taylor approximation terms
The learned neural terms are useful for
Prediction
Learning what specifically is wrong about the original model and how it may be improved
Sparse regression summarizes the dynamics of the neural network
Regression is interpretable and can be used by scientists to improve the model
Debugging Science!
Examples:
Modeling COVID spread
Baseline: neural net trained on 21 days of data
Then took SEIRD model and replaced some terms with neural nets while
Leaving intact the model’s overall structure and most terms
Putting in constraints on the outputs of the individual neural nets
Can ensure that certain constraints like positivity and population conservation are maintained
LIGO Black hole dynamics from gravitational wave data
Augmented a basic ODE system with neural terms
5 digits of accuracy
Much better than basic neural networks
Full simulations are 12 digits but are a lot more expensive
Earthquake-safe building simulation
Researchers made a many-body dynamics model of a building
Model was inaccurate so it was augmented with neural terms to improve accuracy
Inference of car position from sensors
Using a physical model of the car as the base + neural net terms was more accurate than a Gaussian Process regression
Challenge: what if the base ODE was mis-specified?
Its been found that small mis-specifications are ok; not great but still better than not using a base ODE
No real guarantees though
Suggests a hierarchical learning process:
For a given ODE structure can use gradient descent and neural networks to search for functions around the ODE
Then propose additional ODEs based on this search
Space of ODEs is non-differentiable (symbolic regression) but the neural net space is differentiable and easier to search
SciML organization
Suite of tools for putting neural networks into mathematical models
Focused on the Julia language ecosystem
Dynamics -> Non-linear mixed effects modeling
Covariates of the units in the study used to train a structural model
Structural model used to learn a dynamic model of the population
Model includes both fixed and variable(random) effects, need to fit parameters to ensure that the expected value of variable effects is 0 (zero-mean noise)
Pumas (PharmacUtical Modeling and Simulation) model used to design dosing designs, regulatory submissions
Improving the scientific workflow
Use neural net as structural model
Train it and then learn dynamics that produce accurate results
Use sparse regression to find interpretable model that approximates it
Find explanations from prior literature that explains the term of the sparse regression
This needs to be augmented with additional experiments and use of optimal experimental design to expand the available data based on the best way to refine a model
Example: high fidelity surrogates for ocean columns for climate models
Improves efficiency of simulating these dynamics
Simplest option: ODE that is 1D (depth) approximation to 3D model, with a non-linear neural net
Challenge: the neural net is controlling the model’s derivative, which induces drift over time
Alternative: have the neural net control the integral
To do this
We need to evaluate the accuracy of the simulator’s output given the results of the internal neural network
Requires the overall simulation to be differentiable
Automatic differentiation of programs with discrete randomness
Base program is fundamentally discrete, with sharp decisions
Can’t differentiate it at discrete decision points
Idea:
Create a population of runs of the discrete program on a stochastic distribution of inputs/parameters
Derivative is over the probability distribution of program runs, which is smooth
Applied to agent-based models