Diagnosing Zero Inflation in VAE Observation Models

Why Zero Counts Need a Model-Based Check

Single-cell count matrices often contain many zeros. That fact alone does not show that a negative-binomial (NB) observation model is inadequate: an NB distribution can assign substantial probability to zero, especially when its mean is small and its dispersion is large. The relevant question is whether the observed zero pattern is more frequent, or differently distributed across cells and genes, than the fitted decoder predicts.

In a count VAE, the decoder supplies a conditional mean for each cell–gene pair, while the latent representation and gene-specific dispersion help determine the count distribution. Consequently, a single overall zero percentage can be misleading. Cells with small predicted means should have more zeros than cells with large predicted means, even when the same NB observation model is appropriate throughout.

Definition: For cell $i$ and gene $g$, let $X_{ig}$ be the observed count, and let $\mu_{ig}$ be the NB component mean given the decoder output and latent state. With gene-specific dispersion $\theta_g>0$, use the parameterization $$ X_{ig}\mid z_i \sim \operatorname{NB}(\mu_{ig},\theta_g),\qquad \operatorname{Var}(X_{ig}\mid z_i)=\mu_{ig}+\frac{\mu_{ig}^2}{\theta_g}. $$ Its zero probability is $$ p_{ig}^{(0)}=P(X_{ig}=0\mid z_i)=\left(\frac{\theta_g}{\theta_g+\mu_{ig}}\right)^{\theta_g}. $$

A zero-inflated negative-binomial (ZINB) decoder adds a mixture probability $\pi_{ig}$: with probability $\pi_{ig}$ the observation is set to zero, and with probability $1-\pi_{ig}$ it is drawn from the NB component. Thus $$ P(X_{ig}=0\mid z_i)=\pi_{ig}+(1-\pi_{ig})p_{ig}^{(0)}. $$ Here $\mu_{ig}$ is the mean of the NB component, not the mean of the entire mixture. The mixture mean is $(1-\pi_{ig})\mu_{ig}$. The extra zero component is a statistical device; its presence does not by itself establish a distinct biological or technical process that generated those zeros.

Expected Zero Frequencies Under the NB Decoder

For a specified fitted NB decoder, define the zero indicator $I_{ig}=1$ when $X_{ig}=0$, and $I_{ig}=0$ otherwise. Conditional on the latent states and decoder parameters, $I_{ig}$ is Bernoulli with success probability $p_{ig}^{(0)}$. For a set of cell–gene entries, the expected number of zeros and its conditional variance are

E(Z)=\sum_{(i,g)}p_{ig}^{(0)},\qquad \operatorname{Var}(Z)=\sum_{(i,g)}p_{ig}^{(0)}\bigl(1-p_{ig}^{(0)}\bigr), \quad Z=\sum_{(i,g)} I_{ig}.

These expressions assume conditional independence of the entries being summed. A standard count decoder commonly makes that assumption conditional on the latent states and parameters. After latent uncertainty is integrated out, however, entries need not be independent. In that setting, the variance formula is best treated as a conditional reference; simulated replicates can provide a better calibration.

A useful aggregate discrepancy is the observed-minus-expected number of zeros, divided by its conditional standard deviation:

D_0= \frac{\sum_{(i,g)}\bigl(I_{ig}-p_{ig}^{(0)}\bigr)} {\sqrt{\sum_{(i,g)}p_{ig}^{(0)}\bigl(1-p_{ig}^{(0)}\bigr)}}.

A large positive value means more zeros were observed than the conditional NB model expects for the entries examined. This standardized discrepancy is not the sum of individually standardized residuals. Its numerator is the sum of the residual numerators; its denominator is the square root of the sum of their conditional variances. That distinction matters because individual standardization gives each entry a different denominator.

Worked Example: A Zero Count That Is Not Strong Evidence of Inflation

Suppose a hypothetical gene has $\mu_{ig}=2$ and $\theta_g=2$ in each of 50 cells. Assume the fitted decoder's conditional independence model for these entries. The NB zero probability for each cell is $$ p^{(0)}=\left(\frac{2}{2+2}\right)^2=0.25. $$

Suppose 16 of the 50 observed counts are zero. The model expects $50(0.25)=12.5$ zeros, and the conditional variance of the zero count is $$ 50(0.25)(0.75)=9.375. $$ Therefore, $$ D_0=\frac{16-12.5}{\sqrt{9.375}} =\frac{3.5}{3.062}=1.143. $$

This is a modest discrepancy, not compelling evidence on its own that the decoder needs a zero-inflation component. For clarity, the individual zero-indicator residual is $(I_{ig}-0.25)/\sqrt{0.25(0.75)}$: it is $1.732$ for a zero and $-0.5774$ for a nonzero. Their sum is not $D_0$. The aggregate statistic divides the total unstandardized discrepancy, $3.5$, by the square root of the total variance, $9.375$.

