feat: Add JointEmbeddingSCVI model (#3883)

Hi!

I curated the code from "Improving SCVI for low-count cells through
self-supervised augmentation" into a pull request to add the new model
class in scvi-tools.

I am working on a tutorial (that is a bit more visual and intuitive than
the metrics I used in the paper). I'll PR that to the tutorials repo in
a bit.

The summary below was generated by GitHub, and summarizes it nicely:

> Add JointEmbeddingSCVI, a self-supervised SCVI variant that augments
the standard ELBO with a cross-correlation objective (CCO) between the
embedding of the data and a binomially-thinned view of it, encouraging
embeddings that are robust to count dropout/noise. A VICReg-style
variance term guards against dimension collapse.
> 
> The joint-embedding loss is active only during training; evaluation
(get_elbo / get_marginal_ll / get_reconstruction_error) falls back to
the standard reweighted ELBO, and so does any minibatch with fewer than
two cells.
> 
> Contents:
> - src/scvi/module/_joint_embedding_utils.py: binomial thinning,
per-cell thinning-probability sampling, cross-correlation and variance
losses
> - src/scvi/module/_joint_embedding_vae.py: JointEmbeddingVAE module
> - src/scvi/model/_joint_embedding_scvi.py: JointEmbeddingSCVI model
> - tests/model/test_joint_embedding_scvi.py,
tests/module/test_joint_embedding_utils.py
> - docs/api (user + developer) and CHANGELOG entries
> 
>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: ori-kron-wis <ori.kronfeld@weizmann.ac.il>

Author: 3 people

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

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)

docs/api/user.md

@@ -67,6 +67,7 @@ import scvi
    external.SysVI
    external.SCVIVA
    external.DIAGVI
+   external.JointEmbeddingSCVI
    external.Tangram
 ```
 

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}
+}

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

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

docs/user_guide/models/index.md

@@ -12,6 +12,7 @@ decipher
 destvi
 diagvi
 gimvi
+jointembeddingscvi
 linearscvi
 methylanvi
 methylvi

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`
+```

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",

src/scvi/external/joint_embedding_scvi/__init__.py

@@ -0,0 +1,4 @@
+from ._model import JointEmbeddingSCVI
+from ._module import JointEmbeddingVAE
+
+__all__ = ["JointEmbeddingSCVI", "JointEmbeddingVAE"]