Overconfidence in SNE (SNPE/SNRE/SNLE) estimates while comparing to the ground truth estimate by MCMC

[paarth-dudani]

I am working on establishing fidelity of SNE estimates for the following 1D exponential model with noise:

I have used pymc to get the ground truth posterior estimate for this model as shown below (passing convergence diagnostics):

However, when I run SNE algorithms on the same and a slightly different dataset (with a different random instantiation of the same noise level) as the one used for pymc, I get overconfident estimates as shown below:

I would really appreciate some advice in dealing with this issue that I am facing. Following is the code I am using, shown for SNRE:

from torch.distributions import Exponential, MultivariateNormal, Uniform, Normal

import pandas as pd
import matplotlib.pyplot as plt 
import numpy as np
from sbi.analysis import pairplot
from torch.distributions import Exponential, MultivariateNormal, Uniform
import sbi
import torch
from sbi.inference import NRE, DirectPosterior, simulate_for_sbi
import pandas as pd
from sbi.utils import BoxUniform

def simulator_custom(theta, num_sim, input_array, initial_condition):
    output = np.zeros((num_sim, len(input_array)))
    for i in range(num_sim):
        for j in range(len(input_array)):
            output[i,j] = (initial_condition)*np.exp(-input_array[j]*theta[i]) + np.random.normal(0,0.1)
    return torch.Tensor(output)


n = 500
t = np.linspace(0,4,n)
initial_condition = 2.0


theta_true = torch.tensor([1.5])

x_o = simulator_custom(theta_true, 1, t, initial_condition)

plt.scatter(t, x_o, s=40, c='green', edgecolors='black', linewidths=0.5, alpha=0.7)
plt.xlabel('theta',  fontsize=12)
plt.ylabel('Number of iterations',  fontsize=12)
plt.title('Observed data',  fontsize=12)

class WrappedExponential(Exponential):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def log_prob(*args, **kwargs):
        log_probs = Exponential.log_prob(*args, **kwargs)
        return log_probs.squeeze()

exp_prior_mod = WrappedExponential(torch.tensor([2.0]))
uniform_prior = BoxUniform(low=torch.tensor([0.0]), high=torch.tensor([4.0]))



# Simulation conditions
num_dims = 1
num_sim = 500
num_rounds = 3

# Later rounds use the APT loss.

inference = NRE(prior=uniform_prior)
proposal = uniform_prior

def run_SNRE(proposal, x_o, num_rounds, num_sim):
    for _ in range(num_rounds):
        theta = proposal.sample((num_sim,))
        initial_condition = 2.0
        # theta = np.clip(theta, 0, 2.5)
        x = simulator_custom(theta, num_sim, t, initial_condition)
        x = np.clip(x, 0, 8)
        _ = inference.append_simulations(theta, x).train()
        posterior = inference.build_posterior(mcmc_method="slice_np_vectorized",
                                        mcmc_parameters={"num_chains": 10,
                                                        "thin": 5})
    proposal = posterior.set_default_x(x_o)

    return proposal

proposal = run_SNRE(uniform_prior, x_o, num_rounds, num_sim)

Following is the dataset used in for estimation by pymc and SNE algorithms:

[michaeldeistler]

Hi there, thanks for creating this! Do the SNPE/SNLE/SNRE estimates converge to the true solution if you use many more simulations (e.g. 10k)?

Michael

[paarth-dudani]

Hi!
Thanks for the suggestion!

I did it for SNE algorithms and SNRE seems to give a close enough estimate whereas SNLE and SNPE are still overconfident as shown below:

Further complicating the matter is disparity in their validation/training loss graphs, which for SNPE looks like:

but for SNLE looks like:

Would you suggest I use a custom built algorithm using the tutorial on the sbi website? Is there anything else I can do to specifically address overconfidence?

[michaeldeistler]

I think that's tricky, there are a few things that could lead to this poor performance. Your data is quite high dimensional (500), so maybe use an embedding network. I would also suggest trying with less high-dimensional data (e.g. 10) to see if it works there. As a second step, you could then increase n and see at which values it breaks.

[paarth-dudani]

Thanks for your suggestions!

I was also wondering what the sudden spike and change of sign of both validation and training loss means in the case of SNPE as shown above as I keep getting a similar pattern for various sizes of the dataset but only for SNPE and not for SNLE/SNRE.

[paarth-dudani]

I have made some changes to the ground truth as it was misspecified.

As it now turns out, overconfidence is largely resolve and I now face the issue of underconfidence as shown below:

Do you have any suggestions as to how I can tackle this issue of underconfident/conservative estimates by the SNE algorithms?

[janfb]

Good that the misspecification is resolved.

For NPE with MDN one can see how it becomes more confident with more simulations, right (blue vs red)?
So I think it makes sense to just use more simulations, e.g., 1000, or 5000, just to make sure there is no bug.
(Note that for 1-D parameter spaces, MAF is restricted to fitting a Gaussian, which is fine here, but not in general).

[paarth-dudani]

Thanks for the reply!

Do you have any suggestions (apart from more number of simulations) for SNRE as I am consistently getting conservative estimates after using it, whereas for SNPE and SNLE I am able to get good estimates.