Repository: sbi-dev/sbi

Issue:
`force_first_round_loss` attribute in SNPE
When using SNPE, multi-round loss is only triggered when a `proposal` is passed to `append_simulations()`. Even though this is clearly documented, I think (actually, I know) that some users forget about this: they simply use `append_simulations(theta, x)` and expect multi-round loss to happen (but they actually get a cold posterior without knowing it).

We could make this more robust. After `inference.train()` has been called once, the user either **must** pass a proposal or set `train(..., force_first_round_loss=True)` manually.

What do you think @jan-matthis @janfb ?


Discussion:
We should definitely make this more robust. What do you think about raising a warning instead?
Yes, a warning would also work. Not sure which option I prefer, I'll think about it.
I think it would be good to raise an error unless a proposal is passed or a flag is set like @michaeldeistler suggested. 
I would prefer an error over a warning because the result the user gets could be actually incorrect. 

Resolving patch:
diff --git a/sbi/inference/snpe/snpe_a.py b/sbi/inference/snpe/snpe_a.py
--- a/sbi/inference/snpe/snpe_a.py
+++ b/sbi/inference/snpe/snpe_a.py
@@ -107,6 +107,7 @@ def train(
         calibration_kernel: Optional[Callable] = None,
         exclude_invalid_x: bool = True,
         resume_training: bool = False,
+        force_first_round_loss: bool = False,
         retrain_from_scratch: bool = False,
         show_train_summary: bool = False,
         dataloader_kwargs: Optional[Dict] = None,
@@ -143,6 +144,11 @@ def train(
                 cluster. If `True`, the split between train and validation set, the
                 optimizer, the number of epochs, and the best validation log-prob will
                 be restored from the last time `.train()` was called.
+            force_first_round_loss: If `True`, train with maximum likelihood,
+                i.e., potentially ignoring the correction for using a proposal
+                distribution different from the prior.
+            force_first_round_loss: If `True`, train with maximum likelihood,
+                regardless of the proposal distribution.
             retrain_from_scratch: Whether to retrain the conditional density
                 estimator for the posterior from scratch each round. Not supported for
                 SNPE-A.
@@ -431,7 +437,7 @@ def log_prob(self, inputs: Tensor, context: Tensor) -> Tensor:
             )
             return log_prob_proposal_posterior  # \hat{p} from eq (3) in [1]
 
-    def sample(self, num_samples: int, context: Tensor, batch_size: int) -> Tensor:
+    def sample(self, num_samples: int, context: Tensor, batch_size: int = 1) -> Tensor:
         context = context.to(self._device)
 
         if not self._apply_correction:
@@ -461,7 +467,7 @@ def _sample_approx_posterior_mog(
 
         # Compute the precision factors which represent the upper triangular matrix
         # of the cholesky decomposition of the prec_p.
-        prec_factors_p = torch.cholesky(prec_p, upper=True)
+        prec_factors_p = torch.linalg.cholesky(prec_p, upper=True)
 
         assert logits_p.ndim == 2
         assert m_p.ndim == 3
@@ -693,14 +699,14 @@ def _precisions_posterior(self, precisions_pp: Tensor, precisions_d: Tensor):
         # Check if precision matrices are positive definite.
         for batches in precisions_pp:
             for pprior in batches:
-                eig_pprior = torch.symeig(pprior, eigenvectors=False).eigenvalues
+                eig_pprior = torch.linalg.eigvalsh(pprior, UPLO="U")
                 if not (eig_pprior > 0).all():
                     raise AssertionError(
                         "The precision matrix of the proposal is not positive definite!"
                     )
         for batches in precisions_d:
             for d in batches:
-                eig_d = torch.symeig(d, eigenvectors=False).eigenvalues
+                eig_d = torch.linalg.eigvalsh(d, UPLO="U")
                 if not (eig_d > 0).all():
                     raise AssertionError(
                         "The precision matrix of the density estimator is not "
diff --git a/sbi/inference/snpe/snpe_base.py b/sbi/inference/snpe/snpe_base.py
--- a/sbi/inference/snpe/snpe_base.py
+++ b/sbi/inference/snpe/snpe_base.py
@@ -163,6 +163,7 @@ def train(
         calibration_kernel: Optional[Callable] = None,
         exclude_invalid_x: bool = True,
         resume_training: bool = False,
+        force_first_round_loss: bool = False,
         discard_prior_samples: bool = False,
         retrain_from_scratch: bool = False,
         show_train_summary: bool = False,
@@ -189,6 +190,9 @@ def train(
                 cluster. If `True`, the split between train and validation set, the
                 optimizer, the number of epochs, and the best validation log-prob will
                 be restored from the last time `.train()` was called.
+            force_first_round_loss: If `True`, train with maximum likelihood,
+                i.e., potentially ignoring the correction for using a proposal
+                distribution different from the prior.
             discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
                 from the prior. Training may be sped up by ignoring such less targeted
                 samples.
@@ -202,6 +206,20 @@ def train(
         Returns:
             Density estimator that approximates the distribution $p(\theta|x)$.
         """
+        if self._round == 0 and self._neural_net is not None:
+            assert force_first_round_loss, (
+                "You have already trained this neural network. After you had trained "
+                "the network, you again appended simulations with `append_simulations"
+                "(theta, x)`, but you did not provide a proposal. If the new "
+                "simulations are sampled from the prior, you can set "
+                "`.train(..., force_first_round_loss=True`). However, if the new "
+                "simulations were not sampled from the prior, you should pass the "
+                "proposal, i.e. `append_simulations(theta, x, proposal)`. If "
+                "your samples are not sampled from the prior and you do not pass a "
+                "proposal and you set `force_first_round_loss=True`, the result of "
+                "SNPE will not be the true posterior. Instead, it will be the proposal "
+                "posterior, which (usually) is more narrow than the true posterior."
+            )
 
         # Calibration kernels proposed in Lueckmann, Gonçalves et al., 2017.
         if calibration_kernel is None:
@@ -490,7 +508,6 @@ def _loss(
         Returns:
             Calibration kernel-weighted negative log prob.
         """
-
         if self._round == 0:
             # Use posterior log prob (without proposal correction) for first round.
             log_prob = self._neural_net.log_prob(theta, x)
@@ -502,6 +519,7 @@ def _loss(
     def _check_proposal(self, proposal):
         """
         Check for validity of the provided proposal distribution.
+
         If the proposal is a `NeuralPosterior`, we check if the default_x is set.
         If the proposal is **not** a `NeuralPosterior`, we warn since it is likely that
         the user simply passed the prior, but this would still trigger atomic loss.
diff --git a/sbi/inference/snpe/snpe_c.py b/sbi/inference/snpe/snpe_c.py
--- a/sbi/inference/snpe/snpe_c.py
+++ b/sbi/inference/snpe/snpe_c.py
@@ -95,6 +95,7 @@ def train(
         calibration_kernel: Optional[Callable] = None,
         exclude_invalid_x: bool = True,
         resume_training: bool = False,
+        force_first_round_loss: bool = False,
         discard_prior_samples: bool = False,
         use_combined_loss: bool = False,
         retrain_from_scratch: bool = False,
@@ -123,6 +124,9 @@ def train(
                 cluster. If `True`, the split between train and validation set, the
                 optimizer, the number of epochs, and the best validation log-prob will
                 be restored from the last time `.train()` was called.
+            force_first_round_loss: If `True`, train with maximum likelihood,
+                i.e., potentially ignoring the correction for using a proposal
+                distribution different from the prior.
             discard_prior_samples: Whether to discard samples simulated in round 1, i.e.
                 from the prior. Training may be sped up by ignoring such less targeted
                 samples.


Test patch:
diff --git a/tests/linearGaussian_snpe_test.py b/tests/linearGaussian_snpe_test.py
--- a/tests/linearGaussian_snpe_test.py
+++ b/tests/linearGaussian_snpe_test.py
@@ -192,7 +192,9 @@ def test_c2st_snpe_on_linearGaussian_different_dims():
     posterior_estimator = inference.train(
         max_num_epochs=10
     )  # Test whether we can stop and resume.
-    posterior_estimator = inference.train(resume_training=True)
+    posterior_estimator = inference.train(
+        resume_training=True, force_first_round_loss=True
+    )
     posterior = DirectPosterior(
         prior=prior, posterior_estimator=posterior_estimator
     ).set_default_x(x_o)