CRNDiff: Count-Native Diffusion Framework via Chemical Reaction Networks

Yuxuan Qiu; Praful Gagrani; Tetsuya J Kobayashi

cs.LGSep 25, 2026

CRNDiff: Count-Native Diffusion Framework via Chemical Reaction Networks

Authors: Yuxuan Qiu, Praful Gagrani, Tetsuya J Kobayashi

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

Abstract

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.

Figures & tables

Figure 1

Endothelial abundant , $80{,}463$ training cells
Method	Purity $\uparrow$	$W_{1}$ $\downarrow$	MMD 2 $\downarrow$	PCC $\uparrow$
Real cells	$0.944$	$0.011$	$0.0001$	$1.000$
CRNDiff	$\mathbf{0.906}_{\pm.005}$	$\mathbf{0.064}_{\pm.003}$	$\mathbf{0.0011}_{\pm.0001}$	$\mathbf{0.998}_{\pm.000}$
MDLM ( Sahoo et al., 2024 )	$0.898_{\pm.026}$	$0.078_{\pm.029}$	$0.0058_{\pm.0010}$	$0.991_{\pm.003}$
CFGen ( Palma et al., 2025 )	$0.797_{\pm.009}$	$0.102_{\pm.001}$	$0.0036_{\pm.0000}$	$0.993_{\pm.000}$
scANVI ( Xu et al., 2021 )	$0.857_{\pm.016}$	$0.079_{\pm.007}$	$0.0016_{\pm.0001}$	$0.997_{\pm.000}$

Table 2: Conditional fidelity for endothelial, myeloid, and neuronal targets. Distributional metrics use held-out target cells from the test set as the real-data reference; purity uses fixed CellTypist target-label mappings. Boldface indicates the best generative-model mean.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Terminal-time illustration. (a) The estimated excess-covariance signal reaches the empirical Poisson resolution floor at

T_{O}=3.09

. (b) reconstructed marginals obtained from three terminal times of the forward process are compared with the data distribution in gray.

Figure 3: Two-dimensional count mixture and matched one-dimensional marginals. Components $\{c_{1},c_{5}\}$ define the target distribution, and $\{c_{3},c_{7}\}$ define an alternative distribution with identical coordinate marginals. The component pairs occupy different regions of the count space.

Figure 4: Reverse-chain snapshots for marginal tilt, FK steering, and tilted FK steering. Gray markers show training samples; orange markers show generated particles. In the intermediate columns, orange-marker area scales with normalized particle weight. In the output column, one marker per initial ancestor is placed at its last descendant in output-array order, with area proportional to that ancestor’s number of output descendants. Red frames mark resampling events. FK steering also resamples at $k=24$ , which is not shown, and a final resampling is forced at $k=0$ before output sampling.

Figure 5: Unweighted first-coordinate histograms before and after output sampling for FK steering and tilted FK steering, with the exact target marginal as a reference. Before output sampling, histograms use the endpoint particles; Output denotes the final samples drawn using the terminal weights. The comparison measures marginal agreement; the target and alternative component pairs have identical one-dimensional marginals.

Setting	Value
Generator
Architecture	Transformer encoder; self-attention across genes
Input encoding	Scaled/log-count features, gene embeddings, time MLP
Layers / width / attention heads	3 / 128 / 4
Output head	Categorical over counts $0,\ldots,512$
Optimizer	AdamW

Appendix

Table 3: Generator, sampling, residual-discriminator, and downstream-classifier settings used in the single-cell experiments.

Figure 6: Representative generated cells in the fixed atlas PCA basis, with one column per generative model and one row per target population in decreasing training abundance. Gray points show all test cells, with test cells from the target population highlighted in green; orange points show generated cells. Generated cells are subsampled to $2{,}500$ cells per panel; the PCA basis is applied without refitting.

Figure 7: CRNDiff-generated cells in the fixed atlas UMAP embedding. (a) Test cells ( $n=45{,}145$ ); (b)–(d) $2{,}500$ generated cells for endothelial, myeloid, and neuronal target populations. Mesothelial cells appear only in the test-set panel.

Figure 8: Three sampling methods across the three target populations in the fixed atlas PCA basis: (a) marginal tilt, (b) FK steering, and (c) tilted FK steering. Colors follow Figure 6 . Each panel shows $2{,}500$ generated cells subsampled from the combined outputs of three sampling runs.

