Choosing Priors for Single-Cell Latent Spaces

Why the Latent Prior Matters

In a count VAE, the prior \(p(z)\) specifies which latent representations are favored before a cell’s observed counts are considered. The encoder’s approximate posterior \(q_\phi(z\mid x)\) is then regularized toward that prior through the evidence lower bound. Consequently, a prior is not merely a neutral starting point: its support and density can affect which latent configurations are easy or costly for the model to use.

This tutorial compares three choices: a standard normal prior for an unconstrained embedding, a logistic-normal distribution for a bounded or compositional representation, and a mixture prior when several latent modes are plausible. The appropriate choice depends first on what \(z\) represents. A simplex-valued vector of proportions is not interchangeable with an unconstrained coordinate vector, even if both are informally called an embedding.

Three Prior Families and Their Assumptions

For a \(d\)-dimensional unconstrained embedding, the standard normal prior is

where \(I_d\) is the \(d\times d\) identity matrix. This assumes independent coordinates, each with mean zero and variance one, before observing a cell. Its density decreases smoothly with squared distance from the origin. The prior is easy to sample and has a simple density; it does not assert that biological cell populations are actually one spherical Gaussian distribution.

A logistic-normal prior is appropriate when the latent quantity is constrained to a simplex, such as a vector of nonnegative proportions summing to one. Let \(u\in\mathbb{R}^{K-1}\) have a multivariate normal distribution, and define

Then \(y\) has positive entries summing to one. The normal distribution is placed on log-ratios, since \(u_j=\log(y_j/y_K)\), not directly on the proportions. This construction respects the simplex constraint while allowing flexible dependence among log-ratios. It does not describe an unconstrained Euclidean embedding unless a transformation to a bounded or compositional representation is intended.

A Gaussian mixture prior allows multiple regions of high density. With \(c\) a latent component label,

where \(\pi_k\geq 0\), \(\sum_k\pi_k=1\), and each covariance matrix \(\Sigma_k\) is positive definite. The mixture can represent separated modes, but it introduces assumptions about the number of components, their weights, and their locations and shapes. A component is not automatically a known cell type: its meaning must be assessed rather than presumed.

What the Prior Penalizes

The negative log prior density, \(-\log p(z)\), measures how strongly the prior disfavors a particular location, up to constants shared across locations. For a standard normal coordinate \(z\), the location-dependent part is \(z^2/2\). Thus, values far from zero receive a larger penalty than values near zero. For a mixture, the density is a sum of component densities; it is not generally correct to assign a point only to its nearest component.

For example, consider the one-dimensional equal-weight mixture

Each component has standard deviation \(0.5\). At \(z=0\), each component mean is 2 units away, so zero is four standard deviations from either mean. The small density between the modes follows directly from the Gaussian tails.

Logistic-Normal Geometry

The logistic-normal transformation changes both the support and the interpretation of a latent variable. In a three-part composition, choose \(u_1=\log(y_1/y_3)\) and \(u_2=\log(y_2/y_3)\). If \(u=(\log 2,0)\), the inverse transformation gives

The change of variables also affects density calculations. For \(K\) proportions represented by \(K-1\) log-ratios, the absolute Jacobian determinant for the map from \(y_1,\ldots,y_{K-1}\) to \(u\) is \(1/\prod_{j=1}^{K}y_j\). Thus, if \(u\) has density \(f_U\), the density on the simplex’s first \(K-1\) coordinates is

The Jacobian is essential when evaluating a density on the proportion scale. A common logistic-normal model instead works directly with normal log-ratios in its objective; the correct density or probability calculation depends on which scale is being modeled.

Prior Choice as a Modeling Decision

A standard normal prior is a reasonable baseline when the representation is unconstrained and a single centered reference distribution is acceptable. It is simple, but it can exert substantial pressure toward a central region. One way to see the scale of that pressure is the KL divergence for a one-dimensional Gaussian approximate posterior \(q(z)=N(m,s^2)\) against a standard normal prior:

This follows by taking the expectation under \(q\) of \(\log q(z)-\log p(z)\), substituting the two Gaussian log densities, and using \(E_q[(z-m)^2]=s^2\) and \(E_q[z^2]=m^2+s^2\). The expression is nonnegative and is zero when \(m=0\) and \(s^2=1\).

A mixture is an alternative when one smooth central mode seems too restrictive and separated regions are scientifically or computationally plausible. It can reduce the density penalty near multiple component centers, but fitting or specifying component locations and weights adds complexity. Components can overlap, and exchanging their labels leaves the mixture density unchanged; the labels therefore do not by themselves identify distinct biological populations.

A logistic-normal prior is an alternative based on support rather than multimodality. It is suitable when the latent quantity is a composition, but it does not solve the problem of representing multiple cell states in an unconstrained embedding. Conversely, a Gaussian mixture on unconstrained coordinates does not automatically respect compositional constraints. These choices address different modeling needs and should not be treated as interchangeable upgrades.

Diagnostics, Failure Modes, and Limitations

Prior checks should focus on implications, not merely on whether optimization completes. Generate hypothetical \(z\) values from the prior and inspect their range, spread, and shape on the scale that will be interpreted. For a logistic-normal prior, transform draws to proportions and check that they are valid and plausible for the intended quantity. For a mixture, examine whether the weights and component spreads imply modes that are too close together, too narrow, or effectively unused.

Then compare the inferred latent distribution with these prior-predictive implications, keeping in mind that posterior embeddings are informed by counts and are not expected to be identical to prior draws. Warning signs include many inferred points in regions assigned extremely low prior density, a mixture component that receives negligible weight, or a logistic-normal prior used for coordinates that have no compositional meaning. Such patterns do not prove a model is invalid, but they motivate checking the prior specification and the resulting fit.

A frequent mistake is to assume that a flexible mixture is always safer because it can represent more shapes. The correction is that flexibility brings additional choices and can make components difficult to interpret; the mixture should be justified and checked, not selected solely for its capacity. Another mistake is to use a logistic-normal distribution just because an embedding is bounded in some implementation. The correction is to define the latent quantity and its support explicitly, then choose a transformation that matches that meaning.

All three priors are modeling assumptions, not direct descriptions of cell identity. A prior-predictive pattern is not evidence that a biological population exists, and a component label is not a cell-type annotation. The prior is also only one part of the generative model: a suitable latent prior does not guarantee that the count model or inference procedure is adequate.