Backport PR #3883 on branch 1.4.x (feat: Add JointEmbeddingSCVI model) (#3892)

meeseeksmachine · vals · web-flow · commit cebe7a5c8f1d · 2026-06-29T13:21:05.000+03:00
Backport PR #3883: feat: Add JointEmbeddingSCVI model Co-authored-by: Valentine Svensson <v@nxn.se>
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,8 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
     enabling memory-efficient training on large-scale datasets stored as sharded Zarr collections,
     with support for batch covariates, {pr}`3620`.
 - Add support for rapids-singlecell, {pr}`3811`.
+- Add {class}`scvi.external.JointEmbeddingSCVI`, a self-supervised SCVI variant using binomial
+    thinning and a cross-correlation objective (CCO) for robust embeddings, {pr}`3883`.
 - Add {class}`scvi.external.CYTOVI` KNN imputation backend option to be cuML, {pr}`3821`.
 
 #### Fixed
diff --git a/docs/api/developer.md b/docs/api/developer.md
@@ -197,7 +197,7 @@ Module classes in the external API with respective generative and inference proc
    external.scviva.NicheLossOutput
    external.sysvi.SysVAE
    external.diagvi.DIAGVAE
-
+   external.JointEmbeddingVAE
 ```
 
 ## Module (Base)
diff --git a/docs/api/user.md b/docs/api/user.md
@@ -67,6 +67,7 @@ import scvi
    external.SysVI
    external.SCVIVA
    external.DIAGVI
+   external.JointEmbeddingSCVI
    external.Tangram
 ```
 
diff --git a/docs/references.bib b/docs/references.bib
@@ -386,3 +386,12 @@ @article{Weinberger26
   month = apr,
   publisher={Nature Publishing Group}
 }
+
+@article{Svensson26,
+  title={Improving SCVI for low-count cells through self-supervised augmentation},
+  author={Valentine Svensson},
+  journal={bioRxiv},
+  year={2026},
+  publisher={Cold Spring Harbor Laboratory},
+  doi = {10.64898/2026.02.11.705441}
+}
diff --git a/docs/tutorials/index_scrna.md b/docs/tutorials/index_scrna.md
@@ -18,6 +18,7 @@ notebooks/scrna/AutoZI_tutorial
 notebooks/scrna/sysVI
 notebooks/scrna/decipher_tutorial
 notebooks/scrna/velovi
+notebooks/scrna/JointEmbeddingSCVI_tutorial
 notebooks/scrna/Tahoe100_mrVI
 ```
 
@@ -126,6 +127,13 @@ Use Decipher to jointly analyze samples from distinct conditions.
 Use VeloVI to estimate RNA velocity.
 ```
 
+```{customcard}
+:path: notebooks/scrna/JointEmbeddingSCVI_tutorial
+:tags: Analysis
+
+Use JointEmbeddingSCVI for improving SCVI for low-count cells through self-supervised augmentation
+```
+
 ```{customcard}
 :path: notebooks/scrna/Tahoe100_mrVI
 :tags: Analysis, Differential-comparison, Dimensionality-reduction, Removal-of-variance
diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md
@@ -54,6 +54,9 @@ scvi-tools is composed of models that can perform one or many analysis tasks. In
    * - :doc:`/user_guide/models/velovi`
      - Deep generative modeling of transcriptional dynamics for RNA velocity analysis in single cells
      - :cite:p:`GayosoWeiler23`
+   * - :doc:`/user_guide/models/jointembeddingscvi`
+     - Improving SCVI for low-count cells through self-supervised augmentation
+     - :cite:p:`Svensson26`
 ```
 
 ## ATAC-seq analysis
diff --git a/docs/user_guide/models/index.md b/docs/user_guide/models/index.md
@@ -12,6 +12,7 @@ decipher
 destvi
 diagvi
 gimvi
+jointembeddingscvi
 linearscvi
 methylanvi
 methylvi
