Question/feature: perform posterior inference

[itsdfish]

Hello,

I am interested in performing posterior interference with various simulation models, such as agent based models. I was wondering whether that is possible with your package. One example in Python is BayesFlow. If this is not possible currently, I think it would be a very useful feature to add.

As a simple example, suppose I have a Gaussian model and I want to update the prior on mu and sigma after observing 50 observations. Would that be possible? Here is partial code:

using Distributions 

function sample_prior()
    μ = rand(Normal(0, 1))
    σ = rand(Uniform(0, 10))
    return [μ,σ]
end

function generate_data()
    θ = sample_prior()
    y = rand(Normal(θ...))
    return [θ...,y] 
end

# rows: μ, σ, y
training_data = mapreduce(x -> generate_data(), hcat, 1:10000)

# train some data

observed_data = rand(Normal(0, 1), 50)

# obtain posterior distribution of μ and σ

[rafaelorozco]

Yes, our main use for this library is posterior inference and it currently works very well by training on samples on the joint distribution as in your setup.

Im in the middle of travel, but tomorrow I will send a script on how to train on the training pairs you made.

[itsdfish]

Awesome. I would very much appreciate an example.

Also, does it support multivariate distributions?

[lstagner]

An example would be great. I've been looking into doing Variational Inference with Normalizing Flows but I have yet to find a complete Julia example.

[rafaelorozco]

Started a PR so that we can add some more examples. Started with the setup that @itsdfish gave. The basic script is here https://github.com/slimgroup/InvertibleNetworks.jl/pull/84/files#diff-231dfb51cc0355965e85e0fedb61956dd6e52e86eb974de512ede26feb55d8b6

using Distributions

function sample_prior()
    μ = rand(Normal(0, 1))
    σ = rand(Uniform(0, 10))
    return [μ,σ]
end

function generate_data(num_obs)
    θ = sample_prior()
    y = rand(Normal(θ...), num_obs)
    return [θ...,y...] 
end

# train using samples from joint distribution x,y ~ p(x,y) where x=[μ, σ] -> y = N(μ, σ)
# rows: μ, σ, y
num_params = 2; num_obs = 50; n_train = 10000
training_data = mapreduce(x -> generate_data(num_obs), hcat, 1:n_train);

############ train some data
X_train = reshape(training_data[1:num_params,:], (1,1,num_params,:));
Y_train = reshape(training_data[(num_params+1):end,:], (1,1,num_obs,:));

n_epochs   = 2
batch_size = 200
n_batches = div(n_train,batch_size)

# make conditional normalizing flow
using InvertibleNetworks, LinearAlgebra, Flux

L = 3 # RealNVP multiscale layers
K = 4 # Coupling layers per scale
n_hidden = 32 # Hidden channels in coupling layers' neural network
G = NetworkConditionalGlow(num_params, num_obs, n_hidden,  L, K;);
opt = ADAM(4f-3)

# Training logs 
loss_l2   = []; logdet_train = [];

for e=1:n_epochs # epoch loop
    idx_e = reshape(1:n_train, batch_size, n_batches) 
    for b = 1:n_batches # batch loop
        X = X_train[:, :, :, idx_e[:,b]];
        Y = Y_train[:, :, :, idx_e[:,b]];
     
        # Forward pass of normalizing flow
        Zx, Zy, lgdet = G.forward(X, Y)

        # Loss function is l2 norm - logdet
        append!(loss_l2, norm(Zx)^2 / prod(size(X)))  # normalize by image size and batch size
        append!(logdet_train, -lgdet / prod(size(X)[1:end-1])) # logdet is already normalized by batch size

        # Set gradients of flow
        G.backward(Zx / batch_size, Zx, Zy)

        # Update parameters of flow
        for p in get_params(G) 
          Flux.update!(opt,p.data,p.grad)
        end; 
        clear_grad!(G)

        print("Iter: epoch=", e, "/", n_epochs, ", batch=", b, "/", n_batches, "; f l2 = ",  loss_l2[end], 
            "; lgdet = ", logdet_train[end], "; full objective = ", loss_l2[end] + logdet_train[end], "\n")
    end
end

# posterior inference on unseen observation
x_ground_truth = [1,1] # mu=1, sigma=1
observed_data  =  reshape(rand(Normal(x_ground_truth[1], x_ground_truth[2]), num_obs), (1,1,num_obs,:))

