Count Bridges enable Modeling and Deconvolving Transcriptomic Data

Nic Fishman; Gokul Gowri; Tanush Kumar; Jiaqi Lu; Valentin de Bortoli; Jonathan S. Gootenberg; Omar Abudayyeh

cs.LGMar 5, 2026

Count Bridges enable Modeling and Deconvolving Transcriptomic Data

Authors: Nic Fishman, Gokul Gowri, Tanush Kumar, Jiaqi Lu, Valentin de Bortoli, Jonathan S. Gootenberg, Omar Abudayyeh

Organizations: Harvard University · MIT · CNRS · Beth Israel Deaconess Medical Center ·

Abstract

Many modern biological assays, including RNA sequencing, yield integer-valued counts that reflect the number of molecules detected. These measurements are often not at the desired resolution: while the unit of interest is typically a single cell, many measurement technologies produce counts aggregated over sets of cells. Although recent generative frameworks such as diffusion and flow matching have been extended to non-Euclidean and discrete settings, it remains unclear how best to model integer-valued data or how to systematically deconvolve aggregated observations. We introduce Count Bridges, a stochastic bridge process on the integers that provides an exact, tractable analogue of diffusion-style models for count data, with closed-form conditionals for efficient training and sampling. We extend this framework to enable direct training from aggregated measurements via an Expectation-Maximization-style approach that treats unit-level counts as latent variables. We demonstrate state-of-the-art performance on integer distribution matching benchmarks, comparing against flow matching and discrete flow matching baselines across various metrics. We then apply Count Bridges to two large-scale problems in biology: modeling single-cell gene expression data at the nucleotide resolution, with applications to deconvolving bulk RNA-seq, and resolving multicellular spatial transcriptomic spots into single-cell count profiles. Our methods offer a principled foundation for generative modeling and deconvolution of biological count data across scales and modalities.

Figures & tables

Figure 1: Left: Sample paths for several endpoint gaps $d_{1}$ (top). Fixing the prefix $[0,t]$ resample $(t,1]$ by the recursive kernel (bottom). Middle: Bessel slack posteriors at initial and intermediate times. The slack $M_{t}$ concentrates near $0$ as $|d|$ grows. Right: ECDFs of $X_{s}$ from a one–step kernel $(1\!\to\!s)$ and a two–step kernel $(1\!\to\!t\!\to\!s)$ are indistinguishable, confirming composition.

Figure 2

Figure 3

Figure 4: Deconvolution of the low-rank Gaussian mixture across different group sizes and levels of between-group heterogeneity.

Table 7

Method	JSD	RMSE	Spearman
CIBERSORTx	0.194	0.109	0.079
MuSiC	0.313	0.140	0.186
Count Bridge	0.113 $\pm 0.001$	0.073 $\pm 0.000$	0.267 $\pm 0.005$

Table 3: Cell-type proportion deconvolution error for nucleotide level bulk RNA sequencing data.

Method	JSD	RMSE	Spearman
STDeconvolve	0.288	0.177	0.255
Count Bridge	0.231 $\pm 0.002$	0.110 $\pm 0.001$	0.332 $\pm 0.001$

Table 4: Cell-type proportion deconvolution error for spatial transcriptomics.

Comparison	MMD	$\mathbf{W_{2}}$	Energy
Spot mean	0.409	0.030	41.717
Count Bridge	0.203 $\pm 0.000$	0.017 $\pm 0.000$	8.903 $\pm 0.014$

Table 5: Gene expression count profile deconvolution error for spatial transcriptomics

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Method	MMD	$W_{2}$	Energy
CFM	0.065 ± 0.019	0.049 ± 0.008	0.874 ± 0.246
DFM	0.010 ± 0.002	0.010 ± 0.002	0.035 ± 0.014
Count Bridge (CE)	0.0065 ± 0.0023	0.0080 ± 0.0009	0.026 ± 0.004
Count Bridge (ES)	0.0044 ± 0.0018	0.0052 ± 0.0007	0.0098 ± 0.0029
Count Bridge (MSE)	0.030 ± 0.000	0.033 ± 0.001	0.366 ± 0.015

Appendix

Table 6: Discrete Moons Results: Noise → Two Moons

$\lambda_{+}=\lambda_{-}$	MMD	$W_{2}$	Energy
0	0.0038 ± 0.0012	0.0046 ± 0.0002	0.0075 ± 0.0011
8	0.0039 ± 0.0015	0.0045 ± 0.0006	0.0080 ± 0.0022
16	0.0049 ± 0.0003	0.0049 ± 0.0003	0.0095 ± 0.0004
32	0.0052 ± 0.0024	0.0055 ± 0.0015	0.011 ± 0.006
256	0.0064 ± 0.0020	0.0063 ± 0.0009	0.015 ± 0.004

Appendix

Table 7: Count Bridge (Energy Score) Results Across Different

\lambda_{+}=\lambda_{-}

Values

Method	MMD	$W_{2}$	Energy
Count Bridge (ES)	0.0044 ± 0.0018	0.0052 ± 0.0007	0.0098 ± 0.0029
Gaussian Diffusion	0.024 $\pm$ 0.009	0.017 $\pm$ 0.006	0.118 $\pm$ 0.064
Discrete Diffusion	0.017 $\pm$ 0.010	0.013 $\pm$ 0.006	0.072 $\pm$ 0.055

Appendix

