Flow Matching for Count Data

Abstract: High-dimensional count data arise in applications such as single-cell RNA sequencing and neural spike trains, where mapping 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 achieves better sample quality than representative baselines 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.

• The paper introduces count-FM, a flow matching framework that models count data directly using a continuous-time birth-death process to preserve discrete structures.
• It employs conditional binomial bridges to derive a simulation-free training objective while achieving parameter efficiency with only 2d outputs per state.
• Empirical results on scRNA-seq and neural spike data demonstrate superior sample quality, efficiency, and interpretability compared to existing discrete methods.

Flow Matching for Count Data: An Expert Technical Commentary

Introduction and Context

High-dimensional count data, central to domains such as single-cell RNA sequencing (scRNA-seq) and neural spike-train analysis, present nontrivial modeling challenges due to their discrete, sparse, and correlated structure. The proliferation of diffusion and flow-based deep generative models in continuous domains has prompted the development of analogous methods in the discrete setting. However, existing approaches for count data either recast counts as categorical variables—leading to infeasible parameter scaling in large-count regimes—or dequantize to a continuous space, undermining the semantics and statistical structure of integer-valued observations.

This work introduces count-FM, an efficient flow-matching framework specifically designed for count data. The core methodological advance is the direct modeling of distributions in count space via a continuous-time birth-death process with local (±1) jumps, combined with a simulation-free training objective grounded in conditional binomial bridges. This formalism preserves the discrete structure, improves parameter efficiency, supports interpretable transport between distributions, and generalizes across unconditional, conditional, and transport scenarios.

Methodological Innovations

Birth-Death Flow in Count Space

The proposed count-FM parameterizes transitions between count-valued distributions as a continuous-time Markov jump process (CTMC) on $\mathbb{N}_0^d$. At each coordinate $i$, jump rates $\lambda_{t,i}(x)$ and $\mu_{t,i}(x)$ encode stochastic "birth" and "death" events (unit increments or decrements), with time- and state-dependence captured by a neural network. Unlike categorical-state parameterizations that incur a combinatorial explosion in parameters, count-FM only requires $2d$ outputs per state—scaling linearly with dimension and invariant to the maximum count.

Conditional Binomial Bridges and Path-Space KL

A tractable and differentiable training objective is derived using conditional binomial bridges, which define trajectories respecting the geometry of count increments. The derivation yields closed-form coordinate-wise bridge transition rates and a path-space KL objective. Explicitly, for endpoints $(x_0, x_1)$, the bridge at time $t$ has coordinate:

