Deriving the Evidence Lower Bound for a Count VAE

Why a Count VAE Uses a Lower Bound

A latent-variable model describes a cell’s observed gene-count vector $x=(x_1,\ldots,x_G)$ through a latent variable $z$. It specifies a prior $p_\theta(z)$ and a conditional count model $p_\theta(x\mid z)$. The marginal likelihood averages the conditional probability over the prior:

p_\theta(x)=\int p_\theta(x\mid z)p_\theta(z)\,dz =\int p_\theta(x,z)\,dz.

The marginal log likelihood is a natural measure of how well the model assigns probability to observed counts. But the integral over $z$ is generally difficult to calculate. A count VAE instead uses an approximate posterior $q_\phi(z\mid x)$ and maximizes a computable lower bound on $\log p_\theta(x)$. The derivation below makes precise when this bound holds and what its terms mean.

For a model with conditionally independent genes, the observation model may factor as $p_\theta(x\mid z)=\prod_{g=1}^G p_\theta(x_g\mid z)$, using count likelihoods such as Poisson or negative binomial. The resulting reconstruction term is a log probability for the observed count vector, not a squared-error score. The precise likelihood form is determined by the count model.

Deriving the Evidence Lower Bound

Definition: For an observed count vector $x$, the evidence lower bound (ELBO) for approximate posterior $q_\phi(z\mid x)$ is $$ \mathcal{L}(\theta,\phi;x) = \mathbb{E}_{q_\phi(z\mid x)} \left[\log p_\theta(x\mid z)\right] - \mathrm{KL}\!\left(q_\phi(z\mid x)\,\|\,p_\theta(z)\right). $$ When the terms are well-defined, this is also $$ \mathcal{L}(\theta,\phi;x) = \mathbb{E}_{q_\phi(z\mid x)} \left[\log p_\theta(x,z)-\log q_\phi(z\mid x)\right]. $$

Here $\mathrm{KL}(q\|p)=\mathbb{E}_q[\log(q/p)]$ is the Kullback–Leibler divergence. It is nonnegative when $q$ is absolutely continuous with respect to $p$, and is $+\infty$ if $q$ assigns positive probability to a region where $p$ is zero. Expectations and integrals are taken with respect to a common dominating measure; sums replace integrals for discrete $z$.

The derivation starts by restricting attention to the support $S=\{z:q_\phi(z\mid x)>0\}$. Since the joint density or mass is nonnegative, $\int_S p_\theta(x,z)\,dz\leq p_\theta(x)$. If the restricted integral is positive, and the logarithms and expectation below are well-defined, write the integral as an expectation under $q_\phi$. Jensen’s inequality for the concave logarithm gives

\begin{aligned} \log p_\theta(x) &\geq \log\int_S p_\theta(x,z)\,dz\\ &= \log\mathbb{E}_{q_\phi(z\mid x)} \left[\frac{p_\theta(x,z)}{q_\phi(z\mid x)}\right]\\ &\geq \mathbb{E}_{q_\phi(z\mid x)} \left[\log p_\theta(x,z)-\log q_\phi(z\mid x)\right] =\mathcal{L}(\theta,\phi;x). \end{aligned}

The positive-overlap condition matters: $p_\theta(x)>0$ alone does not guarantee that the integral over $S$ is positive. If that restricted integral is zero, then $p_\theta(x,z)=0$ for $q_\phi$-almost every $z$, and the ELBO is $-\infty$. The inequality still holds in the extended-real sense, with $\log 0=-\infty$. Thus the lower-bound claim does not require pretending that a zero restricted integral has a finite logarithm.

When $p_\theta(x)>0$, the exact posterior is $p_\theta(z\mid x)=p_\theta(x,z)/p_\theta(x)$. If $q_\phi$ is absolutely continuous with respect to this posterior and the relevant expectations are finite, substituting this expression into the KL divergence gives an exact identity:

\begin{aligned} \mathrm{KL}\!\left(q_\phi(z\mid x)\,\|\,p_\theta(z\mid x)\right) &=\mathbb{E}_{q_\phi} \left[\log q_\phi(z\mid x)-\log p_\theta(x,z)+\log p_\theta(x)\right]\\ &=\log p_\theta(x)-\mathcal{L}(\theta,\phi;x). \end{aligned}

This identity shows that the ELBO is tight when the approximate posterior equals the exact posterior, and that its gap is a KL divergence. If $q_\phi$ puts positive mass where the exact posterior is zero, this KL is infinite; correspondingly, the ELBO is $-\infty$. The identity is not a license to subtract infinities: in such cases, use the support condition and extended-real lower-bound statement instead.

Interpreting the Two Terms for Counts

