Guidance on scaling to ~40-100 dimensional problems — JIT compilation time

[tomkimpson]

Hi Joshua,

We're using jaxns for nested sampling in a pulsar timing array (PTA) continuous wave gravitational wave search. The likelihood is a Kalman filter implemented in JAX, and evaluates in ~0.02s on GPU after JIT compilation.

What works well:

With 2 free parameters, jaxns runs end-to-end in ~30 minutes — roughly 25 min JIT compilation and ~3 min actual sampling (130k likelihood evals, 2500 samples,
ESS=314). Evidence computation works perfectly. We're using num_live_points=200, s=3, k=0, max_samples=5000.

Where we struggle:

Our full model has ~107 free parameters. With this dimensionality, jaxns never gets past JIT compilation — we've waited 7.5 hours on a P100 GPU with no progress beyond "Number of Markov-chains set to: 200".

We believe the compilation bottleneck is the XLA trace of the full nested sampling loop including our likelihood, not the sampling itself. The likelihood's computation graph is moderately complex (Kalman filter
with lax.scan), and we suspect the combination of this with 107-dim slice sampling creates a very large XLA program.

Questions:

1. What dimensionality range is jaxns designed to handle efficiently? Is ~100 dims realistic? 50 dims? Or is ~10-20 more typical?
2. Is the compilation time expected to scale steeply with the number of parameters? Any strategies to reduce it (e.g., wrapping the likelihood to limit the trace depth)?
3. Would adjusting s, k, or num_live_points help with compilation time, or is it purely a function of the model graph size?
4. Are there any best practices for using jaxns with complex lax.scan-based likelihoods?

Thanks for a great library — the 2-dim case works beautifully.

[Joshuaalbert]

Hi @tomkimpson glad to hear your use case. My suggestion now is to decorate your likelihood function with jit. Make sure it's defined at global scope so that it's compiled once.

Also, watch for release of JAXNS v3 which will address performance and hugely improve algorithmic efficiency. So you will be able to take advantage of parallel multi machine processing easily, and large dimension problems. An associated paper is being released at the same time which will describe the novel algorithmic aspects in case you're interested to learn more.

[Joshuaalbert]

Ps you may be able to ask an AI to help your improve the efficiency of your likelihood implementation. Things like embedding large static arrays, or inefficient things like sorting, can often be avoided.

Let me know if the above helps.

[nenasedk]

@Joshuaalbert

Just wanted to ask a bit about v3. 1) What's changing in the parallelization to help distribute it across devices? 2) Would you recommend testing out the improvements branch to try out the new features? And 3) Do you have a rough estimate of the timeline for the v3 paper - will it be out this year? Thanks, and looking forward to the update!

[Joshuaalbert]

Hi @nenasedk

1. The key for parallelism is a new algorithmic advance that allows embarassingly parallel calculation, so this means we can start a bunch of workers and dispatch sampling requests to them. I implement distribution with ZMQ. The flow will be two-step: 1) On every node that you want to participate, run a process telling it which devices will participate (can be mixture of CPU and GPU devices, and it's logical so you can have the GPU or CPU be many logical devices), 2) on some compute you'll run your main program that defines the Model. It then outsources likelihood evaluation to workers. This also works on a single node. There will still be the fully JAX non-distributed pathway, which uses pmap or vmap to do parallel replacement on a single machine, but for really big problems you'll be able to use an entire compute cluster with 1000's of workers accelerating the computation.
2. I wouldn't get your hands dirty until I put a stamp of approval on it. Still some things I'm not happy with that need to be addressed
3. Yes, this year I'll put out the paper. It will replace the phantom powered NS paper, which was found with the help of Will Handley and Andrew Fowlie to contain incorrect results. But since then I've resolved these problems and in the process understood something a little deeper in the algorithm.