Zero-inflated negative-binomial likelihood

Your task

Function signature

zero_inflated_negative_binomial_log_prob(
    counts: torch.Tensor,
    mean: torch.Tensor,
    inverse_dispersion: torch.Tensor,
    zero_inflation_logits: torch.Tensor,
    observed_mask: torch.Tensor,
    mean_floor: float = 0.0,
) -> dict[str, torch.Tensor]

Worked example

Two routes to an observed zero

import torch

counts = torch.tensor([[0, 1, 0, 2]], dtype=torch.int64)
mean = torch.tensor([[1.0, 1.0, 1.0, 1.0]], dtype=torch.float64, requires_grad=True)
inverse_dispersion = torch.tensor([1.0, 1.0, 1.0, 1.0], dtype=torch.float64, requires_grad=True)
zero_inflation_logits = torch.tensor([[0.0, 0.0, 0.0, 0.0]], dtype=torch.float64, requires_grad=True)
observed_mask = torch.tensor([[True, True, True, True]], dtype=torch.bool)
mean_floor = 0.0

result = zero_inflated_negative_binomial_log_prob(counts=counts, mean=mean, inverse_dispersion=inverse_dispersion, zero_inflation_logits=zero_inflation_logits, observed_mask=observed_mask, mean_floor=mean_floor)

Expected return

{
    'log_prob': torch.tensor([[-0.2876820724517809, -2.0794415416798357, -0.2876820724517809, -2.772588722239781]], dtype=torch.float64),
    'cell_nll': torch.tensor([5.4273944088231785], dtype=torch.float64),
    'structural_zero_posterior': torch.tensor([[0.6666666666666666, 0.0, 0.6666666666666666, 0.0]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement a normalized zero-inflated negative-binomial (ZINB) distribution in log space. counts is int64 [B,G]; mean and zero_inflation_logits are floating [B,G]; inverse_dispersion is [G] or [B,G]. observed_mask is bool [B,G], True for a measured count INCLUDING zero. B,G \ge 1.
• All supplied tensors are finite. Observed counts are nonnegative integers. Observed inverse dispersion \theta>0. Missing entries can contain finite sentinels and must not affect any output or local derivative. The only key is mean_floor, a finite nonnegative scalar defaulting to 0.
• Effective NB mean m=\operatorname{mean}+\mathrm{mean\_floor}; raw observed means must be nonnegative and m strictly positive. Thus with default floor the means must be positive; a positive floor permits raw mean zero. The floor is ADDED to the NB mean, not used to clamp it or added independently to logarithms.
• This preserves a normalized distribution and is an explicit teaching adaptation, not exact scvi-tools \epsilon parity. Let t=\mathrm{zero\_inflation\_logits} and \pi=\operatorname{sigmoid}(t), the probability of the STRUCTURAL-ZERO component (not the nonzero-count probability). With NB inverse dispersion \theta:
• \log p_{\mathrm{NB}}(x)=\log\Gamma(x+\theta)-\log\Gamma(\theta)-\log\Gamma(x+1)+\theta\log\frac{\theta}{\theta+m}+x\log\frac{m}{\theta+m} . If x=0, \log p(0)=\log\left[\pi+(1-\pi)e^{\log p_{\mathrm{NB}}(0)}\right]. Otherwise \mathrm{log\_prob}=\operatorname{log}(1-\pi)+\mathrm{log\_NB}(x). Detect zero by x=0 exactly.
• For observed zero also return \texttt{structural\_zero\_posterior}=\frac{\pi}{\pi+(1-\pi)p_{\mathrm{NB}}(0)}; this responsibility is zero at positive counts. log_prob and responsibility are exactly zero at missing entries. \mathrm{cell\_nll}[b]=-\sum_{g} log_prob [b,g], so entirely missing cells return zero.
• Work in log space: \operatorname{log}(\pi)=-\operatorname{softplus}(-t), \operatorname{log}(1-\pi)=-\operatorname{softplus}(t), a=\log m-\log\theta, \log p_{\mathrm{NB}}(0)=-\theta\operatorname{softplus}(a).
• Compose stable softplus and two-term logsumexp from elementary operations; do not call sigmoid, softplus, logsumexp, distributions, or a ready-made ZINB implementation. Keep the derivative correct at a=0 and t=0 as well as logits near \pm 1 000. torch.lgamma is the only exceptional numerical operation.
• Reject negative mean_floor, negative observed counts, negative raw means, nonpositive effective means, and nonpositive observed \theta with ValueError. Other shape/finiteness/integrality conditions are preconditions.

Why this matters

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 zero_inflated_negative_binomial_log_prob(
    counts: torch.Tensor,  # (B, G)
    mean: torch.Tensor,  # (B, G)
    inverse_dispersion: torch.Tensor,  # (G,) or (B, G)
    zero_inflation_logits: torch.Tensor,  # (B, G)
    observed_mask: torch.Tensor,  # (B, G)
    mean_floor: float = 0.0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    log_prob: (B, G).
    cell_nll: (B,).
    structural_zero_posterior: (B, G)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement zero_inflated_negative_binomial_log_prob")

Sources and further reading

• www.nature.com
• scverse/scvi-tools · _negative_binomial.py