Single-cell and perturbations

Gaussian reparameterization, KL, and a minibatch ELBO

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

scVIscANVItotalVIMULTIVICPA

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Sample a diagonal-Gaussian latent variable using supplied noise and compute its posterior/prior densities and analytic KL divergence. Combine these with supplied reconstruction terms into a masked minibatch negative ELBO.

Function signature

diagonal_gaussian_vae_terms(
    mean: torch.Tensor,
    log_variance: torch.Tensor,
    noise: torch.Tensor,
    prior_mean: torch.Tensor,
    prior_log_variance: torch.Tensor,
    reconstruction_log_prob: torch.Tensor,
    library_kl: torch.Tensor,
    masks: dict,
    kl_weight: float = 1.0,
) -> dict[str, torch.Tensor]

Worked example

A sample and a distribution penalty are different

Take one cell and one latent dimension: mean=[[1]], log_variance=[[0]], noise=[[1]], prior_mean=[0], prior_log_variance=[0]. Both standard deviations are 1, so z=[[2]]. The analytic KL compares the distributions, while the two densities are evaluated at that sample:

\begin{aligned} \mathrm{KL}(q\parallel p)&=\tfrac12(1+1^2-1)=0.5,\\ \log q(z)&=-\tfrac12\bigl(\log(2\pi)+1\bigr)\approx-1.418939,\\ \log p(z)&=-\tfrac12\bigl(\log(2\pi)+4\bigr)\approx-2.918939. \end{aligned}

The last two quantities are returned as log_qz and log_pz. With reconstruction_log_prob=[[-2]], observed_mask=[[True]], library_kl=[0.25], cell_mask=[True], and kl_weight=2, reconstruction NLL is 2 and both the cell negative ELBO and minibatch negative ELBO are 2 + 2·0.5 + 0.25 = 3.25.

Parameters

Parameter	Shape and meaning
`mean`torch.Tensor	$[B, D]$ $[B,D]$ posterior means.
`log_variance`torch.Tensor	$[B, D]$ Posterior log variances.
`noise`torch.Tensor	$[B, D]$ Supplied fixed Gaussian variates for reparameterization.
`prior_mean`torch.Tensor	$[D]$ or $[B, D]$ Prior mean shared across cells or supplied separately for each cell.
`prior_log_variance`torch.Tensor	$[D]$ or $[B, D]$ Or $[B,D]$ prior log variances.
`reconstruction_log_prob`torch.Tensor	$[B, G]$ Supplied per-gene log likelihoods; finite ignored values allowed at missing entries.
`library_kl`torch.Tensor	$[B]$ Supplied unannealed KL contribution, commonly all zero.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`kl_weight`floatDefault `1.0`	Multiplier on the latent KL term.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Implement a diagonal-Gaussian sample, log densities, analytic KL and minibatch negative ELBO. mean, log_variance and supplied noise have shape $(B,D)$ . prior_mean and prior_log_variance have shape $(D,)$ or $(B,D)$ and broadcast over cells. All tensors are finite and exponentials of tested log-variances are representable.
Log-variance is $\operatorname{log}(\sigma^{2})$ , not a standard deviation or variance. $B$ , $D$ , $G \ge 1$ . noise is already sampled; do not generate randomness. $z = \operatorname{mean} + \operatorname{exp}(0.5\cdot \mathrm{log\_variance})\cdot \mathrm{noise}$ .
$\mathrm{log\_qz}[b] = -0.5\cdot \sum_{d}[\operatorname{log}(2\cdot \pi)+\mathrm{log\_variance}+(z-\operatorname{mean})^{2}\cdot \operatorname{exp}(-\mathrm{log\_variance})].$ log_pz uses prior_mean and prior_log_variance in the same density, evaluated at $z$ .
$\begin{gathered}\mathrm{kl}[b] = 0.5\cdot \sum_{d}[\mathrm{prior\_log\_variance}-\mathrm{log\_variance} + \operatorname{exp}(\mathrm{log\_variance}-\mathrm{prior\_log\_variance})\\{}+(\operatorname{mean}-\mathrm{prior\_mean})^{2}\cdot \operatorname{exp}(-\mathrm{prior\_log\_variance})-1]\end{gathered}.$
The log_qz quadratic can equivalently be $\mathrm{noise}^{2}$ at the supplied reparameterized sample. reconstruction_log_prob is a supplied $[B,G]$ tensor, not recomputed from $z$ here. observed_mask is bool $[B,G]$ : True includes a gene likelihood, False is missing.
$\mathrm{reconstruction\_nll}[b] = -\sum_{g} \operatorname{where}(\mathrm{observed\_mask}, \mathrm{reconstruction\_log\_prob}, 0)$ . A zero-valued log probability is not missing. library_kl is a supplied $[B]$ term, usually zero with observed libraries. The only option is kl_weight (default $1.0$ , nonnegative).
$\mathrm{cell\_negative\_elbo} = \mathrm{reconstruction\_nll} + \mathrm{kl\_weight}\cdot \mathrm{kl} + \mathrm{library\_kl}.$ Only the latent kl is annealed; library_kl is never multiplied by kl_weight. cell_mask is bool $[B]$ : True includes a cell in negative_elbo, which is the mean of cell_negative_elbo over included cells.
If none are included, return exactly zero. A valid cell with no observed genes still contributes its KL terms unless excluded by cell_mask. Other outputs, including per-cell diagnostics, are not cell-masked. Do not clamp the analytic KL or add $\epsilon$ .
Reject a negative kl_weight with ValueError; other declared input conditions are preconditions.

Why this matters

Reparameterization separates supplied randomness from differentiable mean/variance parameters; the ELBO then balances reconstruction and regularization.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def diagonal_gaussian_vae_terms(
    mean: torch.Tensor,  # (B, D)
    log_variance: torch.Tensor,  # (B, D)
    noise: torch.Tensor,  # (B, D)
    prior_mean: torch.Tensor,  # (D,) or (B, D)
    prior_log_variance: torch.Tensor,  # (D,) or (B, D)
    reconstruction_log_prob: torch.Tensor,  # (B, G)
    library_kl: torch.Tensor,  # (B,)
    masks: dict,
    kl_weight: float = 1.0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    z: (B, D).
    log_qz: (B,).
    log_pz: (B,).
    kl: (B,).
    reconstruction_nll: (B,).
    cell_negative_elbo: (B,).
    negative_elbo: Scalar tensor, shape ()."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement diagonal_gaussian_vae_terms")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.