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. We propose Count Flow Map, a finite-time extension of count flow matching (count-FM) for one- or few-step generation directly in count space. Our model learns stochastic transitions that replace many local birth--death updates, using Poisson births and Binomial deaths to preserve nonnegative integer counts without a predefined maximum. These transition models are trained to recover count-FM's local dynamics over short intervals 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. Simulations demonstrate strong one-step sample quality and stability across the evaluated inference budgets, including in a high-dimensional, high-count setting. In single-cell drug perturbation prediction and neural population forecasting, Count Flow Map achieves competitive distributional quality, recovers perturbation effects, and supports high-activity event forecasts with one or a few model evaluations.