$$X_t^{(i)} = x_0^{(i)} + \mathrm{sgn}(x_1^{(i)} - x_0^{(i)})\,B_t^{(i)}, \quad B_t^{(i)} \sim \mathrm{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$

Local rate-matching between model and bridge yields a KL objective equivalent (in gradient) to marginal generator matching, aligning the count-FM framework with recent advances in generator/diffuser matching for Markov processes.

Efficient Forward Sampling and OT-Coupled Transport

Sample generation proceeds via discretized local jumps, ensuring all intermediate states remain integer-valued. The endpoint coupling $\pi(x_0, x_1)$ used for training may be either independent or constructed using minibatch optimal transport (OT), the latter yielding straighter and more interpretable transition paths and often improved sample efficiency. OT-coupled count-FM (count-FM-OT) aligns with methods harnessing transport geometry for more efficient generative flows.

Empirical Evaluation

2D Simulation

In controlled two-dimensional simulations, count-FM and count-FM-OT achieve lower Wasserstein-2 and MMD² scores than discrete flow/diffusion baselines (including D3PM, tauLDR, SEDD, and Dirichlet-FM), as well as the count-specific Poisson-JUMP model. Notably, count-FM attains superior generative fidelity with an order-of-magnitude fewer parameters.

Figure 1: Intermediate samples from the native sampling trajectories of different models in the simulation, highlighting the progressive and smooth transport of count-FM compared to abrupt transitions in categorical-state baselines.

Figure 2: Marginal bridge distributions for both coordinates under a common progress variable $s$; count-FM exhibits gradual evolution in count space, while categorical-state baselines collapse quickly toward the target.

Figure 3: Sampling efficiency comparison showing OT-coupled count-FM requires fewer function evaluations and less wall-clock time to attain the same quality threshold.

Single-cell RNA-seq Generation and Transport

Applying count-FM to the Dentate Gyrus scRNA-seq dataset, the framework demonstrates state-of-the-art performance against specialized single-cell generative models (scVI, scDiffusion, scLDM, DCM, CFGen) for both unconditional sample generation and developmental transport (P12 to P35). The preservation of count-valued intermediates enables biologically meaningful interpolation trajectories, maintaining interpretability in cell fate analysis.

Figure 4: Generated transport trajectories for scRNA-seq from P12 to P35 in PCA space, with count-FM yielding lineage-consistent transitions and interpretable fate proportions.

Figure 5: Training-phase transition snapshots illustrating the gradual movement of generated cells from P12 toward P35 in PCA space.

Figure 6: Testing-phase transition snapshots, confirming generalization and smooth transport for held-out cells.

Conditional Generation for Neural Spike Counts

In conditional spike-train generation on the CRCNS hc-3 dataset, count-FM—with classifier-free guidance—captures mean response, variance, sparsity, and population correlation more faithfully than Poisson MLP and deterministic MLP mean regressors. Conditional guidance weight $i$0 modulates calibration versus sharpness, with $i$1 optimizing the overall trade-off.

Figure 7: Mean response and population correlation on the hc-3 linear-track session; count-FM at $i$2 best matches observed cross-neuron dependence.

Figure 8: Bin-wise mean count comparison across signed position and guidance weights, showing the impact of conditional guidance on spatial tuning.

Figure 9: Population correlation matrices for different models and guidance settings, with count-FM most closely adhering to ground-truth dependencies.

Comprehensive quantitative metrics confirm the superior modeling of mean, variance, sparsity, and covariance by count-FM (Table: hc-3 conditional generation metrics).

Further experiments on respiration- and odor-conditioned piriform cortex spike trains extend these conclusions to other neurophysiological settings.

Figure 10: Neuron-by-time mean responses for held-out odor-3 trials, revealing improved recovery of localized, stimulus-locked patterns by count-FM.

Figure 11: Generated responses for a fixed held-out covariate trajectory, showing increased condition-aligned structure with higher guidance weights, but risks of over-amplification.

Theoretical and Practical Implications

The count-FM framework addresses both statistical and computational limitations of prior discrete generative models for count data by:

• Parameter Efficiency: Transition modeling with only $i$3 parameters enables application to high-dimensional, high-count problems in genomics and neuroscience.
• Sample Quality and Diversity: Explicit rate-matching yields higher-quality samples as measured by Wasserstein, MMD, and distributional calibration.
• Interpretable Transport: The preservation of count-valued intermediates aligns with the requirements of developmental biology, where intermediate states are biologically meaningful.
• Flexible Conditioning: Classifier-free guidance supports nuanced control over the strength of conditioning, generalizing to various downstream tasks.
• Sampling Efficiency: OT-coupling and the compatibility with advanced sampling schemes (e.g., shortcut and flow map methods) improve computational performance and the tractability of large-scale data generation.

Limitations and Future Directions

Critical limitations include the simplicity of the conditional binomial bridge (which may excessively constrain intermediate path variability), the use of first-order samplers, and the need for many discretized steps in generation. The framework is readily extensible—suggested directions include:

• Adoption of more expressive bridges (e.g., beta-binomial or latent stochastic interpolants) to capture richer path distributions.
• Integration of advanced simulation schemes (bridge-aware sampling, Feynman-Kac corrections, shortcut or flow-map-based methods) to increase numerical accuracy and reduce sample steps.
• Handling multimodal, structured, or hierarchical count datasets via further augmentation of the conditional architecture.

Conclusion

Count-FM establishes a theoretically principled and empirically validated approach to generative modeling for count data, bridging the gap between efficient flow-matching methods in the continuous domain and the discrete high-dimensional distributions inherent to modern genomics and neuroscience. The innovation in local birth-death modeling, tractable simulation-free training, and parameter efficiency positions count-FM as an effective baseline and compelling building block for future developments in scientific machine learning for count-based data modalities (2605.07746).