3 ms·
There's a difference from what was done before though. While we are using adjoint sensitivity analysis of the old, we have things setup so that the internal vjp
by ChrisRackauckas 7y ago
There's a difference from what was done before though. While we are using adjoint sensitivity analysis of the old, we have things setup so that the internal vjp (vector-Jacobian product) is using reverse-mode AD that is compatible with the neural network framework to do that calculation in O(1) instead of O(n) f-calls. We have a PR setup that will make it do source-to-source AD on that part as well. So this scaling is distinctly different in terms of scalability from the traditional numerical codes like that found in Sundials. We are not the first to do this, CASADI has been doing this for a bit of time, but we can now demonstrate that we can throw arbitrary Flux.jl neural network frameworks in there and reverse-mode will specialize this, and specialize the whole thing.
And yes, parameter estimation in jump SDEs has been around, but using full solver reverse-mode AD on high order GPU-accelerated adaptive solvers has not been around, and we're talking about many orders of magnitude speedups here when you put these all together.
But yes, in terms of the math, the jump is very small here. But last week at a conference I got some questions about whether an existing tool can actually allow you to just define a PDE or jump diffusion with neural networks in there and take an adjoint that is both efficient in the differential equation and the neural network. I was saying yes, but given the previous DiffEqFlux.jl it would require someone to understand the capabilities of the AD/adjoint system, know all of the differential equations that can be solved through the package, and realize they can write a code to do it. I wanted to make it explicit by sharing code: yes, we do neural jump diffusions and nerual PDEs. It's not something hypothesized in the future, it's something that our package is already doing today with all of the goodies you expect from the ODE support.