Selecting Latent Dimension Without Overfitting

Why Latent Dimension Needs a Selection Rule

A count VAE maps each cell to a latent vector with a chosen number of coordinates, \(K\). Increasing \(K\) gives the model more capacity to represent variation, but that variation need not be reproducible or useful. A model can fit its training cells more closely while predicting new cells no better, or even worse. The selection problem is therefore not “Which dimension gives the smallest training loss?” It is “Which dimension gives adequate predictive fit and stable downstream results without unnecessary complexity?”

For a candidate dimension \(K\), write the generative model as \(p_\theta(x\mid z,c)\), with latent variable \(z\), observed cell counts \(x\), and any conditioning variables \(c\) included in the model. The encoder supplies an approximate posterior \(q_\phi(z\mid x,c)\). The evidence lower bound (ELBO) for a cell is the expected log likelihood minus the Kullback–Leibler divergence from the approximate posterior to the prior. The earlier material on evidence lower bounds and count VAEs supplies this objective; here the focus is how to compare candidate dimensions.

Set Up a Fair Comparison

Compare models on the same count data, genes, conditioning variables, and data splits. Use the same observation model and broadly comparable encoder and decoder design; otherwise, a metric difference may reflect a changed model specification rather than \(K\). Train each candidate from multiple random initializations when feasible, since optimization variability can be substantial. Keep a final test set untouched until the candidate and evaluation procedure have been chosen.

The split should match the intended generalization target. A random split of cells can assess prediction for additional cells drawn from represented samples. If the intended use is prediction for new donors or samples, hold out entire donors or samples instead; random cell splitting may otherwise put closely related cells in both training and validation sets. Fit data-dependent preprocessing, such as feature selection, using training data only. Apply the resulting choices unchanged to validation and test data.

For count data, the likelihood comparison must use the same genes and account consistently for library size and any modeled covariates. Report predictive scores per cell, and, where useful, per observed count or gene, so that totals are not confused with changes in the number of observations. Do not compare a score based on one gene set or count representation with a score based on another as though the values were directly comparable.

Held-Out Likelihood and Reconstruction

The marginal likelihood for a held-out cell is \(p_\theta(x\mid c)=\int p_\theta(x\mid z,c)p(z\mid c)\,dz\). It measures how well the model assigns probability to the observed counts after accounting for latent uncertainty. This integral is usually unavailable in closed form. A practical comparison can use an importance-sampling estimate based on the encoder distribution. For draws \(z^{(m)}\) from \(q_\phi(z\mid x,c)\), define the importance weight \(w_m=p_\theta(x,z^{(m)}\mid c)/q_\phi(z^{(m)}\mid x,c)\), where \(p_\theta(x,z\mid c)=p_\theta(x\mid z,c)p(z\mid c)\). Then estimate the log marginal likelihood by:

Use the same estimation procedure and number of importance samples \(M\) across candidate dimensions, and check that the estimate is stable as \(M\) increases. The finite-\(M\) log estimate is not an exact likelihood value; its accuracy can depend on the quality of the proposal distribution. The ELBO is a lower bound on the log marginal likelihood, not an unbiased estimate of it. ELBO comparisons can still be informative when calculated consistently, but a difference may partly reflect a change in the tightness of the variational bound.

Reconstruction measures offer a complementary view. For example, held-out negative log likelihood or a count deviance can summarize how closely predictions match observed counts. State exactly which reconstruction measure is used and whether smaller or larger is better. Reconstruction measures alone can reward fitting idiosyncrasies, so they should not replace held-out predictive evaluation. Training reconstruction is especially unsuitable as the sole selection criterion: additional capacity can improve it even when generalization does not.

KL Usage: A Diagnostic, Not a Biological Count

For a diagonal Gaussian encoder and a standard normal prior, the KL contribution for coordinate \(j\) in a cell is:

Averaging this quantity over cells gives a per-coordinate summary of how much the approximate posterior for that coordinate differs from its prior. A coordinate with a very small average contribution is often described as weakly used under this fitted model. Summing contributions gives the total KL penalty, in nats per cell. These summaries depend on the prior, parameterization, and fitted solution. In particular, KL usage does not count biological processes, and a low value for one coordinate does not prove that the corresponding variation is irrelevant.

Downstream Stability and a Selection Workflow

A dimension can produce a good average predictive score but unstable results across random seeds. Assess stability for the downstream task that motivates fitting the VAE, without treating stability as proof of biological validity. For example, compare nearest-neighbor sets or cluster assignments across repeated fits. Because latent coordinates can rotate or permute without changing the model’s substantive fit, direct coordinate-by-coordinate comparisons may be misleading. Comparing neighborhood structure or assignments avoids requiring identical coordinate axes.

For two fits, a simple neighborhood comparison for cell \(i\) is the Jaccard similarity, the size of the intersection of its neighbor sets divided by the size of their union. Average this over cells and pairs of fits. For cluster assignments, an adjusted Rand index can summarize agreement while adjusting for agreement expected by chance. These are stability summaries, not universal pass/fail thresholds: their meaning depends on the number of neighbors or clusters and on the downstream question.

A practical workflow is to define candidate dimensions and the target of generalization in advance; fit each candidate with repeated seeds; compare held-out predictive scores and reconstruction using the same evaluation procedure; inspect KL usage across cells and seeds; and assess stability for the planned downstream analysis. Then choose a dimension with strong enough held-out performance and adequate stability, preferring a simpler candidate when larger dimensions offer no material, reproducible improvement. If candidate scores are close, report the ambiguity rather than presenting a precise dimension as uniquely determined.

Assumptions, Diagnostics, and Limits

These comparisons are informative only if the evaluation matches the model and scientific target. Inspect whether count inputs, genes, library-size handling, and covariates are consistent across candidates. Check training convergence and replicate fits: a poorly optimized high-dimensional model may appear inferior for computational reasons, while a single lucky initialization may appear superior. For importance-sampling scores, check sensitivity to the number of draws; for all metrics, quantify variation across held-out cells and training seeds. If the validation set is small or cells are strongly grouped by donor, nominal precision based on treating every cell as independent can be misleading; use splits and uncertainty summaries that respect the sampling structure.

A held-out score can also be optimized repeatedly by trying many architectures, splits, or dimensions against the same validation set. That makes the reported validation performance optimistic for the selected model. Use validation data to make the selection and reserve a test set for a final evaluation, or use an appropriate nested evaluation strategy when the full modeling procedure itself is being assessed. Held-out performance measures prediction under the chosen split and data distribution; it does not certify causal meaning, biological truth, or generalization to populations unlike those represented.

No metric dominates in every application. Held-out likelihood evaluates probabilistic prediction but can be sensitive to likelihood specification and estimator quality. Reconstruction focuses on observed counts but can favor overly detailed fits. KL usage indicates how the fitted approximate posterior uses coordinates but is prior- and solution-dependent. Downstream stability detects sensitivity to fitting choices but stable results can still be systematically wrong. These diagnostics are most useful together, with the intended analysis determining which trade-offs matter.

• Pitfall: choosing the smallest training loss. More latent capacity can improve fit to training cells. Correct this by comparing held-out performance and checking whether gains persist across seeds.
• Pitfall: calling every low-KL coordinate useless. An average can hide subgroup usage, and refitting at a smaller \(K\) can reorganize the representation. Inspect per-cell and repeated-fit summaries.
• Pitfall: treating a stable embedding as validated biology. Stability describes reproducibility under the chosen procedure, not truth. Check predictive fit and the scientific assumptions separately.
• Pitfall: comparing scores from different data or estimators. Changes in genes, preprocessing, splits, or likelihood estimation can explain score differences. Hold these fixed for a fair comparison.