thoglu / jammy_flows

A package to describe amortized (conditional) normalizing-flow PDFs defined jointly on tensor products of manifolds with coverage control. The connection between different manifolds is fixed via an autoregressive structure.
MIT License
42 stars 3 forks source link

transform & jacobian weights #1

Closed JohannesBuchner closed 3 years ago

JohannesBuchner commented 3 years ago

Hi,

I see that the example script learns how to fit jammy_flows to characters (conditional on the character as the label). But I am a bit confused how to do the simpler unconditional case, with Gaussianization flow. The notebook does not seem to use any samples, the visualisations just come after initialisation of the structure.

I'd also like to know how to use the pdf object. Specifically, if I have a coordinate in the "basis" space (where the PDF is a unit gaussian), how do I transform to the target space and obtain the jacobian weight for that location/transformation? If I understand correctly, forward() is the other direction?

Cheers, Johannes

thoglu commented 3 years ago

Hi Johannes, yes that is a good point I will add a script where no conditional distribution is fitted in some future commit. You just leave out the "conditional_input" and it will fit without any conditional input.

Regarding your second question: Forward evaluates the PDF in the target space and automatically takes care of the jacobian factor and calculates all involved quantities. It requries the target space point on which you want to evaluate the PDF, and it then returns the log_pdf of that target-space point, the log_pdf of the corresponding base space point, and the corresponding base space point.

Internally, forward calculates f^-1(x_target), the inverse function. Forward is the usual function associated to any pytorch model. It has nothing to do with the "forward" direction of the injective function of the normalizing flow, and it is synonymous with a direct call of the model as model(x) = model.forward(x).

Is there any reason or use case you would want to have explicitly just the Jacobian weight?

For sampling you call "sample", and it calculates the target and base space point, and also the target and base log_pdfs.

I was hoping the description in the README somehow made that clear, I wanted to add a more thorough documentation once all intended features are there.

Best, Thorsten

JohannesBuchner commented 3 years ago

I want to do something a little weird, which involves running a sampler in the base space, and back-correcting for the space distortions. Essentially this is about using the PDF for proposal efficiency, without using it as a prior. I'll look into the forward function. Thanks!

thoglu commented 3 years ago

Ok, so maybe you can modify the "sample" function, such that it also returns the jacobian factor. Intrenally, the sample function samples in base space, and then calculates f(z_base)=x_target and the jacobian factor - but in the end combines the result such that it returns x_target, z_base, log_pdf(x_target), log_pdf_base(z_base) - by changing the return values you might get something that suits you.

One way which would also work is calculate the log_pdf of the base space sample and subtract that from the third return value of the "sample" function - that should be equivalent to the jacobian factor.

JohannesBuchner commented 3 years ago

Ok, thanks, I will have a look at that. Just to clarify, I want to do the walking on the base space myself externally, I don't want to (importance) sample from the Gaussian. But that information sounds like exactly what I need.