another discretization fix

Author: abhinadduri

src/state/tx/models/state_transition.py

@@ -875,10 +875,12 @@ def _dose_smoothness_loss(self, batch: Dict[str, torch.Tensor], pred: torch.Tens
             return pred.new_tensor(0.0)
         return torch.stack(losses).mean()
 
-    def _compute_distribution_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    def _compute_distribution_loss(
+        self, pred: torch.Tensor, target: torch.Tensor, *, allow_discrete: bool = False
+    ) -> torch.Tensor:
         """Apply the primary distributional loss, optionally chunking feature dimensions for SamplesLoss."""
 
-        if self.discretize:
+        if self.discretize and not allow_discrete:
             raise RuntimeError("Distributional loss is not used when discretize is enabled.")
 
         if isinstance(self.loss_fn, SamplesLoss) and self.mmd_num_chunks > 1:
@@ -905,6 +907,18 @@ def _reshape_logits(self, logits: torch.Tensor, padded: bool) -> torch.Tensor:
         reshaped = self._reshape_sequence(logits, padded, self._project_out_dim)
         return reshaped.view(reshaped.size(0), reshaped.size(1), self._prediction_dim, self.num_expression_buckets)
 
+    def _expected_expression_from_logits(self, logits: torch.Tensor) -> torch.Tensor:
+        """Compute expected expression values by weighing bucket means with predicted probabilities."""
+
+        if self.bucket_boundaries is None:
+            raise RuntimeError("Bucket boundaries must be initialized to compute expectations from logits.")
+
+        means = self._bucket_means().to(logits.device)
+        probs = F.softmax(logits, dim=-1)
+        expanded_means = means.unsqueeze(0).unsqueeze(0)
+        expectation = (probs * expanded_means).sum(dim=-1)
+        return expectation
+
     def _get_target_expression(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
         if "pert_cell_counts" in batch and batch["pert_cell_counts"] is not None:
             return batch["pert_cell_counts"]
@@ -1008,11 +1022,31 @@ def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int, padded=T
             with torch.no_grad():
                 bucket_targets = self._discretize_expression(target_expression)
                 self._update_bucket_statistics(target_expression, bucket_targets)
-            per_token_losses = self._emd_per_token(logits, bucket_targets)
-            per_set_main_losses = per_token_losses.view(per_token_losses.size(0), -1).mean(dim=1)
-            main_loss = per_set_main_losses.mean()
-            self.log("train_loss", main_loss)
-            total_loss = main_loss
+
+            per_token_emd_losses = self._emd_per_token(logits, bucket_targets)
+            emd_per_set = per_token_emd_losses.view(per_token_emd_losses.size(0), -1).mean(dim=1)
+            emd_loss = emd_per_set.mean()
+
+            expected_expression = self._expected_expression_from_logits(logits)
+            per_set_main_losses = self._compute_distribution_loss(
+                expected_expression,
+                target_expression,
+                allow_discrete=True,
+            )
+            main_loss = torch.nanmean(per_set_main_losses)
+
+            if hasattr(self.loss_fn, "sinkhorn_loss") and hasattr(self.loss_fn, "energy_loss"):
+                sinkhorn_component = self.loss_fn.sinkhorn_loss(expected_expression, target_expression).nanmean()
+                energy_component = self.loss_fn.energy_loss(expected_expression, target_expression).nanmean()
+                self.log("train/sinkhorn_loss", sinkhorn_component)
+                self.log("train/energy_loss", energy_component)
+
+            self.log("decoder_loss", emd_loss)
+            self.log("train/emd_loss", emd_loss)
+            self.log("train/mmd_loss", main_loss)
+
+            total_loss = main_loss + emd_loss
+            self.log("train_loss", total_loss)
         else:
             target = batch["pert_cell_emb"]
             pred = self._reshape_sequence(pred, padded, self.output_dim)
@@ -1159,10 +1193,28 @@ def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> Non
             target_expression = self._reshape_sequence(target_expression, padded=True, feature_dim=self._prediction_dim)
             with torch.no_grad():
                 bucket_targets = self._discretize_expression(target_expression)
-            per_token_losses = self._emd_per_token(logits, bucket_targets)
-            per_set_main_losses = per_token_losses.view(per_token_losses.size(0), -1).mean(dim=1)
-            loss = per_set_main_losses.mean()
-            self.log("val_loss", loss)
+            per_token_emd_losses = self._emd_per_token(logits, bucket_targets)
+            emd_per_set = per_token_emd_losses.view(per_token_emd_losses.size(0), -1).mean(dim=1)
+            emd_loss = emd_per_set.mean()
+
+            expected_expression = self._expected_expression_from_logits(logits)
+            per_set_main_losses = self._compute_distribution_loss(
+                expected_expression,
+                target_expression,
+                allow_discrete=True,
+            )
+            main_loss = torch.nanmean(per_set_main_losses)
+
+            if hasattr(self.loss_fn, "sinkhorn_loss") and hasattr(self.loss_fn, "energy_loss"):
+                sinkhorn_component = self.loss_fn.sinkhorn_loss(expected_expression, target_expression).mean()
+                energy_component = self.loss_fn.energy_loss(expected_expression, target_expression).mean()
+                self.log("val/sinkhorn_loss", sinkhorn_component)
+                self.log("val/energy_loss", energy_component)
+
+            total_loss = main_loss + emd_loss
+            self.log("val_loss", total_loss)
+            self.log("val/decoder_loss", emd_loss)
+            self.log("val/mmd_loss", main_loss)
         else:
             pred = self._reshape_sequence(pred, padded=True, feature_dim=self.output_dim)
             target = batch["pert_cell_emb"]
@@ -1232,10 +1284,22 @@ def test_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> None:
             target_expression = self._reshape_sequence(target_expression, padded=False, feature_dim=self._prediction_dim)
             with torch.no_grad():
                 bucket_targets = self._discretize_expression(target_expression)
-            per_token_losses = self._emd_per_token(logits, bucket_targets)
-            per_set_main_losses = per_token_losses.view(per_token_losses.size(0), -1).mean(dim=1)
-            loss = per_set_main_losses.mean()
-            self.log("test_loss", loss)
+            per_token_emd_losses = self._emd_per_token(logits, bucket_targets)
+            emd_per_set = per_token_emd_losses.view(per_token_emd_losses.size(0), -1).mean(dim=1)
+            emd_loss = emd_per_set.mean()
+
+            expected_expression = self._expected_expression_from_logits(logits)
+            per_set_main_losses = self._compute_distribution_loss(
+                expected_expression,
+                target_expression,
+                allow_discrete=True,
+            )
+            main_loss = torch.nanmean(per_set_main_losses)
+
+            total_loss = main_loss + emd_loss
+            self.log("test_loss", total_loss)
+            self.log("test/decoder_loss", emd_loss)
+            self.log("test/mmd_loss", main_loss)
         else:
             target = batch["pert_cell_emb"]
             pred = self._reshape_sequence(pred, padded=False, feature_dim=self.output_dim)