Technical Summary: Flow Matching for Count Data (count-FM)

Problem Statement

High-dimensional count data, prevalent in fields such as single-cell RNA sequencing (scRNA-seq) and neural spike train analysis, presents unique challenges for generative modeling. These data types are characterized by discreteness, sparsity, and complex correlation structures. Existing deep generative approaches often struggle in this domain:

• Categorical State Approaches: Methods treating counts as categorical variables (e.g., tokens) require output dimensions that scale with the maximum count range. This leads to parameter inefficiency and treats adjacent counts as unrelated states, ignoring the inherent geometry of count space.
• Continuous Transformations: Approaches mapping counts to continuous spaces via dequantization or latent representations replace discrete probability mass with continuous densities, potentially blurring the underlying discrete structure.
• Limitations of Existing Count Models: While some count-specific jump models exist (e.g., Poisson-JUMP), they are often not designed for transport between arbitrary count distributions or lack the efficient training and generation properties of modern flow-based methods.

The paper addresses the need for a generative framework that operates directly in count space, respects the discrete geometry of the data, and enables efficient transport between arbitrary source and target distributions without the parameter overhead of categorical-state models.

Methodology: count-FM

The authors propose count-FM, a flow-matching framework tailored for count data based on a continuous-time birth-death process with local unit jumps.

1. Continuous-Time Birth-Death Process

The model defines the transition between a source distribution $p_0$ and a target distribution $p_1$ as a continuous-time Markov jump process (CTMC) on the space of count vectors $\mathbb{N}_0^d$. The process evolves via local unit births ($+1$) and deaths ($-1$) for each coordinate $i$:

• Birth Rate: $\lambda_{t,i}(x)$
• Death Rate: $\mu_{t,i}(x)$
• Dynamics: Over a small interval $h$, the probability of transitioning from state $x$ to $x + e_i$ is $h \lambda_{t,i}(x)$, and to $x - e_i$ is $h \mu_{t,i}(x)$. Death is disallowed when a count is zero.

This parameterization is highly efficient: instead of predicting a full categorical distribution over all possible count levels (which scales with the maximum count), the model only predicts two rates per dimension ($2d$ outputs), regardless of the count range.

2. Conditional Binomial Bridge and Training

To train the model, the authors define a conditional bridge between fixed endpoints $(x_0, x_1)$. For each coordinate $i$, the bridge is constructed using a Binomial distribution:
$$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$
This bridge ensures that the conditional mean moves linearly from $x_0$ to $x_1$ while maintaining integer-valued intermediate states.

By enforcing mass preservation (the Kolmogorov forward equation) on this conditional bridge, the authors derive closed-form expressions for the conditional birth and death rates:
$$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$$
where $(\cdot)^+$ denotes positive rectification.

3. Training Objective

The model is trained by minimizing a path-space Kullback-Leibler (KL) divergence between the conditional bridge process and the model's learned process. This decomposes into a local rate-matching objective:
$$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$$
where $\ell(u, v) = v - u \log v$ is a generalized KL term. This allows for simulation-free training of the marginal transition rates by learning the identifiable conditional rates.

4. Sampling and Coupling

• Sampling: Samples are generated by simulating the learned birth-death process forward from $t=\epsilon_t$ to $t=1-\epsilon_t$ using a first-order local-jump discretization.
• Endpoint Coupling: The paper explores two coupling strategies for pairing source and target samples:
  • Independent Coupling: Random pairing ($p_0(x_0)p_1(x_1)$).
  • Optimal Transport (OT) Coupling: Minibatch OT pairing using a symmetric Poisson cost. This tends to induce transitions with lower curvature, improving interpretability and sampling efficiency.

5. Conditional Generation

For conditional tasks, the model incorporates covariates $y$ and utilizes Classifier-Free Guidance (CFG). During training, conditioning variables are randomly dropped. At sampling, the conditional and unconditional rates are combined via a guidance scale $w$ to control the strength of the condition.

Key Contributions

1. count-FM Framework: A flow-matching framework for count data based on local birth-death dynamics, enabling transport directly in count space with intermediate states remaining count-valued.
2. Tractable Training Scheme: A method to learn marginal transition rates via a conditional binomial bridge, retaining the efficiency of flow matching while avoiding the complexity of categorical-state modeling.
3. Parameter Efficiency: A formulation that predicts only local birth and death rates ($2d$ outputs) rather than full categorical distributions, making it suitable for high-dimensional data with large count ranges.

Experimental Results

1. Simulation (2D Gamma-Poisson Mixture)

• Performance: count-FM and count-FM-OT achieved the best sample quality (lowest 2-Wasserstein distance and MMD$^2$) compared to baselines including D3PM, discrete-FM, tauLDR, SEDD, and Poisson-JUMP.
• Efficiency: count-FM achieved these results with substantially fewer trainable parameters (2,372) compared to categorical-state baselines (14,000–67,000+).
• Geometry: Intermediate states evolved smoothly in count space. The OT-coupled version followed a visibly straighter transition path and reached target quality with fewer function evaluations (NFEs).

2. Single-Cell RNA-seq (Dentate Gyrus)

• Unconditional Generation: On a dataset of 2,930 cells and 13,913 genes, count-FM outperformed latent-variable baselines (scVI) and latent diffusion models (scDiffusion, scLDM) in sample quality (W2 and MMD$^2$) while using fewer parameters (approx. 9.8M vs. 29M–39M for competitors).
• Transport (P12 to P35): The model successfully modeled developmental transport from postnatal day 12 to day 35. Using lineage-restricted OT coupling, the transport paths respected known developmental structures (e.g., neuroblast to granule cell progression), producing interpretable intermediate distributions that remained in count space.

3. Neural Spike Trains (Hippocampal/Entorhinal)

• Conditional Generation: The model was applied to generate spike counts conditioned on linearized position and running direction.
• Comparison: count-FM with guidance scale $w=1$ outperformed a deterministic MLP regressor and a Poisson MLP baseline.
• Correlation Structure: Unlike the Poisson MLP, which underestimated population correlations, count-FM better preserved the cross-neuron dependence structure.
• Guidance Tradeoff: $w=1$ provided the best balance between spatial tuning and distributional calibration. $w=2$ sharpened place fields but distorted amplitude calibration.

Significance and Claims

The paper claims that count-FM provides an effective framework for modeling high-dimensional count data by enabling generation and transport directly in the count space. Its primary significance lies in:

• Geometric Fidelity: By modeling transport via local unit jumps, the method respects the discrete geometry of count data, avoiding the blurring effects of continuous approximations.
• Scalability: The parameter efficiency makes it viable for high-dimensional applications (e.g., thousands of genes) where categorical-state methods become computationally prohibitive.
• Interpretability: The intermediate states in the transport path are valid count samples, allowing for the inspection of meaningful transitional distributions in applications like developmental biology.
• Versatility: The framework supports unconditional generation, transport between arbitrary distributions, and conditional generation with controllable guidance.

The authors acknowledge limitations, noting that the conditional binomial bridge is relatively simple and may not explore all plausible intermediate paths, and that the current first-order discretization could be improved with more accurate simulation schemes. However, the empirical results across simulation and biomedical datasets demonstrate that count-FM achieves high sample quality with greater modeling efficiency than competing approaches.