Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Stateful Controls #559

Draft
wants to merge 16 commits into
base: dev
Choose a base branch
from
Draft

Conversation

lockwo
Copy link
Contributor

@lockwo lockwo commented Jan 5, 2025

This is the v1 of the stateful control PR.

This PR introduces a suite of internally breaking changes which allow for stateful paths. To the non-advanced user, this will likely have no impact on the way they interact with diffrax. It should not impact ODE users at all. SDEs users may be impacted as they might want to try the new "UnsafeBrownianPath" (or as I am calling it "DirectBrownianPath" since it isn't really that unsafe any more), but this should just allow for faster code. I tried several approaches to this problem, but I think based on the conversations in other issues (e.g. #548), that this is at least someone an aligned vision.

Functionaly, this was just a lot of adding arguments to things to support an extra path_state variable. Hence the internal breaking changes. However, these breaking changes come with some decent motivation:

  • Flexibility and speed: this allows for SDE users to see meaningful speedups (on smaller problems where the VF isn't something massive). VBT is necessarily (and unecessarily VBT vs brownian path slowdown #489) slower than UBP, and in theory is only necessary for adaptive solvers that reject steps (I think? Also maybe stateful paths allow some computations to carry brownian bridge states to work with rejecting steps, I think Frank had some jax code that did that but wasn't in diffrax). Thus for anyone running non adaptive solvers (or running solvers that use pre-visible step size adaptation, not currently in diffrax but we have several of those), VBT is not necessary. In addition to the speedups users get from being able to use UBP where before they had to use VBT, UBP is also faster. UBP is faster in both of its new modes, one in which it just splits the key as it goes along and (much faster) in the case where we precompute brownian motion. This addresses Performance issue with SDE solver #517 and Why can't UBP be backproped? #490.
  • Generalization: paths can be stateful, and this is a more complete and generalizable picture that will likely unlock new functionality (and previously, the UBP relied on a bit of a hack to split the keys to not be stateful). Being applied at the AbstractPath level also checks off one point of Tidy-up to-dos #217.

Currently, this PR is missing a few things though.

  • Polish (there's some naming/typing inconsistencies as I made it)
  • More analysis on speed testing
  • Direct adjoint fix: Bounded while loop equinox#923
  • Backsolve adjoint (maybe) fix: I'm not very knowledgeable on this, so I'm not sure how stateful paths could fit in, but I think they can work (they don't error currently, but I get pretty high errors, I assume there needs to be some path tracking alignment that I don't currently do)
  • There's some comments/questions I left for myself (and reviewers) that I still need to address (I can extract some of the main ones to be sure to get feedback)
  • Probably could use a few more tests

@lockwo lockwo marked this pull request as draft January 5, 2025 18:21
@lockwo lockwo changed the base branch from main to dev January 6, 2025 00:22
@patrick-kidger
Copy link
Owner

I can see this is a draft, LMK when you want a review! Happy to take a look on specific points if you want direct feedback earlier.

(Although FWIW I think it'd be cleanest if #484 is the next 'big' thing we merge, just to avoid us having too many 'big' things in flight at the same time.)

@lockwo
Copy link
Contributor Author

lockwo commented Jan 8, 2025

(Although FWIW I think it'd be cleanest if #484 is the next 'big' thing we merge, just to avoid us having too many 'big' things in flight at the same time.)

I'm good with that, although I think that's in @andyElking's control ;)

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.

Bounded while loop
2 participants