High weak order solvers and adjoint sensitivity analysis for stochastic differential equations

GSoC 2020: The Julia Language – Final report.

Project summary

In this project, we have implemented new promising tools within the SciML organization which are relevant for tasks such as optimal control or parameter estimation for stochastic differential equations. The high weak order solvers will allow for massive performance advantages for fitting expectations of equations. Instead of automatic differentiation (AD) through the operations of an SDE solver, which scales poorly in memory, one can now use efficient stochastic adjoint sensitivity methods.

Blog posts

The following posts describe the work during the entire period in more detail:

  1. GSoC 2020: High weak order SDE solvers and their utility in neural SDEs
  2. High weak order SDE solvers
  3. Adjoint sensitivity methods – To be uploaded soon.

Docs

The documentation of the solvers is available here. Docs with respect to the adjoint sensitivity tools will be available here.

Achievements

Please find below a list of the PRs carried out during GSoC in the different repositories in chronological order within the SciML ecosystem.

StochasticDiffEq.jl

Merged:

Open:

DiffEqSensitivity.jl

Merged:

Open:

DiffEqNoiseProcess.jl

Merged:

DiffEqGPU.jl

Merged:

ModelingToolkit.jl

Merged:

DiffEqDevTools.jl

Merged:

DiffEqBase.jl

Merged:

Future work

There is still a lot that we’d like to do, e.g.,

  • Writing up more docs and examples
  • Implementing drift-implicit weak stochastic Runge-Kutta solvers
  • Finishing the SDE adjoints for the Ito sense
  • Implementing a virtual Brownian tree to store the noise processes in O(1) memory
  • Setting up an OptimalControl library that allows for easy usage of the new tools within a symbolic interface
  • Benchmarking of the new solvers and adjoints

Contributions, suggestions & comments are always welcome! You might like to join our slac channels #diffeq-bridged and #neuralsde to get in touch.

Acknowledgement

I would like to thank my mentors Chris Rackauckas, Moritz Schauer, and Yingbo Ma for their amazing support during this project. It was a great opportunity to work in such an inspiring collaboration and I highly appreciate their detailed feedback. I would also like to thank Christoph Bruder, Niels Lörch, Martin Koppenhöfer, and Michal Kloc for helpful comments on my blog posts. Many thanks to the very supportive julia community and to Google’s open source program for funding this experience!

Frank Schäfer
Frank Schäfer
PhD candidate in Physics

My research interests include many-body physics, machine learning, and differentiable programming.

Related