scverse / scvi-tools

Deep probabilistic analysis of single-cell and spatial omics data
http://scvi-tools.org/
BSD 3-Clause "New" or "Revised" License
1.21k stars 344 forks source link

interoperability with pyro to simplify implementing complex hierarchical priors? #862

Closed vitkl closed 3 years ago

vitkl commented 3 years ago

Hi scVI team!

I am wondering if you are planning any interoperability with pyro to simplify implementing complex hierarchical priors? Maybe you have other ideas of simplifying that? I would like to try to adding Variational Autoencoding posteriors to cell2location (https://github.com/BayraktarLab/cell2location, https://www.biorxiv.org/content/10.1101/2020.11.15.378125v1.full) - would be cool if this is possible by integrating with scVI.

Thanks

Vitalii

adamgayoso commented 3 years ago

Hi Vitalii,

We don't currently support pyro -- we could after we finish the current model changes (move to pytorch lightning for training and it seems pyro can work with pytorch lightning).

Since you're considering amortized VI, do you have a sketch already of the model and posteriors? I see your model has many latent variables, but if you're using a mean field approximation, manually coding the ELBO maybe could not be so challenging?

I'm happy to continue this discussion, but it would be very helpful for us to continue at https://discourse.scvi-tools.org/ -- which we are hoping to use for topics like this.

vitkl commented 3 years ago

Hi Adam

Support for pyro would be quite amazing. Sound like it is better to wait until you finish the changes. I will re-initiate the discussion on the forum.

With this many parameters, my concern is that it would be necessary to reinvent something similar to pyro to automatically collect KL contributions from parameters. I struggle a bit with wrapping my head around the minimal number of steps needed to add this to scVI (I tried inferring this from the AutoZiVAE model but it has extra bits like reparametrisation of Beta distr. which make it harder to understand what is minimally needed).

adamgayoso commented 3 years ago

@vitkl now that we have redone our backend, we are going to explore interoperability with Pyro. It actually seems like it can be easily done. Do you have a version of your model in pyro now that I could take a look at?

If this all works out and you're interested, you'd be able to have a model potentially in scvi.external, or you can use our new template repository (https://github.com/YosefLab/scvi-tools-skeleton)

vitkl commented 3 years ago

Hi Adam

This sounds very exciting! Keep me posted.

You can take this model for example: https://github.com/BayraktarLab/cell2location/blob/a03b9dc03c97f326f384c15fa0b276273d9d38a9/cell2location/models/LocationModelLinearDependentWPyro.py#L179 It relies on Pyro AutoNormal guide / Variational distribution. Ideally it would be possible to combine AutoNormal guide for global parameters and amortized inference encoder NN on local parameters for each observation.

Thanks for pointers - we will think about the choices

On Thu, 28 Jan 2021, 19:02 Adam Gayoso, notifications@github.com wrote:

@vitkl https://github.com/vitkl now that we have redone our backend, we are going to explore interoperability with Pyro. It actually seems like it can be easily done. Do you have a version of your model in pyro now that I could take a look at?

If this all works out and you're interested, you'd be able to have a model potentially in scvi.external, or you can use our new template repository ( https://github.com/YosefLab/scvi-tools-skeleton)

— You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub https://github.com/YosefLab/scvi-tools/issues/862#issuecomment-769304291, or unsubscribe https://github.com/notifications/unsubscribe-auth/AFMFTV7IIFE7LWZFU7PW6GLS4GYDZANCNFSM4T67WIDQ .

adamgayoso commented 3 years ago

Thanks for sharing. I'm hoping the interface will be generic enough such that your code would just basically have to be copy pasted into different places -- structured in such a way that it will interface well with our higher-level models and black box train methods.