feat: Add JointEmbeddingSCVI model

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>

Author: vals

CHANGELOG.md

@@ -15,6 +15,7 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
     with support for batch covariates, {pr}`3620`.
 - Add support for rapids-singlecell, {pr}`3811`.
 - Add {class}`scvi.external.CytoVI` KNN imputation backend option to be cuML, {pr}`3821`.
+- Add {class}`scvi.model.JointEmbeddingSCVI`, a self-supervised SCVI variant using binomial thinning and a cross-correlation objective (CCO) for robust embeddings, {pr}`XXXX`.
 
 #### Fixed
 

docs/api/developer.md

@@ -152,6 +152,7 @@ Existing module classes with respective generative and inference procedures.
 
    module.AutoZIVAE
    module.Classifier
+   module.JointEmbeddingVAE
    module.LDVAE
    module.MRDeconv
    module.PEAKVAE

docs/api/user.md

@@ -26,6 +26,7 @@ import scvi
    model.AUTOZI
    model.CondSCVI
    model.DestVI
+   model.JointEmbeddingSCVI
    model.LinearSCVI
    model.PEAKVI
    model.SCANVI

src/scvi/model/__init__.py

@@ -8,6 +8,7 @@
 from ._autozi import AUTOZI
 from ._condscvi import CondSCVI
 from ._destvi import DestVI
+from ._joint_embedding_scvi import JointEmbeddingSCVI
 from ._linear_scvi import LinearSCVI
 from ._multivi import MULTIVI
 from ._peakvi import PEAKVI
@@ -25,6 +26,7 @@
     "PEAKVI",
     "CondSCVI",
     "DestVI",
+    "JointEmbeddingSCVI",
     "MULTIVI",
     "AmortizedLDA",
     "utils",

src/scvi/model/_joint_embedding_scvi.py

