Low Storage Methods

[luke-a-thompson]

Hi,

I've been working on low-storage Runge-Kutta methods recently, in particular Williamson 2N methods.

These methods compute the ODE flow via a recursive scheme that only uses two registers, so they can be useful in memory-constrained settings like PDE solves. The tradeoff is usually that they need more stages for a given order.

I put together an implementation here: diffrax-lowstorage.

At the moment it includes:

• a 2N solver class with substantially lower memory use at the same order,
• embedded error estimates,
• and two well-known methods.

@patrick-kidger, would there be any interest in merging something like this into Diffrax? I'd be happy to open a PR and make whatever changes you think would be needed.

As far as I know, the only other mainstream library with this kind of support is Diffeq.jl LowStorageRK, so it could be a nice addition.

Thanks

[patrick-kidger]

Hey there! This is pretty awesome to see.

I think realistically this isn't something I can commit to maintaining myself (at least right now), but I'd love to be able to advertise this where appropriate. Perhaps we can add them as a section to this page under a new section that contains a link to your library? Also a link on the Equinox awesome list would be appropriate.

I'd be happy to take a PR on those :)

That aside, one technical question – IIUC then the memory savings won't help you under autodiff, when basically everything is saved anyway? I imagine you could resolve this by wrapping things in a jax.checkpoint, or something else of similar character. Just curious how you're thinking about this.

[luke-a-thompson]

Hi Patrick,

That's understandable, but I'd love to add section to the ODE solver docs on these methods. I will open the PRs when I've made some more progress on the library/docs.

I ran some memory benchmarks with Compiled.memory_analysis() and there was was still some solid memory savings under each different adjoint.

This is quite surprising to me, as I was intending to use #727 for the very reason highlighted in your message. I'm no jax expert so please let me know if I'm not testing memory usage in the right way.

Thanks,
Luke

tests/test_bench_lowstorage.py
compiled-memory-bytes (size=8192):
  bosh3: 266048 bytes  (vs bosh3: 1.000x)
  heun: 167232 bytes  (vs bosh3: 0.629x)
  bwrrk33: 137473 bytes  (vs bosh3: 0.517x)
  bwrrk53: 137473 bytes  (vs bosh3: 0.517x)

compiled-reverse-mode-memory-bytes-by-adjoint (size=8192):
  adjoint=recursive_checkpoint:
    bosh3: 11981136 bytes  (vs bosh3: 1.000x)
    heun: 6018096 bytes  (vs bosh3: 0.502x)
    bwrrk33: 6019368 bytes  (vs bosh3: 0.502x)
    bwrrk53: 6019880 bytes  (vs bosh3: 0.502x)
  adjoint=direct:
    bosh3: 5020144 bytes  (vs bosh3: 1.000x)
    heun: 2525400 bytes  (vs bosh3: 0.503x)
    bwrrk33: 2912257 bytes  (vs bosh3: 0.580x)
    bwrrk53: 4059905 bytes  (vs bosh3: 0.809x)
  adjoint=backsolve:
    bosh3: 499328 bytes  (vs bosh3: 1.000x)
    heun: 301416 bytes  (vs bosh3: 0.604x)
    bwrrk33: 206849 bytes  (vs bosh3: 0.414x)
    bwrrk53: 207105 bytes  (vs bosh3: 0.415x)

runtime (size=8192):
  bosh3: 49.906 ms  (vs bosh3: 1.000x)
  heun: 12.556 ms  (vs bosh3: 0.252x)
  bwrrk33: 10.358 ms  (vs bosh3: 0.208x)
  bwrrk53: 11.600 ms  (vs bosh3: 0.232x)

[patrick-kidger]

Interesting! I don't know. Diffrax has some checkpointing internally, maybe that is helping :)