Gene-Specific Dispersion in Negative Binomial VAEs

Why Dispersion Matters in a Count Decoder

A count decoder must describe not only the expected number of molecules for each gene in each cell, but also how much counts vary around that expectation. A Poisson likelihood ties the variance to the mean. Single-cell counts can vary more than that, even after accounting for the cell-specific library size and the decoder’s expected expression. The negative binomial (NB) likelihood allows this extra variability through a dispersion parameter.

The choice of dispersion structure matters. A single shared value is economical but may be too restrictive when genes differ in variability. Giving each gene its own value is more flexible, but requires enough information to estimate many parameters. Allowing dispersion to depend on a cell covariate can represent systematic changes in variability, but it also adds modeling assumptions and can be difficult to distinguish from an incorrectly specified mean.

Let $Y_{ig}$ be the observed count for cell $i$ and gene $g$. Conditional on the cell’s latent representation and any modeled library-size term, the decoder specifies a mean $\mu_{ig}>0$. This tutorial uses the NB size parameter $\theta$, with larger $\theta$ meaning less overdispersion. This convention is equivalent to using an overdispersion coefficient $\phi=1/\theta$.

Definition: Under the mean–size parameterization, $Y\sim NB(\mu,\theta)$, for $\mu>0$ and $\theta>0$, has probability mass function $$ P(Y=y)=\frac{\Gamma(y+\theta)}{\Gamma(\theta)y!} \left(\frac{\theta}{\theta+\mu}\right)^\theta \left(\frac{\mu}{\theta+\mu}\right)^y,\qquad y=0,1,2,\ldots. $$ Its mean is $\mu$, and its variance is $\mu+\mu^2/\theta$. Equivalently, with $\phi=1/\theta$, the variance is $\mu+\phi\mu^2$.

The decoder mean $\mu_{ig}$ and the dispersion $\theta$ play distinct roles: $\mu_{ig}$ sets the conditional expected count, while $\theta$ controls the additional variance beyond the Poisson variance $\mu_{ig}$. For a fixed mean, reducing $\theta$ increases the variance. As $\theta$ grows without bound, the variance approaches the Poisson value $\mu$.

Deriving the Mean–Size Parameterization

A useful justification for this parameterization is the Gamma–Poisson mixture. Suppose the count is Poisson conditional on an unobserved rate $\Lambda$, while that rate varies across observations according to a Gamma distribution. Use a Gamma distribution with shape $\theta$ and scale $\mu/\theta$, so that $E(\Lambda)=\mu$ and $\operatorname{Var}(\Lambda)=\mu^2/\theta$. Conditional on $\Lambda=\lambda$, let $Y\mid\Lambda=\lambda\sim\operatorname{Poisson}(\lambda)$.

Integrating out $\Lambda$ gives the marginal count probability. The Gamma density has rate $\theta/\mu$, and the Poisson probability is $e^{-\lambda}\lambda^y/y!$. Multiplying these densities and integrating yields:

\begin{aligned} P(Y=y) &=\int_0^\infty \frac{e^{-\lambda}\lambda^y}{y!} \frac{(\theta/\mu)^\theta}{\Gamma(\theta)} \lambda^{\theta-1}e^{-(\theta/\mu)\lambda}\,d\lambda\\ &=\frac{(\theta/\mu)^\theta}{\Gamma(\theta)y!} \int_0^\infty \lambda^{y+\theta-1} e^{-(1+\theta/\mu)\lambda}\,d\lambda\\ &=\frac{\Gamma(y+\theta)}{\Gamma(\theta)y!} \left(\frac{\theta}{\theta+\mu}\right)^\theta \left(\frac{\mu}{\theta+\mu}\right)^y. \end{aligned}

The last equality uses the Gamma integral $\int_0^\infty x^{a-1}e^{-bx}\,dx=\Gamma(a)/b^a$, for $a,b>0$. The mixture also gives the moments directly. By iterated expectation, $E(Y)=E\{E(Y\mid\Lambda)\}=E(\Lambda)=\mu$. By the law of total variance,

\operatorname{Var}(Y) =E\{\operatorname{Var}(Y\mid\Lambda)\} +\operatorname{Var}\{E(Y\mid\Lambda)\} =E(\Lambda)+\operatorname{Var}(\Lambda) =\mu+\frac{\mu^2}{\theta}.

This derivation also gives an interpretation: NB overdispersion represents rate heterogeneity in the mixture. In a VAE, it is a conditional observation model for counts given the decoder output and latent variables; it does not establish that a literal Gamma-distributed biological rate generated each observed cell.

Worked Example: Computing a Negative Binomial Probability

Suppose a hypothetical decoder assigns a count for one gene in one cell a mean of $\mu=4$ molecules and a size of $\theta=2$. Find the probability of observing $y=3$, and calculate the variance.

Substitute into the mass function:

