Modeling Library Size Within a Count VAE

Why Library Size Is a Modeling Choice

A cell with more recorded molecules will often have larger counts for many genes, even if its relative expression profile resembles that of a cell with fewer molecules. A count VAE therefore needs to distinguish the overall scale of a cell’s counts from how that cell’s counts are allocated across genes. The distinction is consequential: it determines whether a cell’s total count is treated as fixed, as arising around a known exposure, or as a random quantity to infer.

Let \(x_{ig}\) be the observed count for gene \(g\) in cell \(i\), and let \(z_i\) be that cell’s biological latent representation. Write \(C_i=\sum_g x_{ig}\) for the observed total. Suppose the decoder maps \(z_i\) to gene proportions \(\rho_{ig}(z_i)\), where \(\rho_{ig}\geq 0\) and \(\sum_g\rho_{ig}=1\). The library-size component supplies the scale multiplying those proportions.

These approaches are not merely three notations for the same procedure. They differ in whether the total is modeled as random, whether it contributes information to the likelihood, and whether uncertainty about the scale is represented. The introductory material on generative latent-variable models and count VAE likelihoods provides the general setting; here the focus is how library size enters the decoder.

Three Ways to Represent the Scale

A common fixed-offset model uses a known positive scale \(O_i\) and assumes conditionally independent Poisson counts:

The conditional expected count is \(O_i\rho_{ig}(z_i)\), and the expected total is \(O_i\), because the proportions sum to one. The offset is treated as given, not estimated from the count likelihood. It may be based on a library-size measure or a chosen exposure. If the offset is expressed through a log link, \(\log O_i\) is added to the gene-specific decoder log mean. The Poisson assumption also implies conditional independence across genes and equality of each conditional variance and mean.

Alternatively, one can condition directly on the observed total. Under the same gene proportions, the conditional model is

Its expected count for gene \(g\), conditional on the observed total, is \(C_i\rho_{ig}(z_i)\). The total is fixed by construction; this likelihood models how the recorded molecules are allocated, not why the cell has \(C_i\) molecules. It therefore cannot use variation in totals to learn a cell-specific scale.

A latent-scale model instead uses a positive random variable \(L_i\):

The lognormal prior is one standard choice, not a requirement. A VAE can use an encoder distribution \(q_\phi(\log L_i\mid x_i)\) to approximate the posterior. Conditional on \(L_i\) and \(z_i\), the expected count is \(L_i\rho_{ig}(z_i)\). Marginally, if \(L_i\) and \(z_i\) are independent under the distribution being averaged, the expectation is \(E[L_i]E[\rho_{ig}(z_i)]\). If they are dependent, the expectation is \(E[L_i\rho_{ig}(z_i)]\), which need not factor into those two terms.

Why the Poisson and Multinomial Views Are Related

The connection between an offset model and a conditional-total model can be derived directly. Suppose counts are conditionally independent Poisson with means \(L_i\rho_{ig}\), and the proportions sum to one. For one cell, write \(C_i=\sum_g x_{ig}\). The joint probability of the gene counts is

The first bracket is the probability mass function of \(C_i\sim\operatorname{Poisson}(L_i)\). The second is the multinomial probability of the observed allocation given that total. Thus independent Poisson counts imply a Poisson total and a multinomial allocation conditional on the total. This factorization requires independent Poisson counts and proportions summing to one; it should not be assumed automatically for other count likelihoods.

This result clarifies the modeling distinction. Conditioning on \(C_i\) removes the Poisson-total factor. A fixed Poisson offset retains that factor, although if the chosen offset equals the observed total, its value is constant with respect to the decoder proportions. A latent \(L_i\) makes the total informative about the scale and combines that information with the prior and encoder approximation.

Working With a Latent Library Size

A latent scale lets a model represent cell-to-cell differences in total count while propagating uncertainty in that scale. For example, with \(q_\phi(\log L_i\mid x_i)=N(m_i,v_i)\), the expected scale under this approximate posterior is \(\exp(m_i+v_i/2)\). This follows from the mean of a lognormal random variable; using \(\exp(m_i)\) instead gives the median, not the mean, unless \(v_i=0\).

The latent-scale contribution to a VAE objective is an expectation of the count log likelihood under the approximate posterior, together with a penalty for departing from the prior. In a model with both \(z_i\) and \(\log L_i\), the relevant terms have the form

Here \(D_{\mathrm{KL}}\) is the Kullback–Leibler divergence, which penalizes a variational distribution that differs from its prior. The exact decomposition of the divergence depends on the chosen prior and whether the latent variables are modeled independently. The important point for library size is that the count likelihood and prior both affect its inferred distribution; it is not simply replaced by the observed total.

Choosing Among the Approaches

The table assumes the same normalized gene-proportion decoder for clarity. Other parameterizations are possible, but a decoder that produces unnormalized rates must define how those rates and the library scale interact. If both can change freely in overall magnitude, the model may have difficulty distinguishing a larger library scale from uniformly larger gene rates. Constraining the proportions to sum to one makes the intended separation explicit.

A fixed offset is computationally simple, but conclusions can be sensitive to how that offset is constructed. Conditioning on the observed total avoids modeling total-count variation, but cannot explain systematic differences in totals. A latent scale is more flexible and expresses uncertainty, but introduces a prior and an additional inference problem. It does not automatically correct technical biases, and a poorly chosen prior can shrink unusual cells toward typical scales.

Diagnostics, Assumptions, and Common Pitfalls

The assumptions should be checked in relation to the scientific target. For the Poisson derivation, counts must be conditionally independent given the cell’s scale and decoder proportions, and the proportions must sum to one. The offset must be treated as fixed and appropriate for the role assigned to it. The multinomial alternative requires conditioning on a known total and interpreting the target as gene allocation. The latent model requires a defensible prior and an encoder capable of representing uncertainty in the scale.

Warning signs include observed totals that are systematically inconsistent with fixed offsets; latent posterior scales that are repeatedly pushed against the prior’s extreme tail; or cells whose large totals are explained only by implausible shifts in the gene proportions. Compare observed totals with the model-implied distribution of totals when totals are modeled. For a latent Poisson model with proportions summing to one, \(C_i\mid L_i\sim\operatorname{Poisson}(L_i)\); for an observed-total multinomial model, such a total check is not meaningful because the total is conditioned upon. Sensitivity to plausible offset definitions or prior scales is also informative.

A common mistake is to say, “The observed total is 200, so the Poisson model treats the total as fixed at 200.” That is not generally correct. Under independent Poisson counts with fixed offset \(O_i=200\), the total is random with distribution \(\operatorname{Poisson}(200)\); it has expectation 200, but can take other values. Under the multinomial model conditional on \(C_i=200\), the total is fixed. The correction is to state explicitly whether the likelihood models the total or conditions it away.