Custom_vjp in jax.experimental.ode odeint

[agobL]

Hello,

I use jax.experimental.ode odeint function in my work. You implemented the custom_vjp functionnality for this function (which derivates the ODE system and solves it).
What is the mathematical theory you used for implementing that custom derivative? Is there a scientific article that tells about this technique ?

[mattjj]

Thanks for the question!

I wrote up my own derivation for differentiation of ODEs here: https://implicit-layers-tutorial.org/implicit_functions/, specifically this subsection.

I like that derivation because it's simple and mechanical: no steps require creative leaps and it doesn't involve importing extraneous machinery (like Lagrange multipliers). Also unlike any other derivation I've seen it allows decomposition of reverse-mode differentiation into linearization, partial evaluation, and transposition; that makes it a natural fit for JAX's autodiff system, though our implementation hasn't yet caught up.

Some other resources:

1. Sec 2.2 of these notes or Sec 2.1.1 of this tutorial or Sec 2.1 of this paper, but those all require embedding the system in an optimization problem so as to use Lagrange multipliers (where the first link above does not), or
2. Appendix B of the Neural ODEs paper, but that requires too much discretization and calculation for my tastes, and also mainly checks the form of the answer.

There are probably lots of other resources with which I'm not familiar.

WDYT?

[patrick-kidger]

Shameless self promotion. By far the simplest derivation I know of is that of Section 5.1.2.1 of On Neural Differential Equations, which gives a derivation in just 5 lines of elementary calculus.

---

If you're curious to know more: Theorem 5.2 on the previous page gives the statement being proved; Appendic C.3.1 gives an ever-so-slightly-longer derivation with full mathematical rigor; finally Diffrax is a JAX-based faster and more featureful alternative to jax.experimental.ode.odeint.