Count flow map models integer data transitions in few steps
Stochastic Flow Map for Count Data
Machine Learning
Summary
High-dimensional count data, like numbers of cells or events, are hard to model with traditional tools that expect continuous or categorical data. The authors created Count Flow Map, a model that directly handles integer counts by simulating random births and deaths to reflect changes over time in just a few steps. This method works well even with large counts and can predict how biological systems respond to drugs or how neural activity changes. It allows quick and accurate generation of realistic count data without many repeated calculations.
What this means in practice
- •For bioinformatics teams: Predict outcomes of drug treatments on single-cell data with fast and accurate count-based models.
- •For neural engineers: Forecast large-scale neural population activity during high-activity events using efficient count modeling.
Authors
Ganchao Wei
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.