KL Divergence for Latent SDEs - #463
Conversation
patrick-kidger
left a comment
There was a problem hiding this comment.
Okay, I really like the example here!
I'm afraid this might still take a bit more iteration to get to something clean enough to merge, though -- see my comments. :)
There was a problem hiding this comment.
As per this comment:
#402 (review)
and also the updated term docs:
https://docs.kidger.site/diffrax/api/terms/
then this outer MultiTerm isn't really in-keeping. We're not adding all of these extra terms on to the same evolving state.
Bearing in mind that the rest of Diffrax has to see this as just another SDE solve.
I think this one might take a bit more iteration to get to something that's obeying the abstractions in the way they're designed, I'm afraid.
There was a problem hiding this comment.
I think I get what you're saying, multiterm implies a single differential equation "unit". So the composed multi terms is bad form. However, I'm not sure I see the difficulty going forward, I can replace it with tuple (multiterm, multiterm) or even tuple (multiterm, ode term). Which seems to adhere to this principle of multiterm = sde unit, since we are integrating two simultaneous SDEs, while also falling in line with other solvers (such as implicit Euler as you remarked).
On the terms vs solver approach, I am open to both. I think in my many iterations/experimentations I found the solver approach more in line with my thinking about the nature of the problem, specifically the original idea of (terms, kl_term) I didn't see as appealing since the KL_term relies on information from the other term and I didn't see a clean way to do that. However, having terms with a term wrapper is very doable (but may not mesh with the repo as well).
There was a problem hiding this comment.
Maybe introducing a totally new term is ok (given the remarks in #453), in which case the approach of a KLTerm (rather than a solver), is doable. Given the restricted nature of terms so far, I originally thought that wasn't in line with the package
Happy to iterate on cleaning it, I think the biggest question is the design one (on solvers, terms, and how to represent the problem in diffrax). Once that is resolved, I can iterate quickly to get the rest in :) |

A continuation of #402 with the new 0.6.0 lineax changes.
Relevant origin issue: #401