Logging metrics during an ODE solve

Hello @patrick-kidger,

thank you for open-sourcing this nice library! I was going to resume work on my own small ODE lib, but since this is much more elaborate than what I came up with so far, I am inclined to use this instead for a small project in the future.

One question that came up to me when reading the source code: Is there currently a way to compute step-wise metrics during the solve? (Think logging step sizes, Jacobian eigenvalues, etc.)

This would presumably happen in the integrate method. Could I e.g. use the solver_state pytree for this in, say, overridden solver classes? Thank you for your consideration.

[patrick-kidger]

So these are an example of an “integral loss” (or regularisation term): a loss of the form \int k(y(t)) dt for some function k. Indeed it’s a known trick to avoid calculating these on the forward pass, and only calculate their gradients on the backward pass.

(In practice this is pretty much only ever possible when using a custom backward pass like BacksolveAdjoint though. I don’t think any autodifferentiation framework, JAX or otherwise, is currently smart enough to elide the forward pass if using the “standard” RecursiveCheckpointAdjoint.)

Right now the best way to tackle this in Diffrax is indeed just to append these values during the forward calculation. Even with other libraries, I think people have often had to really know what they’re doing and then hack something together, since it’s such an edge-case in terms of autodifferentiation frameworks.

I’ve opened #69 to track this though, as this would be nice to have. CC @jacobjinkelly CC @Zymrael as I recall you have something similar in torchdyn?

[pharringtonp19]

@patrick-kidger – Awesome library

In this context of “logging-metrics”, I was wondering how one might go about implementing Learning Differential Equations that are Easy to Solve. My naive implementation “appended” the regularizing term to the state. However, from discussions with one of the authors, they suggest a more efficient way to implement it.

we actually only need to integrate the regularization term on the backwards pass - on the forwards pass we can just integrate a placeholder value of 0. This is essentially because we only care about the gradient with respect to the regularization term, and the gradient of the regularization can be constructed from the value of the rest of the state (everything but the regularization) and its gradient.

Do you have any suggestions on how to tackle this with diffrax? – happy to open a separate issue if you would prefer.