NeuroFlow: Joint Categorical Flow Matching for Neural Population Activity

Leading models of neural population activity commonly predict time-varying firing rates and connect them to observed spikes through a conditionally independent count emission model. We introduce NeuroFlow, which instead places a coupled generative distribution directly over binned neuron–time spike counts. NeuroFlow represents counts as vertices of categorical probability simplices and learns joint transport from simplex-valued noise using complementary Dirichlet or Fisher–Rao flow-matching paths. Position-wise flow times and visibility masks allow the same model to condition on arbitrary subsets of neural and behavioral observations, supporting neuron imputation, temporal forecasting, region-level generation, and generation without observed neural context. Beyond conventional prediction, the resulting probability flow provides joint samples and a dequantization-based likelihood bound for complete held-out spike sequences. Using a Poisson baseline with matched region encoders, behavioral pathway, and Transformer backbone, we find that the evaluated Dirichlet model remains competitive on marginal prediction and better matches the empirical inter-spike-interval distribution. A proof of concept shows how this interface can test a theory-derived invariance: neuron-wise temporal shifts preserve individual-neuron statistics while disrupting population alignment, and NeuroFlow assigns consistently lower likelihood to the resulting surrogates. More broadly, NeuroFlow provides a common probabilistic framework for comparing which hypothesized forms of population structure are supported across large neural datasets.