feat: Add DRVI

[moinfar]

Add DRVI

This PR adds DRVI as scvi.external.DRVI.

This is a rewrite of #3699, reusing scvi-tools functionality and removing extras as much as possible.

Design:

DRVI now subclasses scvi.model.SCVI directly and only changes what is actually different:

• Decoder is replaced by a splitting + additive decoder + pooling (DecoderDRVI extends DecoderSCVI, SplitFCLayers extends FCLayers). Disentanglement is induced in the decoder — the latent is reshaped into n_split_latent independent groups, each decoded separately, then aggregated.
• Adds log-space likelihood for negative binomial.
• Adds latent stats and interpretability mixins to interpret latent dimensions (plotting functions are removed).
• Reuses everything else from SCVI unchanged: encoder, prior, covariate/batch handling (including batch_representation="embedding"), RNASeqMixin, size factors, minified mode, scArches transfer, training via a datamodule, and setup_anndata.

To enable the decoder swap without forking FCLayers, scvi.nn.FCLayers gained a few small overridable seams (_build_layer, _apply_layer, _apply_batch_norm). These are behavior-preserving for existing models (SCVI, ...) —SplitFCLayers overrides them to use per-split weights.

Likelihoods

In addition to scvi's nb, zinb and poisson, DRVI adds log-space nb:

