Scalable Single-Cell Gene Expression Generation with Latent Diffusion Models

AI summary

Overview

• Research area: Generative machine learning applied to single-cell genomics (stat.ML / computational biology), sitting at the intersection of variational autoencoders, diffusion/flow-matching models, and single-cell transcriptomics.
• Technical level: Advanced. The paper assumes familiarity with VAEs, the ELBO, attention/transformers, diffusion models, flow matching, and classifier-free guidance. The biological motivation (gene exchangeability, sparsity of scRNA-seq) is explained clearly, but the methodology is dense.
• Scope: The paper introduces scLDM, a fully transformer-based latent diffusion model for single-cell gene expression that respects the permutation-invariant (exchangeable) nature of gene sets, and demonstrates it on reconstruction, unconditional/conditional generation, perturbation, and downstream classification tasks.

What This Paper Is About

Modeling single-cell gene expression with generative models is hard because expression data are counts with complex dependencies among genes, and because the set of genes has no meaningful order. Existing methods either impose a fixed gene ordering (tying each input dimension to one gene, which prevents transferring across tissues/species without retraining), operate on a restricted panel of highly variable genes, or rely on shallow architectures such as MLPs and GANs with known training instabilities. This paper's goal is a scalable, exchangeable generative model that generates realistic expression profiles across thousands of genes, and that can be steered with multiple conditions (e.g., cell type plus perturbation).

Key Contributions

1. A fully transformer-based VAE for exchangeable data. The model uses a fixed-size, permutation-invariant set of latent variables rather than per-gene latents. A single Multi-head Cross-Attention Block (MCAB) serves two roles: permutation-invariant pooling in the encoder and permutation-equivariant unpooling in the decoder, eliminating the need for separate pooling/unpooling modules as in SetTransformer and SetVAE.

2. A latent diffusion model as the prior. The standard Gaussian prior is replaced with a diffusion model trained using linear interpolants and the flow matching loss (Scalable Interpolant Transformers / SiT), with the denoiser parameterized by Diffusion Transformers (DiT). This better matches the aggregated posterior and enables controlled generation.

3. Joint multi-attribute classifier-free guidance. Instead of the additive conditioning used by CFGen, scLDM supports guidance over simultaneous combinations of attributes, which the paper shows outperforms additive guidance on perturbational benchmarks.

4. A sparsity-aware input scheme. Because 70%+ of gene–cell entries are zero, the encoder filters to expressed genes only and pads to a fixed context length, reducing compute without limiting the decoder's ability to model zero counts (the decoder conditions on the latent and can place mass at zero via the Negative Binomial likelihood).

Main Findings

• Reconstruction improvements are large and consistent. Across three datasets (Dentate Gyrus, Tabula Muris, HLCA), scLDM with a Negative Binomial likelihood achieves the lowest reconstruction error, highest Pearson correlation, and lowest MSE compared to scVI and CFGen. On Tabula Muris, PCC rises from 0.221 (scVI) and 0.136 (CFGen) to 0.376; on HLCA, MSE drops from 0.238 (scVI) and 0.117 (CFGen) to 0.095.

• Negative Binomial beats Gaussian. scLDM (NB) substantially outperforms scLDM (Gauss) on generation benchmarks (e.g., Wasserstein-2 of 7.267 vs. 14.670 on Tabula Muris unconditional), confirming that respecting the count nature of the data matters. Gaussian variants degrade sharply in precision and recall.

• Strong generation on observational data. In both unconditional and conditional settings, scLDM (NB) attains the best Wasserstein-2 and MMD² (RBF) scores in nearly all dataset/setting combinations, and its 1-NN score is closer to the ideal 0.5 than either scDiffusion or CFGen.

• Precision–recall trade-off. Some baselines (notably scDiffusion and CFGen in certain settings) achieve higher precision, while scLDM (NB) tends to achieve higher recall — suggesting it covers more of the true expression distribution, at some cost in sample sharpness. On unconditional HLCA, scDiffusion reaches precision 0.788 versus scLDM's 0.540, while scLDM's recall is far higher (0.622 vs. 0.008).

• Downstream utility. The autoencoder's learned embeddings are useful for cell-level classification, and the framework's advantages extend to perturbational single-cell data, where the joint conditioning formulation outperforms additive multi-attribute guidance.

Methodology in Plain English

The authors build the model in two stages.

