KL Divergence for Latent SDEs by lockwo · Pull Request #463 · patrick-kidger/diffrax · GitHub
Skip to content

KL Divergence for Latent SDEs - #463

Open
lockwo wants to merge 28 commits into
patrick-kidger:mainfrom
lockwo:kl-pr
Open

KL Divergence for Latent SDEs#463
lockwo wants to merge 28 commits into
patrick-kidger:mainfrom
lockwo:kl-pr

Conversation

@lockwo

@lockwo lockwo commented Jul 12, 2024

Copy link
Copy Markdown
Contributor

A continuation of #402 with the new 0.6.0 lineax changes.

Relevant origin issue: #401

@lockwo
lockwo marked this pull request as ready for review July 12, 2024 18:14

@patrick-kidger patrick-kidger left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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. :)

Comment thread diffrax/_integrate.py Outdated
Comment thread diffrax/_solver/kl.py Outdated
Comment thread diffrax/_solver/kl.py Outdated
Comment thread diffrax/_solver/kl.py Outdated
Comment thread diffrax/_solver/kl.py Outdated
Comment thread diffrax/_solver/kl.py Outdated
Comment on lines +132 to +135

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread diffrax/_solver/kl.py Outdated
Comment thread diffrax/_solver/kl.py Outdated
Comment thread diffrax/_solver/kl.py Outdated
@lockwo

lockwo commented Jul 22, 2024

Copy link
Copy Markdown
Contributor Author

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. :)

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 :)

@lockwo

lockwo commented Aug 15, 2024

Copy link
Copy Markdown
Contributor Author

@lockwo
lockwo requested a review from patrick-kidger August 15, 2024 19:21
@lockwo lockwo mentioned this pull request Aug 15, 2024
@lockwo lockwo mentioned this pull request Jun 3, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants