Fix: grouped-label classification when use_labels_groups=True need to use logits =false. (#3805)

close #3804

Author: ori-kron-wis

CHANGELOG.md

@@ -27,6 +27,9 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
 
 #### Removed
 
+- Removed grouped-label classification legacy code in {class}`scvi.model.SCANVI`,
+    {class}`scvi.external.TOTALANVI`, and {class}`scvi.external.METHYLANVI`, {pr}`3805`.
+
 ### 1.4.2 (2026-02-26)
 
 #### Added

src/scvi/external/methylvi/_methylanvi_module.py

@@ -55,8 +55,6 @@ class METHYLANVAE(SupervisedModuleClass, METHYLVAE, BSSeqModuleMixin):
         If None, initialized to uniform probability over cell types
     labels_groups
         Label group designations
-    use_labels_groups
-        Whether to use the label groups
     linear_classifier
         If `True`, uses a single linear layer for classification instead of a
         multi-layer perceptron.
@@ -86,7 +84,6 @@ def __init__(
         dispersion: Literal["region", "region-cell"] = "region",
         y_prior=None,
         labels_groups: Sequence[int] = None,
-        use_labels_groups: bool = False,
         linear_classifier: bool = False,
         classifier_parameters: dict | None = None,
         use_batch_norm: Literal["encoder", "decoder", "none", "both"] = "both",
@@ -156,30 +153,7 @@ def __init__(
             y_prior if y_prior is not None else (1 / n_labels) * torch.ones(1, n_labels),
             requires_grad=False,
         )
-        self.use_labels_groups = use_labels_groups
         self.labels_groups = np.array(labels_groups) if labels_groups is not None else None
-        if self.use_labels_groups:
-            if labels_groups is None:
-                raise ValueError("Specify label groups")
-            unique_groups = np.unique(self.labels_groups)
-            self.n_groups = len(unique_groups)
-            if not (unique_groups == np.arange(self.n_groups)).all():
-                raise ValueError()
-            self.classifier_groups = Classifier(
-                n_latent, n_hidden, self.n_groups, n_layers, dropout_rate
-            )
-            self.groups_index = torch.nn.ParameterList(
-                [
-                    torch.nn.Parameter(
-                        torch.tensor(
-                            (self.labels_groups == i).astype(np.uint8),
-                            dtype=torch.uint8,
-                        ),
-                        requires_grad=False,
-                    )
-                    for i in range(self.n_groups)
-                ]
-            )
 
     @auto_move_data
     def classify(

src/scvi/external/totalanvi/_module.py

@@ -143,7 +143,6 @@ def __init__(
         extra_decoder_kwargs: dict | None = None,
         y_prior=None,
         labels_groups: Sequence[int] = None,
-        use_labels_groups: bool = False,
         linear_classifier: bool = False,
         classifier_parameters: dict | None = None,
     ):
@@ -234,30 +233,7 @@ def __init__(
             y_prior if y_prior is not None else (1 / n_labels) * torch.ones(1, n_labels),
             requires_grad=False,
         )
-        self.use_labels_groups = use_labels_groups
         self.labels_groups = np.array(labels_groups) if labels_groups is not None else None
-        if self.use_labels_groups:
-            if labels_groups is None:
-                raise ValueError("Specify label groups")
-            unique_groups = np.unique(self.labels_groups)
-            self.n_groups = len(unique_groups)
-            if not (unique_groups == np.arange(self.n_groups)).all():
-                raise ValueError()
-            self.classifier_groups = Classifier(
-                n_latent, n_hidden, self.n_groups, n_layers_encoder, dropout_rate_encoder
-            )
-            self.groups_index = torch.nn.ParameterList(
-                [
-                    torch.nn.Parameter(
-                        torch.tensor(
-                            (self.labels_groups == i).astype(np.uint8),
-                            dtype=torch.uint8,
-                        ),
-                        requires_grad=False,
-                    )
-                    for i in range(self.n_groups)
-                ]
-            )
 
     @auto_move_data
     def classify(

src/scvi/module/_scanvae.py

@@ -70,8 +70,6 @@ class SCANVAE(SupervisedModuleClass, VAE):
         If None, initialized to uniform probability over cell types
     labels_groups
         Label group designations
-    use_labels_groups
-        Whether to use the label groups
     linear_classifier
         If `True`, uses a single linear layer for classification instead of a
         multi-layer perceptron.
@@ -102,7 +100,6 @@ def __init__(
         use_observed_lib_size: bool = True,
         y_prior: torch.Tensor | None = None,
         labels_groups: Sequence[int] = None,
-        use_labels_groups: bool = False,
         linear_classifier: bool = False,
         classifier_parameters: dict | None = None,
         use_batch_norm: Literal["encoder", "decoder", "none", "both"] = "both",
@@ -176,30 +173,7 @@ def __init__(
             y_prior if y_prior is not None else (1 / n_labels) * torch.ones(1, n_labels),
             requires_grad=False,
         )
-        self.use_labels_groups = use_labels_groups
         self.labels_groups = np.array(labels_groups) if labels_groups is not None else None
-        if self.use_labels_groups:
-            if labels_groups is None:
-                raise ValueError("Specify label groups")
-            unique_groups = np.unique(self.labels_groups)
-            self.n_groups = len(unique_groups)
-            if not (unique_groups == np.arange(self.n_groups)).all():
-                raise ValueError()
-            self.classifier_groups = Classifier(
-                n_latent, n_hidden, self.n_groups, n_layers, dropout_rate
-            )
-            self.groups_index = torch.nn.ParameterList(
-                [
-                    torch.nn.Parameter(
-                        torch.tensor(
-                            (self.labels_groups == i).astype(np.uint8),
-                            dtype=torch.uint8,
-                        ),
-                        requires_grad=False,
-                    )
-                    for i in range(self.n_groups)
-                ]
-            )
 
     def loss(
         self,

src/scvi/module/base/_base_module.py

@@ -778,16 +778,7 @@ class SupervisedModuleClass:
 
     @auto_move_data
     def classify_helper(self, z):
-        if self.use_labels_groups:
-            w_g = self.classifier_groups(z)
-            unw_y = self.classifier(z)
-            w_y = torch.zeros_like(unw_y)
-            for i, group_index in enumerate(self.groups_index):
-                unw_y_g = unw_y[:, group_index]
-                w_y[:, group_index] = unw_y_g / (unw_y_g.sum(dim=-1, keepdim=True) + 1e-8)
-                w_y[:, group_index] *= w_g[:, [i]]
-        else:
-            w_y = self.classifier(z)
+        w_y = self.classifier(z)
         return w_y
 
     @auto_move_data

tests/model/test_scanvi.py

@@ -167,10 +167,8 @@ def test_scanvi():
     assert scanvi_pxr is not scvi_pxr
     scanvi_model.train(1)
 
-    # Test without label groups
-    scanvi_model = SCANVI.from_scvi_model(
-        m, "label_0", labels_key="labels", use_labels_groups=False
-    )
+    # Test without label groups (default)
+    scanvi_model = SCANVI.from_scvi_model(m, "label_0", labels_key="labels")
     scanvi_model.train(1)
 
     # test from_scvi_model with size_factor
@@ -283,10 +281,8 @@ def test_scanvi_with_external_indices():
     assert scanvi_pxr is not scvi_pxr
     scanvi_model.train(1)
 
-    # Test without label groups
-    scanvi_model = SCANVI.from_scvi_model(
-        m, "label_0", labels_key="labels", use_labels_groups=False
-    )
+    # Test without label groups (default)
+    scanvi_model = SCANVI.from_scvi_model(m, "label_0", labels_key="labels")
     scanvi_model.train(1)
 
     # test from_scvi_model with size_factor