# posterior sampling with conditional normalizing flow
num_post_samps = 1000
ZX_noise = randn(1,1,num_params,num_post_samps) 
Y_forward = repeat(observed_data,  1, 1, 1, num_post_samps) 

_, Zy_fixed_train, _ = G.forward(ZX_noise, Y_forward); #needed to set the proper transforms on inverse

X_post = G.inverse(ZX_noise, Zy_fixed_train );

# plot results
using PyPlot

X_prior = mapreduce(x -> sample_prior(), hcat, 1:num_post_samps)
X_prior = reshape(X_prior[1:num_params,:], (1,1,num_params,:))

fig = figure()
subplot(1,2,1)
hist(X_prior[1,1,1,:];alpha=0.7,density=true,label="Prior")
hist(X_post[1,1,1,:];alpha=0.7,density=true,label="Posterior")
axvline(x_ground_truth[1], color="k", linewidth=1,label="Ground truth")
xlabel(L"\mu"); ylabel("Density"); 
legend()

subplot(1,2,2)
hist(X_prior[1,1,2,:]; alpha=0.7,density=true,label="Prior")
hist(X_post[1,1,2,:]; alpha=0.7,density=true,label="Posterior")
axvline(x_ground_truth[2], color="k", linewidth=1,label="Ground truth")
xlabel(L"\sigma"); ylabel("Density");
legend()
tight_layout()

And after two epochs, the posterior samples look like this for mu=1 std=1

@itsdfish yes it is quite good for multivariate distributions. I actually think that the networks we have shine best on posterior distributions over images because they are currently mostly conv net based. Although, we are working on dense network for 1d inputs in #77

Will be happy to help anyone get these working for their specific applications and then we can mold them to nice examples.

[itsdfish]

Thank you for providing an example. This is very helpful. While looking through the code, I thought of a few high level questions. I hope you don't mind me asking some basic questions (machine learning is outside of my expertise).

1. Is it possible to perform inference on a data set of any size, as is possible in BayesFlow? In the example you provided, the training size and the observed size for inference is num_obs. An error is thrown if they are different. Unfortunately, this limits the usefulness of the network.
2. Why do you use 4 dimensional arrays (e.g., X_train etc)? What are the first two dimensions?
3. Can this example serve as a template for models with more parameters and data dimensions?
4. What kind of neural network are you using? Is it a conditional auto encoder?
5. Are there any types of models that this package is not suitable for? For example, time series data, or non iid?

Thanks again!

[rafaelorozco]

Dont mind at all! thank you for the interest and good questions:

1.) Yes it is possible. In general you need a summary statistic that transorms arbitrary sized observations to a single fixed size. Either learned summary statistic (summary network) like in BayesFlow or a transformation related to the forward operator like in:
https://doi.org/10.1117/12.2651691
and
https://openreview.net/pdf?id=LoJG-lUIlk

The learned summary networks will be merged soon in #82

2.) This isn't anything inherent. The 4d arrays are just a result of our lab applications being mostly on structured images that have width height dimension and 5d tensors for volume inputs. So to hack in 1d inputs with the current code we put the data into the channel dimension and the rest are singletons. It looks like the community is also interested in data with 1 dimension so we are adding dense layers and logic to handle 1 dimension in the PR of dense layers #77

3.) Yes, currently it should work for any parameter and data dimension. I played around with a few without errors but if you run into anything let me know and I might be able to point to a solution quicker.

4.) I cant think of anything that doesnt work off of the top of my head. There is some technical details on the limits of normalizing flows for learning distributions. It has to do with there existing a diffeomorphism between the data distribution and the base distribution but in my experience this is mostly a technical detail and they have been practically used on essentially every type of data application. so AFAIK, the package should work for any marginal distribution and any conditional distribution. The marginal you just need samples x~p(x). But the conditional distributions can be trained with different types of samples depending on the scenario. For the above example (amortized posterior) you need samples from the joint distribution. But you can also do non-amortized Variational inference that only requires a single observation and access to the forward model and its adjoint (if that was part of the process used to make the observation). This method is used in tandem with traditional Amortized variational inference in one of the papers in this package readme:
https://arxiv.org/abs/2207.11640
https://github.com/slimgroup/ReliableAVI.jl

[itsdfish]

Awesome! Thanks for your answers. The approaches you described in your answer to my first question will be very useful to the Julia community because there are so many types of models for which Bayesian inference is challenging.