diff --git a/docs/user_guide/models/jointembeddingscvi.md b/docs/user_guide/models/jointembeddingscvi.md
@@ -0,0 +1,20 @@
+# Joint-Embedding SCVI
+
+:::{note}
+This page is under construction.
+:::
+
+**JointEmbeddingSCVI** {cite:p}`Svensson26` (Python class {class}`~scvi.external.JointEmbeddingSCVI`) is a ...
+
+The advantages of JointEmbeddingSCVI are:
+
+-   ...
+
+The limitations of JointEmbeddingSCVI include:
+
+-   ...
+
+```{topic} Tutorials:
+
+-   {doc}`/tutorials/notebooks/scrna/JointEmbeddingSCVI_tutorial`
+```
diff --git a/src/scvi/external/__init__.py b/src/scvi/external/__init__.py
@@ -4,6 +4,7 @@
 from .decipher import Decipher
 from .diagvi import DIAGVI
 from .gimvi import GIMVI
+from .joint_embedding_scvi import JointEmbeddingSCVI, JointEmbeddingVAE
 from .methylvi import METHYLANVI, METHYLVI
 from .mrvi import MRVI
 from .poissonvi import POISSONVI
@@ -22,6 +23,8 @@
     "SCAR",
     "SOLO",
     "GIMVI",
+    "JointEmbeddingSCVI",
+    "JointEmbeddingVAE",
     "Decipher",
     "RNAStereoscope",
     "SpatialStereoscope",
