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

Renzo G. Soatto; Anders Hoel; Greycen Ren; Shorna Alam; Stephen Bates; Nikolaos P. Daskalakis; Caroline Uhler; Maria Skoularidou

cs.LGApr 4, 2026

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

Authors: Renzo G. Soatto, Anders Hoel, Greycen Ren, Shorna Alam, Stephen Bates, Nikolaos P. Daskalakis, Caroline Uhler, Maria Skoularidou

Organizations: Massachusetts Institute of Technology (MIT), Cambridge, MA, USA · Laboratory for Information and Decision Systems (LIDS), MIT, Cambridge, MA, USA · Eric and Wendy Schmidt Center, Broad Institute, Cambridge, MA, USA · Boston University, Boston, MA, USA

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.

Figures & tables

Figure 1: Visualization of CountsDiff’s forward corruption process (top) and reverse sampling process (bottom). The top diagram depicts the progression of a

p

-schedule, a pure death process. The bottom shows a single step of the generalized, birth-death sampling process.

Figure 2: Histogram of model-generated samples versus ground truth and distributional distance metrics (top); variance statistics (bottom) for a subset of dimensions. Existing diffusion models exhibit failure cases even in a simple toy dataset: Gaussian diffusion suffers from mode collapse, while masked diffusion overfits outliers (inflated variance). Full results for all ten dimensions can be found in Appendix E.1 .

Figure 3: 5 images guided by each Cifar10 class sampled from CountsDiff with guidance scale $2.0$ and $\eta_{\text{rescale}}=0.005$ .

Figure 4: Nine images drawn from CountsDiff trained on CIFAR-10 with $\eta_{\text{rescale}}$ attrition schedule for varying levels of $\eta_{\text{{rescale}}}$

$p$ -schedule	$\gamma$	$\eta_{\mathrm{rescale}}$	FID $\downarrow$	IS $\uparrow$
FI Discrete	uncond	none	$5.73$	$9.12\pm 0.05$
FI Continuous	uncond	none	$5.44$	$9.09\pm 0.16$
FI Continuous	$1.0$	$0.01$	5.20	$9.64\pm 0.17$
FI Continuous	$2.0$	$0.02$	$9.71$	$9.77\pm 0.11$
Cosine	uncond	none	$5.76$	$9.29\pm 0.18$
Cosine	$1.0$	$0.01$	$5.26$	$9.85\pm 0.08$

Table 1: FID and IS of 50k images sampled with from CountsDiff trained on CIFAR-10 for a selected set of sampling hyperparameters. FI Discrete with unconditional generation and no attrition is equivalent to Blackout ( Santos et al., 2023 )

	Pointwise			Distributional
Method	Spearman $\uparrow$	RMSE $\downarrow$	Bias	ED $\downarrow$	log(scFID) $\downarrow$	log(MMD) $\downarrow$	SWD $\downarrow$
Zero imputation	N/A	1.91(0.01)	-1.31(0.00)	1.44(0.00)	-2.35(0.00)	-4.08(0.00)	0.17(0.00)
Mean imputation	0.17(0.00)	1.31(0.04)	0.00(0.00)	0.17(0.00)	-5.01(0.01)	-7.03(0.00)	0.12(0.01)
Conditional Mean	0.20(0.00)	1.12(0.01)	0.00(0.00)	0.15(0.00)	-6.23(0.01)	-7.49(0.01)	0.08(0.00)
MAGIC	0.21(0.00)	1.88(0.01)	-1.31(0.00)	1.44(0.00)	-2.35(0.00)	-4.07(0.00)	0.17(0.00)
scIDPMs, 1-sample	0.08(0.00)	2.25(0.02)	0.86(0.00)	0.68(0.00)	-3.04(0.01)	-3.39(0.00)	0.20(0.00)

Table 2: Benchmarking on scRNA-seq imputation, human fetus cell atlas with 50% MCAR. Mean (standard error). Methods are grouped into three categories: naive baseline (top), imputation/scRNAseq-specific (middle), and general generative (bottom). Best performance in each category for each metric is bolded, and second best is italicized.

	Pointwise			Distributional