I would like to take you up on your offer to develop some tutorial examples. If you would be willing to adapt your gaussian example to use the summary statistic technique once #82 is merged, I can use it to develop tutorials for different types of models. Here are some initial thoughts:

1. Multivariate Gaussian as a simple demonstration of multivariate distributions

2. A decision making model called the Linear Ballistic Accumulator, which describes how people choose between multiple options. The distribution is multivariate where one dimension (choice) is discrete, and the other is continuous (response time). The LBA has a tractable likelihood function, which would allow us to potentially compare the invertible network to results from Turing.jl. However, there are many decision making models that are only simulation based. So this would be of broad interest to that community.

3. An agent based model from Agents.jl. The developers of Agents.jl have several example models available. I think this could become very popular with that modeling framework. In fact, there was recently a question on discourse which prompted me to look for possible solutions.

4. I think it would be a good idea to demonstrate a few things, such as saving trained networks (presumably with BSON.jl), how to generate posterior predictive distributions, and perform quality checks. I would be happy to help with these.

When all is said and done, it might be worth wrapping these procedures into useful functions and putting them in a separate package in the slimgroup organization. If that turns out to be valuable, I can also assist with that.

[itsdfish]

@rafaelorozco, I would like to make an tutorial example for model of discrete data. The constructor for NetworkConditionalGlow thew an "inexact" error. So I was wondering if it is even possible. I added the code here in case its possible and you would be willing to have a look.

Update

It seems like the problem that NetworkConditionalGlow fails when num_params = 1. If I duplicate the parameter theta so that there are two parameters, it seems to work correctly and agree with the analytic posterior distribution of theta.

[rafaelorozco]

I came to answer the question but you arrived to exactly the correct conclusion in the update! It is also cool that you came up with the hack that I typically use when doing invertible layers on arrays of length 1.

Thank you so much for the interest and the cool ideas for examples. Im working with @mloubout to merge the summary network PR soon. We have been using them successfully for a while, just want to make the interface a little bit more user friendly.

The network saving and loading is super important so I will make that example right now for the PR.

[itsdfish]

Awesome! I look forward to digging in and making some tutorial examples once the summary network is complete.

[itsdfish]

One of my goals is to define a user-friendly interface for training and sampling. Maybe with a bit more thought, it can be flexible, general and extendable. Here is the first iteration. See train! and sample_posterior, which wrap around the code you provided. The training info is displayed with a progress bar using ProgressMeter.jl

[rafaelorozco]

I spent some time looking at the code but had a hard time coming up with suggestions. Looks good! I think it is a matter of increasing complexity of examples and seeing where things break and need to be more generalizable.
I added an example of loading and saving networks here in the PR:

https://github.com/slimgroup/InvertibleNetworks.jl/blob/3c5dddd24afb6ca737f5921cc8a0305cb4598e1a/examples/utils/save_load_network.jl

There was an error with BSON so you need to be on that PR for it to work.

[itsdfish]

Awesome. Thanks for adding functionality for saving and loading trained networks. This will be very useful!

Thanks for looking over the code. I did encounter an issue while trying to develop a multivariate gaussian model. The constructor for NetworkConditionalGlow threw a similar inexact error as before. It appears that it throws an error anytime n_parms is odd. Perhaps the dimensions of arrays must be 2^n, or even. Is this something were we need some logic to add a redundant parameter and remove it later?

I do have another question. When extending the y_train to a multivariate gaussian with 2 variables, are the dimensions supposed to be (1,2,n_obs,n_train)? That's my best guess, but I anticipate that might be the next problem.

update

I forgot to add a link above to the training data: https://github.com/itsdfish/NormalizingFlowsTutorials.jl/blob/066c32a2030bc8c7a67f6d434a1b92a57d1b15f6/sandbox/mv_gaussian.jl#L46

[rafaelorozco]

aah great catch thank you for bringing that up. I thought I had fixed it but I missed it in that layer. It is a simple problem related to the invertible coupling layer. It is an easy fix but I forgot to also do it for the conditional layers.

With regards to y_train, we will need to go to proper 1dim implementation that looks like (nx, n_chan, n_train). Where nx is the main size of the variable (in your case n_obs). and n_chan would be used for multiple channels or 2 variables in this case. Again, this is working for non-conditional layer, just need to port it to conditional layers. Should be able to get it done by tomorrow.

Thank you for the very useful input!