The first term, $\mathbb{E}_{q_\phi}[\log p_\theta(x\mid z)]$, is the expected log probability that the decoder assigns to the observed counts under latent values drawn from $q_\phi(z\mid x)$. It is commonly called the reconstruction term. Higher values indicate that the observed count vector receives greater conditional probability under the decoder, averaged over the approximate posterior. Its scale depends on the observation model and the number of genes; it is not a percentage of counts reconstructed correctly.

The second term, $\mathrm{KL}(q_\phi(z\mid x)\|p_\theta(z))$, measures how far the cell-specific approximate posterior departs from the latent prior. It is subtracted. A large departure therefore reduces the ELBO, all else equal. The prior is not a target that must be matched exactly for every cell: the likelihood can justify a departure when it improves the reconstruction term enough. Both terms are measured in nats when natural logarithms are used.

The ELBO balances these terms, but the balance is not a universal numerical trade-off independent of modeling choices. Changing the count likelihood, gene set, or prior can change their scale and interpretation. In particular, a reconstruction term based on count probabilities is not interchangeable with a squared-error loss unless a specific model justifies that replacement.

Worked Example: A Discrete Latent Variable and the ELBO Gap

Consider one gene with observed count $x=0$ and two possible latent states, $A$ and $B$. Suppose the prior probabilities are $p(A)=0.25$, $p(B)=0.75$, and the count model gives $p(x=0\mid A)=0.20$, $p(x=0\mid B)=0.60$. Let the approximate posterior assign $q(A\mid x)=0.40$, $q(B\mid x)=0.60$. These are hypothetical probabilities under a count model.

The joint probabilities are $p(x,A)=0.25(0.20)=0.05$ and $p(x,B)=0.75(0.60)=0.45$. Thus $p(x)=0.05+0.45=0.50$, so $\log p(x)=\log(0.50)=-0.6931$. The exact posterior probabilities are $p(A\mid x)=0.05/0.50=0.10$ and $p(B\mid x)=0.45/0.50=0.90$.

Calculate the ELBO using the joint form:

\begin{aligned} \mathcal{L} &=0.40\log(0.05/0.40)+0.60\log(0.45/0.60)\\ &=0.40\log(0.125)+0.60\log(0.75)\\ &=0.40(-2.0794)+0.60(-0.2877)\\ &=-1.0044. \end{aligned}

The posterior KL divergence is $0.40\log(0.40/0.10)+0.60\log(0.60/0.90)=0.3112$. The identity checks numerically: $-0.6931-(-1.0044)=0.3113$, with the small discrepancy due to rounding. The ELBO is below the log marginal likelihood, and the gap is the divergence between the chosen approximation and the exact posterior.

Worked Example: Separating Count Reconstruction From Prior KL

For a hypothetical cell with counts $x=(2,0)$, use two equally likely latent states, so $p(A)=p(B)=0.50$, and let $q(A\mid x)=0.70$, $q(B\mid x)=0.30$. Suppose the decoder uses independent Poisson counts conditional on the latent state. Its rates for the two genes are $(2,1)$ in state $A$ and $(1,0.5)$ in state $B$.

For a Poisson count $x_g$ with rate $\lambda_g$, the log probability is $x_g\log\lambda_g-\lambda_g-\log(x_g!)$. In state $A$, the log probability for gene 1 is $2\log 2-2-\log(2!)=-1.3069$; for gene 2 it is $0\log 1-1-\log(0!)=-1.0000$. The total is $-2.3069$. In state $B$, the corresponding values are $2\log 1-1-\log(2!)=-1.6931$ and $0\log(0.5)-0.5-\log(0!)=-0.5000$, for a total of $-2.1931$.

The expected reconstruction term is $0.70(-2.3069)+0.30(-2.1931)=-2.2727$. The prior KL is $0.70\log(0.70/0.50)+0.30\log(0.30/0.50)=0.0823$. Therefore $\mathcal{L}=-2.2727-0.0823=-2.3550$ nats. This calculation checks the conditional independence assumption used to add the two gene log probabilities, and it shows that the prior penalty is separate from the expected count log likelihood.

For comparison, the exact marginal probability is $0.50(2e^{-3})+0.50(e^{-1.5}/2)=0.1056$, giving $\log p(x)\approx-2.2484$. Thus the ELBO gap is approximately $-2.2484-(-2.3550)=0.1066$ nats. The bound is valid, and this particular approximate posterior does not make it tight.

Support, Diagnostics, and Failure Modes

The lower bound depends on probability mass where the approximate posterior places mass. A decoder that assigns zero probability to the observed counts at those latent values can make the reconstruction expectation $-\infty$. Even if the full marginal likelihood is positive because other latent values explain the observation, that fact does not rescue an approximation concentrated entirely on incompatible values.

Worked Example: Positive Marginal Likelihood but a Negative-Infinite ELBO