Why Cell-Specific Predictions Matter

A gene's decoder mean varies across cells, so its zero probability usually varies as well. The expected zero count must be computed by summing those cell-specific probabilities. Substituting an average mean into the zero-probability formula generally does not give the average zero probability: the NB zero-probability function is nonlinear.

Worked Example: Heterogeneous Means Change the Expected Number of Zeros

Consider a hypothetical set of 50 cells for one gene, all with $\theta_g=2$. The decoder gives $\mu_{ig}=0.5$ for 20 cells and $\mu_{ig}=2$ for 30 cells. For the first group, $$ p_{\mathrm{low}}^{(0)}=\left(\frac{2}{2+0.5}\right)^2=0.64, $$ so the expected number of zeros is $20(0.64)=12.8$. For the second group, $$ p_{\mathrm{high}}^{(0)}=\left(\frac{2}{2+2}\right)^2=0.25, $$ giving $30(0.25)=7.5$ expected zeros. The total expectation is $12.8+7.5=20.3$, and the conditional variance is $$ 20(0.64)(0.36)+30(0.25)(0.75)=4.608+5.625=10.233. $$

Suppose 24 counts are zero. The standardized discrepancy is $$ D_0=\frac{24-20.3}{\sqrt{10.233}} =\frac{3.7}{3.199}=1.157. $$ By contrast, replacing the cell-specific means by their average, $1.4$, gives an NB zero probability of $(2/3.4)^2=0.3460$ and an expected total of only $17.30$ zeros. That pooled calculation would misrepresent the decoder's predictions. The observed total is not especially surprising relative to the correctly summed cell-specific expectation.

Residual Patterns and Targeted Predictive Checks

A total zero count can conceal where a model is failing. Define the zero-indicator Pearson residual $$ r_{ig}^{(0)}=\frac{I_{ig}-p_{ig}^{(0)}}{\sqrt{p_{ig}^{(0)}(1-p_{ig}^{(0)})}}. $$ Inspect residuals in meaningful groups: by gene, by bins of predicted mean, by library size, or by relevant cell groups. A pattern of positive residuals concentrated among cells with moderate or high predicted means may indicate a different problem from positive residuals confined to low-mean entries. The relevant grouping should be chosen to reveal plausible systematic misfit, not repeatedly searched until a striking pattern appears.

Residuals are diagnostics, not independent observations to be treated automatically as standard normal. Their distributions are discrete and can be unstable when $p_{ig}^{(0)}$ is near zero or one. Moreover, shared latent variables and estimated parameters can induce dependence or additional uncertainty. For those reasons, a simulation-based reference is often more informative than judging one residual or one aggregate discrepancy against a normal cutoff.

Fit the NB decoder.
Obtain the decoder means and gene-specific dispersions under the NB observation model. Check that the entries compared use the same conditioning information and data-processing conventions.

Compute observed summaries.
For each planned gene or cell group, record the number or proportion of zeros and, if useful, the residual pattern against predicted mean.

Generate comparable replicates.
Draw latent states from the chosen approximate posterior or other specified latent distribution, then generate counts from the fitted NB decoder. Recompute the same zero summaries for every replicate.

Compare the observed summaries with the replicates.
Assess whether the observed zero frequency and its distribution across the planned groups fall in the range of replicated values. A persistent upper-tail discrepancy is evidence of a predictive mismatch, not a direct proof of a biological zero-generating mechanism.

The simulation scheme must match the diagnostic question. Holding latent states fixed checks the conditional observation distribution. Drawing latent states from an approximate posterior propagates some latent uncertainty, but if those states were inferred from the same counts being checked, the comparison is conditional on data that may already have adapted to the observations. A held-out set of entries or cells, when appropriate to the analysis, can reduce that reuse. Parameter uncertainty and the fitting procedure can also matter; a check based only on point estimates does not automatically include them.

Worked Example: Reading a Posterior Predictive Zero Check

Continue with the hypothetical heterogeneous-mean setting above, but suppose the observed number of zeros is 27 rather than 24. Generate 500 replicated data sets from the fitted NB decoder using the specified latent-state procedure, and calculate the number of zeros in the same 50 entries each time. Suppose the simulated central 95% range is 14 to 26 zeros, and 7 of the 500 replicates have at least 27 zeros.

The observed count, 27, lies just above the stated central range. The simulated upper-tail proportion with the common plus-one correction is $$ \frac{7+1}{500+1}=\frac{8}{501}=0.01597. $$ This indicates that a count at least as large as 27 was uncommon in these particular replicates. It is a diagnostic of potential excess zeros relative to this fitted model, not a posterior probability that the NB model is false and not a calibrated hypothesis-test p-value unless the full simulation and fitting procedure has been designed to provide that calibration. The result warrants checking whether the discrepancy recurs across relevant genes or groups and whether other model misspecification could explain it.

