search

Joshuaalbert / jaxns

Probabilistic Programming and Nested sampling in JAX
https://jaxns.readthedocs.io/
Other
134 stars 9 forks source link

Implementation of Gradient-guided nested sampling #136

Open juehang opened 6 months ago

juehang commented 6 months ago

Is your feature request related to a problem? Please describe. While nested sampling is implemented in jax in this package, it does not (to my knowledge) make use of jax autograd. There is a recent paper describing a gradient-guided nested sampling method that actually cites this package! https://arxiv.org/pdf/2312.03911.pdf

Describe the solution you'd like I would like an implementation in this package, ideally so that I can use it in numpyro. I do not have time to work full-time on an implementation myself, but I would be interested in a joint-effort as well if that is preferred/needed.

Additional context There is an example pytorch implementation here: https://github.com/Pablo-Lemos/GGNS/tree/main

Joshuaalbert commented 6 months ago

Hi thanks for the idea. I've considered this before. I could introduce this as an experimental feature. It would need to be assessed will before becoming part of the core engine.

The focus of jaxns is fast easy general purpose probabilistic programming. I'm trying to keep the amount implementation experimentation needed by a user to solve their problem to a minimum so they can focus on building better models. To that end I'm not looking to have a library of nested sampling implementations. But I can explore this from an experimental perspective. If the engine can be improved and still remain general purpose and robust with good defaults that solve most problems then it could be an option. This year at the MaxEnt conference I'll discuss the latest nested sampling research with others in the field and plan to update jaxns accordingly.

juehang commented 6 months ago

I see, thank you for the update! My initial impression from the paper is that this method allows for nested sampling to be extended to higher dimensional models. My current research interest is in applying bayesian methods to simulation-based inference problems which often have quite large latent spaces, and so I was hoping this would be the best of both worlds between NUTS/HMC (better scaling with large models) and nested sampling (marginal likelihoods and evidence).

I will keep posted for now, then; feel free to reach out if you ever need an early tester!

Joshuaalbert commented 6 months ago

You should try out an experimental feature of jaxns: parametrised models. In jaxns you can let some variables be Bayesian and some be pointwise estimates of a prior. When you solve the model it finds the parameter values that maximizes the evidence and gives you the posterior of the rest. This can be really powerful, since a full Bayesian treatment with huge latent spaces is impractical for large dimension. You can use this to have Bayesian hierarchical models that perform well.On Jan 20, 2024 06:53, Juehang Qin @.***> wrote: I see, thank you for the update! My initial impression from the paper is that this method allows for nested sampling to be extended to higher dimensional models. My current research interest is in applying bayesian methods to simulation-based inference problems which often have quite large latent spaces, and so I was hoping this would be the best of both worlds between NUTS/HMC (better scaling with large models) and nested sampling (marginal likelihoods and evidence). I will keep posted for now, then; feel free to reach out if you ever need an early tester!

—Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you commented.Message ID: @.***>

juehang commented 6 months ago

Thank you for the suggestion; do you have any papers on this topic? I couldn't find much about this kind of method with a cursory search; I suppose it is similar to profile likelihoods but I was wondering if there are any papers exploring the statistical properties of this procedure, what the limitations are, etc.

Joshuaalbert commented 6 months ago

It's novel, and I will write a paper after I answer some of these questions myself. Are you interested in being a part of the paper? If your example use case is a good fit then it could be the example for the paper.On Jan 20, 2024 23:53, Juehang Qin @.***> wrote: Thank you for the suggestion; do you have any papers on this topic? I couldn't find much about this kind of method with a cursory search; I suppose it is similar to profile likelihoods but I was wondering if there are any papers exploring the statistical properties of this procedure, what the limitations are, etc.

—Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you commented.Message ID: @.***>

juehang commented 6 months ago

I would love to, and I do think I have a model that would be a reasonable use-case (there are 1e2~1e5 latent parameters, depending on dataset, but because those latent parameters correspond to individual events we don't necessarily care about knowing these parameters better than a best-fit estimate).

On top of that, model comparison is something I would be interested in with this model, so nested sampling would be perfect. Would you like to talk about this over zoom or something? Otherwise perhaps we can move this conversation off of the github issue for now 🙂

From: Joshua George Albert @.> Sent: Sunday, January 21, 2024 9:58:35 AM To: Joshuaalbert/jaxns @.> Cc: Juehang Qin @.>; Author @.> Subject: Re: [Joshuaalbert/jaxns] Implementation of Gradient-guided nested sampling (Issue #136)

It's novel, and I will write a paper after I answer some of these questions myself. Are you interested in being a part of the paper? If your example use case is a good fit then it could be the example for the paper.On Jan 20, 2024 23:53, Juehang Qin @.***> wrote: Thank you for the suggestion; do you have any papers on this topic? I couldn't find much about this kind of method with a cursory search; I suppose it is similar to profile likelihoods but I was wondering if there are any papers exploring the statistical properties of this procedure, what the limitations are, etc.

—Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you commented.Message ID: @.***>

— Reply to this email directly, view it on GitHubhttps://github.com/Joshuaalbert/jaxns/issues/136#issuecomment-1902678833, or unsubscribehttps://github.com/notifications/unsubscribe-auth/AJLCK34SGYY7IQCXXO2CF3TYPU3KXAVCNFSM6AAAAABB7SJNZOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMBSGY3TQOBTGM. You are receiving this because you authored the thread.Message ID: @.***>

Joshuaalbert commented 6 months ago

Yes, sounds like a good idea. Do you have LinkedIn? This is my profile https://www.linkedin.com/in/joshua-g-albert-66734638 and easy to connect there and move to email/zoom.