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

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 natively 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 varying 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 a fetal cell and heart cell atlas. Remarkably, we find that even this simple instantiation matches or surpasses the performance of a state-of-the-art discrete generative model and leading RNA-seq imputation methods, while leaving substantial headroom for further gains through optimized design choices in future work.

• The paper introduces a principled diffusion process over natural numbers using birth-death processes that accurately models count-based data.
• It generalizes noise schedules from Gaussian diffusion to unordered discrete domains, enhancing sample quality and distribution matching.
• CountsDiff demonstrates robust performance in imputation tasks for scRNA-seq, reducing outliers while preserving key statistical properties.

CountsDiff: Diffusion Modeling on the Natural Numbers for Count-Based Data Generation and Imputation

Introduction

The paper "CountsDiff: A Diffusion Model on the Natural Numbers for Generation and Imputation of Count-Based Data" (2604.03779) introduces a substantial advancement in the methodology of generative modeling for count-based data. While diffusion models have become the dominant paradigm for generative modeling in both continuous and categorical domains, effective extensions to ordered, unbounded discrete domains such as $\mathbb{N}_0$ have remained elusive. Traditional approaches are fundamentally limited: categorical diffusion models incur heavy computational burdens and ignore ordinal structure, while continuous relaxations (e.g., mapping counts to $\mathbb{R}$ and later quantizing) fail to meaningfully capture the inherent discreteness and often yield suboptimal sample quality for sparse counts.

CountsDiff formalizes a diffusion process over the natural numbers, grounded in the theory of birth-death processes, and introduces a flexible, extensible framework with forward and reverse dynamics that admit a principled suite of algorithmic augmentations paralleling the contemporary continuous and categorical diffusion literature.

The CountsDiff Framework

CountsDiff builds on birth-death Markov processes to define a forward process that monotonically "thins" counts, and a reverse process that probabilistically "unthins" towards the data manifold. The forward process is parameterized directly via the survival probability schedule $p(t)$, corresponding to the probability that a count survives up to time $t$. This parameterization is an explicit generalization of Blackout Diffusion, but with inclusion of continuous-time training, advanced loss weighting strategies, classifier-free guidance, and new reverse process modifications (attrition or "remasking") allowing nonmonotone denoising steps.

Figure 1: CountsDiff's forward process (pure-death, top) and generalized birth-death reverse process (bottom).

The forward marginal at time $t$ is:

$$q(x_t \mid x_0) = \binom{x_0}{x_t} p(t)^{x_t} (1-p(t))^{x_0-x_t}$$

These binomial marginals are analytically tractable and allow tight theoretical guarantees on the existence and expressiveness of the process over $\mathbb{N}_0$. Remarkably, all noise schedules from Gaussian diffusion have explicit analogs in CountsDiff, providing direct transfer of decades of empirical findings.

Reverse Process, Attrition, and Guidance

CountsDiff's reverse process reproduces the birth process, but is extended to include "attrition," i.e., a controlled death rate during sampling, generalizing beyond strictly monotonic "densification." The model framework proves that, irrespective of the attrition schedule, the final sample marginals remain exact. This insight is especially potent for introducing analogues of churn (continuous) and remasking (discrete) procedures, enabling enhanced diversity, improved sample quality, and mitigation of local mode collapse.

Further, the model incorporates predictor-free classifier guidance by leveraging both unconditional and conditional heads in a shared neural parameterization, following the approach in conditional diffusion literature but generalized to count-valued domains. The guidance mechanism directly modulates the birth process rates.

Rounding and Implementation Details

Diffusion models universally face the problem of reconciling real-valued network outputs with discrete-valued reconstructions at generation time. CountsDiff dispenses with naive (and demonstrably lossy) rounding strategies by employing a randomized rounding scheme that preserves the expectation and prevents the pathological collapse to zero, especially significant in sparse regimes.

Empirical Results

Toy Count Vector Modeling

