Count-native diffusion improves cell-type specific gene expression modeling
CRNDiff: Count-Native Diffusion Framework via Chemical Reaction Networks
Machine Learning
Summary
Many scientific measurements count things, like the number of RNA molecules in a cell, but models often treat these counts as smooth numbers instead. The authors build a new method called CRNDiff that works directly with whole numbers by modeling them with chemical reaction networks. This approach helps generate realistic examples of rare cell types without retraining the model. Tested on human heart data, CRNDiff creates synthetic cells that closely match real ones, especially for rare types, and these generated cells improve classification tasks when used instead of real cells.
What this means in practice
- •For biomedical data scientists: Generate synthetic single-cell RNA datasets matching specific cell types to improve machine learning model training when real data is scarce.
- •For computational genomics teams: Create targeted rare cell population samples without retraining models by using chemical reaction network–based diffusion for better downstream analysis.
Authors
Yuxuan Qiu, Praful Gagrani, Tetsuya J Kobayashi
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.