\begin{aligned} P(Y=3) &=\frac{\Gamma(3+2)}{\Gamma(2)3!} \left(\frac{2}{2+4}\right)^2 \left(\frac{4}{2+4}\right)^3\\ &=\frac{24}{1\cdot6}\left(\frac{1}{3}\right)^2 \left(\frac{2}{3}\right)^3 =4\cdot\frac{1}{9}\cdot\frac{8}{27} =\frac{32}{243}\approx0.1317. \end{aligned}

The variance is $4+4^2/2=12$ molecules squared. The mean is 4 molecules, so the variance is three times the mean. The probability calculation is conditional on the stated decoder mean and size; it is not a probability that the model parameters are correct.

Three Dispersion Structures

For counts $Y_{ig}$ in a VAE, the mean $\mu_{ig}$ may vary over both cells and genes. The dispersion structure specifies which observations share a size parameter. The main options differ in how many parameters they introduce and what variation they can represent.

Structure	Parameterization	Strength	Main limitation
Shared	$\theta_{ig}=\theta$	Few parameters; stable when data are limited	Assumes the same relative overdispersion for every gene and cell
Gene-specific	$\theta_{ig}=\theta_g$	Allows genes to have different conditional variability	Many parameters; estimates for sparse genes can be unstable
Covariate-dependent	$\log\theta_{ig}=w_i^\top\gamma_g$, for example	Allows dispersion to change systematically with measured covariates	Requires adequate covariate coverage and careful separation from mean effects

The logarithm in the covariate-dependent form ensures $\theta_{ig}>0$. Here $w_i$ is a vector of observed cell-level covariates, such as a technical quality measure, and $\gamma_g$ is a coefficient vector for gene $g$. A model may instead share some or all coefficients across genes; that choice determines how much pooling is imposed. Covariates that affect the expected count belong in the mean model as well as, or instead of, the dispersion model when justified. A dispersion term is not a substitute for a missing mean effect.

A shared dispersion model is not necessarily a claim that all genes have the same raw variance: their means may differ, and the variance includes both $\mu_{ig}$ and $\mu_{ig}^2/\theta$. It is the relative overdispersion at a given mean that is constrained. Gene-specific dispersion relaxes that constraint. Covariate-dependent dispersion additionally permits the degree of overdispersion to change with $w_i$.

Worked Example: Comparing Gene-Specific Dispersion Estimates

Consider hypothetical counts from five cells for each of two genes. Gene A has counts $0,1,2,3,4$, and gene B has counts $0,0,2,4,4$. Both have sample mean 2 counts per cell. Use the sample variance and the NB moment relationship to obtain rough gene-specific size estimates.

For gene A, the squared deviations from 2 sum to $4+1+0+1+4=10$. Thus $s_A^2=10/(5-1)=2.5$ counts squared. For gene B, the squared deviations sum to $4+4+0+4+4=16$, so $s_B^2=16/4=4$ counts squared.

Setting the sample variance approximately equal to $\bar y+\bar y^2/\theta$ gives $\hat\theta=\bar y^2/(s^2-\bar y)$, when $s^2>\bar y$. Therefore:

\hat\theta_A=\frac{2^2}{2.5-2}=8,\qquad \hat\theta_B=\frac{2^2}{4-2}=2.

The estimates suggest that gene B is more overdispersed: its estimated size is smaller, and its estimated overdispersion coefficient is $1/2$, compared with $1/8$ for gene A. These are moment estimates from only five cells, not precise fitted parameters. The example illustrates why a shared size could misrepresent one of the genes even when their sample means are equal.

Covariate-Dependent Dispersion in a Small Example

A covariate-dependent model is useful only if the covariate has a defensible role in the observation process and there are observations across its relevant range. To illustrate the calculation, suppose a hypothetical technical covariate $x_i$ has two levels, low and high. Counts for one gene are $0,2,4,6,8$ at the low level and $0,0,4,8,8$ at the high level. The two groups have the same sample mean, so the example isolates a difference in observed variability.

Worked Example: Translating a Covariate Effect Into Dispersion

For the low-$x$ group, the mean is 4 and the squared deviations sum to $16+4+0+4+16=40$, giving sample variance $40/4=10$. For the high-$x$ group, the squared deviations sum to $16+16+0+16+16=64$, giving sample variance $64/4=16$.

The method-of-moments sizes are:

\hat\theta_{\mathrm{low}}=\frac{4^2}{10-4}=\frac{8}{3}, \qquad \hat\theta_{\mathrm{high}}=\frac{4^2}{16-4}=\frac{4}{3}.

If $x=0$ denotes low and $x=1$ high, the two estimates correspond approximately to $\log\theta_i=\gamma_0+\gamma_1x_i$, with $\gamma_0=\log(8/3)\approx0.9808$ and $\gamma_1=\log\{(4/3)/(8/3)\}=\log(1/2)\approx-0.6931$. The negative slope means the estimated size is lower, hence overdispersion is higher, at the high covariate level. With only five cells per group, these values are illustrative rather than reliable evidence of a covariate effect.

