Posterior Predictive Checks for Count VAEs

Why Check a Fitted Count VAE?

A count VAE can fit its training objective well and still fail to reproduce important features of the observed count matrix. For example, it may capture typical expression levels while generating too few zero counts, or match gene-wise means while missing variability across cells. A posterior predictive check asks a direct question: if the fitted model were a reasonable description of these cells, what data sets would it generate?

The check is a comparison, not a new fitting method. It can reveal which features of the data are not reproduced by the model, but it does not by itself identify the cause of a discrepancy or establish that one model is better than another.

Definition: Let \(x_i=(x_{i1},\ldots,x_{iG})\) be the observed counts for cell \(i\), with \(n\) cells and \(G\) genes. A posterior predictive replicate \(x_i^{\mathrm{rep}}\) is generated from the fitted observation model after drawing latent variables from an approximate posterior informed by the observed data for that cell. A replicated count matrix has the same dimensions as the observed matrix and is generated under the same specified covariates and sampling conditions.

Write the fitted model's parameters as \(\hat{\theta}\), and let \(c_i\) denote any observed cell-level covariates used by the model. The variational encoder supplies an approximate posterior \(q_{\hat{\phi}}(z_i\mid x_i,c_i)\) for the latent state. A basic simulation step is

z_i^{\mathrm{rep}}\sim q_{\hat{\phi}}(z_i\mid x_i,c_i), \qquad x_i^{\mathrm{rep}}\sim p_{\hat{\theta}}(x_i\mid z_i^{\mathrm{rep}},c_i).

Repeat this for each cell to obtain one replicated matrix, and repeat the procedure \(M\) times to obtain \(M\) matrices. For a count VAE with conditionally independent gene counts, the second draw is made gene by gene using the fitted decoder's Poisson or negative binomial distributions and their fitted means and dispersion parameters. Any modeled library-size variable must also be drawn as specified by that model.

This procedure is an approximate, plug-in posterior predictive check: it uses fitted parameters and an approximate latent posterior. It does not generally integrate over uncertainty in all model parameters. Nor is it the same as generating entirely new cells from the prior over \(z_i\). The latter is a prior predictive check; the posterior predictive check conditions on the observed cells through their approximate latent posteriors.

Summaries to Compare

Choose summaries that target features relevant to the model's intended use. For each observed or replicated matrix, useful summaries include the library size \(L_i=\sum_{g=1}^{G}x_{ig}\), each gene's mean across cells, each gene's sample variance across cells, and each gene's zero fraction. These summaries describe different aspects of the data; matching one does not guarantee matching the others.

L_i=\sum_{g=1}^{G}x_{ig}, \qquad \bar{x}_g=\frac{1}{n}\sum_{i=1}^{n}x_{ig}, \qquad s_g^2=\frac{1}{n-1}\sum_{i=1}^{n}(x_{ig}-\bar{x}_g)^2, \qquad r_g=\frac{1}{n}\sum_{i=1}^{n}\mathbf{1}(x_{ig}=0).

Here \(s_g^2\) is the sample variance across cells for gene \(g\), and \(r_g\) is that gene's observed zero fraction. Calculate the same summaries for every replicate. Then compare the observed value with the distribution of replicated values, using a histogram, an overlay, or a predictive interval. An interval from the central 90% of replicated summaries describes the range containing the middle 90% of those simulations; it is not a 90% confidence interval for a fixed model parameter.

Worked Example: Comparing a Small Observed Matrix With One Replicate

Suppose a hypothetical data set contains four cells and three genes. Its observed count matrix is

X= \begin{pmatrix} 4&0&1\\ 2&1&0\\ 0&3&1\\ 1&0&0 \end{pmatrix}.

The row sums are \(5,3,4,1\), so the four observed library sizes are \(L=(5,3,4,1)\). The gene-wise summaries are calculated across the four rows.

Gene 1. Its counts are \(4,2,0,1\), giving \(\bar{x}_1=(4+2+0+1)/4=1.75\). Its sample variance is \(s_1^2=[(4-1.75)^2+(2-1.75)^2+(0-1.75)^2+(1-1.75)^2]/3=8.75/3=2.917\). There are no zeros, so \(r_1=0/4=0\).

Gene 2. Its counts are \(0,1,3,0\), giving \(\bar{x}_2=4/4=1\). Its sample variance is \([1+0+4+1]/3=2\), and its zero fraction is \(2/4=0.5\).

Gene 3. Its counts are \(1,0,1,0\), giving \(\bar{x}_3=2/4=0.5\). Its sample variance is \([0.25+0.25+0.25+0.25]/3=0.333\), and its zero fraction is \(2/4=0.5\).

One hypothetical posterior predictive replicate is

X^{\mathrm{rep}}= \begin{pmatrix} 2&0&1\\ 1&2&0\\ 0&1&1\\ 1&0&0 \end{pmatrix}.

