fix: support validation set in DestVI training and raise clear errors for unsupported validation in RESOLVI (#3881)

Author: ori-kron-wis

.pre-commit-config.yaml

@@ -1,6 +1,7 @@
 fail_fast: false
 default_language_version:
   python: python3
+  node: 22.20.0
 default_stages:
   - pre-commit
   - pre-push

CHANGELOG.md

@@ -25,6 +25,8 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
 - Changed {class}`scvi.external.Tangram` backend to be in Pytorch, {pr}`3786`.
 - Consolidate parts of the training and data loading between {class}`~scvi.external.GIMVI`
     and {class}`scvi.external.DIAGVI`, {pr}`3830`.
+- support validation set in DestVI training and raise clear errors for unsupported validation
+    in RESOLVI, {pr}`3881`.
 
 #### Removed
 

src/scvi/external/resolvi/_model.py

@@ -229,7 +229,25 @@ def train(
             List of parameters to train if running model in Arches mode.
         **kwargs
             Other keyword args for :class:`~scvi.train.Trainer`.
+
+        Notes
+        -----
+        RESOLVI trains with Pyro SVI and maintains per-cell global parameters, so it does not
+        support a held-out validation set. ``train_size`` must be ``1.0`` and ``early_stopping``
+        is not available.
         """
+        train_size = kwargs.pop("train_size", 1.0)
+        if train_size != 1.0:
+            raise ValueError(
+                "RESOLVI does not support a validation set: it uses Pyro SVI with per-cell "
+                f"global parameters, so `train_size` must be 1.0 (got {train_size})."
+            )
+        if kwargs.pop("early_stopping", False):
+            raise ValueError(
+                "RESOLVI does not support `early_stopping` because it trains without a "
+                "validation set (`train_size` must be 1.0)."
+            )
+
         blocked = self._block_parameters.copy()
         for name, param in self.module.named_parameters():
             if not param.requires_grad:

src/scvi/external/resolvi/_utils.py

@@ -172,9 +172,13 @@ def get_normalized_expression_importance(
             indices = np.arange(adata.n_obs)
         scdl = self._make_data_loader(adata=adata, indices=indices, batch_size=batch_size)
 
-        transform_batch = _get_batch_code_from_category(
-            self.get_anndata_manager(adata, required=True), transform_batch
-        )
+        if transform_batch is not None:
+            warnings.warn(
+                "`transform_batch` is not supported by "
+                "`get_normalized_expression_importance` and will be ignored.",
+                UserWarning,
+                stacklevel=settings.warnings_stacklevel,
+            )
 
         gene_mask = slice(None) if gene_list is None else adata.var_names.isin(gene_list)
 
@@ -231,6 +235,7 @@ def get_normalized_expression_importance(
             else:
                 exprs.append(samples[1, ...].cpu())
         exprs = torch.cat(exprs, axis=1).numpy()
+        exprs = exprs[..., gene_mask]
         if return_mean:
             exprs = exprs.mean(0)
         weighting = torch.cat(weighting, axis=0).numpy()

src/scvi/model/_destvi.py

@@ -623,6 +623,26 @@ def train(
         **kwargs
             Other keyword args for :class:`~scvi.train.Trainer`.
         """
+        # A validation set is only meaningful when every spot-specific parameter is amortized.
+        # With amortization != "both", the proportions (`V`) and/or latents (`gamma`) are stored
+        # per spot and only optimized for the spots seen during training, so held-out validation
+        # spots keep their random initialization and the validation loss is not meaningful.
+        validation_requested = (
+            train_size < 1.0
+            or bool(validation_size)
+            or kwargs.get("early_stopping", False)
+            or kwargs.get("check_val_every_n_epoch", None) is not None
+        )
+        if validation_requested and self.module.amortization != "both":
+            raise ValueError(
+                "DestVI only supports a validation set (`train_size < 1.0`, `validation_size`, "
+                "`early_stopping`, or `check_val_every_n_epoch`) when `amortization='both'`. With "
+                f"`amortization={self.module.amortization!r}`, the per-spot parameters (`V` and/or"
+                "`gamma`) are not amortized and are trained only on the training spots, leaving "
+                "validation spots at their random initialization. Use `amortization='both'`, or "
+                "train on all spots with `train_size=1.0` and no early stopping / validation."
+            )
+
         update_dict = {
             "lr": lr,
             "n_epochs_kl_warmup": n_epochs_kl_warmup,
@@ -659,7 +679,10 @@ def setup_anndata(
         %(param_adata)s
         %(param_layer)s
         smoothed_layer
-            param that...
+            When provided, the model adds a spatial-smoothing regularization that encourages a
+            spot's inferred cell-type proportions to agree with those inferred from its
+            neighborhood, yielding spatially smoother deconvolution.
+            Optional; if ``None`` the regularization is disabled.
         %(param_batch_key)s
         """
         setup_method_args = cls._get_setup_method_args(**locals())

src/scvi/module/_mrdeconv.py

@@ -520,7 +520,10 @@ def loss(
             n_samples = self.n_samples_augmentation + 1
         m = x_augmented.shape[1]
 
-        if self.augmentation:
+        # Augmentation tensors (prior_sampled, ratio_augmentation, ratios_ct_augmentation) are
+        # only produced during training (see `_get_inference_input`, gated on `self.training`).
+        # During validation they are `None`, so fall back to the static (non-augmented) prior.
+        if self.augmentation and self.training:
             prior_sampled = inference_outputs["prior_sampled"].reshape(
                 n_samples, 1, x_augmented.shape[1], self.n_labels, self.n_latent
             )
@@ -565,7 +568,7 @@ def loss(
         # beta loss
         glo_neg_log_likelihood_prior = -beta_reg * Normal(mean, scale).log_prob(self.beta).sum()
         loss_augmentation = torch.tensor(0.0, device=x_augmented.device)
-        if self.augmentation:
+        if self.augmentation and self.training:
             expected_proportions = torch.cat(
                 [
                     ratios_ct_augmentation,

tests/external/resolvi/test_resolvi.py

@@ -30,6 +30,21 @@ def test_resolvi_train(adata):
     )
 
 
+def test_resolvi_train_validation_unsupported(adata):
+    # RESOLVI trains with Pyro SVI and per-cell global parameters, so it does not support a
+    # validation set. `train_size != 1.0` previously raised a cryptic TypeError (collision with
+    # the hardcoded `train_size=1.0`), and `early_stopping` has no validation set to monitor.
+    # Both must now raise a clear ValueError, while the supported settings still train.
+    RESOLVI.setup_anndata(adata)
+    model = RESOLVI(adata)
+    with pytest.raises(ValueError, match="train_size"):
+        model.train(max_epochs=2, train_size=0.8)
+    with pytest.raises(ValueError, match="early_stopping"):
+        model.train(max_epochs=2, early_stopping=True)
+    # explicit train_size=1.0 must not collide, and the default path still works
+    model.train(max_epochs=2, train_size=1.0)
+
+
 def test_resolvi_train_size_factor(adata):
     RESOLVI.setup_anndata(adata, batch_key="batch", size_factor_key="cell_area")
     model = RESOLVI(adata, size_scaling=True)
@@ -95,6 +110,30 @@ def test_resolvi_downstream(adata):
     )
 
 
+def test_resolvi_normalized_expression_gene_list(adata):
+    RESOLVI.setup_anndata(adata, size_factor_key="cell_area")
+    model = RESOLVI(adata)
+    model.train(
+        max_epochs=2,
+    )
+    gene_list = adata.var_names[:3].tolist()
+
+    # both functions must honor `gene_list` and return only the requested subset
+    expr = model.get_normalized_expression(n_samples=2, gene_list=gene_list)
+    assert list(expr.columns) == gene_list
+    assert expr.shape == (adata.n_obs, len(gene_list))
+
+    expr_imp = model.get_normalized_expression_importance(n_samples=30, gene_list=gene_list)
+    assert list(expr_imp.columns) == gene_list
+    assert expr_imp.shape == (adata.n_obs, len(gene_list))
+
+    # `transform_batch` is not supported by the importance estimator: it must be ignored
+    # (not error) and warn the user, so that `differential_expression(weights="importance")`
+    # can still call it.
+    with pytest.warns(UserWarning, match="transform_batch.*ignored"):
+        model.get_normalized_expression_importance(n_samples=30, transform_batch=0)
+
+
 def test_resolvi_downstream_size_scaling(adata):
     RESOLVI.setup_anndata(adata, size_factor_key="cell_area")
     model = RESOLVI(adata, size_scaling=True)

tests/model/test_destvi.py

@@ -61,6 +61,44 @@ def test_destvi():
             spatial_model.get_normalized_expression()
 
 
+def test_destvi_validation():
+    # DestVI must be trainable with a validation set: `early_stopping`,
+    # `check_val_every_n_epoch` and `train_size < 1.0` previously crashed in the loss because the
+    # augmentation branch was not gated on `self.training` (validation produces no augmentation
+    # tensors). See `MRDeconv.loss`.
+    n_latent = 2
+    n_labels = 5
+    dataset = synthetic_iid(n_labels=n_labels)
+    CondSCVI.setup_anndata(dataset, labels_key="labels", batch_key="batch")
+    sc_model = CondSCVI(dataset, n_latent=n_latent, n_layers=2, prior="mog", num_classes_mog=10)
+    sc_model.train(1, train_size=0.9)
+
+    DestVI.setup_anndata(dataset, layer=None)
+
+    # train_size < 1.0 with validation evaluated every epoch
+    model = DestVI.from_rna_model(dataset, sc_model, amortization="both")
+    model.train(max_epochs=2, train_size=0.9, check_val_every_n_epoch=1)
+    assert "validation_loss" in model.history
+    assert not np.isnan(model.history["validation_loss"].values[0][0])
+
+    # early stopping (requires a working validation pass)
+    model = DestVI.from_rna_model(dataset, sc_model, amortization="both")
+    model.train(max_epochs=2, train_size=0.9, early_stopping=True)
+    assert "validation_loss" in model.history
+
+    # Non-amortized spot parameters (here `V`, since amortization="latent") are only trained on the
+    # training spots, so a validation set would be evaluated on randomly-initialized parameters.
+    # Reject it. ("proportion"/"none" are not exercised here because mog priors require amortized
+    # latents, but they hit the same `amortization != 'both'` guard.)
+    model = DestVI.from_rna_model(dataset, sc_model, amortization="latent")
+    with pytest.raises(ValueError, match="amortization='both'"):
+        model.train(max_epochs=2, train_size=0.9)
+    with pytest.raises(ValueError, match="amortization='both'"):
+        model.train(max_epochs=2, early_stopping=True)
+    # train_size=1.0 (all spots trained) must still work
+    model.train(max_epochs=1)
+
+
 @pytest.mark.internet
 def test_destvi_new(save_path: str):
     CELL_TYPE_ID = "broad_cell_types"