Repository: sbi-dev/sbi

Issue:
Calibration kernels not applied in first round.
calibration kernels are not applied in the first round, but they should be applied, no?
https://github.com/mackelab/sbi/blob/09349ac01b018c880c930f3a27938ba12c8f0731/sbi/inference/snpe/snpe_base.py#L372-L380

# Fix:
- apply the calibration kernel in first round as well.

Calibration kernels not applied in first round.
calibration kernels are not applied in the first round, but they should be applied, no?
https://github.com/mackelab/sbi/blob/09349ac01b018c880c930f3a27938ba12c8f0731/sbi/inference/snpe/snpe_base.py#L372-L380

# Fix:
- apply the calibration kernel in first round as well.



Discussion:

Resolving patch:
diff --git a/sbi/inference/snpe/snpe_b.py b/sbi/inference/snpe/snpe_b.py
--- a/sbi/inference/snpe/snpe_b.py
+++ b/sbi/inference/snpe/snpe_b.py
@@ -66,7 +66,7 @@ def __init__(
             logging_level=logging_level,
         )
 
-    def _get_log_prob_proposal_posterior(
+    def _log_prob_proposal_posterior(
         self, theta: Tensor, x: Tensor, masks: Tensor
     ) -> Tensor:
         """
@@ -97,11 +97,7 @@ def _get_log_prob_proposal_posterior(
         self._assert_all_finite(log_prob_proposal, "proposal posterior eval")
 
         # Compute log prob with importance weights.
-        log_prob = (
-            self.calibration_kernel(x)
-            * torch.exp(log_prob_prior - log_prob_proposal)
-            * log_prob_posterior
-        )
+        log_prob = torch.exp(log_prob_prior - log_prob_proposal) * log_prob_posterior
 
         # XXX This evaluates the posterior on _all_ prior samples.
         if self._use_combined_loss:
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
@@ -202,7 +202,7 @@ def __call__(
         self._neural_posterior._num_trained_rounds = num_rounds
         return self._neural_posterior
 
-    def _get_log_prob_proposal_posterior(self, theta: Tensor, x: Tensor, masks: Tensor):
+    def _log_prob_proposal_posterior(self, theta: Tensor, x: Tensor, masks: Tensor) -> Tensor:
         """
         Return the log-probability used for the loss.
 
@@ -363,23 +363,16 @@ def _train(
             self._neural_posterior.neural_net.train()
             for batch in train_loader:
                 optimizer.zero_grad()
-                theta_batch, x_batch, masks = (
+                theta_batch, x_batch, masks_batch = (
                     batch[0].to(self._device),
                     batch[1].to(self._device),
                     batch[2].to(self._device),
                 )
 
-                # Just do maximum likelihood in the first round.
-                if round_ == 0:
-                    log_prob = self._neural_posterior.neural_net.log_prob(
-                        theta_batch, x_batch
-                    )
-                else:  # or call the SnpeC loss
-                    log_prob = self._get_log_prob_proposal_posterior(
-                        theta_batch, x_batch, masks
-                    )
-                loss = -torch.mean(log_prob)
-                loss.backward()
+                batch_loss = torch.mean(
+                    self._loss(round_, theta_batch, x_batch, masks_batch)
+                )
+                batch_loss.backward()
                 if clip_max_norm is not None:
                     clip_grad_norm_(
                         self._neural_posterior.neural_net.parameters(),
@@ -394,21 +387,16 @@ def _train(
             log_prob_sum = 0
             with torch.no_grad():
                 for batch in val_loader:
-                    theta_batch, x_batch, masks = (
+                    theta_batch, x_batch, masks_batch = (
                         batch[0].to(self._device),
                         batch[1].to(self._device),
                         batch[2].to(self._device),
                     )
-                    # Just do maximum likelihood in the first round.
-                    if round_ == 0:
-                        log_prob = self._neural_posterior.neural_net.log_prob(
-                            theta_batch, x_batch
-                        )
-                    else:
-                        log_prob = self._get_log_prob_proposal_posterior(
-                            theta_batch, x_batch, masks
-                        )
-                    log_prob_sum += log_prob.sum().item()
+                    # Take negative loss here to get validation log_prob.
+                    batch_log_prob = -self._loss(
+                        round_, theta_batch, x_batch, masks_batch
+                    )
+                    log_prob_sum += batch_log_prob.sum().item()
 
             self._val_log_prob = log_prob_sum / num_validation_examples
 
@@ -435,6 +423,25 @@ def _train(
         self._summary["epochs"].append(epoch)
         self._summary["best_validation_log_probs"].append(self._best_val_log_prob)
 
+    def _loss(self, round_idx: int, theta: Tensor, x: Tensor, masks: Tensor) -> Tensor:
+        """Return the right calibrated loss, depending on the round.
+
+        The loss is the negative log prob. Irrespective of the round or SNPE method
+        (A, B, or C) it is weighted with a calibration kernel.
+
+        Returns:
+            Calibration kernel-weighted negative log prob.
+        """
+
+        if round_idx == 0:
+            # Use posterior log prob (without proposal correction) for first round.
+            log_prob = self._neural_posterior.neural_net.log_prob(theta, x)
+        else:
+            # Use proposal posterior log prob tailored to snpe version (B, C).
+            log_prob = self._log_prob_proposal_posterior(theta, x, masks)
+
+        return -(self.calibration_kernel(x) * log_prob)
+
 
 class PotentialFunctionProvider:
     """
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
@@ -76,7 +76,7 @@ def __init__(
             logging_level=logging_level,
         )
 
-    def _get_log_prob_proposal_posterior(
+    def _log_prob_proposal_posterior(
         self, theta: Tensor, x: Tensor, masks: Tensor
     ) -> Tensor:
         """
@@ -138,13 +138,11 @@ def _get_log_prob_proposal_posterior(
         self._assert_all_finite(log_prob_prior, "prior eval")
 
         # Compute unnormalized proposal posterior.
-        unnormalized_log_prob_proposal_posterior = log_prob_posterior - log_prob_prior
+        unnormalized_log_prob = log_prob_posterior - log_prob_prior
 
         # Normalize proposal posterior across discrete set of atoms.
-        log_prob_proposal_posterior = self.calibration_kernel(
-            x
-        ) * unnormalized_log_prob_proposal_posterior[:, 0] - torch.logsumexp(
-            unnormalized_log_prob_proposal_posterior, dim=-1
+        log_prob_proposal_posterior = unnormalized_log_prob[:, 0] - torch.logsumexp(
+            unnormalized_log_prob, dim=-1
         )
         self._assert_all_finite(log_prob_proposal_posterior, "proposal posterior eval")
 
@@ -160,8 +158,6 @@ def _get_log_prob_proposal_posterior(
 
         return log_prob_proposal_posterior
 
-    def _get_log_prob_proposal_MoG(
-        self, theta: Tensor, x: Tensor, masks: Tensor
-    ) -> Tensor:
+    def _log_prob_proposal_MoG(self, theta: Tensor, x: Tensor, masks: Tensor) -> Tensor:
 
         raise NotImplementedError


Test patch: