WFR Flow Matching (WFR-FM)

• WFR-FM is a simulation-free algorithm that unifies flow matching with dynamic unbalanced optimal transport under the WFR metric to model evolving state and mass.
• It jointly regresses a displacement vector field and a scalar mass change function, ensuring exact geodesic recovery and principled dynamic modeling.
• Empirical results demonstrate state-of-the-art performance in single-cell biology and high-dimensional synthetic benchmarks with near-optimal interpolation metrics.

WFR Flow Matching (WFR-FM) is a simulation-free training algorithm unifying flow matching and dynamic unbalanced optimal transport (OT) under the Wasserstein–Fisher–Rao (WFR) metric. It allows simultaneous regression of a displacement vector field and a scalar mass change function (birth–death dynamics) to model systems where both state and mass evolve over time. WFR-FM provides theoretical guarantees for exact geodesic recovery under the WFR geometry and demonstrates state-of-the-art empirical performance, particularly in single-cell biology, where imbalanced snapshots with proliferating and apoptotic dynamics are prevalent (Peng et al., 11 Jan 2026).

1. Theoretical Foundation: The WFR Metric

The WFR metric extends the Benamou–Brenier dynamic formulation of OT to accommodate unbalanced measures, where total mass can change via birth–death processes. Given two nonnegative measures $\rho_0, \rho_1$ on $\mathbb{R}^d$, the dynamic WFR formulation seeks trajectories $(\rho(t,x), v(t,x), r(t,x))$ solving

$$\partial_t \rho(t,x) + \nabla \cdot (\rho(t,x) v(t,x)) = r(t,x) \rho(t,x), \quad t \in [0,1], x \in \mathbb{R}^d,$$

with $\rho(0,\cdot) = \rho_0$, $\rho(1,\cdot) = \rho_1$. The objective is to minimize the action functional:

$$\mathrm{WFR}^2(\rho_0, \rho_1) = \inf_{\substack{\rho,v,r\ \partial_t \rho + \nabla\cdot(\rho v)=r\rho}} \int_0^1 \int_{\mathbb{R}^d} \bigl(\|v(t,x)\|^2 + \alpha\,r(t,x)^2\bigr) \rho(t,x) \, dx\,dt,$$

where $\alpha > 0$ controls the transport versus mass-change penalty. Setting $r \equiv 0$ recovers classical balanced OT. The mass growth term $r(t,x)$ captures local proliferation and apoptosis, enabling a principled metric for unbalanced data (Peng et al., 11 Jan 2026).

2. WFR-FM Objective and Regression Loss

WFR-FM parameterizes a vector velocity field $v_\theta(t, x)$ and a scalar mass growth rate $r_\theta(t, x)$ using neural networks. Training is conducted using a weighted mean-square regression loss against closed-form targets $(v^*(t,x), r^*(t,x))$ derived from an analytic “conditional path”: a traveling Gaussian whose mean and mass interpolate along the two-Dirac WFR geodesic. The loss is:

$$L_{\rm WFRFM}(\theta) = \mathbb{E}_{t \sim U[0,1], x \sim p_t} \left[ \|v_\theta(t,x) - v^*(t,x)\|^2 + \beta |r_\theta(t,x) - r^*(t,x)|^2 \right],$$

where $p_t$ denotes the induced marginal at time $t$, and $\beta$ (typically set to $\alpha$) regulates the growth regression term. This framework generalizes classical flow matching, which regresses only $v(t,x)$, by explicitly learning both transport and local mass change (Peng et al., 11 Jan 2026).

3. Recovery of WFR Geodesics

The theoretical underpinning of WFR-FM is conditional flow matching (CFM), which guarantees that regression against the analytic targets $(v^*, r^*)$ yields a learned flow matching the true WFR geodesic. Specifically:

• The analytic traveling Gaussian path has marginals that converge to $\rho_0, \rho_1$ at endpoints.
• Its mean and mass exactly follow the two-Dirac WFR geodesic.
• The induced $(u^*, g^*)$ satisfy the dynamic continuity equation, and the path realizes a constant-speed geodesic with action equalling $\mathrm{WFR}(\rho_0, \rho_1)$.

As established by Theorem 4.2 in (Peng et al., 11 Jan 2026), minimizing the WFR-FM loss recovers the unique WFR geodesic, due to identical gradients between the regression loss and the intractable true flow loss. Regressing $(v_\theta, r_\theta)$ to $(u^*, g^*)$ ensures exact dynamic OT (Peng et al., 11 Jan 2026).

4. Algorithm Structure and Simulation-Free Training

The WFR-FM algorithm proceeds as follows:

1. For each snapshot pair $(\rho_k, \rho_{k+1})$, solve a local mini-batch WFR OET (static entropy–regularized OT) to obtain a semi-coupling $y^{(k)}(x_0, x_1)$.
2. Build the traveling Gaussian conditional path $p_t(x|x_0, x_1)$ with mean $\mu_t$ and mass $m_t$ given by the closed-form two-Dirac WFR formulas.
3. Iteratively sample pairs $(x_0, x_1)$, time $t$, and points $x \sim \mathcal{N}(\mu_t, \sigma^2 I)$. Compute targets $v^* = \frac{\mu_t - \mu_{t-dt}}{dt}$ and $r^* = \partial_t \log m_t$.
4. Evaluate the regression loss weighted by $m_t$ and update neural network parameters.
5. Repeat until convergence.

No ODE simulations are required during training. Only Gaussian sampling and analytic target computation are used, ensuring scalability for high-dimensional and large-scale data (Peng et al., 11 Jan 2026).

5. Empirical Performance and Benchmarking

WFR-FM demonstrates competitive or superior empirical results on both synthetic and biological datasets:

• On synthetic benchmarks (2D Gene, 5D Dyngen, 1000D Gaussian), WFR-FM achieves the lowest $W_1$ (1–Wasserstein error) and near-zero Relative Mass Error (RME), outperforming balanced FM and ODE-based unbalanced OT baselines.
• The learned trajectories display path actions within 1–2% of static WFR references, confirming geodesic optimality.
• For four real single-cell RNA-seq datasets (EMT, EB, CITE, Mouse), WFR-FM attains the best hold-out interpolation metrics (unseen-timepoint $W_1$ and RME).
• Scalability is established by linear scaling of per-epoch runtime with the number of samples (up to 16,000 cells in 100D), outperforming both FM-only and ODE-based alternatives.
• On a synthetic gene proliferation model, $r_\theta$ recovers the ground-truth proliferation rate with Pearson correlation $>0.99$ (Peng et al., 11 Jan 2026).

6. Hyperparameters, Implementation, and Computational Complexity

WFR-FM utilizes 5-layer MLPs (256 hidden units, LeakyReLU activations) for $v_\theta$ and $r_\theta$. Growth penalty $\alpha$ and regression weight $\beta$ are typically set in $[1,3]$; empirical ablations indicate robustness in $[1,2]$. The conditional Gaussian bandwidth $\sigma$ matches the end-time spread of samples following balanced FM conventions. Typical batch sizes are $b=512$ to $2048$ (FM) and $B=1024$ to $2048$ (OET). Per-gradient-step computational complexity is $O(b d^2)$ (sampling) plus $O(B^2)$ (local OET), with practical scalability for mini-batch OT. Implementation leverages PyTorch and the POT library for Sinkhorn, and does not require ODE backpropagation (Peng et al., 11 Jan 2026).

7. Summary and Significance

WFR Flow Matching (WFR-FM) establishes a unified and efficient paradigm for modeling dynamical systems from unbalanced snapshot data. By leveraging the WFR metric and simulation-free analytic regression, it provides (i) principled dynamic unbalanced OT under WFR, (ii) closed-form, highly scalable training routines, (iii) theoretical guarantees for exact geodesic recovery, and (iv) state-of-the-art accuracy and stability in single-cell trajectory inference and generative modeling with evolving mass (Peng et al., 11 Jan 2026).