Let the latent state be $A$ or $B$, each with prior probability $0.50$, and observe count $x=1$. Under a Poisson decoder, set the rate to $0$ in state $A$ and $1$ in state $B$. A Poisson distribution with rate zero assigns probability zero to count 1; with rate 1, it assigns probability $e^{-1}$. Hence $p(x=1)=0.50(0)+0.50e^{-1}=0.1839>0$, and $\log p(x=1)$ is finite.

Now choose $q(A\mid x)=1$. Its support contains only $A$, where $p(x,A)=0.50(0)=0$. The restricted integral over the support of $q$ is zero, and the ELBO is $\log 0=-\infty$. It remains true that $\log p(x=1)\geq-\infty$, but an ordinary finite-log Jensen calculation over this support is unavailable. This boundary case demonstrates why positive full marginal likelihood is not enough to guarantee positive overlap.

In applications, inspect the reconstruction and KL terms separately, and check that the likelihood calculation is finite for observed counts under the latent values being considered. A zero or non-finite likelihood may signal an invalid parameter value or a support mismatch. A very small KL is not, by itself, evidence of a good count model: it can accompany a weak reconstruction term. Conversely, a strong reconstruction term can coexist with a large posterior departure from the prior.

The ELBO also has limits as a model assessment. It is a lower bound, not the exact marginal log likelihood unless the approximation matches the posterior. Comparing bound values across models can be misleading when they use different likelihoods, gene sets, or data scales. And a high bound does not establish that latent coordinates have a unique biological interpretation; the bound evaluates the specified probabilistic model and approximation, not every scientific claim one might make from its latent variables.

How the ELBO Differs From Alternatives

Quantity or approach	What it measures	Key distinction
Exact marginal log likelihood	$\log\int p_\theta(x,z)\,dz$	The target evidence; often difficult to evaluate exactly.
ELBO	Expected conditional log likelihood minus prior KL	Tractable lower bound; its gap is the posterior KL when the identity’s conditions hold.
Reconstruction term alone	Expected $\log p_\theta(x\mid z)$	Omits the prior KL penalty and is not, by itself, the ELBO.
Plug-in latent likelihood	Conditional likelihood at a selected latent value	Does not average over posterior uncertainty and is not generally the marginal evidence.

The ELBO is useful because it connects a difficult marginal target to expectations under a chosen approximate posterior. Its mathematical validity does not guarantee that the approximation is accurate or that optimization finds a useful solution. The support and finiteness checks are part of interpreting the objective, not merely technical details.

Specify the count model.
Write the prior $p_\theta(z)$ and conditional likelihood $p_\theta(x\mid z)$, including any conditional independence assumptions used to factor the count likelihood.

Choose an approximate posterior.
Identify $q_\phi(z\mid x)$ and check whether it places mass only where the joint model can assign positive probability to the observed counts.

Evaluate both ELBO terms.
Compute the expected count log likelihood and subtract $\mathrm{KL}(q_\phi(z\mid x)\|p_\theta(z))$, retaining the model’s stated count units and likelihood.

Interpret the result as a bound.
When the posterior KL identity applies, the ELBO gap is that KL; otherwise use the support-aware extended-real inequality and do not assume a finite gap.

Check Your Understanding

Answer each question, then compare your reasoning with the solutions.

For a count VAE, state the two terms in the ELBO and explain why the second is subtracted.
If the exact log marginal likelihood is $-1.20$ nats and the ELBO is $-1.55$ nats, what is the posterior KL gap, assuming the identity’s conditions hold?
Can a positive marginal likelihood ensure that the ELBO is finite? Explain using support.
For counts $x=(1,0)$, a decoder gives independent Poisson rates $(1,2)$ at a fixed latent state. Calculate the conditional log likelihood at that state.

Solutions. (1) The ELBO is the expected conditional log likelihood minus $\mathrm{KL}(q_\phi(z\mid x)\|p_\theta(z))$. The KL is nonnegative and penalizes departure of the approximate posterior from the prior. (2) The gap is $-1.20-(-1.55)=0.35$ nats. (3) No. The marginal may be positive because of latent values outside the support of $q_\phi$, while the joint probability can be zero throughout the support of $q_\phi$, making the ELBO $-\infty$. (4) The log likelihood is $[1\log 1-1-\log(1!)]+[0\log 2-2-\log(0!)]=-1-2=-3.000$ nats. The calculation uses the stated conditional independence and Poisson likelihood.

Quantity or approach	What it measures	Key distinction
Exact marginal log likelihood	\(\log\int p_\theta(x,z)\,dz\)	The target evidence; often difficult to evaluate exactly.
ELBO	Expected conditional log likelihood minus prior KL	Tractable lower bound; its gap is the posterior KL when the identity’s conditions hold.
Reconstruction term alone	Expected \(\log p_\theta(x\mid z)\)	Omits the prior KL penalty and is not, by itself, the ELBO.
Plug-in latent likelihood	Conditional likelihood at a selected latent value	Does not average over posterior uncertainty and is not generally the marginal evidence.