Stochastic Flow Map for Count Data
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. Here, we propose Count Flow Map, a generative model that learns finite-time transitions directly in count space for one- or few-step generation and transport between count distributions. 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, 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.
Related
- Equivalence testing and run-count design for stochastic machine-learning experiments
- Flow Matching for Count Data
- Flow decomposition and stochastic completeness of weighted graphs
- Numerical approximation of probabilistically strong solutions to the stochastic harmonic map heat flow
- Stochastic Riccati equations with jumps for indefinite linear-quadratic control: A stochastic value-flow approach
BibTeX
@misc{wei2026stochastic,
title = {{Stochastic Flow Map for Count Data}},
author = {Ganchao Wei},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2609.23290},
eprint = {2609.23290},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.23290},
}
RIS
TY - UNPB
TI - Stochastic Flow Map for Count Data
AU - Ganchao Wei
PY - 2026
DA - 2026/10/28
PB - arXiv
DO - 10.48550/arxiv.2609.23290
UR - https://arxiv.org/abs/2609.23290
AB - 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. Here, we propose Count Flow Map, a generative model that learns finite-time transitions directly in count space for one- or few-step generation and transport between count distributions. 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, 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.
ER -