Connect with AdvancedVI.jl

[Red-Portal]

Hi all,

Now that AdvancedVI@0.3.0 is shaping up, shall we consider connecting it with NormalizingFlows? I think most of the work would be about setting up documentation and tests matching AdvancedVI. Is there anything else that needs to be done? Also, what are the current experiences with GPUs?

[zuhengxu]

Hi @Red-Portal, thanks for bringing this up and great work on AdvancedVI.jl! I'm thinking to first sync the API of NormalizingFlows.jl with AdvancedVI.jl (e.g., enable close form potential, stickandlanding, etc.). The rest of it will just be including some examples. What do you think?

I can't comment too much on the GPU side. As long as Distributions.jl and Bijectors.jl play nicely with CuArray, we should be good to go. This was not case a year ago, but could have changed now. @torfjelde, what is your current opinion on it?

[Red-Portal]

Sounds all good to me! I vaguely remember that you didn't experience much performance gains from using GPUs. Any idea what was the bottleneck?

[zuhengxu]

Sounds all good to me! I vaguely remember that you didn't experience much performance gains from using GPUs. Any idea what was the bottleneck?

Yes! I didn't do a proper investigation on the bottleneck, so I can't comment too much on it. If I have to take an educated guess, the current implementation of ELBO est is not ``batched''; it uses map function to apply the `elbo_single_sample` to a collection of samples (see this line). In my perspective, this could be one of the biggest performance lose when working with GPUs.

If I recall correctly, the reason why I chose this way is because many Bijectors.jl layers did not support batched input. (The situation might have changed now.)

[Red-Portal]

That makes a lot of sense. I wonder if it would be hard to bypass bijectors and make a batch version locally until upstream catches up. Or maybe re-use as much upstream component while making a local batch interface. I could then make a NomalizingFlows.jl extension in AdvancedVI that matches that.

[Red-Portal]

Actually that isn't even necessary. It would suffice to define rand(rng, q_flow, n_samples here without even going through the Bijectors interface.

[zuhengxu]

Actually that isn't even necessary. It would suffice to define rand(rng, q_flow, n_samples here without even going through the Bijectors interface.

Ahh I think I didn't explain this part clearly. the rand or logpdf interface of Bijections.jl is not the issue here; it is more relavent to whether the bijective transformation b is compatible with the batchwise operation. To do rand(rng, q_flow, n_sample), ideally we want to dorand from reference to get bunch of x0s and transform it usingb(x0s), which I believe was not fully supported. (similarly for density evaluation, which requires the batched operation on binv).

[Red-Portal]

@zuhengxu @yebai I just tried using NormalizingFlows.jl in AdvancedVI today ,and it just works perfectly. I am planning to add a tutorial showing how to do this once AdvancedVI@0.5 is released and the Turing interface is updated accordingly. Let's think about what else is there to do to connect the two. Though it seems NormalizingFlows.jl itself doesn't provide interfaces for easily constructing flows. Is this planned for the future? or is it intended? For example, the boilerplate for the neural spline flow looks like it should belong under src/.

[zuhengxu]

@Red-Portal Glad to hear that the NF.jl is compatible with AdvancedVI.jl functionality! I assume you just construct the flow with ADVI references and some NF flow layers, and train it with those utility functions defined in AdvancedVI.jl?

itself doesn't provide interfaces for easily constructing flows. Is this planned for the future? or is it intended?

It was intended as the point was to allow users to freely define their flows. Though I've been thinking lately to
include a constructor for at least RealNVP. I'm pretty confident to the RealNVP architecture provided in example/, and its probably the most widely used flow in literature.

[zuhengxu]

Actually, it maybe the right time to discuss more about how to connect NF and AVI. @Red-Portal There are a few things that I'd like to hear your thoughts. Two things that I've been wanting to implement but haven't got the chance to are:

1. Using more structured reference distribution---such full-rank Gaussian---for flows (though I'm not sure if its really beneficial in practice given those flexible flow transformations). One remark is that some parameters in the structured family (e.g., covaraince) would require specific operators such as the proximal operator or gradclip to ensure pd, while the rest params just get updated normally. How do we identify these structured params (e.g., cov) from the flattened param array?

2. The STL gradient estimator for NF is bit tricky due to non-close form entropy, which typically require a stop_grad operator. I wasn't super sure about how to do it in an AD-backend-agnostic way. Is ChainRuleCore.jl the best choice here?

[Red-Portal]

It was intended as the point was to allow users to freely define their flows

I think it would be nice to have an interface that provides varying levels of control. For Turing, for example, we probably need a well-crafted default constructor that only takes the capacity of the flow. Something like RealNVP(dim, depth, width). Here width could be a scalar for the simple-minded or a vector for those who want to customize the width-per-layer, for example. (I would like to have something like this for all of the flows 😄 )

Using more structured reference distribution---such full-rank Gaussian---for flows (though I'm not sure if its really beneficial in practice given those flexible flow transformations). One remark is that some parameters in the structured family (e.g., covaraince) would require specific operators such as the proximal operator or gradclip to ensure pd, while the rest params just get updated normally. How do we identify these structured params (e.g., cov) from the flattened param array?

AdvancedVI provides an interface called operators where we can unflatten the parameters, modify some values, and re-flatten them back. This way, custom projection operators can be easily implemented. But for flows, since everything is non-convex anyway, I think it would be okay to use nonlinear parametrizations (softplus the diagonal of the scale matrix and such.)

The STL gradient estimator for NF is bit tricky due to non-close form entropy, which typically require a stop_grad operator. I wasn't super sure about how to do it in an AD-backend-agnostic way. Is ChainRuleCore.jl the best choice here?

This works out of the box with AdvancedVI as long as logpdf(q, x) is implemented and differentiable. For the way it is implemented, you can check it out here. Instead of relying on ChainsRulesCore, we just construct a copy of the current q and pass it as a DifferentiationInterface.Constant. This works across all AD frameworks supported by DifferentiationInterface.

[zuhengxu]

I think it would be nice to have an interface that provides varying levels of control. For Turing, for example, we probably need a well-crafted default constructor that only takes the capacity of the flow. Something like Something like RealNVP(dim, depth, width)...

Yeah, for RealNVP such interface make sense. But for other flows such as NSF, each layer will have too many hyperparamter to control (e.g., bandwidth, number of nodes, NN architectures). I'm not a huge fan of having a NeuralSplineFlow(...) with 10 arguments. I'll make a PR shortly to include the interface.

Instead of relying on ChainsRulesCore, we just construct a copy of the current q and pass it as a DifferentiationInterface.Constant

This is smart, thanks! I'll give it a shot.

[Red-Portal]

Yeah, for RealNVP such interface make sense. But for other flows such as NSF, each layer will have too many hyperparamter to control (e.g., bandwidth, number of nodes, NN architectures). I'm not a huge fan of having a NeuralSplineFlow(...) with 10 arguments. I'll make a PR shortly to include the interface.

Can this stuff be moved to src/ so it can be called from other packages?

[zuhengxu]

Can this stuff be moved to src/ so it can be called from other packages?

I'm working on this in #53.