Designing Encoders and Decoders for Count Data

From Network Architecture to a Count Model

In a count VAE, the encoder and decoder are not just generic neural networks. Their output layers must correspond to the probability model for the observed counts. The encoder produces a distribution over a cell's latent vector; the decoder uses that latent vector, together with any modeled exposure, to produce parameters of a count distribution. An architecture can be expressive and still be unsuitable if its outputs do not obey the constraints of that distribution.

Here, \(x_i=(x_{i1},\ldots,x_{iG})\) denotes the observed nonnegative integer counts for cell \(i\), and \(z_i\) is its \(K\)-dimensional latent vector. Write \(\ell_i>0\) for the cell's exposure or library-size factor. The discussion uses a diagonal-Gaussian encoder and a negative-binomial decoder, with gene proportions summing to one and expected counts scaled by \(\ell_i\). These are architecture choices, not the only possible count-VAE design.

The usual likelihood and latent-variable setup is assumed: conditional on \(z_i\) and \(\ell_i\), gene counts are modeled with the chosen decoder distribution, often independently across genes; the encoder approximates the latent posterior; and the network weights are shared across cells. Any independence here is a modeling assumption, not a claim that genes are biologically independent. For negative-binomial counts, this tutorial uses mean \(\mu\) and dispersion \(\theta>0\), so the variance is \(\mu+\mu^2/\theta\).

Encoder Outputs: A Valid Latent Distribution

A standard encoder first transforms the count vector to a numerical input, for example \(u_{ig}=\log(1+x_{ig})\), and maps it through one or more hidden layers. A hidden activation such as ReLU may be used between affine layers. The final encoder layers return a mean vector \(m_i\) and a log-variance vector \(a_i\). Returning log variance rather than variance allows the network to produce any real-valued output while a deterministic transformation guarantees positive variance.

Here \(v_{ik}\) is the variance, not the standard deviation, of latent coordinate \(k\); the corresponding standard deviation is \(\sqrt{v_{ik}}\). The reparameterized expression generates a draw from the encoder distribution using noise \(\epsilon_i\) independent of the network inputs. This construction lets gradients pass through the deterministic mapping from \(m_i\) and \(v_i\) while sampling the noise separately.

The log-variance output can become extremely large or small, leading to unstable computations or nearly deterministic latent coordinates. Implementations may constrain it, use a numerically stable equivalent, or monitor its range. Such choices alter the parameterization or optimization behavior and should be documented; they do not change the requirement that the resulting variances be positive.

Decoder Outputs: Valid Means and Dispersions

A compositional decoder can map \(z_i\) to one real-valued logit per gene, \(r_{ig}\). Real-valued logits are unconstrained; applying softmax converts them into positive proportions summing to one. Scaling those proportions by a positive exposure gives nonnegative expected counts. In particular, the resulting means are strictly positive when \(\ell_i>0\) and all logits are finite.

The negative-binomial decoder can then use \(\mu_{ig}\) and a positive dispersion \(\theta_g\) (or a cell- and gene-specific positive dispersion, if that model is intended). A common output transformation for a real-valued raw dispersion \(d_g\) is \(\theta_g=\operatorname{softplus}(d_g)+\varepsilon\), where \(\operatorname{softplus}(d)=\log(1+\exp(d))\) and \(\varepsilon>0\) is a small numerical floor. An exponential transformation is another option, though it can be more sensitive to extreme raw outputs.

Justification, Diagnostics, and Alternatives

The output transformations enforce the relevant constraints directly. For finite logits, \(\exp(r_{ig})>0\), and the softmax denominator is the sum of positive terms, so every \(\rho_{ig}>0\). Summing the proportions cancels the common denominator:

Since \(\ell_i>0\), \(\mu_{ig}=\ell_i\rho_{ig}>0\) and its units are expected counts when \(\ell_i\) is a count-scale exposure. Similarly, softplus is positive for every finite input, so adding a positive floor guarantees \(\theta_g>0\). These facts ensure valid parameter values; they do not establish that the model fits the data well.

Useful checks include verifying that encoder means and log variances are finite, that transformed variances and dispersions are positive, and that each softmax vector sums to one within numerical tolerance. Inspect the distributions of \(m_i\), \(a_i\), \(\theta_g\), and predicted \(\mu_{ig}\) across cells and genes. Saturated hidden units, extreme log variances, near-zero proportions for many genes, or means systematically at odds with observed count scales can signal poor scaling, an unsuitable architecture, or an optimization problem. These are warning signs rather than formal tests of model adequacy.

The compositional decoder above allocates a specified exposure across genes. A common alternative is to have the decoder output a separate raw value per gene and transform it, for example with an exponential or softplus, to obtain a positive mean directly. That alternative does not force gene means to sum to \(\ell_i\), so it represents a different mean structure; it is not merely a different way to calculate the same proportions. Likewise, a shared dispersion, a gene-specific dispersion, and a cell-specific dispersion encode different assumptions. Dispersion output transformations should match the chosen parameterization and should not be confused with choosing which dispersion is shared.

A frequent mistake is to apply a sigmoid independently to each gene logit and call the results proportions. Sigmoids produce values between zero and one, but their sum need not equal one, so multiplying them by an exposure does not allocate that exposure as a composition. Use softmax when the model requires proportions summing to one; use independent positive means only when that is the intended model. Another mistake is to interpret a decoder mean as a denoised observed count or as a guarantee of a particular count. It is the expectation of the specified conditional distribution, with uncertainty represented by that distribution.