Figure 9: Number of distinct initial ancestors represented along the reverse chain, normalized by output sample size. On logarithmic

y

-axes, curves show means and standard-deviation bands over three sampling runs. Dashed lines mark chain-internal resampling events shared across the three sampling runs, excluding the terminal weighted output selection. Final values correspond to Table 7 .

Endothelial abundant , $80{,}463$ training cells; $n_{\mathrm{out}}=10{,}057$
Conditioning	Purity $\uparrow$	$W_{1}$ $\downarrow$	MMD 2 $\downarrow$	PCC $\uparrow$
Best-of- $N$ , $1\times$	0.830	0.046	0.0012	0.998
Best-of- $N$ , $2\times$	0.949	0.051	0.0033	0.998
Best-of- $N$ , $4\times$	0.958	0.078	0.0079	0.996
[3pt/2pt]Value-guided, $\gamma=1$	0.711	0.119	0.0117	0.992
Value-guided, $\gamma=2$	0.889	0.029	0.0006	0.999

Appendix

Table 10: Conditioning strategies for the CRNDiff generator, using the metrics of Table 2 . Entries are means over three sampling runs. The frozen-generator methods use the reference CRNDiff generator;

\dagger

denotes its conditionally fine-tuned counterpart. Candidate-pool sizes and guidance strengths do not imply equal computational cost.

Diffusion Models

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

See all 1,169 papers

Jul 31, 2026

Improving scDiffusion with Sparsity-Biased Classifier-Free Guidance

Single-cell RNA sequencing (scRNA-seq) has become an essential tool in modern cellular biology, and generating accurate synthetic scRNA-seq data is becoming increasingly important. Although diffusion models have achieved promising results in conditional scRNA-seq generation, existing guidance strategies, including classifier guidance and classifier-free guidance (CFG), rely on an unconditional branch trained to approximate the true marginal distribution, which may retain substantial gene-specific structure and limit guidance effectiveness. Inspired by recent work showing that diffusion models can be effectively guided using intentionally degraded references, we propose a sparsity-biased classifier-free guidance (SB-CFG) strategy for scRNA-seq generation. Rather than approximating the assumed "neutral" marginal distribution, SB-CFG introduces a deliberately under-informative sparse reference for the unconditional branch, removing gene identity while preserving only coarse sparsity statistics. This "bad" reference amplifies the contrast between conditional and unconditional predictions, leading to stronger and more effective guidance during sampling. We evaluated SB-CFG as a training-free sampling modification on five publicly available scRNA-seq datasets. Experimental results demonstrate consistent improvements over standard CFG-based sampling in terms of marker gene expression fidelity, cell-type consistency, and sparsity preservation, indicating that SB-CFG better captures biologically meaningful gene expression patterns.

Yu Song, Hao Sun, Ikuko Nishikawa +1

College of Information Science and Engineering, Ritsumeikan University, Osaka, Japan

Classifier-Free Guidance Diffusion Sampling Single-Cell RNA Sequencing

Ask PDF

Apr 4, 2026

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

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.

Renzo G. Soatto, Anders Hoel, Greycen Ren +5

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 +1

Diffusion Models Synthetic Data Generation Incomplete Data Imputation

Ask PDF

May 8, 2026

Flow Matching for Count Data

High-dimensional count data arise in applications such as single-cell RNA sequencing and neural spike trains, where mappings between distributions across successive batches or time points form critical components of data analysis. The recent success of diffusion- and flow-based deep generative models for images, video, and text motivates extending these ideas to count-valued settings, but many existing methods either treat each count as a categorical state or transform counts into a continuous space, neither of which is natural or efficient when the count range is large. We propose count-FM, a flow-matching framework for count data based on a continuous-time birth-death process with local unit jumps. Count-FM learns marginal transitions efficiently in count space through simulation-free training of conditional transition rates, allowing transport between arbitrary count-distributed source and target populations. In simulation, count-FM variants achieve strong sample quality while using substantially fewer parameters. We further apply count-FM to scRNA-seq and neural spike-train data for unconditional generation, transport, and conditional generation. Across these tasks, count-FM yields improved sample quality, greater modeling efficiency, and interpretable transport paths.

Ganchao Wei, John Pearson

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

Flow Matching Synthetic Data Generation Flow-Based Generative Modeling

Ask PDF

Abstract

Diffusion Models

Explore similar work

Improving scDiffusion with Sparsity-Biased Classifier-Free Guidance

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

Flow Matching for Count Data