Method	Spearman $\uparrow$	RMSE $\downarrow$	Bias	ED $\downarrow$	log(scFID) $\downarrow$	log(MMD) $\downarrow$	SWD $\downarrow$
Zero imputation	N/A	1.00(0.00)	-0.99(0.00)	1.41(0.00)	-5.29(0.01)	-6.64(0.00)	0.03(0.00)
Mean imputation	0.26(0.00)	0.51(0.00)	0.32(0.00)	0.12(0.00)	-4.81(0.01)	-7.99(0.01)	0.02(0.00)
Conditional Mean	0.36(0.00)	0.47(0.00)	0.28(0.00)	0.12(0.00)	-5.82(0.01)	-8.69(0.01)	0.01(0.00)
MAGIC	0.09(0.02)	1.00(0.00)	-0.99(0.00)	1.41(0.00)	-5.28(0.00)	-6.63(0.00)	0.03(0.00)
scIDPMs, 1-sample	0.03(0.00)	2.22(0.01)	1.20(0.00)	0.98(0.00)	-3.61(0.00)	-4.20(0.00)	0.14(0.00)

Table 3: Benchmarking on scRNA-seq imputation, human fetus cell atlas with 25% low-biased missingness (MNAR). Mean (standard error). Methods are grouped into three categories: naive baseline (top), imputation/scRNAseq-specific (middle), and general generative (bottom). Best performance in each category for each metric is bolded, and second best is italicized.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Converted

p

-schedule from Blackout Diffusion (see B.2 ) versus cosine

p

-schedule

Figure 6: Weights from Blackout Diffusion versus proposed

p

-schedule

	Real Data	CountsDiff	Gaussian	Masked
Dimension	(Ground Truth)	( Ours )
Dim 0	0.78	0.55	0.19	3.06
Dim 1	4.71	1.99	0.46	9.22
Dim 2	0.10	0.07	0.01	1.89
Dim 3	0.12	0.08	0.02	2.49
Dim 4	0.28	0.21	0.05	7.27

Appendix

Table 4: Variance of generated samples per dimension. CountsDiff (Ours) consistently maintains variance closer to the Real Data, whereas Gaussian Diffusion collapses (near-zero variance) and masked diffusion suffers from extreme outliers (high variance).

Panel of Figure 7 — Figure 7: Histograms of marginals of dimensions 2-9

Panel of Figure 8 — Figure 8: Joint Kernel Density Estimate (KDE) plots between dimensions 0 and 1 of real data (blue contours) versus model-generated samples (red dashed contours). Gaussian Diffusion suffers from mode collapse, resulting in a much tighter, less diverse distribution. Masked Diffusion exhibits a broader, more diffuse distribution with ’leaked’ probability mass and slightly less correlation between the dimensions, indicating outliers and overfitting to the marginals. CountsDiff (Ours) more closely aligns with the true data distribution.

Panel of Figure 9 — Figure 9: Comparison of Binomial sampling standard rounding, Poisson with no rounding, and Binomial sampling with stochastic rounding across two different counts regimes.

Panel of Figure 10 — Figure 10: Train losses and validation metrics throughout training

Table 5: FID and IS of 5k images sampled with from CountsDiff with the Continuous FI (left) and Cosine (right)

p

-schedules at various guidance scales with

\eta_{\mathrm{rescale}}=0.01

Attrition Rate Strategy	FID $\downarrow$	IS $\uparrow$
None	$9.666$	$9.463\pm 0.402$
$\eta_{rescale}=0.005$	$9.562$	$9.504\pm 0.382$
$\eta_{rescale}=0.01$	$9.507$	$9.561\pm 0.368$
$\eta_{rescale}=0.02$	$10.400$	$9.653\pm 0.269$
$\eta_{rescale}=0.05$	$12.727$	$9.565\pm 0.268$

Appendix

Table 6: FID and IS of 5k images sampled with from CountsDiff various attrition rate strategies

$p$ -schedule	$\gamma$	attrition schedule	FID $\downarrow$
FI Continuous	$1.0$	$\eta_{\text{rescale}}=0.01$	$9.844$
Cosine Continuous	$1.0$	$\eta_{\text{rescale}}=0.01$	$\mathbf{7.637}$

Appendix

Table 7: FID of 10k images sampled with from

30

M-parameter CountsDiff model trained on CelebA for

500

k steps

$p$ -schedule	$\gamma$	attrition schedule	FID $\downarrow$
Cosine Continuous	$1.5$	$\eta_{\text{rescale}}=0.005$	$4.948$

Appendix

Table 8: FID of 50k images sampled with from

60

M-parameter CountsDiff model trained on CelebA for

1

M steps

$p$ -schedule	$\gamma$	attrition schedule	FID $\downarrow$
Cosine Continuous	$1.0$	$\eta_{\text{rescale}}=0.005$	$7.580$
Cosine Continuous	$1.0$	$\eta_{\text{rescale}}=0.01$	$7.541$
Cosine Continuous	$1.5$	$\eta_{\text{rescale}}=0.005$	$7.217$
Cosine Continuous	$1.5$	$\eta_{\text{rescale}}=0.01$	$7.483$

