Why Condition on Batch?
Single-cell count data can vary across batches because of technical differences in processing, sequencing, or measurement. The same data can also vary because the batches contain different cell types or biological states. A conditional variational autoencoder (VAE) supplies an observed batch label to part of the model so that the model can represent batch-associated variation explicitly. This changes the model’s assumptions; it does not, by itself, determine whether a difference is technical or biological.
The earlier material on identifiability and disentanglement in VAE latents motivates an important caution: a latent representation is not guaranteed to correspond to distinct biological and technical factors just because the model has separate inputs for them. Here the goal is narrower and practical: specify how a batch covariate enters a count VAE, understand what the resulting objective says, and recognize when the data cannot support the desired separation.
Model Specification and Conditioning Choices
A useful generic specification has prior \(p(z_i)\), often \(N(0,I)\), and an approximate posterior \(q_\phi(z_i\mid x_i,b_i)\). A count decoder can specify \(p_\theta(x_i\mid z_i,b_i,l_i)\). For example, it can produce gene logits \(a_i=f_\theta(z_i,b_i)\), convert them to proportions \(\rho_{ig}\) with a softmax, and set the count mean to \(\mu_{ig}=l_i\rho_{ig}\). A negative binomial observation model then has mean \(\mu_{ig}\) and gene-specific dispersion \(r_g\), with variance \(\mu_{ig}+\mu_{ig}^2/r_g\).
This is one general construction, not a requirement that every conditional count VAE use exactly this decoder. The observation model must be selected for the counts and its parameterization stated. For the example above, counts are conditionally independent across genes given \(z_i,b_i,l_i\), and each gene follows the specified negative binomial distribution. These are modeling assumptions, not facts guaranteed by the count matrix. Batch labels are assumed known and correctly assigned. The library-size offset is treated as given in this formulation.
There are two distinct places to supply the batch label. Conditioning the encoder gives \(q_\phi(z_i\mid x_i,b_i)\): the inference network may use batch when approximating the latent posterior. Conditioning the decoder gives \(p_\theta(x_i\mid z_i,b_i,l_i)\): the generative model may produce different count distributions at the same \(z_i\) in different batches. A model may use either route or both. The choice changes the assumptions and the information available to each network.
| Where batch enters | What the model can do | What it does not guarantee |
|---|---|---|
| Encoder only | Use batch when approximating the posterior over \(z_i\). | Represent a batch-specific count distribution in the decoder; the decoder is still shared across batches. |
| Decoder only | Allow the count distribution to depend on batch at a given \(z_i\). | Ensure that \(z_i\) contains no batch information. |
| Both | Use batch in inference and allow batch-conditional count generation. | Separate batch effects from biology when the design or model does not identify that separation. |
For a decoder-conditioned model, the joint distribution conditional on the observed batch and library size is \(p_\theta(x_i,z_i\mid b_i,l_i)=p(z_i)p_\theta(x_i\mid z_i,b_i,l_i)\). If the encoder is also conditioned, it supplies the approximation \(q_\phi(z_i\mid x_i,b_i)\) used for inference. A common per-cell objective is the evidence lower bound (ELBO):
The first term rewards count distributions that assign high probability to the observed counts. The Kullback–Leibler (KL) divergence penalizes a posterior that departs from the prior. The batch label appears explicitly in the likelihood and, when the encoder is conditioned, in the approximate posterior. It is not a separate penalty that forces latent variables to be independent of batch.
Why the ELBO Applies
For fixed observed \(x_i,b_i,l_i\), suppose \(q_\phi(z_i\mid x_i,b_i)\) has support wherever the joint density \(p_\theta(x_i,z_i\mid b_i,l_i)\) is positive. Multiply and divide the integral defining the marginal likelihood by \(q_\phi\), then apply Jensen’s inequality to the concave logarithm:
The last equality uses the factorization \(p_\theta(x_i,z_i\mid b_i,l_i)=p(z_i)p_\theta(x_i\mid z_i,b_i,l_i)\). Thus the ELBO is a lower bound on the conditional log marginal likelihood under this model. This derivation does not establish that the model is correctly specified, that optimization finds a good solution, or that the latent coordinates have a unique biological interpretation.
Worked Example: A Batch-Conditioned Encoder
Consider a hypothetical two-dimensional latent variable with prior \(p(z)=N(0,I)\). A count encoder receives the same observed count vector in two hypothetical inputs, but the batch label differs. Suppose its outputs are \(q(z\mid x,b=1)=N((0.2,-0.4),\operatorname{diag}(0.5,1))\) and \(q(z\mid x,b=2)=N((0.6,-0.4),\operatorname{diag}(0.5,1))\). The means are vectors and the diagonal entries are variances.
For a diagonal Gaussian \(q=N(m,\operatorname{diag}(v))\) and standard normal prior, the KL divergence is \(\frac{1}{2}\sum_j(m_j^2+v_j-1-\log v_j)\). For batch 1, the first coordinate contributes \(\frac{1}{2}(0.2^2+0.5-1-\log(0.5))=0.1166\), and the second contributes \(\frac{1}{2}((-0.4)^2+1-1-\log(1))=0.0800\). The total is \(0.1966\) nats.
For batch 2, the first-coordinate contribution is \(\frac{1}{2}(0.6^2+0.5-1-\log(0.5))=0.2766\); the second remains \(0.0800\). The total is \(0.3566\) nats. These values show how a batch-conditioned encoder can produce different approximate posteriors and KL costs. They do not show that either posterior isolates biology: that interpretation requires evidence beyond the fact that batch was an encoder input.
Worked Example: Batch-Conditional Decoder Means
Suppose a hypothetical decoder has two genes and, at a fixed latent value \(z\), returns logits \((\log 3,0)\) for batch 1 and \((0,\log 3)\) for batch 2. Let the library-size offset be \(l_i=100\) counts. The softmax proportions for batch 1 are \((3/(3+1),1/(3+1))=(0.75,0.25)\). Its mean counts are therefore \((100(0.75),100(0.25))=(75,25)\) counts.
For batch 2, the proportions are \((0.25,0.75)\), giving mean counts \((25,75)\) counts. At the same \(z\) and library size, the decoder allows a batch-dependent expected composition. These are model means, not a claim that every cell has exactly those observed counts. If the observation model is negative binomial, the realized counts vary around the means according to the specified dispersions as well.
Confounding: What Conditioning Cannot Resolve
A batch covariate tells the model which batch generated an observation. It does not create comparisons absent from the experimental design. Consider a hypothetical study in which all cells of type A were measured in batch 1 and all cells of type B in batch 2. Suppose the average count for one gene is 10 in type A and 20 in type B. The observed groups differ, but cell type and batch change together.
Worked Example: Two Explanations Fit a Confounded Design
For illustration, consider a log-mean model for this gene, \(\log\mu=\alpha+\beta I(\text{type B})+\gamma I(\text{batch 2})\). Here \(I(\cdot)\) is 1 when the stated condition holds and 0 otherwise. In the hypothetical data, the only observed combinations are type A in batch 1, with mean 10, and type B in batch 2, with mean 20.
One explanation sets \(\alpha=\log(10)\), \(\beta=\log(2)\), and \(\gamma=0\). It predicts a mean of \(\exp(\log(10))=10\) for type A in batch 1 and \(\exp(\log(10)+\log(2))=20\) for type B in batch 2. A second explanation sets \(\alpha=\log(10)\), \(\beta=0\), and \(\gamma=\log(2)\). It predicts the same two observed means. The data therefore do not distinguish a cell-type difference from a batch difference under this model.
No increase in the number of cells from these same two combinations supplies the missing comparisons. To estimate these effects separately, the design needs overlap—for example, both cell types represented in both batches—or other justified information. A conditional VAE can fit a model to the observed counts, but batch labels alone cannot identify which of these explanations is correct.
Diagnostics, Alternatives, and Limitations
Evaluate a conditional model against its stated purpose, rather than assuming that supplying batch labels has succeeded in removing unwanted variation. Useful checks include whether observed and generated count distributions resemble one another within each batch, whether reconstruction errors or residual patterns differ systematically by batch, and whether predictions are implausible for batch–cell-type combinations represented in the data. A model that fits pooled counts but misses a batch-specific pattern may have an inadequate conditional decoder or observation model. Conversely, apparently similar latent distributions across batches do not prove that biological variation has been preserved.
Inspect the design before fitting: cross-tabulate batch against known cell types, conditions, and other important covariates. Empty or nearly empty combinations reveal weak or absent support for comparisons. Also check label quality, library-size handling, and whether batch is being supplied to the intended network. If adding batch to the encoder substantially changes posterior summaries, investigate whether it improves inference or instead changes the representation in a way that weakens the biological question. These checks diagnose concerns; none can prove disentanglement from observational data alone.
The main alternatives are not interchangeable. An encoder-only model uses batch to help infer \(z_i\), but its shared decoder cannot express a direct batch-specific count distribution at fixed \(z_i\). A decoder-conditioned model can express such differences, while its encoder need not use batch explicitly. Conditioning both networks permits both functions but adds flexibility and does not impose batch invariance on \(z_i\). A model that omits batch makes a stronger shared-distribution assumption; it may be suitable when that assumption is plausible, but can misattribute batch variation to latent structure or leave systematic misfit. The appropriate choice depends on the scientific target and design, not a universal rule.
A further limitation is that a batch effect may interact with cell state or gene: one shared additive adjustment may not represent all differences. More flexible batch-conditioned networks can fit complex patterns, but may also learn poorly supported combinations or overfit. Finally, batch can be associated with genuine biological conditions. Removing every signal associated with batch can therefore remove biology when batch and condition are confounded. Conditional generation is a modeling tool, not a causal correction or a guarantee of a batch-free embedding.
Worked Example: Calculating a Simple ELBO
Take a hypothetical one-gene Poisson decoder whose rate is fixed at \(\lambda=3\) counts for a particular batch and latent value, independent of \(z\). The observed count is \(x=4\). Let the approximate posterior be \(q(z\mid x,b)=N(0.5,0.25)\), where 0.25 is the variance, and let the prior be \(N(0,1)\). Because the decoder rate is independent of \(z\), its expected log likelihood under \(q\) is just its value at rate 3.
The Poisson log likelihood is \(\log p(x=4\mid\lambda=3)=4\log(3)-3-\log(4!)=4(1.098612)-3-\log(24)=-1.7836\). The one-dimensional Gaussian KL divergence is \(\frac{1}{2}(0.5^2+0.25-1-\log(0.25))=\frac{1}{2}(0.25+0.25-1+1.386294)=0.4431\) nats. Thus the ELBO is \(-1.7836-0.4431=-2.2267\) nats. It is a lower bound on the conditional log likelihood for this observation under the stated model; it is not itself a probability.
Check Your Understanding
Answer each question, then compare your reasoning with the solutions.
- In a model with \(q_\phi(z\mid x,b)\) and \(p_\theta(x\mid z,l)\), which part receives batch, and can the decoder represent a direct batch-specific count distribution at fixed \(z\)?
- A two-gene decoder returns logits \((0,\log 4)\) and \(l=80\). What are the proportions and mean counts?
- Why can observing only type A in batch 1 and type B in batch 2 prevent separate estimation of cell-type and batch effects?
- Does decoder conditioning guarantee that the latent variable is batch-independent? Give the reason.
Solutions. (1) Batch enters the encoder only. The decoder has no batch input, so it cannot represent a direct batch-specific distribution at fixed \(z\) and \(l\). (2) The softmax proportions are \((1/5,4/5)=(0.2,0.8)\); the mean counts are \((80(0.2),80(0.8))=(16,64)\). (3) Cell type and batch are perfectly aligned in the observed combinations, so distinct cell-type and batch effects can produce the same fitted means; the missing combinations provide no comparison to distinguish them. (4) No. Conditioning the decoder permits batch-dependent generation, but does not impose independence of the encoder’s latent representation from batch or guarantee identifiability.