Table 8: Discrete Moons Results: Eight Gaussians → Two Moons

Dim	Method	NFE	MMD	$W_{2}$	EMD
4	CFM	8	0.027 ± 0.014	0.019 ± 0.002	0.716 ± 0.420
		32	0.026 ± 0.017	0.015 ± 0.004	0.584 ± 0.460
		128	0.026 ± 0.018	0.015 ± 0.004	0.565 ± 0.469
	DFM	8	0.025 ± 0.004	0.011 ± 0.001	0.245 ± 0.053
		32	0.035 ± 0.003	0.014 ± 0.001	0.458 ± 0.078
		128	0.046 ± 0.005	0.018 ± 0.002	0.759 ± 0.144

Appendix

Table 9: Performance Comparison Across Dimensions and Methods

Group Size (n)	$\alpha={1}$	$\alpha={10}$	$\alpha={1000}$
4	0.0091 ± 0.0005	0.010 ± 0.000	0.011 ± 0.000
8	0.011 ± 0.001	0.014 ± 0.001	0.016 ± 0.000
32	0.020 ± 0.002	0.023 ± 0.002	0.025 ± 0.002
128	0.050 ± 0.008	0.053 ± 0.006	0.057 ± 0.002

Appendix

Table 10: Deconvolution Performance:

W_{2}

vs Group Size and Dirichlet Concentration

Group Size (n)	$\alpha={1}$	$\alpha={10}$	$\alpha={1000}$
4	0.130 ± 0.006	0.195 ± 0.025	0.207 ± 0.040
8	0.286 ± 0.023	0.363 ± 0.037	0.483 ± 0.010
32	0.530 ± 0.051	0.657 ± 0.095	0.921 ± 0.222
128	2.24 ± 0.58	2.22 ± 0.54	2.63 ± 0.14

Appendix

Table 11: Deconvolution Performance: EMD vs Group Size and Dirichlet Concentration

Group Size (n)	$\alpha={1}$	$\alpha={10}$	$\alpha={1000}$
4	0.011 ± 0.001	0.011 ± 0.002	0.012 ± 0.004
8	0.016 ± 0.001	0.017 ± 0.002	0.021 ± 0.001
32	0.014 ± 0.003	0.020 ± 0.003	0.025 ± 0.010
128	0.037 ± 0.005	0.045 ± 0.007	0.044 ± 0.001

Appendix

Table 12: Deconvolution Performance: MMD vs Group Size and Dirichlet Concentration

Method	MMD	$W_{2}$	EMD
exact	$\mathbf{0.037\pm 0.005}$	$\mathbf{0.050\pm 0.008}$	$\mathbf{2.24\pm 0.58}$
randomized	$0.038\pm 0.000$	$0.051\pm 0.007$	$2.33\pm 0.43$
round	$0.038\pm 0.003$	$0.052\pm 0.007$	$2.32\pm 0.44$

Appendix

Table 13: Performance comparison of different rounding methods with standard errors.

Figure 5: UMAP of true single-cell and CB-deconvolved nucleotide-level expression, aggregated to the gene level and hued by cell type (as assigned by Yazar et al. (2022) ).

Method	JSD	RMSE	Spearman
STDeconvolve	0.288	0.177	0.255
RCTD	0.161	0.113	0.580
Count Bridge	0.229	0.110	0.332

Appendix

Table 14: Cell-type deconvolution error for spatial transcriptomic data. Note that RCTD requires a single-cell reference dataset for deconvolution.

Panel of Figure 6 — Figure 6: Performance of RCTD, STDeconvolve and Count Bridge on MERFISH deconvolution across number of cell types using (a) Jensen Shannon Divergence (JSD), (B) RMSE and (C) Spearman Correlation

Figure 7: UMAP of true vs deconvolved cell profiles using the Count Bridge.

Figure 8: Abundance of putative cell types (automatically annotated by celltypist) in synthetic unit-level data generated by Count Bridges through deconvolution, vs. abundance of putative cell types in true unit-level data.

Generative Modeling

Generative modeling learns data distributions to create new samples, from images and text to molecules, simulations, and structured records.

See all 493 papers

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

Sep 20, 2026

Stochastic Flow Map for Count Data

High-dimensional count data are common in scientific applications, but most diffusion and flow models are designed for continuous or categorical data, and generation often requires many sequential model evaluations. We propose Count Flow Map, a generative model that learns finite-time transitions directly in count space for one- or few-step generation. Our model directly learns stochastic transitions over finite time intervals, using Poisson births and Binomial deaths to preserve nonnegative integer counts without a predefined maximum. These transition models are trained to match the underlying local birth--death dynamics and to maintain consistency across step sizes. We characterize the connection between local dynamics and finite-time transition consistency and derive a bound on the generation error. After validating Count Flow Map in several simulations, including a high-dimensional, high-count setting, we apply it to single-cell drug perturbation prediction and neural population forecasting, where it captures perturbation responses and supports forecasts of high-activity events with only one or a few model evaluations. Together, these experiments demonstrate that Count Flow Map enables high-quality generation directly in count space across inference budgets, from one-step to few-step generation, using a single trained model.

Ganchao Wei

Flatiron Institute New York, NY 10010, USA

Synthetic Data Generation One-Step Generative Modeling Flow Map Learning Flow-Based Generative Modeling

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

Abstract

Generative Modeling

Explore similar work

Flow Matching for Count Data

Stochastic Flow Map for Count Data

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