diff --git a/src/scvi/external/joint_embedding_scvi/__init__.py b/src/scvi/external/joint_embedding_scvi/__init__.py
@@ -0,0 +1,4 @@
+from ._model import JointEmbeddingSCVI
+from ._module import JointEmbeddingVAE
+
+__all__ = ["JointEmbeddingSCVI", "JointEmbeddingVAE"]
diff --git a/src/scvi/external/joint_embedding_scvi/_model.py b/src/scvi/external/joint_embedding_scvi/_model.py
@@ -0,0 +1,157 @@
+"""Joint Embedding SCVI model."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from scvi.external.joint_embedding_scvi._module import JointEmbeddingVAE
+from scvi.model._scvi import SCVI
+
+if TYPE_CHECKING:
+    from typing import Literal
+
+    from anndata import AnnData
+
+
+class JointEmbeddingSCVI(SCVI):
+    """SCVI with joint embedding loss using binomial thinning and CCO :cite:p:`Svensson26`.
+
+    This model extends the standard SCVI with a cross-correlation objective (CCO)
+    loss that encourages the embedding of a thinned view to match the embedding
+    of the original data. This promotes robustness to count dropout/noise.
+
+    Thinning probabilities are dynamically sampled per cell to produce target
+    library sizes that are log-uniform between min_library_size and the observed
+    library size. This matches realistic library size variation in single-cell data.
+
+    Parameters
+    ----------
+    adata
+        AnnData object that has been registered via :meth:`~scvi.model.SCVI.setup_anndata`.
+    registry
+        Registry of the datamodule used to train the model, passed through when training
+        without an :class:`~anndata.AnnData` object (see :class:`~scvi.model.SCVI`).
+    joint_embedding_weight
+        Weight for the CCO loss. Default is 100.0.
+    lambda_off_diag
+        Off-diagonal penalty in CCO loss. Default is 0.01.
+    min_library_size
+        Minimum target library size for thinning. Default is 10.
+        Thinned library sizes are sampled log-uniformly between this
+        value and the observed library size.
+    reconstruction_weight
+        Weight for reconstruction loss. Default is 1.0.
+        Set to 0.0 for pure self-supervised training with only CCO loss.
+    variance_weight
+        Weight for variance regularization loss (VICReg-style). Default is 0.0.
+        Set to positive value (e.g., 1.0) to prevent dimension collapse
+        in self-supervised training.
+    use_joint_embedding
+        Whether to use joint embedding loss. Default is True.
+        Set to False to train as standard SCVI.
+    n_hidden
+        Number of nodes per hidden layer.
+    n_latent
+        Dimensionality of the latent space.
+    n_layers
+        Number of hidden layers used for encoder and decoder NNs.
+    dropout_rate
+        Dropout rate for neural networks.
+    dispersion
+        One of the following:
+
+        * ``'gene'`` - dispersion parameter of NB is constant per gene across cells
+        * ``'gene-batch'`` - dispersion can differ between different batches
+        * ``'gene-label'`` - dispersion can differ between different labels
+        * ``'gene-cell'`` - dispersion can differ for every gene in every cell
+    gene_likelihood
+        One of:
+
+        * ``'nb'`` - Negative binomial distribution
+        * ``'zinb'`` - Zero-inflated negative binomial distribution
+        * ``'poisson'`` - Poisson distribution
+        * ``'normal'`` - ``EXPERIMENTAL`` Normal distribution
+    use_observed_lib_size
+        If ``True``, use the observed library size for RNA as the scaling factor in the mean of the
+        conditional distribution.
+    latent_distribution
+        One of:
+
+        * ``'normal'`` - Normal distribution
+        * ``'ln'`` - Logistic normal distribution (Normal(0, I) transformed by softmax)
+    **kwargs
+        Additional keyword arguments for :class:`~scvi.external.JointEmbeddingVAE`.
+
+    Examples
+    --------
+    >>> adata = anndata.read_h5ad(path_to_anndata)
+    >>> scvi.external.JointEmbeddingSCVI.setup_anndata(adata, batch_key="batch")
+    >>> model = scvi.external.JointEmbeddingSCVI(
+    ...     adata,
+    ...     joint_embedding_weight=1.0,
+    ...     lambda_off_diag=0.01,
+    ... )
+    >>> model.train()
+    >>> latent = model.get_latent_representation()
+
+    See Also
+    --------
+    :class:`~scvi.model.SCVI`
+    :class:`~scvi.external.JointEmbeddingVAE`
+    """
+
+    _module_cls = JointEmbeddingVAE
+
+    def __init__(
+        self,
+        adata: AnnData | None = None,
+        registry: dict | None = None,
+        joint_embedding_weight: float = 100.0,
+        lambda_off_diag: float = 0.01,
+        min_library_size: float = 10.0,
+        reconstruction_weight: float = 1.0,
+        variance_weight: float = 0.0,
+        use_joint_embedding: bool = True,
+        n_hidden: int = 128,
+        n_latent: int = 10,
+        n_layers: int = 1,
+        dropout_rate: float = 0.1,
+        dispersion: Literal["gene", "gene-batch", "gene-label", "gene-cell"] = "gene",
+        gene_likelihood: Literal["zinb", "nb", "poisson", "normal"] = "zinb",
+        use_observed_lib_size: bool = True,
+        latent_distribution: Literal["normal", "ln"] = "normal",
+        **kwargs,
+    ):
+        # Pass joint embedding params through kwargs to module
+        super().__init__(
+            adata=adata,
+            registry=registry,
+            n_hidden=n_hidden,
+            n_latent=n_latent,
+            n_layers=n_layers,
+            dropout_rate=dropout_rate,
+            dispersion=dispersion,
+            gene_likelihood=gene_likelihood,
+            use_observed_lib_size=use_observed_lib_size,
+            latent_distribution=latent_distribution,
+            joint_embedding_weight=joint_embedding_weight,
+            lambda_off_diag=lambda_off_diag,
+            min_library_size=min_library_size,
+            reconstruction_weight=reconstruction_weight,
+            variance_weight=variance_weight,
+            use_joint_embedding=use_joint_embedding,
+            **kwargs,
+        )
+
+        # Update model summary string
+        self._model_summary_string = (
+            "JointEmbeddingSCVI model with the following parameters: \n"
+            f"n_hidden: {n_hidden}, n_latent: {n_latent}, n_layers: {n_layers}, "
+            f"dropout_rate: {dropout_rate}, dispersion: {dispersion}, "
+            f"gene_likelihood: {gene_likelihood}, latent_distribution: {latent_distribution}, "
+            f"joint_embedding_weight: {joint_embedding_weight}, "
+            f"lambda_off_diag: {lambda_off_diag}, min_library_size: {min_library_size}, "
+            f"reconstruction_weight: {reconstruction_weight}, "
+            f"variance_weight: {variance_weight}, "
+            f"use_joint_embedding: {use_joint_embedding}."
+        )
diff --git a/src/scvi/external/joint_embedding_scvi/_module.py b/src/scvi/external/joint_embedding_scvi/_module.py
diff --git a/src/scvi/external/joint_embedding_scvi/_utils.py b/src/scvi/external/joint_embedding_scvi/_utils.py
diff --git a/tests/external/joint_embedding_scvi/test_joint_embedding_scvi.py b/tests/external/joint_embedding_scvi/test_joint_embedding_scvi.py
diff --git a/tests/external/joint_embedding_scvi/test_joint_embedding_utils.py b/tests/external/joint_embedding_scvi/test_joint_embedding_utils.py