Its library sizes are \(3,3,2,1\), and the replicated gene means are \(1,0.75,0.5\). The replicated zero fractions are \(0.25,0.5,0.5\). This one draw suggests a lower total count than observed, but it cannot show whether that difference is unusual: the distribution across many replicates is needed. It also happens to reproduce the zero fractions for genes 2 and 3, which does not establish that the model adequately fits their zero behavior.

Why the Replicated Distribution Is the Reference

A posterior predictive check compares a chosen statistic \(T\) calculated on the observed matrix with that statistic calculated on simulated matrices. For example, \(T(X)\) might be the sample variance of library sizes, or the zero fraction for a prespecified gene. Under the fitted model, the predictive distribution of \(T(X^{\mathrm{rep}})\) is obtained by simulation.

This simulation approximates the following integration over latent states and replicated counts, with fitted parameters held fixed:

p(x^{\mathrm{rep}}\mid X,\hat{\theta}) \approx \prod_{i=1}^{n} \int p_{\hat{\theta}}(x_i^{\mathrm{rep}}\mid z_i,c_i) q_{\hat{\phi}}(z_i\mid x_i,c_i)\,dz_i.

The approximation arises because the encoder distribution is used in place of the exact latent posterior. To see why simulation estimates this distribution, define \(T_m=T(X_m^{\mathrm{rep}})\) for replicate \(m\). Each \(T_m\) is generated by first drawing the latent states and then drawing counts from the decoder. Repeating these two draws therefore samples the integrated distribution above. A simulated quantile, such as the 95th percentile of \(T_1,\ldots,T_M\), estimates the corresponding quantile of that predictive distribution as \(M\) grows.

A commonly reported descriptive tail area is

\hat{p}_{\mathrm{pred}} = \frac{1}{M}\sum_{m=1}^{M} \mathbf{1}\{T(X_m^{\mathrm{rep}})\geq T(X)\}.

This is the fraction of simulations whose statistic is at least as large as the observed statistic. It is not a conventional null-hypothesis p-value: posterior predictive checks reuse the observed data to infer the latent states, and the resulting tail areas need not be uniformly distributed when the model is correct. Treat the tail area as a diagnostic summary, alongside plots and the size of the discrepancy, rather than as a threshold-based proof of model failure.

Worked Example: Checking Whether Library-Size Variability Is Reproduced

Continue with the hypothetical observed library sizes \(5,3,4,1\). Their mean is \(3.25\), and their sample variance is

T(X)=\frac{(5-3.25)^2+(3-3.25)^2+(4-3.25)^2+(1-3.25)^2}{3} =\frac{8.75}{3}=2.917.

Suppose 500 posterior predictive matrices are generated. In this hypothetical run, the 5th and 95th percentiles of the replicated library-size variances are 1.10 and 6.40, and 310 replicates have variance at least 2.917. The observed variance lies inside the central 90% predictive interval. The estimated upper-tail area is \(310/500=0.620\). These results provide no striking evidence that the observed library-size variance is unusually high under this check; they do not prove that the model is correct.

Now suppose instead that only 8 of the 500 variances were at least 2.917. The estimated upper-tail area would be \(8/500=0.016\), suggesting that the observed variability is in the upper tail of the fitted model's predictions. That would motivate investigation of the size and pattern of the mismatch, rather than an automatic declaration that the model has failed. The numerical quantiles and counts here are hypothetical summaries, not results from a biological study.

Library Size Requires a Model-Specific Check

Whether a replicated matrix can diagnose library-size fit depends on how the VAE handles library size. If the model generates a cell-specific library-size variable, simulate it from the fitted model along with the latent state and counts. Then compare replicated and observed distributions of \(L_i\), including their spread and, where relevant, their behavior across covariate groups.

Some models condition on the observed library size or use it as a fixed size factor in the decoder. In that case, the total count in a replicate is partly or wholly constrained by the input. A close match of library sizes is then not independent evidence that the model explains library-size variation. If the total is fixed exactly, \(L_i^{\mathrm{rep}}=L_i\) by construction, so this check cannot diagnose the distribution of library sizes. State clearly which quantity is generated and which is held fixed.

Worked Example: Detecting a Library-Size Check That Is Built In

Consider four hypothetical cells with observed library sizes \(5,3,4,1\). A model generates counts conditional on each cell's observed total. A simulation returns totals \(5,3,4,1\) in every replicate. The mean observed library size is \((5+3+4+1)/4=3.25\), and every replicated mean is also \(3.25\).

This exact match is a consequence of conditioning, not a successful test of whether the model could generate the observed variation in total counts. The observed sample variance is \(2.917\), as calculated above, but the replicate variance is always \(2.917\) too. Since the totals were fixed, their replicated variance contains no new information about model fit.

The appropriate interpretation is to omit library size as a diagnostic target for this conditional simulation, or to use a model and simulation design that generates library size if its distribution is a target of interest. Other summaries, such as gene-specific means, variances, and zero fractions conditional on the totals, can still be informative.

Gene Means, Variances, and Zero Rates

Gene-wise summaries help localize mismatches that a global count total can hide. For each gene, plot its observed mean against the distribution of replicated means, and similarly compare sample variances and zero fractions. A useful display may show the observed statistic against a predictive interval for each gene, or compare observed and replicated distributions of these statistics across genes.

