Stochastic Flow Map for Count Data

Ganchao Wei

stat.MLSep 20, 2026

Stochastic Flow Map for Count Data

Authors: Ganchao Wei

Organizations: Flatiron Institute New York, NY 10010, USA

Abstract

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.

Figures & tables

Figure 1: 2-D quality–efficiency trade-offs over the full NFE grid. Left and center show endpoint TV and

W_{2}

versus NFE. Right shows

W_{2}

versus measured generation time. Count Flow Map gives the strongest one-step endpoint quality and remains stable across the NFE grid. Discrete Flow Maps becomes competitive by NFE=4 but is worse at NFE=1, while Count-FM, CountsDiff, and D3PM require more inference steps to approach their best endpoint quality.

Method	Coupling	NFE	SW ${}_{2}\downarrow$	logFC $r\uparrow$	Top overlap $\uparrow$	Time (s) $\downarrow$
Count Flow Map	Indep.	1	$0.448\pm 0.006$	$0.779\pm 0.003$	$0.362\pm 0.011$	$0.21\pm 0.00$
Count Flow Map		16	$0.429\pm 0.013$	$0.771\pm 0.006$	$0.386\pm 0.008$	$1.62\pm 0.00$
Count-FM (unit jump)		128–512	$0.424\pm 0.006$	$0.782\pm 0.005$	$\mathbf{0.396\pm 0.003}$	$10.85\pm 7.05$
Count-FM ( $\tau$ -leap)		128	$0.430\pm 0.007$	$\mathbf{0.786\pm 0.006}$	$0.394\pm 0.005$	$5.47\pm 0.08$
Count Flow Map	OT	1	$\mathbf{0.394\pm 0.002}$	$0.779\pm 0.006$	$0.389\pm 0.006$	$0.20\pm 0.00$
Count Flow Map		16	$0.423\pm 0.003$	$0.764\pm 0.006$	$0.381\pm 0.006$	$1.60\pm 0.01$

Table 1: Held-out perturbation prediction. Quality metrics are averaged equally across 24 test conditions. For neural models, we report mean

\pm

sample standard deviation across three training seeds. These repeated runs assess neural-network training variability, while the non-neural Linear and Sinkhorn baselines are evaluated once. Count-FM’s NFE is selected on validation for each seed. Ranges show the lowest and highest selected NFE. Time is the number of seconds needed to generate 9,600 cells, excluding setup and metric computation. Evaluation and timing details can be found in Appendix C.1 .

Figure 2: Held-out single-cell prediction and generation cost. A. Sliced

W_{2}

versus NFE. B. Sliced

W_{2}

versus generation time for 9,600 cells. C. Observed versus generated perturbation strengths at 1 NFE under both couplings, using one fixed training seed. Each point represents one of the 24 test conditions. The dashed line indicates perfect agreement.

r

is Pearson correlation across condition magnitudes, and the slope is fitted through the origin. In A–B, error bars show one sample standard deviation across three training seeds, colors identify the model or sampler, and solid and dashed lines indicate OT and independence coupling, respectively. Effect definitions are provided in Appendix C.3 .

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Method	NFE	TV $\downarrow$	$W_{2}\downarrow$	Time (s) $\downarrow$
Count Flow Map	1	0.153 $\pm$ 0.005	2.50 $\pm$ 0.34	0.008 $\pm$ 0.000
	4	0.145 $\pm$ 0.007	2.74 $\pm$ 0.50	0.028 $\pm$ 0.001
	128	0.129 $\pm$ 0.003	2.71 $\pm$ 0.34	0.739 $\pm$ 0.039
	256	0.126 $\pm$ 0.002	2.54 $\pm$ 0.49	1.478 $\pm$ 0.076
Count-FM + unit jump	1	0.742 $\pm$ 0.002	30.10 $\pm$ 0.34	0.004 $\pm$ 0.000
	4	0.698 $\pm$ 0.001	26.99 $\pm$ 0.43	0.014 $\pm$ 0.001

Appendix

Table 2: 2-D simulation endpoint quality at selected inference budgets. Entries are mean

\pm

standard deviation over five seeds. Lower is better for all metrics. The best quality metric at each NFE is shown in bold.

