🧬 biology

Scalable Single-Cell Gene Expression Generation with Latent Diffusion Models

This paper introduces scLDM, a scalable latent diffusion model that leverages a permutation-invariant Multi-head Cross-Attention Block and Diffusion Transformers to generate realistic, high-quality single-cell gene expression profiles while respecting the data's fundamental exchangeability property.

Original authors: Giovanni Palla, Sudarshan Babu, Payam Dibaeinia, James D. Pearce, Donghui Li, Aly A. Khan, Theofanis Karaletsos, Jakub M. Tomczak

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Giovanni Palla, Sudarshan Babu, Payam Dibaeinia, James D. Pearce, Donghui Li, Aly A. Khan, Theofanis Karaletsos, Jakub M. Tomczak

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are trying to teach a computer to understand the "personality" of a single cell. In biology, a cell's personality is written in its gene expression: a massive list of numbers showing how active thousands of different genes are.

The problem is that this list is messy. It's like a grocery list where the order of items doesn't matter (you can list "milk" before "eggs" or vice versa, and it's still the same list), but most computer programs insist on a strict order. Also, the list is full of zeros (genes that aren't active), and the numbers are whole counts, not smooth decimals.

The paper introduces a new AI model called scLDM (single-cell Latent Diffusion Model) that solves these problems. Here is how it works, using simple analogies:

1. The "Indifferent Chef" (Handling Order)

Most AI models treat genes like a sentence where word order matters. If you swap "cat" and "dog," the meaning changes. But in a cell, swapping the order of genes doesn't change the cell's identity.

The Old Way: Imagine a chef who insists you list ingredients in a specific order (Apples, then Bananas, then Cherries). If you give them Cherries, then Apples, they get confused or break.
The scLDM Way: This model is like a chef who doesn't care about the order. They just look at the set of ingredients. Whether you hand them the list as "Apples, Bananas" or "Bananas, Apples," they understand it's the same recipe. They use a special tool called a Multi-head Cross-Attention Block to mix the ingredients together without worrying about which one came first.

2. The "Compressed Sketch" (The Latent Space)

Cells have thousands of genes, which is too much data for a computer to process all at once efficiently.

The Analogy: Imagine trying to memorize a 10,000-page encyclopedia. Instead, you create a "compressed sketch" or a summary note that captures the essence of the book.
How it works: The scLDM first reads the cell's gene list and compresses it into a small, fixed-size "summary token" (the latent variable). It doesn't matter if the original list had 10 genes or 10,000; the summary is always the same size. This makes the model scalable (it can handle huge datasets without crashing).

3. The "Denoising Artist" (Diffusion)

Once the model has the "summary sketch," it needs to generate new, realistic cells.

The Analogy: Imagine a sculpture made of clay.
- Old Models: Tried to build the sculpture from scratch, often resulting in lumpy, unrealistic shapes.
- scLDM (Diffusion): Starts with a block of pure, chaotic noise (like static on an old TV). It then slowly "denoises" this block, step-by-step, refining the static into a perfect, realistic sculpture.
- The Twist: Instead of doing this on the raw, messy gene data (which is slow and hard), it does the "denoising" on the compressed sketch (the latent space). This is like sculpting a small, manageable clay ball first, then expanding it into a huge statue. This is much faster and produces higher-quality results.

4. The "Custom Request" (Conditional Generation)

Scientists often want to generate cells that have specific traits, like "a lung cell that has been exposed to a virus."

The Analogy: Think of the model as a custom furniture maker.
- Old Models: Could make a chair, but if you asked for a "red chair with a broken leg," it might get confused or make a chair that looked like a table.
- scLDM: Uses a technique called Classifier-Free Guidance. You can tell it, "Make me a cell that is a Lung Cell AND has a Virus." The model learns to blend these instructions perfectly. The paper claims this "joint" way of listening to instructions works better than trying to add them up one by one.

What Did They Prove?

The authors tested scLDM against other top models (like scVI, scDiffusion, and CFGen) on real biological data:

Reconstruction: When asked to "remember" a cell it had seen before, scLDM did a better job than the others, capturing the details more accurately.
Generation: When asked to create new fake cells that look real, scLDM produced data that was statistically closer to real cells than its competitors. It didn't just copy-paste; it created new, valid variations.
Prediction: When used as a tool to classify cells (e.g., "Is this cell infected with COVID-19?"), the "summary sketches" created by scLDM helped the computer make more accurate decisions than older methods.

The Bottom Line

The paper presents a new, flexible, and powerful way to teach computers to understand and create single-cell data. By respecting the fact that gene order doesn't matter and by using a "sculpting from noise" technique on compressed data, scLDM creates more realistic biological simulations than previous methods.

Note: The paper focuses entirely on the computational model's ability to generate and reconstruct data. It does not claim that this model is currently being used to cure diseases or make clinical decisions in hospitals, but rather that it is a superior tool for simulating and understanding cellular biology in a computer.

Technical Summary: Scalable Single-Cell Gene Expression Generation with Latent Diffusion Models

Problem Statement

Computational modeling of single-cell gene expression is essential for understanding cellular heterogeneity, differentiation, and responses to perturbations. However, generating realistic expression profiles remains a significant challenge due to two primary factors: the count-based nature of gene expression data and the complex latent dependencies among genes.

Existing generative models face critical limitations:

Violation of Exchangeability: Many methods impose artificial fixed orderings of genes or rely on restricted subsets of highly variable genes (HVGs). This contradicts the biological reality that gene expression profiles are exchangeable sets where the order of genes carries no intrinsic meaning. Fixed-order architectures tie input dimensions to specific genes, preventing the model from handling varying gene vocabularies across tissues or species without retraining.
Architectural and Training Instabilities: Approaches based on Generative Adversarial Networks (GANs) suffer from training instabilities and mode collapse. Meanwhile, standard Variational Auto-Encoders (VAEs) often rely on shallow architectures or require complex hierarchical latent structures to handle set sizes, which may be unnecessary if parameterization is handled correctly.

Methodology: scLDM

The authors introduce scLDM (scalable Latent Diffusion Model), a fully transformer-based deep generative model designed to respect the fundamental exchangeability property of single-cell data. The framework consists of two main stages:

1. Transformer-Based VAE with Permutation-Invariant/Equivariant Design

The core of scLDM is a VAE architecture that utilizes a Multi-head Cross-Attention Block (MCAB) to handle the unordered nature of gene sets.

Input Processing: To manage computational complexity given the high sparsity of single-cell RNA-seq data (often >70% zeros), the model filters inputs to include only expressed genes and pads the context to a fixed length $d$ . This focuses the model on biologically meaningful signals without restricting the decoder's ability to model zero counts.
Encoder (Permutation-Invariant): The encoder processes the set of gene counts and IDs to produce fixed-size latent variables $Z$ . It employs MCAB where the "queries" are learnable pseudo-inputs (inducing points), and the "keys/values" are derived from the gene embeddings. This design ensures the latent representation is permutation-invariant regardless of the input gene order.
Decoder (Permutation-Equivariant): The decoder reconstructs gene expression counts for a requested set of gene IDs. It uses the same MCAB architecture but modifies the queries to be gene-specific embeddings indexed by the requested gene IDs. This ensures the output is permutation-equivariant, meaning permuting the input gene IDs results in a corresponding permutation of the output parameters.
Likelihood: The model parameterizes the conditional likelihood using either a Negative Binomial (NB) or Gaussian distribution, effectively capturing the overdispersion and count nature of the data.

2. Latent Diffusion Model (LDM) with Flow Matching

Instead of a standard Gaussian prior, the authors replace the prior with a Latent Diffusion Model trained in the latent space of the VAE.

Architecture: The diffusion process is modeled using Diffusion Transformers (DiT).
Training Objective: The model is trained using Flow Matching with linear interpolants, optimizing a loss function that learns a vector field to transport noise to the data distribution.
Controlled Generation: The framework supports multi-conditional generation (e.g., cell type + perturbation) via Classifier-Free Guidance (CFG). The authors propose a joint conditioning strategy where a single token represents a unique combination of attributes, contrasting with additive conditioning approaches used in prior work.

Key Contributions

Exchangeable Architecture: The paper proposes a novel, fully transformer-based VAE that uses a unified MCAB architecture to achieve permutation-invariant pooling in the encoder and permutation-equivariant unpooling in the decoder. This eliminates the need for separate pooling/unpooling components or fixed gene orderings.
Latent Diffusion with Flow Matching: The integration of a Diffusion Transformer-based LDM trained with flow matching and linear interpolants allows for high-quality generation and controlled sampling via classifier-free guidance.
Joint Conditioning: The authors demonstrate that a joint conditioning formulation for classifier-free guidance (treating attribute combinations as single tokens) outperforms additive conditioning strategies in multi-attribute perturbation benchmarks.

Experimental Results

The authors evaluated scLDM on observational datasets (Dentate Gyrus, Tabula Muris, HLCA) and perturbational datasets (Parse 1M, Replogle), comparing against state-of-the-art baselines including scVI, CFGen, scDiffusion, CPA, and scGPT.

Reconstruction Performance: scLDM achieved superior reconstruction metrics (lower Reconstruction Error, higher Pearson Correlation Coefficient, and lower MSE) across all observational datasets compared to scVI and CFGen. For instance, on the HLCA dataset, scLDM achieved a PCC of 0.310 compared to 0.146 for CFGen.
Generative Quality: In unconditional and conditional generation tasks, scLDM demonstrated better distributional matching (lower Wasserstein-2 and MMD2 distances) and a more balanced precision-recall trade-off, avoiding the mode collapse observed in some baselines.
Perturbation Modeling: On the Parse 1M and Replogle datasets, scLDM showed competitive or superior performance in generating gene expression profiles for unseen cell-perturbation combinations. It outperformed baselines in perturbation prediction metrics (e.g., DEG overlap, Spearman correlation of log-fold changes).
Downstream Classification: Embeddings learned by the scLDM-VAE were evaluated on classification tasks (COVID-19 infection detection and cell type classification in Tabula Sapiens 2.0). The 270M parameter scLDM model achieved state-of-the-art F1 scores (0.820) on the COVID-19 task, outperforming other foundation models like TranscriptFormer and scGPT.
Ablation Studies: Experiments confirmed that the zero-padding input strategy is computationally efficient and maintains reconstruction quality. Furthermore, joint conditioning was shown to be superior to additive conditioning for multi-attribute generation.

Significance and Claims

The paper claims that scLDM provides a principled framework for learning exchangeable models of single-cell data, moving beyond the artificial constraints of fixed gene orderings. By combining a permutation-invariant/equivariant transformer architecture with latent diffusion, the model achieves state-of-the-art performance in both reconstructing and generating single-cell expression profiles.

The authors emphasize that their approach naturally supports flexible conditioning and generalization across unseen contexts, enabling faithful modeling of complex cellular responses. They position this work as a step toward foundational models for exchangeable biological data, suggesting the design is applicable beyond transcriptomics to multi-omics and multi-modal data. The work aims to accelerate therapeutic development by enabling in silico perturbation experiments, though the authors note that experimental validation remains necessary for clinical applications.