Appendix

Table 9: FID of 5k images sampled with from

60

M-parameter CountsDiff model trained on CelebA for

1

M steps

Figure 11: 25 images drawn from CountsDiff trained on CelebA with increasing $\eta_{rescale}$ .

	Pointwise			Distributional
Method	Spearman $\uparrow$	RMSE $\downarrow$	Bias	ED $\downarrow$	log(scFID) $\downarrow$	log(MMD) $\downarrow$	SWD $\downarrow$
Zero imputation	N/A	9.27(0.13)	-3.49(0.01)	1.73(0.00)	-0.32(0.00)	-4.50(0.00)	1.43(0.03)
Mean imputation	0.31(0.00)	8.32(0.17)	0.03(0.01)	0.87(0.00)	-2.89(0.00)	-5.19(0.01)	1.25(0.04)
Conditional Mean	0.49(0.00)	7.62(0.19)	-0.01(0.01)	0.39(0.00)	-4.29(0.01)	-6.65(0.01)	1.11(0.06)
MAGIC	0.39(0.00)	9.24(0.18)	-3.45(0.01)	1.68(0.00)	-0.38(0.00)	-4.50(0.00)	1.47(0.05)
scIDPMs, 1-sample	0.31(0.00)	8.89(0.11)	1.55(0.01)	0.76(0.00)	-3.16(0.01)	-5.11(0.01)	0.88(0.04)

Appendix

Table 10: Benchmarking on scRNA-seq imputation, heart with 50% MCAR. Mean (standard error). Methods are grouped into three categories: naive baseline (top), scRNAseq/imputation-specific (middle), and general generative (bottom). Best performance in each category for each metric is bolded, and second best is italicized.

Table 11: Attrition sweep across schedules. Metrics are reported as mean (standard error)

Table 12: Guidance sweep across schedules. Metrics are reported as mean (standard error)

CountsDiff Sampling Steps	scFID $\times 10^{4}$ $\downarrow$	ED $\downarrow$
30	$2.73(0.19)$	$0.01(0.00)$
20	$2.55(0.28)$	$0.01(0.00)$
15	$2.32(0.18)$	$0.01(0.00)$
10	$2.35(0.15)$	$0.02(0.00)$
7	$2.58(0.18)$	$0.02(0.00)$
5	$2.97(0.20)$	$0.03(0.00)$

Appendix

Table 13: Effect of number of sampling steps on generation quality, measured by scFID and Energy distance (ED) between imputed and ground-truth test set samples. Metrics are reported as mean (standard error)

	Sample-level						Distributional
Method	R ${}^{2}\uparrow$	RMSE $\downarrow$	MAE $\downarrow$	Bias	Spearman $\uparrow$	Pearson $\uparrow$	ED $\downarrow$	log(scFID) $\downarrow$	log(MMD) $\downarrow$	SWD $\downarrow$
Zero imputation	-7.84(0.01)	1.91(0.01)	1.31(0.00)	-1.31(0.00)	N/A	N/A	1.44(0.00)	-2.35(0.00)	-4.08(0.00)	0.17(0.00)
Mean imputation	-0.86(0.00)	1.31(0.04)	0.48(0.00)	0.00(0.00)	0.17(0.00)	0.25(0.00)	0.17(0.00)	-5.01(0.01)	-7.03(0.00)	0.12(0.01)
Conditional Mean	-0.75(0.01)	1.12(0.01)	0.45(0.00)	0.00(0.00)	0.20(0.00)	0.31(0.00)	0.15(0.00)	-6.23(0.01)	-7.49(0.01)	0.08(0.00)
MAGIC	-7.82(0.01)	1.88(0.01)	1.31(0.00)	-1.31(0.00)	0.21(0.00)	0.52(0.01)	1.44(0.00)	-2.35(0.00)	-4.07(0.00)	0.17(0.00)
scIDPMs, 1-sample	-22.26(0.16)	2.25(0.02)	1.18(0.00)	0.86(0.00)	0.08(0.00)	0.13(0.00)	0.68(0.00)	-3.04(0.01)	-3.39(0.00)	0.20(0.00)

Appendix

Table 14: Full metrics (fetus, 50% MCAR).

Model Method	Accuracy $\uparrow$	F1 $\uparrow$
Zero imputation	$0.82(0.00)$	$0.53(0.01)$
Mean imputation	$0.81(0.00)$	$0.49(0.01)$
Conditional Mean	$0.82(0.00)$	$0.55(0.01)$
MAGIC	$0.62(0.00)$	$0.27(0.01)$
ReMDM, 1-sample	$0.82(0.00)$	$0.53(0.01)$
ReMDM, 5-samples	$0.82(0.00)$	$0.53(0.01)$