Interpret each summary in relation to the others. A gene whose mean is well reproduced but whose variance is consistently above the predictive range may have more cell-to-cell heterogeneity than the model generates. An observed zero fraction above its replicated range indicates that the model tends to generate too few zeros for that gene under the specified check. Neither pattern uniquely identifies a cause: possibilities include a misspecified observation distribution, inadequate latent representation, unmodeled covariates, or a difference between the data used for fitting and the data being checked.

Worked Example: Interpreting a Gene-Level Zero-Rate Check

For Gene 2 in the hypothetical matrix, the observed zero fraction is \(2/4=0.50\). Suppose 500 replicated matrices give a median zero fraction of 0.50 and a central 90% predictive interval from 0.25 to 0.75. The observed value is within that interval, so this check does not show an unusual zero fraction for Gene 2 in this small example.

For a different hypothetical gene, suppose the observed zero fraction is \(0.80\), while its replicated central 90% interval is 0.30 to 0.60. The observed value exceeds the interval's upper endpoint. This points to a mismatch in zero frequency for that gene under the fitted posterior predictive procedure. It does not, by itself, prove that a zero-inflated likelihood is needed: the discrepancy could also reflect covariates, latent structure, or another feature the fitted model does not capture. With many genes, isolated extremes may occur by chance, so prioritize patterns across genes and checks motivated in advance.

Diagnostics, Alternatives, and Limitations

A credible check requires a faithful simulation of the fitted model. Use the same gene set, cell covariates, offsets, and preprocessing conventions as in the observed data. Confirm that generated counts are nonnegative integers, that all intended random components are sampled, and that the decoder parameters used in simulation correspond to the fitted model. If the model includes a dispersion parameter, use the fitted dispersion in the count draw rather than generating only from the conditional mean.

Look beyond a single aggregate number. Compare distributions, stratify by relevant covariates when the model includes them, and inspect whether discrepancies are systematic across genes or cells. If a check is based on a statistic with a wide predictive distribution, an observed value inside that range may still be scientifically important; conversely, with a very large number of genes, some individual summaries may fall outside nominal predictive ranges even when no single systematic failure is present.

Approach	What it answers	Important distinction
Posterior predictive check	Does the fitted model reproduce selected features of these observed cells?	Replicates use approximate latent posteriors informed by the observed cells.
Prior predictive check	What data does the model generate before conditioning on observed cells?	Latents come from the prior rather than the cell-specific approximate posterior.
Held-out predictive assessment	How well does the model predict data not used for the relevant fit?	It targets predictive performance on held-out observations, not only reproduction of training-data summaries.

Posterior predictive checks have several limitations. The fitted parameters are often treated as fixed, and uncertainty in parameter estimates is therefore not fully represented. The approximate variational posterior can also differ from the exact posterior. Because each observed cell helps determine its approximate latent state, a flexible model may reproduce some features of its training data particularly well, making certain checks less sensitive. Finally, a check can only detect discrepancies targeted by its summaries: matching library sizes, means, variances, and zero rates does not imply that every aspect of the joint distribution is correct.

A frequent mistake is to say, “The observed zero rate is inside the 90% predictive interval, so there is a 90% probability the model is correct.” The interval describes simulated values of a chosen statistic under the fitted procedure, not the probability that the model is true. The correction is to report the statistic, its predictive distribution, and the feature of the data being assessed, then describe the result as evidence that is or is not conspicuous under that particular check.

Key takeaway: Generate replicated counts according to the fitted model's actual conditional structure, then compare targeted summaries of those replicates with the observed summaries. A posterior predictive check can identify features the model does not reproduce; it cannot certify overall correctness.

Check Your Understanding

Answer each question using the definitions and procedures above. Short solutions follow the questions.

What is the difference between drawing a latent state from \(q_{\hat{\phi}}(z_i\mid x_i,c_i)\) and drawing one from the prior?
For observed library sizes \(2,4,6\), calculate the mean library size and sample variance.
A model fixes each replicated total to the corresponding observed total. Can its replicated library-size distribution assess whether the model generates realistic totals? Explain.
A gene's observed zero fraction is 0.70, while its replicated central 90% interval is 0.20 to 0.50. State a careful interpretation.

Solutions. (1) The approximate posterior draw is informed by the observed cell and is used for a posterior predictive check; the prior draw is not conditioned on that cell's observed counts and belongs to a prior predictive simulation. (2) The mean is \((2+4+6)/3=4\). The sample variance is \([(2-4)^2+(4-4)^2+(6-4)^2]/(3-1)=(4+0+4)/2=4\), in squared count units. (3) No. The matching totals are imposed by conditioning, so the simulation cannot diagnose whether the model generates the observed total-count variation. (4) The observed zero fraction is above the central 90% replicated interval, suggesting that the fitted procedure tends to generate too few zeros for this gene under this check. This is a diagnostic discrepancy, not proof of a particular model defect or a conventional hypothesis-test conclusion.