Assumptions, Diagnostics, and Model Choice

The NB likelihood assumes nonnegative integer counts and a correctly specified conditional mean. Given the modeled latent variables, library-size term, and covariates, observations are treated as conditionally independent under the likelihood unless the model explicitly represents dependence. The dispersion must be positive, and the chosen structure must be adequate: a shared value, gene-specific values, or a covariate relationship should be supported by the data and design. As with any fitted model, the likelihood describes a statistical approximation; the Gamma–Poisson derivation does not prove the biological mechanism.

Several checks help identify problems. Compare observed variability with the fitted variance $\mu_{ig}+\mu_{ig}^2/\theta_{ig}$ across ranges of fitted means and genes. Inspect whether systematic groups of cells or genes show residual variation the model consistently misses. Check whether gene-specific estimates are extreme or poorly determined, especially for genes with few nonzero observations. For a covariate-dependent model, verify that observations cover the covariate range and that the estimated trend is not driven by a small number of cells. These checks diagnose fit; they do not by themselves establish which biological or technical mechanism caused a discrepancy.

A key failure mode is confusing mean misspecification with overdispersion. If a group of cells has a systematically different expected count but the decoder omits that group effect, the resulting extra variability can be absorbed by a smaller estimated $\theta$. Likewise, estimates of gene-specific dispersion can become noisy when the gene is rarely observed. A method-of-moments denominator $s^2-\bar y$ that is zero or negative does not yield a finite positive overdispersion estimate; it is not evidence for a negative NB size. In practice, dispersion is typically estimated within a fitted likelihood and may be stabilized or constrained, rather than inferred independently from a handful of counts.

The principal alternative is the Poisson likelihood, which has variance equal to its mean. It is simpler and is the limiting case as $\theta$ tends to infinity, but can understate variability when counts are overdispersed. Among NB models, shared dispersion trades flexibility for stability, gene-specific dispersion captures gene-to-gene differences at the cost of more parameters, and covariate-dependent dispersion represents systematic variability patterns at still greater modeling cost. Selection should account for fit, parameter uncertainty, and whether the covariate relationship is scientifically and statistically plausible; flexibility alone is not evidence of a better model.

Dispersion modeling matters because the observation likelihood contributes to the VAE objective and influences what deviations from the decoder mean are considered plausible. Too little modeled variability can make ordinary count fluctuations appear surprising; too much can make meaningful structure harder to distinguish from noise. The earlier material on modeling library size addresses variation in expected count due to cell-level scale. Dispersion addresses residual conditional variation around that expected count, so the two components should not be treated as interchangeable.

Key takeaway: State the NB convention before interpreting a fitted parameter. With the size convention used here, $\operatorname{Var}(Y)=\mu+\mu^2/\theta$: smaller $\theta$ means greater overdispersion. Choose shared, gene-specific, or covariate-dependent values according to the variability the data can support, and check that mean structure is not being misrepresented as dispersion.

Check Your Understanding

Answer each question, then compare with the solutions below.

For $\mu=5$ and $\theta=10$, what are the NB variance and the overdispersion coefficient $\phi$?
At a fixed mean, which is more overdispersed: $\theta=3$ or $\theta=12$? Explain using the variance formula.
Why can a shared dispersion model fit poorly even if it allows each cell and gene to have a different mean?
A gene has sample mean 3 and sample variance 3. What does the moment formula for $\theta$ indicate, and why should it not be interpreted as a negative size?
What additional modeling risk arises when dispersion depends on a cell covariate?

Solutions. (1) $\operatorname{Var}(Y)=5+25/10=7.5$ counts squared, and $\phi=1/10=0.1$. (2) $\theta=3$ is more overdispersed: at mean $\mu$, the added variance $\mu^2/\theta$ is larger for the smaller size. (3) Different means do not remove the shared constraint on relative overdispersion; genes can have different conditional variances beyond what their means imply. (4) The denominator $s^2-\bar y$ is zero, so the moment estimate has no finite value and corresponds to the Poisson boundary in this rough calculation. A negative denominator, if present, would likewise not imply a valid negative size; the NB size must be positive. (5) The covariate effect may be weakly identified or confounded with a missing mean effect, and extra coefficients can overfit when covariate coverage or sample size is limited.

Structure	Parameterization	Strength	Main limitation
Shared	\(\theta_{ig}=\theta\)	Few parameters; stable when data are limited	Assumes the same relative overdispersion for every gene and cell
Gene-specific	\(\theta_{ig}=\theta_g\)	Allows genes to have different conditional variability	Many parameters; estimates for sparse genes can be unstable
Covariate-dependent	\(\log\theta_{ig}=w_i^\top\gamma_g\), for example	Allows dispersion to change systematically with measured covariates	Requires adequate covariate coverage and careful separation from mean effects