• pnb (default of DRVI) — parametrized NB, same mean as nb but evaluated in log space via the new
  LogNegativeBinomial distribution (should be numerically more stable but equivalent to scvi's NegativeBinomial).
• normal / normal_unit_var — Gaussian with the mean modeled directly (no library/softmax), for
  continuous / log-normalized input.

Interpretability

Because the decoder is additive over splits, each latent dimension's contribution to every gene can be measured directly. The included methods — set_latent_dimension_stats, calculate_interpretability_scores, get_interpretability_scores — quantify these contributions and extract top genes per dimension.

Dependencies

Adds stacked-linear>=0.2.0 (used for the parallel per-split linear layers). That package requires no dependencies other than PyTorch. Also, the same logic can be included as a snippet instead.

[moinfar]

The code is ready for review.
Please let me know if you need any further explanation of the code's functionality.

Here are some TODOs in my list that will be issues in the next few days:

• A tutorial is prepared. A PR for that will be created.
• I am still not 100% sure on scArches compatibility. It runs, but extra tests will be added.

After the first iteration of review/changes, I will finally run the DRVI benchmark to make sure it performs the same.

[ori-kron-wis]

overall looks very good. several few points to think about

[cvduffy]

Are there any notable differences between this implementation and the drvi-py standalone one? I trained models with drvi-py and am considering whether to re-train with this implementation once it's merged/released. Thanks!

[moinfar]

Hi @cvduffy,

I'm not changing any of DRVI's logic for the scvi-tools port, which means no retraining will be needed!

I'll also do my best to add a function/snippet in drvi-py for backward compatibility, making it easy to move already-trained models over to the new scvi-tools version (but this may take a bit of time).

[ori-kron-wis]

@moinfar I merged #3880 and #3879 after some additions, added some considerations here for dropout rate and size_factor, and uncommented the scarches test function.
Perhaps you will want to take a look at it before we finally merge it.

Besides that, the tutorial is all that is missing now.

[moinfar]

@ori-kron-wis
Thanks for your efforts and for merging the required PRs.
I made a PR for a minimal tutorial: scverse/scvi-tutorials#539
All your changes look good to me. I only left one comment.

Only two things from my side:

1. I will run a complete benchmark (should be finished by the end of the week) to make sure this port indeed produces the same output as drvi-py.
2. There is a decision that for DRVI, we anneal KL for max_epoch epochs (compared to the default 400 epochs in scvi-tools). I applied this behavior in the tutorial only (when calling model.train) to keep the changes in scvi-tools minimal. Would you suggest this approach or applying it as below?

class DRVI(SCVI, GenerativeMixin, InterpretabilityMixin):
    ...
    _training_plan_cls = DRVITrainingPlan
    ...
    

class DRVITrainingPlan(TrainingPlan):
    """:class:`~scvi.train.TrainingPlan` that anneals the KL weight over the whole training run.

    DRVI's disentanglement objective benefits from spreading the KL warmup across the entire run,
    especially when the number of epochs is small (compare to the default of ``400``).
    With the default ``n_epochs_kl_warmup="auto"``, the warmup length is derived from the
    trainer's ``max_epochs`` at train time, so it matches the requested number of epochs.
    Passing an explicit ``n_epochs_kl_warmup`` (an ``int`` or ``None``) restores the standard
    behavior.
    """

    def __init__(self, *args, n_epochs_kl_warmup: int | str | None = "auto", **kwargs):
        self._auto_kl_warmup = n_epochs_kl_warmup == "auto"
        super().__init__(
            *args,
            n_epochs_kl_warmup=None if self._auto_kl_warmup else n_epochs_kl_warmup,
            **kwargs,
        )

    @property
    def kl_weight(self):
        """Scaling factor on KL divergence, warming up over ``max_epochs`` when ``"auto"``."""
        n_epochs_kl_warmup = self.n_epochs_kl_warmup
        if self._auto_kl_warmup:
            try:
                n_epochs_kl_warmup = self.trainer.max_epochs
            except RuntimeError:
                # Not yet attached to a trainer (``kl_weight`` is read once during __init__).
                # The exact warmup length is irrelevant here since ``current_epoch`` is 0; any
                # positive value yields ``min_kl_weight``.
                n_epochs_kl_warmup = 1
        klw = _compute_kl_weight(
            self.current_epoch,
            self.global_step,
            n_epochs_kl_warmup,
            self.n_steps_kl_warmup,
            self.max_kl_weight,
            self.min_kl_weight,
        )
        return torch.tensor(klw).to(self.device)

[ori-kron-wis]

@ori-kron-wis Thanks for your efforts and for merging the required PRs. I made a PR for a minimal tutorial: scverse/scvi-tutorials#539 All your changes look good to me. I only left one comment.

Only two things from my side:

1. I will run a complete benchmark (should be finished by the end of the week) to make sure this port indeed produces the same output as drvi-py.
2. There is a decision that for DRVI, we anneal KL for max_epoch epochs (compared to the default 400 epochs in scvi-tools). I applied this behavior in the tutorial only (when calling model.train) to keep the changes in scvi-tools minimal. Would you suggest this approach or applying it as below?

class DRVI(SCVI, GenerativeMixin, InterpretabilityMixin):
    ...
    _training_plan_cls = DRVITrainingPlan
    ...
    

class DRVITrainingPlan(TrainingPlan):
    """:class:`~scvi.train.TrainingPlan` that anneals the KL weight over the whole training run.

    DRVI's disentanglement objective benefits from spreading the KL warmup across the entire run,
    especially when the number of epochs is small (compare to the default of ``400``).
    With the default ``n_epochs_kl_warmup="auto"``, the warmup length is derived from the
    trainer's ``max_epochs`` at train time, so it matches the requested number of epochs.
    Passing an explicit ``n_epochs_kl_warmup`` (an ``int`` or ``None``) restores the standard
    behavior.
    """

    def __init__(self, *args, n_epochs_kl_warmup: int | str | None = "auto", **kwargs):
        self._auto_kl_warmup = n_epochs_kl_warmup == "auto"
        super().__init__(
            *args,
            n_epochs_kl_warmup=None if self._auto_kl_warmup else n_epochs_kl_warmup,
            **kwargs,
        )

    @property
    def kl_weight(self):
        """Scaling factor on KL divergence, warming up over ``max_epochs`` when ``"auto"``."""
        n_epochs_kl_warmup = self.n_epochs_kl_warmup
        if self._auto_kl_warmup:
            try:
                n_epochs_kl_warmup = self.trainer.max_epochs
            except RuntimeError:
                # Not yet attached to a trainer (``kl_weight`` is read once during __init__).
                # The exact warmup length is irrelevant here since ``current_epoch`` is 0; any
                # positive value yields ``min_kl_weight``.
                n_epochs_kl_warmup = 1
        klw = _compute_kl_weight(
            self.current_epoch,
            self.global_step,
            n_epochs_kl_warmup,
            self.n_steps_kl_warmup,
            self.max_kl_weight,
            self.min_kl_weight,
        )
        return torch.tensor(klw).to(self.device)

@moinfar Lets not change scvi-tools' n_epochs_kl_warmup right now to be reproducible (perhaps on a new branch) and add the additional DRVITrainingPlan class to DRVI

[moinfar]

Thanks for the comments. It makes sense to keep changes minimal for now and set n_epochs_kl_warmup manually.

One update: I benchmarked this port against DRVI on drvi-py. It looked good, but the results were not identical. So, in this commit (632ee48), I fixed one difference that I missed: decoder should not have dropout. There is still one minor difference in the dispersion initialization that I prefer not to change to be more similar to scvi. But if the results were significantly different, I would change that as well and benchmark again.

I am running the benchmark again and will get back to you soon.

[ori-kron-wis]

I am running the benchmark again and will get back to you soon.

Do you think we can push it today or tomorrow? As we want to have it in the new release which is ready.

[moinfar]

Yes. The branch can be merged.
Thanks for your efforts on this PR.

Details:
The benchmarking has finished. Now, there is no significant difference between scvi.external.DRVI and drvi.model.DRVI.
After this merge, I will publish a porting utility in drvi-py so users can port their already trained models to this implementation and deprecate the previous one.

[ori-kron-wis]

@moinfar What about DRVITrainingPlan then? We keep it manually for now?

[moinfar]

Yes. I think you also suggested keeping it manual. Right? I am totally fine with it.
Or did I misunderstand?

If you prefer the DRVITrainingPlan, I can create it in minutes (already a git patch).

[ori-kron-wis]

Yes. I think you also suggested keeping it manual. Right? I am totally fine with it. Or did I misunderstand?

If you prefer the DRVITrainingPlan, I can create it in minutes (already a git patch).

Sorry, yes, there was a misunderstanding - let's add DRVITrainingPlan as its a good practice to have it and it guarantees the correct usage as you described.
I thought you meant an additional change in base scvi, but I guess not.

[moinfar]

No worries. Thanks for the clarification.
I pushed the patch (22ee41a).

[moinfar]

@cvduffy
This will address your comment: theislab/DRVI#130