When Is a Zero-Inflated Decoder Warranted?

A ZINB decoder is a plausible alternative when the NB decoder repeatedly underpredicts zeros in a way that is not explained by its cell-specific means, gene-specific dispersion, or other systematic structure. In a ZINB model, the predicted zero probability is increased by $\pi_{ig}$, while the positive-count distribution is also affected through the mixture. Fit and assess the alternative using comparable predictive checks; an improved in-sample zero fit alone can be achieved simply by adding flexibility.

Approach	What it assumes or targets	Important caution
Negative binomial	All counts, including zeros, arise from one overdispersed count distribution conditional on the decoder.	Excess zeros may signal a mismatch, but the model already allows many zeros.
Zero-inflated negative binomial	Counts arise from a mixture of an extra-zero component and an NB component.	The mixture probability can be confounded with a low mean or high dispersion; its interpretation need not be mechanistic.
Hurdle model	A separate process models whether a count is zero; a truncated count distribution models positive values.	This changes the positive-count distribution too, so it is not interchangeable with zero inflation.

The mixture probability and NB parameters can be weakly distinguished when many entries have low means, because both a small mean and a high extra-zero probability can produce more zeros. A fitted $\pi_{ig}$ should therefore not be treated as a definitive estimate of a technical dropout rate. If formal nested-model testing is attempted, the null value $\pi_{ig}=0$ lies on the boundary of its allowed range; a routine likelihood-ratio reference based on a standard chi-square distribution may not apply. Predictive comparison or a simulation-calibrated procedure is safer than assuming that ordinary large-sample test conditions hold.

Common Failure Modes and Interpretation

Calling any large zero fraction “zero inflation.” The correction is to calculate expected zeros from the fitted NB probabilities, not from an intuitive benchmark or a Poisson model when NB dispersion is part of the decoder.
Using one average mean for all cells. The correction is to calculate $p_{ig}^{(0)}$ for each entry and sum those probabilities, because the zero-probability function is nonlinear.
Adding standardized residuals to obtain an aggregate statistic. The correction is to sum residual numerators and divide by the square root of the sum of variances, as in $D_0$. The sum of individually standardized residuals is a different quantity.
Treating a simulated tail proportion as a conventional p-value. The correction is to describe it as a model-based diagnostic unless the null calibration, parameter estimation, and simulation design justify a formal test interpretation.
Reading the ZINB mixture as a proven biological state. The correction is to distinguish predictive adequacy from causal interpretation. Zero inflation may absorb unmodeled heterogeneity, decoder misspecification, or data-processing effects.

Persistent discrepancies should prompt checks of the model inputs and grouping as well as the observation distribution. Incorrect library-size handling, omitted batch or cell-group structure, poorly estimated dispersion, latent states that fail to represent relevant heterogeneity, and data filtering can all create zero patterns the decoder does not reproduce. A zero-inflated decoder is not a substitute for diagnosing these issues. Conversely, if NB simulations reproduce the observed zero frequencies and their patterns across relevant groups, adding a mixture component may not be justified even when the raw data matrix is sparse.

Key takeaway: Evidence for zero inflation is evidence that the fitted NB decoder fails a targeted zero-pattern check, after accounting for its cell-specific predictions and uncertainty. It does not, by itself, identify why the zeros occurred or establish that a separate zero-generating mechanism exists.

Check Your Understanding

Answer each question using the conditional NB framework above. Brief solutions follow.

For $\mu=1$ and $\theta=2$, calculate the NB zero probability.
Why should expected zeros be summed across cell-specific probabilities rather than calculated using the average decoder mean?
A dataset has an observed zero count slightly above its NB expectation, but the count lies well within the simulated predictive range. What conclusion is justified?
Explain why the sum of individual Pearson residuals is not generally the aggregate discrepancy $D_0$.
Does a fitted ZINB mixture probability prove that the corresponding zeros arose from dropout? Explain.

Solutions. (1) $p^{(0)}=(2/3)^2=4/9=0.4444$. (2) The zero-probability formula is nonlinear in the mean, so the probability at the average mean need not equal the average of the cell-specific probabilities. (3) The observed count is compatible with the fitted model under this diagnostic; the small excess alone does not support adding zero inflation. (4) Individual Pearson residuals divide each numerator by its own standard deviation, whereas $D_0$ divides the sum of numerators by the square root of the sum of variances. (5) No. The mixture is a flexible statistical representation and can capture unmodeled heterogeneity or misspecification; it does not establish a causal dropout mechanism.