Appendix

Table 17: Cell classifier results for scRNA-seq imputation on human fetal cell atlas with 50% MCAR. Metrics are reported as mean (standard error)

Model Method	Accuracy $\uparrow$	F1 $\uparrow$
Zero imputation	$0.82(0.00)$	$0.54(0.01)$
Mean imputation	$0.81(0.00)$	$0.51(0.01)$
Conditional Mean	$0.82(0.00)$	$0.53(0.01)$
MAGIC	$0.76(0.00)$	$0.47(0.01)$
ReMDM, 1-sample	$0.82(0.00)$	$0.54(0.01)$
ReMDM, 5-samples	$0.82(0.00)$	$0.54(0.01)$

Appendix

Table 18: Cell classifier results for scRNA-seq imputation on human fetal cell atlas with 25% low-biased MNAR. Metrics are reported as mean (standard error).

Model Method	Accuracy $\uparrow$	F1 $\uparrow$
Zero imputation	$0.99(0.00)$	$0.94(0.01)$
Mean imputation	$0.99(0.00)$	$0.93(0.01)$
Conditional Mean	$0.99(0.00)$	$0.95(0.01)$
MAGIC	$0.93(0.00)$	$0.79(0.01)$
ReMDM, 1-sample	$0.99(0.00)$	$0.94(0.01)$
ReMDM, 5-samples	$0.99(0.00)$	$0.940.01)$

Appendix

Table 19: Cell classifier results for scRNA-seq imputation on human heart cell atlas with 50% MCAR. Metrics are reported as mean (standard error)

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

CRNDiff: Count-Native Diffusion Framework via Chemical Reaction Networks
Sep 25, 2026Yuxuan Qiu, Praful Gagrani, Tetsuya J KobayashiDiffusion Models Conditional Diffusion Models
Institute of Industrial Science The University of Tokyo Tokyo 153-8505, Japan
Flow Matching for Count Data
May 8, 2026Ganchao Wei, John PearsonFlow Matching Synthetic Data Generation
Department of Neurobiology Department of Statistical Science Duke University Durham, NC, USA · Department of Neurobiology Department of Electrical and Computer Engineering Duke University Durham, NC, USA
Jumping up and down: Denoiser diffusion models for discrete ordinal data
Oct 2, 2026Yair Shenfeld, Ricardo Baptista, Stefano PeluchettiDiffusion Models
Division of Applied Mathematics Brown University Providence, RI, USA 02912 · Department of Statistical Sciences University of Toronto Toronto, ON, Canada M5G 1X6 · Sakana AI Tokyo, Japan

Sep 25, 2026J/K move · Enter open · S save

CRNDiff: Count-Native Diffusion Framework via Chemical Reaction Networks

Yuxuan Qiu, Praful Gagrani, Tetsuya J Kobayashi

Institute of Industrial Science The University of Tokyo Tokyo 153-8505, Japan

Diffusion Models Conditional Diffusion Models Diffusion Model Guidance Single-Cell RNA Sequencing

Open paper AskPDF arXiv

Scientific measurements such as single-cell RNA (scRNA) sequencing often take the form of nonnegative integer counts, whereas continuous-state diffusion models approximate this discrete structure using continuous coordinates. Building on stochastic chemical reaction networks (CRNs), a class of count-native Markov jump processes, we introduce CRNDiff, a structured framework that combines count-space diffusion with inference-time conditioning on rare subpopulations. An independent birth--death instantiation yields a closed-form transition kernel for forward noising. This kernel enables reverse sampling via forward-filtering backward-sampling (FFBS) and supports data-driven selection of the terminal noising time, eliminating the need for a validation sweep. This tractability also lets us introduce tilted Feynman--Kac (FK) steering, a method for sampling target subpopulations from a frozen generator without retraining. By tilting posterior marginals before FK particle correction, steering mitigates importance-weight concentration when the target population is rare. Using scRNA-seq data from the human heart cell atlas, we test the ability of CRNDiff to generate cell-type-specific distributions. Across the three evaluated target populations, CRNDiff achieves the highest conditional fidelity among the evaluated generative models, with larger mean purity margins for rarer target populations. Generated cells preserve marker-level differential-expression structure. Replacing real training cells for the target classes with generated cells yields downstream classification performance approaching that of the real-data reference.

Abstract

Diffusion Models

Explore similar work

CRNDiff: Count-Native Diffusion Framework via Chemical Reaction Networks

Flow Matching for Count Data

Jumping up and down: Denoiser diffusion models for discrete ordinal data