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.

翻译：单细胞基因表达的计算建模对理解细胞过程至关重要，但生成逼真的表达谱仍是一项重大挑战。该难点源于基因表达数据的计数特性以及基因之间复杂的潜在依赖关系。现有生成模型往往强加人工基因排序或依赖浅层神经网络架构。我们提出了一种针对单细胞基因表达数据的可扩展潜在扩散模型，称为scLDM，该模型尊重数据的基本可交换性属性。我们的变分自编码器采用固定大小的潜在变量，利用统一多头交叉注意力模块架构，该架构具有双重作用：编码器中的置换不变池化与解码器中的置换等变反池化。我们通过使用扩散变换器和线性插值方法将高斯先验替换为潜在扩散模型来增强此框架，从而支持多条件无分类器引导的高质量生成。我们展示了该模型在观测和扰动单细胞数据的各类实验以及细胞级分类等下游任务中的优越性能。