arXiv:2604.03779cs.LGcs.AI2026-04

用扩散模型生成和补全计数数据,突破传统方法局限。

CountsDiff: A Diffusion Model on the Natural Numbers for Generation and Imputation of Count-Based Data

论文配图:CountsDiff: A Diffusion Model on the Natural Numbers for Generation and Imputation of Count-Based Data
图 1 · 摘自论文原文
  • 基于生存概率调度直接建模自然数分布,设计更灵活。
  • 在图像和单细胞测序数据上表现优于现有方法。
  • 适合生物计数数据补全,如胎儿与心脏细胞图谱分析。

扩散模型在连续和基于符号的数据生成中表现优异,但在离散序数数据上的应用仍不充分。我们提出CountsDiff,一种针对自然数分布的扩散框架。该模型通过直接参数化生存概率调度并显式加权损失,简化了Blackout扩散框架,引入可调控的设计参数,与现有扩散模型保持一致。此外,CountsDiff融入现代扩散模型特性:连续时间训练、无分类器引导及允许非单调反向轨迹的扰动/掩码反向动态。我们初步实现该框架,在CIFAR-10和CelebA等自然图像数据集上验证设计参数影响,展示其在复杂且可解释数据域中的有效性。随后,以单细胞RNA-seq计数数据补全为典型应用场景,评估其在胎儿与心脏细胞图谱上的表现。结果表明,即使当前简单版本已达到或超过先进离散生成模型及主流scRNA-seq补全方法性能,未来优化设计仍有显著提升空间。

原文摘要 · Abstract (English)

Diffusion models have excelled at generative tasks for both continuous and token-based domains, but their application to discrete ordinal data remains underdeveloped. We present CountsDiff, a diffusion framework designed to model distributions on the natural numbers. CountsDiff extends the Blackout diffusion framework by simplifying its formulation through a direct parameterization in terms of a survival probability schedule and an explicit loss weighting. This introduces flexibility through design parameters with direct analogues in existing diffusion modeling frameworks. Beyond this reparameterization, CountsDiff introduces features from modern diffusion models, previously absent in counts-based domains, including continuous-time training, classifier-free guidance, and churn/remasking reverse dynamics that allow non-monotone reverse trajectories. We propose an initial instantiation of CountsDiff and validate it on natural image datasets (CIFAR-10, CelebA), exploring the effects of the introduced design parameters in a complex, well-studied, and interpretable data domain. We then highlight biological count assays as a natural use case, evaluating CountsDiff on single-cell RNA-seq imputation in fetal and heart cell atlases. Remarkably, we find that even this simple instantiation matches or surpasses the performance of a state-of-the-art discrete generative model and leading scRNA-seq imputation methods, while leaving substantial headroom for further gains through optimized design choices in future work.

扩散模型计数数据单细胞生成模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。