Figure 3: 2-D endpoint samples at NFE $\in\{1,4,128,256\}$ . Columns show the target and the six compared methods, and rows correspond to different inference budgets.

Figure 4: Generation paths for the 2-D simulation at normalized progress $q$ . Count Flow Map and Discrete Flow Maps use NFE=4 finite-time compositions. Count-FM, CountsDiff, and D3PM are shown at matched progress along 256-step sampling paths.

Method	NFE	SW ${}_{2}\downarrow$	MMD ${}^{2}_{\rm RBF}\downarrow$	Time (s) $\downarrow$
Count Flow Map	1	0.556 $\pm$ 0.134	0.0007 $\pm$ 0.0004	0.005 $\pm$ 0.000
	4	0.724 $\pm$ 0.134	0.0015 $\pm$ 0.0009	0.012 $\pm$ 0.001
	128	0.518 $\pm$ 0.048	0.0005 $\pm$ 0.0001	0.308 $\pm$ 0.005
	256	0.555 $\pm$ 0.087	0.0005 $\pm$ 0.0002	0.617 $\pm$ 0.008
Count-FM + unit jump	1	10.898 $\pm$ 0.490	0.1717 $\pm$ 0.0029	0.003 $\pm$ 0.000
	4	9.388 $\pm$ 0.284	0.1762 $\pm$ 0.0018	0.005 $\pm$ 0.000

Appendix

Table 3: 32-D simulation endpoint quality at selected inference budgets. Entries are mean

\pm

standard deviation over five seeds. Lower is better for all metrics. The best quality metric at each NFE is shown in bold.

Figure 5: 32-D quality–efficiency curves. Left and center show sliced

W_{2}

and RBF MMD 2 versus NFE. Right shows sliced

W_{2}

versus measured generation time. Count Flow Map achieves strong quality at NFE=1 and retains the lowest sliced

W_{2}

across the full NFE grid. CountsDiff only slightly improves MMD

{}^{2}_{\mathrm{RBF}}

at the largest inference budget.

Figure 6: Condition-level perturbation recovery with OT coupling. A. Observed and generated gene effects at 1 NFE projected onto training-only PCA axes. B. Corresponding projections at 16 NFE. C. Pearson correlations between observed and generated effects across all 2,000 retained genes. Paired markers compare the two sampling budgets for each of the 24 test conditions, ordered by decreasing observed effect magnitude. Both budgets use the same Count Flow Map model from one fixed training run. A–B share the same PCA axes, with lines connecting the same condition. The correlations in C differ from the table’s top-200-gene logFC correlation.

Figure 7: Neural forecast quality, computational cost, and predictive dependence. A. Dimension-normalized energy score for the joint count vector versus NFE. B. Population CRPS for mean count per neuron versus NFE. C. Population CRPS versus latency per 64-draw forecast. D. High-activity event Brier score after marginal-preserving neuron shuffling minus the joint-sample score. Positive values favor joint predictions. A–C use the full evaluated grids, covering 1–64 NFE for Count Flow Map and 1–256 NFE for Count-FM. Recording points in D average three training seeds. Aggregate results in A–C and the mean row in D average six recordings equally within each seed. Error bars show one sample standard deviation across the three seed-level means.

Explore similar work

Flow Matching for Count Data
May 8, 2026Ganchao Wei, John PearsonGenerative Flow Networks Single-Cell Rna
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
A Few-Step Generative Model on Cumulative Flow Maps
May 5, 2026Zhiqi Li, Duowen Chen, Yuchen Sun +1Flow Map Flow Models
Georgia Institute of Technology, Atlanta, United States of America
Binomial flows: Denoising and flow matching for discrete ordinal data
May 1, 2026Yair Shenfeld, Ricardo Baptista, Stefano PeluchettiGenerative Models Diffusion Models

May 8, 2026, stat.MLJ/K move · Enter open · S save

Flow Matching for Count Data

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

Generative Flow Networks Single-Cell Rna Counting Twitter-Fmnerg Benchmarks

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Flow Matching for Count Data

A Few-Step Generative Model on Cumulative Flow Maps

Binomial flows: Denoising and flow matching for discrete ordinal data