Sparse Minibatch Computation for Single-Cell VAEs

Why Sparse Computation Matters

A single-cell count matrix usually contains many zeros. For a batch of \(B\) cells and \(G\) genes, a dense representation stores \(BG\) numbers even when only a small fraction of entries are nonzero. This can make memory, rather than the statistical model, the limiting factor in fitting a count VAE. Sparse computation can reduce storage and avoid unnecessary arithmetic, but it must preserve the model’s treatment of every gene, including genes with observed count zero.

Let \(X=(x_{ig})\) denote the observed count matrix, with \(x_{ig}\) the nonnegative integer count for cell \(i\) and gene \(g\). A minibatch selects \(B\) rows, denoted \(X_{\mathcal B}\). Write \(\operatorname{nnz}(X_{\mathcal B})\) for the number of nonzero entries in that batch. In a row-compressed sparse representation, each row stores its nonzero gene indices and corresponding values, rather than \(G\) entries. A zero stored by omission is still an observed zero; it is not an unmeasured value.

Sparse Encoder Inputs

For a first encoder layer that is affine in its input, let \(S_{\mathcal B}\) be the encoder-input matrix, \(W\) a \(G\)-by-\(H\) weight matrix, and \(b\) an \(H\)-vector bias. The layer output is \(S_{\mathcal B}W+b\). If the input is a count transformation applied separately to each entry and maps zero to zero, it can remain sparse. For example, \(\log(1+x_{ig}/s_i)\), with a positive cell-specific scale \(s_i\), is zero when \(x_{ig}=0\). The transformation can be applied to stored values only.

The sparse product can be understood row by row. For cell \(i\), the output is the bias plus a sum over only the genes with nonzero encoder input. This is an exact matrix product, not an approximation that drops low-count genes. The gradient with respect to \(W\) is also a sparse-dense product: each nonzero input contributes its value times the corresponding back-propagated hidden-layer gradient. Ordinary dense activations in later layers can still be used; their size is typically \(B\)-by-\(H\), not \(B\)-by-\(G\).

Sparse inputs are not guaranteed merely because the original counts are sparse. A preprocessing step that adds a nonzero offset to every gene, or an imputation step that fills zeros with positive values, generally creates a dense input. Also, library-size scaling should be defined carefully: multiplying each observed entry by a cell-specific factor preserves zeros, while centering each gene by subtracting a nonzero mean does not. When the encoder needs library size, it can receive that scalar as an additional input rather than expanding it into a dense gene vector.

Exact Count Likelihoods Without Dense Outputs

The decoder specifies a likelihood for each observed gene count. For a Poisson model with mean \(\mu_{ig}>0\), the log likelihood is \(x_{ig}\log\mu_{ig}-\mu_{ig}-\log(x_{ig}!)\). The \(-\mu_{ig}\) term applies to every gene, including zeros. For a negative-binomial model parameterized by mean \(\mu_{ig}\) and inverse dispersion \(\theta_g>0\), use

This parameterization has conditional mean \(\mu_{ig}\) and variance \(\mu_{ig}+\mu_{ig}^{2}/\theta_g\). The likelihood sum over genes can be organized as a zero-count baseline over all genes plus corrections for nonzero observations. For the negative binomial, the baseline for gene \(g\) is \(\theta_g\log(\theta_g/(\theta_g+\mu_{ig}))\). At a nonzero count, add the correction consisting of the gamma-function terms and \(x_{ig}\log(\mu_{ig}/(\theta_g+\mu_{ig}))\). This decomposition accounts for every zero without storing a dense count array.

The baseline still depends on decoder parameters for all genes, so sparse observations do not generally remove the need to evaluate all gene contributions. A practical approach is to generate decoder parameters in gene chunks. For each cell and chunk, compute the relevant means and accumulate the baseline; then use the sparse indices to add the nonzero corrections. If decoder means are defined using a softmax over genes, its normalizing constant must be computed across all genes, not independently within each chunk. A stable implementation accumulates the row-wise log-sum-exp across chunks, then evaluates or recomputes chunks using that global normalization. Chunking limits peak memory; it does not, by itself, reduce the total number of gene computations.

Minibatch Gradients and Objective Scaling

Let \(\ell_i\) denote the per-cell objective contribution for cell \(i\), such as its evidence lower bound, with any latent-variable expectation estimated using the chosen reparameterization procedure. The full-data average objective is \(N^{-1}\sum_{i=1}^{N}\ell_i\), where \(N\) is the number of cells. If a batch of \(B\) cells is sampled uniformly, the batch average \(B^{-1}\sum_{i\in\mathcal B}\ell_i\) is an unbiased estimate of the full-data average. The corresponding statement holds for gradients when differentiation and expectation can be interchanged; in practice this requires a differentiable computation with finite expected gradients.

If instead the target objective is the full-data sum, its minibatch estimate must be scaled by \(N/B\). The sparse matrix operations determine how each per-cell term and its gradient are computed; they do not change this sampling factor. This distinction matters whether the encoder input uses sparse multiplication or the decoder likelihood is accumulated in chunks.

The same linearity argument applies to the per-cell gradients. Reparameterization noise from sampling continuous latent variables adds Monte Carlo variability, but does not alter the minibatch inclusion calculation.

Implementation Choices, Diagnostics, and Limits

Row-compressed storage is convenient when minibatches are selected by cell, because their nonzero entries can be accessed row by row. Sparse-dense multiplication is a natural encoder operation. On the decoder side, streaming gene chunks avoids holding all \(B G\) outputs at once, while retaining exact gene-wise likelihood terms. These approaches are complementary: sparse storage exploits zeros in the input, and chunking controls the memory used by the generally dense decoder outputs.

Compared with densifying the batch, sparse operations can reduce input memory from order \(BG\) to order \(\operatorname{nnz}(X_{\mathcal B})\), plus index storage. The dense route can be faster for small matrices or hardware and libraries optimized for dense arithmetic, so it may be reasonable when memory is sufficient. Sampling only a subset of genes can reduce decoder computation further, but it is not the exact full likelihood unless the estimator and its correction are specifically designed for that purpose. Simply ignoring genes not observed as nonzero is wrong: the zero-count terms depend on decoder parameters.

Useful checks include comparing sparse and dense calculations on a tiny batch, verifying that the likelihood includes all genes, checking that chunked and unchunked results agree within numerical tolerance, and confirming whether the optimized objective is a sum or an average. Monitor peak memory as well as runtime: accidentally converting the count batch or decoder output to a dense \(B\)-by-\(G\) array defeats the main memory benefit. Check that transformed sparse inputs preserve zero and that stored values have not been confused with missing observations.

Sparse computation has limits. The decoder may still require order \(BG\) arithmetic per batch, even if chunks cap its peak output memory. Very small chunks can add overhead, and floating-point summation order can cause small differences from a dense calculation. The likelihood decomposition is exact only when the baseline and every nonzero correction correspond to the stated observation model. It does not make a misspecified likelihood correct, nor does it eliminate stochastic variation from minibatch selection or latent-variable sampling.

Common Pitfall

A tempting but incorrect argument is: “The sparse matrix stores only nonzero counts, so I can sum the likelihood over its stored entries.” Sparse storage omits zero values, not the likelihood contributions of observed zeros. For example, in the negative-binomial calculation above, the zero at gene 2 contributes \(-0.8109\) to the log likelihood. The correction is to sum the zero baseline over every modeled gene and then add the appropriate corrections at nonzero entries. Similarly, do not apply the \(N/B\) multiplier to a batch mean when the target is the full-data average objective.