Amortized Inference With Neural Encoders

From Counts to Approximate Posterior Parameters

A latent-variable model describes each cell’s observed gene-count vector \(x_i\) through an unobserved representation \(z_i\). Inference asks what values of \(z_i\) are plausible after observing \(x_i\). Directly calculating the exact posterior \(p(z_i\mid x_i)\) is often impractical, so a variational autoencoder uses a tractable approximation, such as \(q_\phi(z_i\mid x_i)\). A neural encoder supplies the parameters of that distribution.

The key idea is that the encoder’s weights are shared across cells, while its output depends on the input cell. Thus, one learned function can produce different approximate posterior distributions for different count vectors. The encoder does not itself observe the latent state, nor does it generally return a single certain value for that state. It returns parameters describing a distribution from which latent values can be sampled.

Here \(x_i\) is the count vector for cell \(i\); \(\phi\) denotes encoder parameters; and \(\sigma_\phi^2(x_i)\) is obtained by exponentiating the log variance. Exponentiation ensures positive variances. For a latent vector of dimension \(d\), the encoder typically returns \(d\) means and \(d\) log variances. The diagonal covariance assumption means that the approximate posterior treats latent coordinates as conditionally independent given the cell. It is a modeling choice, not a property guaranteed by the data.

The encoder’s input representation must match the one used when fitting it. In the examples below, the inputs are hypothetical, already-specified numerical features derived from counts; the calculations do not prescribe a general transformation of count data. The generative model specifies what counts mean, while the encoder provides a convenient approximation to the posterior under that model.

Calculating an Encoder Output

For a Gaussian encoder, the network’s outputs are distribution parameters, not latent samples. Once the mean and variance have been computed, the reparameterization trick expresses a draw as \(z_i=\mu_\phi(x_i)+\sigma_\phi(x_i)\odot\epsilon\), where \(\epsilon\sim N(0,I)\) and \(\odot\) denotes coordinatewise multiplication. This separates the random noise from the learned parameters and allows gradients to be computed through the sampling step.

The output’s dependence on the input is essential: another cell passed through the same trained encoder can receive different posterior parameters. The parameters \(\phi\), however, are common to all cells. During training, the model adjusts those shared parameters so that the encoder’s distributions support a useful approximation to the posterior across the training data.

Amortized Inference and Per-Cell Optimization

To state the distinction precisely, fix the generative-model parameters and let \(\mathcal{Q}\) be a specified family of candidate posterior distributions. For cell \(i\), let \(\mathcal{L}_i(q)\) denote the evidence lower bound (ELBO) evaluated using candidate distribution \(q\). Per-cell variational inference chooses a separate \(q_i\in\mathcal{Q}\) for each cell to maximize that cell’s ELBO. Amortized inference instead restricts the choices to distributions produced by one shared encoder, \(q_\phi(z_i\mid x_i)\), and learns \(\phi\) using the cells’ combined objective.

These objectives clarify the trade-off. Separate optimization can adapt each distribution without requiring the same function to produce all of them, but it requires an optimization problem for each cell. An encoder shares statistical strength and can produce approximate posterior parameters for a new input with a forward pass. That efficiency comes with a constraint: the encoder must represent many cell-specific distributions using the same learned mapping.

This definition separates two issues. Even the best distribution in \(\mathcal{Q}\) may be an imperfect approximation to the exact posterior; that is a limitation of the variational family. The amortization gap instead compares the encoder’s output with the best choice available within that same family. In practice, a per-cell optimization that has not converged is not necessarily \(q_i^*\), so a measured difference from it may also reflect incomplete optimization.

Why Sharing Parameters Creates a Trade-Off

A simple quadratic objective shows how a shared mapping can compromise across cells. This is a mathematical illustration, not a count-VAE likelihood. Suppose each cell’s variational parameters are summarized by a pair \((m,r)\), where \(m\) is the mean and \(r\) is the log variance. Assume its objective has the form

The maximum for cell \(i\) is at its target \((m_i^*,r_i^*)\), because the subtracted sum of squares is smallest, namely zero, there. If the encoder is restricted to output the same pair for two cells, it must choose one compromise pair. Differentiating the summed objective, or equivalently minimizing the summed squared distances, shows that this compromise is the coordinatewise average of the two targets.

In a neural encoder, outputs are not usually forced to be constant. Its input-dependent layers allow different cells to receive different distributions. Nevertheless, limited model capacity, an uninformative input representation, or a difficult mapping can leave a residual gap. Sharing parameters is therefore neither automatically harmful nor a guarantee of accurate per-cell inference; its performance depends on the data, the encoder, and the variational family.

How Shared Parameters Learn From Multiple Cells

The shared objective also explains why a cell’s encoder output can be influenced indirectly by other cells: all cells contribute to updates of \(\phi\). The following small calculation illustrates the shared trade-off for a single scalar parameter. It is again a quadratic surrogate rather than a complete VAE training objective.

For a neural encoder, gradient updates aggregate contributions from many cells and adjust many shared parameters. The calculation is a simplified justification for this principle: the learned mapping seeks performance across cells, rather than solving an isolated optimization for each one. It also shows why inspecting only an average training objective can conceal cells whose individual posterior approximations fit poorly.

Diagnostics, Alternatives, and Limitations

Useful checks ask whether the encoder is producing valid, responsive, and adequately fitted distributions. Inspect the ranges of predicted means and log variances, and verify that exponentiated variances are positive and finite. Extremely large or small log variances, non-finite outputs, or nearly identical outputs across clearly different inputs can signal numerical problems, a weakly responsive mapping, or an issue with the input representation. Similar outputs are not by themselves proof of failure: distinct cells can legitimately have similar posterior distributions.

Compare objective values and posterior summaries across cells rather than relying only on an aggregate. When feasible, per-cell optimization within the same variational family provides a reference for assessing whether the encoder is leaving a substantial amortization gap. Such a comparison must use the same generative model and family, and the per-cell optimization must be sufficiently converged to be informative. A gap can also be investigated by checking whether a more flexible inference function or a better-suited variational family changes the results, while recognizing that these changes address different limitations.

These alternatives do not remove all uncertainty or approximation. Per-cell variational optimization remains restricted to its chosen family and can be sensitive to optimization. A point estimate discards distributional uncertainty. An encoder can also behave poorly on inputs unlike those used to fit it; producing numerical posterior parameters for a new cell does not establish that those parameters are accurate. The encoder’s posterior is an approximation conditional on the specified model and its inputs, not a guarantee of a uniquely determined biological interpretation.

A common mistake is to say, “The encoder returns the latent value for the cell.” In a Gaussian encoder, it returns parameters such as \(\mu_\phi(x_i)\) and \(\ell_\phi(x_i)\); a latent draw is obtained only after applying the reparameterization with a noise draw. Another mistake is to assume that sharing the encoder means every cell has the same posterior. The correction is that the weights are shared but the outputs depend on each cell’s input. Conversely, distinct outputs do not prove that the approximation is optimal: a shared mapping can still have an amortization gap.