Stage 1 — Learn a compact, order-free representation. A transformer encoder reads a cell's expressed genes as a set: each gene is identified by an embedding vector indexed by its gene ID, not by its position in the input. A cross-attention block with a fixed number of learned "pseudo-input" queries aggregates all genes into a fixed-size set of latent tokens. Because the queries are fixed and the aggregation is attention-based, scrambling the input gene order does not change the latent — the model is permutation-invariant. The decoder reverses this: it takes the latent tokens plus a requested list of gene IDs, uses those gene embeddings as queries, and produces per-gene distribution parameters (mean/variance for Gaussian, or mean/dispersion for Negative Binomial). Permuting the requested gene IDs permutes the outputs in the same way — permutation-equivariance — so the full conditional likelihood is exchangeable.

Stage 2 — Generate latents with diffusion. The VAE is frozen, and a Diffusion Transformer is trained in the latent space to model the distribution of latent tokens, using flow matching with linear interpolants and guidance-enabled sampling (SiT). Conditioning is handled with classifier-free guidance: during training, conditions are randomly dropped with probability ρ, and at sampling time the vector field is extrapolated between an unconditional and a conditional prediction. Unlike CFGen's additive scheme, the guidance here is defined over joint attribute combinations.

Training and sampling. The VAE is trained with a β-weighted ELBO (β=0 recovers a pure autoencoder); the diffusion stage is trained separately. To generate a cell, one samples a latent from the diffusion model and passes it through the frozen decoder for the desired gene IDs.

Why This Matters

• Impact on research. The permutation-invariant design removes a structural constraint that has limited prior single-cell models: no fixed gene ordering, no retraining or weight surgery when the gene vocabulary changes across tissues, donors, or species. Combined with the deeper transformer backbone and diffusion prior, this provides a stronger foundation for virtual cell modeling and perturbation prediction, and the reconstruction and generation metrics set a new bar against scVI, scDiffusion, and CFGen on established benchmarks.

• Real-world applications:

  • Virtual cell simulation: generating synthetic cells with specified identities or perturbation states to explore hypotheses without wet-lab cost.
  • Drug and genetic perturbation modeling: conditional generation of expression profiles under multiple simultaneous perturbations to prioritize candidates.
  • Cross-tissue and cross-species transfer: because genes are addressed by embedding index rather than input position, a single architecture can in principle cover different gene panels.
  • Data augmentation and imputation: synthesizing expression profiles to bolster downstream classifiers or fill in sparse measurements.

• Industry relevance. Pharmaceutical and biotech companies increasingly use single-cell data for target discovery and compound screening; a scalable, controllable generative model of expression profiles fits directly into in-silico screening pipelines. The public release of code (github.com/czi-ai/scldm) lowers the barrier to adoption. The work also has general methodological relevance as a case study in combining permutation-equivariant architectures with latent diffusion for high-dimensional scientific data.

Future Directions

• Extending to genuinely variable gene sets. The current formulation uses fixed-size latents and a fixed context length; handling arbitrary, tissue-specific gene panels without retraining or padding remains an open design problem.

• Closing the precision gap. scLDM (NB) often trails baselines on precision while leading on recall. Future work could explore better likelihood parameterizations or guidance schedules to sharpen generated samples without sacrificing coverage.

• Broader benchmarking. The paper evaluates on three observational datasets and perturbational benchmarks; scaling to larger atlases, multi-modal data (e.g., paired RNA and protein), and cross-species settings would test the claimed flexibility more thoroughly.

• Interpretability and control. The paper mentions cross-attention score analyses and UMAP visualizations in the appendix; turning these into a principled interpretability layer could reveal which genes drive a given latent direction and enable finer-grained biological control.

Target Audience

Machine learning researchers working on generative models for scientific data (particularly set-structured or count-valued data), computational biologists and bioinformaticians building single-cell foundation or virtual-cell models, and practitioners in pharma or biotech who need controllable synthetic expression profiles. Readers will benefit most if they already understand VAEs, diffusion/flow matching, and attention mechanisms, since the paper's contributions are primarily architectural and methodological rather than tutorial in nature.

Authors’ abstract

Computational modeling of single-cell gene expression is crucial for understanding cellular processes, but generating realistic expression profiles remains a major challenge. This difficulty arises from the count nature of gene expression data and complex latent dependencies among genes. Existing generative models often impose artificial gene orderings or rely on shallow neural network architectures. We introduce a scalable latent diffusion model for single-cell gene expression data, which we refer to as scLDM, that respects the fundamental exchangeability property of the data. Our VAE uses fixed-size latent variables leveraging a unified Multi-head Cross-Attention Block (MCAB) architecture, which serves dual roles: permutation-invariant pooling in the encoder and permutation-equivariant unpooling in the decoder. We enhance this framework by replacing the Gaussian prior with a latent diffusion model using Diffusion Transformers and linear interpolants, enabling high-quality generation with multi-conditional classifier-free guidance. We show its superior performance in a variety of experiments for both observational and perturbational single-cell data, as well as downstream tasks like cell-level classification.