CountsDiff: A Diffusion Model on the Natural Numbers for Generation and Imputation of Count-Based Data
CountsDiff models natural number distributions with a novel p(t) schedule and non-monotonic reverse dynamics, excelling in image and scRNA-seq data imputation.
Key Findings
Methodology
CountsDiff builds on Blackout diffusion by introducing a direct parameterization via a survival probability schedule p(t) and explicit loss weighting. It incorporates continuous-time training, classifier-free guidance, and non-monotonic reverse dynamics with death/reset mechanisms, enabling flexible modeling of count distributions. The forward process is a pure death process controlled by p(t), while the reverse process allows non-monotonic trajectories through attrition parameters σt,s. Randomized rounding ensures outputs are natural numbers. Validated on synthetic counts, CIFAR-10, CelebA, and scRNA-seq datasets, it outperforms or matches state-of-the-art discrete generative models.
Key Results
- On synthetic sparse count data, CountsDiff achieved lower MMD and Wasserstein distances compared to Gaussian and masked diffusion models, accurately capturing the distribution's sparsity and correlations.
- In image generation, guidance and non-monotonic reverse dynamics improved FID to as low as 2.8 and IS to 8.4 on CIFAR-10, demonstrating high fidelity and diversity.
- In scRNA-seq imputation tasks, CountsDiff reduced reconstruction errors by over 20%, outperforming existing methods like MAGIC, GAIN, and masked autoencoders, especially for low-expression genes.
Significance
This work advances the application of diffusion models to discrete count data, overcoming the limitations of previous categorical or continuous approaches. By enabling non-monotonic reverse trajectories and flexible scheduling, CountsDiff provides a powerful, interpretable framework for biological data imputation and high-dimensional generative tasks, bridging a critical gap in generative modeling for natural and biological counts.
Technical Contribution
The paper introduces a p(t)-parameterized pure death process for the forward diffusion, continuous-time training, and non-monotonic reverse dynamics with attrition. It employs randomized rounding to maintain integer outputs, and guidance techniques adapted for natural numbers. These innovations collectively enhance modeling flexibility, stability, and performance in count-based domains, expanding the scope of diffusion models beyond traditional continuous and categorical data.
Novelty
CountsDiff is the first to model natural number distributions using a p(t)-scheduled pure death process with non-monotonic reverse trajectories, integrating modern diffusion techniques like continuous training and guidance. Its non-monotonic reverse process, combined with randomized rounding, sets it apart from prior discrete models, enabling more accurate and flexible count data generation.
Limitations
- The model's performance degrades in extremely sparse or high-noise scenarios, and training remains computationally intensive. Its generalization to very large counts or multi-modal distributions needs further exploration.
- Parameter tuning for attrition and guidance requires empirical adjustment, limiting ease of use. Future work should automate hyperparameter selection and improve sampling speed.
- Theoretical guarantees for convergence and stability under all settings are still under development, especially for complex biological data.
Future Work
Future directions include optimizing training and sampling efficiency, extending the framework to multi-modal and larger count ranges, and integrating multi-omics data for comprehensive biological modeling. Developing theoretical guarantees and interpretability tools will further enhance practical deployment in biomedical research and high-dimensional data synthesis.
AI Executive Summary
CountsDiff represents a significant leap in diffusion modeling for discrete count data, especially natural numbers. Traditional diffusion models excel in continuous or categorical spaces but struggle with sparse, ordinal, and high-dimensional count data common in biology and image synthesis. CountsDiff addresses these challenges by introducing a p(t)-controlled pure death process as the forward diffusion, enabling a flexible and interpretable framework.
The core innovation lies in the non-monotonic reverse process, which incorporates death and reset mechanisms, allowing trajectories to overshoot or correct errors dynamically. This flexibility is achieved through attrition parameters and guided by modern techniques like classifier-free guidance, resulting in high-quality, diverse samples. The model's continuous-time training further enhances stability and scalability.
Extensive experiments validate CountsDiff’s effectiveness. On synthetic datasets, it captures sparse distributions better than Gaussian or categorical models. In image generation tasks, it surpasses baselines with FID scores as low as 2.8 and high diversity. Most notably, in biological applications such as scRNA-seq imputation, it outperforms state-of-the-art methods, reducing error rates significantly.
These results demonstrate the model’s broad applicability, from high-fidelity image synthesis to accurate biological data reconstruction. Its interpretability and flexibility open new avenues for research in generative modeling of natural and biological counts. Future work will focus on efficiency, scalability, and multi-modal integration, promising transformative impacts across AI and biomedical fields.
Deep Analysis
Background
扩散模型在连续空间和类别空间已取得显著成功，代表算法如Ho et al. (2020)的DDPM和Austin et al. (2021)的掩码扩散，推动了图像、音频等多模态生成技术的发展。然而，生物计数数据（如scRNA-seq）具有稀疏、离散、非负的特性，传统模型难以有效建模。近年来，Blackout Diffusion等尝试在离散空间引入噪声，但存在单调逆动态限制，限制了模型的表达能力。当前研究亟需一种既符合自然数特性，又能处理稀疏性和高维复杂性的扩散框架，成为该领域的研究热点。
Core Problem
核心问题在于如何在自然数空间中有效建模稀疏、离散、非负的计数数据。现有方法多采用类别离散化或连续空间逼近，存在信息损失或效率低下的问题。类别模型难以捕获序性，连续模型在离散空间表现不佳。尤其在高维生物数据中，模型需兼顾稀疏性、多样性和生成质量。设计一个既符合自然数特性，又具有高效训练和采样能力的扩散框架，成为亟待解决的难题。
Innovation
本研究的创新点包括：1）引入p(t)调度的纯死亡过程，简化模型表达并增强可解释性；2）结合连续时间训练，提升模型的稳定性和泛化能力；3）逆过程引入非单调性（死亡/复位机制），允许轨迹非单调变化，避免模式崩溃；4）随机舍入策略，确保输出为自然数且避免模式崩溃。这些创新共同赋予模型更强的表达能力和灵活性，特别适合稀疏离散数据。
Methodology
- �� 设计p(t)调度的纯死亡过程作为前向扩散，定义生存概率p(t)控制信息丢失；
- �� 采用连续时间训练，利用Kolmogorov前向方程优化模型参数；
- �� 逆过程引入死亡和复位（非单调逆动态），通过参数σt,s调节逆轨迹的非单调性；
- �� 利用神经网络预测剩余计数，结合随机舍入确保输出为自然数；
- �� 引入无分类引导机制，实现条件采样；
- �� 采用重绘（RePaint）策略进行缺失值插补，保持模型在实际任务中的适用性。
Experiments
在合成稀疏计数数据上验证模型的生成能力，比较MMD和Wasserstein距离，显示CountsDiff优于Gaussian和掩码模型。在图像任务中，使用CIFAR-10和CelebA，调节引导和逆动态参数，获得FID最低2.8，IS最高8.4。生物数据中，评估scRNA-seq缺失值插补，误差明显低于传统方法和其他扩散模型，验证其在实际应用中的潜力。
Results
CountsDiff在合成数据中成功捕获稀疏分布，指标优于对比模型。在图像生成中，结合引导和逆动态调节，显著提升FID和IS。在生物数据中，模型实现了高精度的基因表达重建，误差降低20%以上，优于现有最优方法。模型还展现出良好的泛化能力和调节灵活性，验证了其广泛适用性。
Applications
模型适用于生物计数数据的缺失值插补、基因表达模拟、稀疏数据生成等场景。其灵活的逆动态设计也适合高维图像、视频等复杂数据的生成任务。未来可结合多模态信息，推动个性化医疗、虚拟现实等行业的发展，改善数据缺失和生成效率。
Limitations & Outlook
模型在极端稀疏或高噪声环境下表现仍有限，训练速度较慢，逆动态参数调优依赖经验。未来需优化算法，提高效率，增强泛化能力，扩展到更大规模和多模态数据中。
Plain Language Accessible to non-experts
想象你在一家工厂里，工厂每天都在生产不同的商品。每次生产的商品数量可能很少，甚至为零。传统的工厂模型就像用一个固定的流程生产商品，容易出现生产过多或过少的问题。CountsDiff就像给工厂装上了智能调节器，它可以根据需要灵活调整生产流程，确保每次生产的商品数量合理。它用一种特别的方法，逐步“打乱”生产过程，然后再“逆转”回来，确保每个商品的数量都符合实际需求。这个过程可以用在画画、基因分析等很多场景，就像工厂可以生产不同的商品一样。它还能在数据缺失时帮忙补全，就像工厂在缺料时自动补充。这个方法让数据生成变得更智能、更灵活，也更贴近实际应用。
ELI14 Explained like you're 14
想象你在玩一个猜糖果的游戏，你的朋友藏了多少糖果。你不知道确切的数字，但可以通过一些线索逐步缩小范围。CountsDiff就像一个聪明的助手，它可以帮你逐步猜出正确的糖果数。它会先“打乱”糖果的数量，让你看不清，然后再慢慢“还原”到真实的数字。这个过程就像拼图游戏，先打乱碎片，再拼回完整的样子。它还能在你不知道的情况下帮你补全缺失的糖果，就像朋友偷偷藏了一些。这个助手不仅在猜糖果时厉害，还能帮你在画画或分析基因时做类似的事情。它让复杂的数字变得简单又有趣，就像魔法一样！
Abstract
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.