The ability of CountsDiff to model sparse, high-dimensional count data is demonstrated in comparison with Gaussian and masked (categorical) diffusion on synthetic negative binomial data. CountsDiff faithfully preserves joint and marginal distributions, with significant improvements over categorical models in sliced Wasserstein distance and avoidance of excessive variance/“hallucinated” outliers, while Gaussian diffusion collapses and fails to capture discreteness.

Figure 2: Histogram and distributional statistics of model-generated samples versus ground truth for simulated counts.

Conditional Image Generation

Trained on CIFAR-10 and CelebA, CountsDiff achieves FID and IS scores commensurate with, or exceeding, state-of-the-art discrete generative models (e.g., Blackout Diffusion, ReMDM), while maintaining the scaling and expressivity necessary for high-dimensional image modeling. The effectiveness of class-conditional guidance is visually and quantitatively evident, with optimal guidance and attrition hyperparameters improving both diversity and fidelity.

Figure 3: Class-conditional samples from CountsDiff on CIFAR-10 with guidance scaling.

Figure 4: Unconditional samples illustrating the impact of increasing attrition/churn on output smoothness and texture.

scRNA-seq Imputation

CountsDiff demonstrates surprising strength and stability as a generative imputation method for single-cell RNA sequencing data under both MCAR and MNAR missing data regimes. The model surpasses or matches leading methods, including ReMDM and scIDPMs, in log(scFID), RMSE, bias, and rank correlation, with lower propensity to produce outlier artifacts. Notably, CountsDiff achieves these results with fewer parameters and improved stability, suggesting overfitting and model collapse are alleviated relative to alternatives.

Theoretical and Practical Implications

CountsDiff elucidates the alignment between noise schedule, weighting strategies, and reverse process flexibility—mirroring and generalizing the design space of continuous and categorical diffusion approaches. This enables systematic transfer of algorithmic advancements and theoretical tools from one domain to another (e.g., direct translation of Gaussian noise schedules to count-space via signal-to-noise calibration). The explicit demonstration that all modifications to the birth-death reverse process (i.e., any attrition schedule) preserve final sample marginals directly motivates further exploration of adaptive reverse process manipulation and meta-learning.

Practically, the model is immediately applicable to biological count data, including large-scale genomics (e.g., scRNA-seq, ATAC-seq) and offers a foundation for future extensions into structured counts (e.g., spatial transcriptomics, multiomics). The design space—encompassing $p$-scheduling, adaptive guidance, randomized rounding, and attrition—invites further empirical optimization and theoretical analysis.

Future Directions

Several avenues for future work are made apparent:

• Extension to Count-Transformed Nonnegative Reals: By integrating the Poisson data randomization technique from JUMP, CountsDiff can accommodate continuous, nonnegative domains with exact NLL optimization.
• Value-Dependent Attrition: Flexible, possibly data-dependent attrition schedules could adaptively improve mode coverage and mitigate outlier generation in scientific domains where overfitting outliers is especially damaging.
• Fast Solvers and Continuous-Time ODE/SDE Sampling: The continuous parameterization unlocks compatibility with adaptive solvers, enabling rapid and high-fidelity (potentially deterministic) sample generation.
• Hierarchical/Bayesian Extensions: The analytic tractability of CountsDiff opens space for integration with full probabilistic modeling frameworks, possibly facilitating Bayesian uncertainty quantification in imputation and generation.
• Foundation Models for Cellular Data: As evidenced in the single-cell setting, CountsDiff's architecture and representational capacity position it as a candidate backbone for foundation modeling in systems biology.

Conclusion

CountsDiff presents a complete framework for diffusion modeling directly in the space of natural numbers. The elucidation of a flexible, interpretable, and theoretically principled design space unifies previously disparate approaches and overcomes structural failures of both categorical and real-valued adaptations. The empirical results across synthetic, imaging, and biological data domains support its status as a competitive and robust model for count-based generation and imputation. The CountsDiff framework stands poised as a foundation for further developments in generative modeling for discrete, ordinal, and count-valued data.