@@ -0,0 +1,157 @@
+"""Joint Embedding SCVI model."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from scvi.model._scvi import SCVI
+from scvi.module import JointEmbeddingVAE
+
+if TYPE_CHECKING:
+    from typing import Literal
+
+    from anndata import AnnData
+
+
+class JointEmbeddingSCVI(SCVI):
+    """SCVI with joint embedding loss using binomial thinning and CCO.
+
+    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.module.JointEmbeddingVAE`.
+
+    Examples
+    --------
+    >>> adata = anndata.read_h5ad(path_to_anndata)
+    >>> scvi.model.JointEmbeddingSCVI.setup_anndata(adata, batch_key="batch")
+    >>> model = scvi.model.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.module.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}."
+        )

src/scvi/module/__init__.py

@@ -6,6 +6,7 @@
 from ._amortizedlda import AmortizedLDAPyroModule
 from ._autozivae import AutoZIVAE
 from ._classifier import Classifier
+from ._joint_embedding_vae import JointEmbeddingVAE
 from ._mrdeconv import MRDeconv
 from ._multivae import MULTIVAE
 from ._peakvae import PEAKVAE
@@ -21,6 +22,7 @@
     "AutoZIVAE",
     "SCANVAE",
     "Classifier",
+    "JointEmbeddingVAE",
     "PEAKVAE",
     "VAEC",
     "MRDeconv",

src/scvi/module/_joint_embedding_utils.py

@@ -0,0 +1,159 @@
+"""Utility functions for joint embedding loss."""
+
+from __future__ import annotations
+
+import torch
+
+
+def variance_loss(z: torch.Tensor, target_std: float = 1.0) -> torch.Tensor:
+    """Compute variance regularization loss (VICReg-style).
+
+    Penalizes dimensions with standard deviation below target_std.
+    This prevents collapse where all samples map to similar embeddings.
+
+    Parameters
+    ----------
+    z
+        Embedding tensor of shape (batch_size, n_latent).
+    target_std
+        Target standard deviation for each dimension. Default is 1.0.
+
+    Returns
+    -------
+    torch.Tensor
+        Scalar loss value (mean of hinge losses across dimensions).
+    """
+    # Use the biased estimator so the loss is well-defined (0, not NaN) for
+    # single-sample minibatches, which the default data splitter can emit.
+    std = z.std(dim=0, unbiased=False)
+    return torch.relu(target_std - std).mean()
+
+
+def binomial_split(
+    x: torch.Tensor, p: float | torch.Tensor = 0.5
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split counts via binomial thinning.
+
+    Parameters
+    ----------
+    x
+        Input count tensor of shape (batch_size, n_features).
+    p
+        Probability for binomial split. Can be a scalar or a tensor of shape
+        (batch_size, 1) for per-cell probabilities. Default is 0.5.
+
+    Returns
+    -------
+    tuple
+        Two tensors (x1, x2) where x1 ~ Binomial(x, p) and x2 = x - x1.
+        Guarantees x1 + x2 = x and both are non-negative.
+
+    Notes
+    -----
+    Binomial thinning requires non-negative integer counts. Inputs are rounded
+    and clamped to ``>= 0`` so that data stored as integer-valued floats (the
+    usual single-cell convention) is handled, and non-count inputs do not raise
+    mid-training.
+    """
+    counts = x.round().clamp(min=0)
+    x1 = torch.distributions.Binomial(total_count=counts, probs=p).sample()
+    x2 = counts - x1
+    return x1, x2
+
+
+def sample_thinning_probs(x: torch.Tensor, min_library_size: float = 10.0) -> torch.Tensor:
+    """Sample per-cell thinning probabilities for realistic library size variation.
+
+    Samples target library sizes from a log-uniform distribution between
+    min_library_size and the observed library size for each cell. This
+    matches the variation typically observed in real single-cell data.
+
+    Parameters
+    ----------
+    x
+        Input count tensor of shape (batch_size, n_features).
+    min_library_size
+        Minimum target library size. Default is 10.
+
+    Returns
+    -------
+    torch.Tensor
+        Per-cell thinning probabilities of shape (batch_size, 1).
+    """
+    # Compute observed library sizes per cell
+    lib_size = x.sum(dim=1, keepdim=True)  # (batch_size, 1)
+
+    # Sample target library size from LogUniform(min_library_size, lib_size)
+    # log(L_target) ~ Uniform(log(min_lib), log(lib_size))
+    log_min = torch.log(torch.tensor(min_library_size, device=x.device, dtype=x.dtype))
+    log_lib = torch.log(lib_size.clamp(min=min_library_size + 1e-6))
+
+    # Sample uniformly in log space
+    log_target = log_min + torch.rand_like(lib_size) * (log_lib - log_min)
+    target_lib_size = torch.exp(log_target)
+
+    # Compute thinning probability: p = target_lib_size / lib_size
+    # Clamp to [0, 1] for safety (handles edge cases where lib_size <= min_library_size)
+    p = (target_lib_size / lib_size.clamp(min=1e-6)).clamp(0.0, 1.0)
+
+    return p
+
+
+def cross_correlation_loss(
+    z1: torch.Tensor,
+    z2: torch.Tensor,
+    lambda_off_diag: float = 0.01,
+    return_components: bool = False,
+) -> torch.Tensor | dict[str, torch.Tensor]:
+    """Compute cross-correlation objective (CCO) loss.
+
+    This loss has two components:
+    - Diagonal (invariance): encourages C_ii = 1, meaning corresponding dimensions
+      of the two views should be correlated. This is the self-supervision signal.
+    - Off-diagonal (redundancy reduction): encourages C_ij = 0 for i != j, meaning
+      different latent dimensions should be decorrelated. This prevents mode collapse.
+
+    Parameters
+    ----------
+    z1
+        First embedding tensor of shape (batch_size, n_latent).
+    z2
+        Second embedding tensor of shape (batch_size, n_latent).
+    lambda_off_diag
+        Weight for off-diagonal penalty. Default is 0.01.
+    return_components
+        If True, return a dict with individual components. Default is False.
+
+    Returns
+    -------
+    torch.Tensor or dict
+        If return_components is False, returns scalar loss value.
+        If True, returns dict with keys: 'total', 'invariance', 'redundancy'.
+    """
+    # Normalize (z-score across batch dimension). The biased std keeps this
+    # finite for single-sample minibatches and makes C a proper correlation matrix.
+    z1_norm = (z1 - z1.mean(0)) / (z1.std(0, unbiased=False) + 1e-8)
+    z2_norm = (z2 - z2.mean(0)) / (z2.std(0, unbiased=False) + 1e-8)
+
+    # Cross-correlation matrix: C[i,j] = corr(z1[:,i], z2[:,j])
+    batch_size = z1.shape[0]
+    C = z1_norm.T @ z2_norm / batch_size
+
+    # Diagonal loss (invariance): (1 - C_ii)^2 - self-supervision signal
+    invariance_loss = ((1 - C.diagonal()) ** 2).sum()
+
+    # Off-diagonal loss (redundancy reduction): sum(C_ij^2 for i != j) - anti-collapse
+    # Clone C before modifying diagonal to avoid in-place operation issues
+    C_off_diag = C.clone()
+    C_off_diag.fill_diagonal_(0)
+    redundancy_loss = (C_off_diag**2).sum()
+
+    total_loss = invariance_loss + lambda_off_diag * redundancy_loss
+
+    if return_components:
+        return {
+            "total": total_loss,
+            "invariance": invariance_loss,
+            "redundancy": redundancy_loss,
+        }
+    return total_loss