Clarification on leakage

[ali-akhavan89]

I have a question on the useful documentation here:

sbi/docs/docs/faq/question_01_leakage.md

Line 27 in 9152e93

_ = inference.append_simulations(theta, x).train(force_first_round_loss=True)

In my workflow (multi-round training using SNPE-C), if I only set force_first_round_loss=True (without applying any proposal truncation), direct sampling goes through very fast and without any problems. I also get good inference, perhaps, with overly tight CIs. I wondered if I am introducing any bias by just using force_first_round_loss=True alone combined with sampling_method='direct' without using the proposal truncation that was suggested in the FAQ. Thanks a lot in advance for your suggestions.

Thanks again,
Ali

[janfb]

Hi @ali-akhavan89

If you set force_first_round_loss=True and then perform multi-round NPE without doing the truncation as described in the FAQ linked above then you are effectively performing SNPE without correction, i.e., you are learning the proposal posterior and not the actual posterior. So yes, this will be biased. For this to be unbiased you would need to either run normal atomic SNPE-C with the automatic correction, or construct the truncated proposal in each round to perform the correction via TSNPE.

In your case, setting force_first_round_loss=True is likely much faster because your prior or proposal is broad or unbounded so that the standard posterior sampling in each round is very fast.

Does this help?

Cheers,
Jan

[ali-akhavan89]

Thank you so much, Jan. This is very helpful. The prior is bounded (uniform distributions), but as you noted, proposal is not truncated, which I am working on it. I have a few follow-up questions:

1. I see that in both rejection and mcmc sampling method a potential_fn is used. Is this the log probability of the density estimator (the neural network model)? So, even if we don't have an explicit likelihood function in our model, will an approximate be created, and will MCMC (or rejection) be done using that?

2. In our workflow, we are dealing with time series data and the training is done on GPU. I'm not sure why, but sample_with='mcmc' seems to engage GPU (or somehow utilize the memory). So, using a few thousand time series would lead to the out-of-memory error. After the density estimator is trained, should we move everything to CPU to resolve this? (please see 3, if yes. FYI, we faced a similar error when using sample_with='rejection')

3. I am in the process of implementing training on GPU but MCMC on CPU, but I was getting the 'device mismatch' error. After tracing back, I've noticed that in_prior_support in class PosteriorBasedPotential(BasePotential) remains on GPU. I used in_prior_support = in_prior_support.to(self.device) and the code goes through. But it seems that it gets stuck in a loop (or maybe I am using a low number of datasets; thus, there is less information to learn, or I haven't waited long enough).

4. I have seen within_support in DirectPosterior class. If the logic is the same, i.e., to draw eligible samples, why do we bother doing MCMC?

I thank you all so much in advance for all the clarifications and support.

All the best,
Ali

[janfb]

You are welcome! Regarding your follow-up questions:

1. the potential function is a function that is often used in the context of MCMC sampling. It takes theta and returns the (unnormalized) log-posterior probability. It is used mostly in MCMC, when we only have the (learned) likelihood and the prior and want to generate samples from the posterior. However, we can also use it in the context of posterior estimation: e.g., when we face the leakage problem, sampling from the leaking posterior can be more efficient by using MCMC instead of the accept-reject scheme. Then we just define the potential function to directly return the posterior log-prob.
   2 + 3. When you run MCMC, it will use the posterior.log_prob to evaluate each sample on the MCMC-chain. Thus, it will perform lots of forward and backward passes through the underlying embedding network and flow. It is therefore probably faster to keep doing everything on the GPU. If you want to do it on CPU, you need to make sure that everything is on CPU, see my comment here: Truncated proposals for SNPE (TSNPE) implementation #1354 (comment)
2. The direct posterior just samples from the underlying posterior estimator and then does simple accept and reject on the prior support, because we don't want to have any samples outside the originally defined prior. MCMC and proper RejectionSampling are just more efficient ways of doing this, which become relevant when we are dealing with leakage.

Does this help?

Cheers,
Jan