flux-longitudinal-flow-matching

name

flux-longitudinal-flow-matching

description

Geometry-aware longitudinal flow matching framework for unpaired biological snapshot data. FLUX learns data-dependent metrics, constructs manifold-aware conditional paths, and uses mixture-of-experts velocity fields for joint transport modeling and unsupervised regime discovery. From Ortega Caro et al. 2026 (arXiv:2605.08648). Use when: modeling unpaired longitudinal biological data, flow matching on manifolds, neural dynamics trajectory reconstruction, cell differentiation modeling, calcium imaging analysis, regime discovery in time series.

FLUX: Geometry-Aware Longitudinal Flow Matching with Mixture of Experts

Framework from "FLUX: Geometry-Aware Longitudinal Flow Matching with Mixture of Experts" (Josue Ortega Caro, Yongxu Zhang, Hannah M. Batchelor, Sizhuang He, Jessica Cardin, Shreya Saxena -- Yale University, arXiv:2605.08648, May 2026) for reconstructing longitudinal trajectories from unpaired population snapshots while discovering latent regime transitions.

Core Problem

Many biological systems are observed only as unpaired population snapshots (not tracked trajectories):

• Widefield calcium imaging: neurons drift between sessions, trial structure not aligned
• Single-cell RNA-seq: measurement destroys the cell, no tracking across timepoints
• Data are sequences of marginal distributions mu_0, mu_1, ..., mu_{T-1} rather than matched trajectories

Two coupled challenges:

1. Geometry: Trajectories must respect curved low-dimensional manifolds embedded in high-dimensional space
2. Regime switching: The transport mechanism itself may change over time (learning phases, developmental stages)

Architecture

Three-Stage Training Pipeline

Stage 1: Metric Learning

• Learn data-dependent metric G(x) from pooled labeled + unlabeled observations
• Geometry model defines low-energy paths on the data manifold
• Supports RBF metrics (low-dim) and RBF-MLP deep-kernel metrics (high-dim)

Stage 2: Bend Network Training

• Train bend network B_psi with frozen metric
• Produces geometry-aware interpolants between adjacent-marginal endpoints
• Conditional paths: z_{k,alpha} = B_psi(x_k, x_{k+1}, alpha; G)
• Tangents: z_dot_{k,alpha} = partial B_psi / partial alpha

Stage 3: Velocity Field Training

• Geometry model and bend network frozen
• Train velocity experts and router on geometry-aware path points and tangents
• L_FM = E[||v_theta(t_{k,alpha}, z_{k,alpha}) - z_dot_{k,alpha}||^2]

Regime-Switching Velocity Fields (MoE)

FLUX decomposes the longitudinal velocity field into M expert vector fields:

• f_m: [0,1] x R^d -> R^d, m = 1,...,M
• Router g: [0,1] x R^d -> R^M maps time+state to expert logits
• Straight-Through Gumbel-Softmax enables differentiable discrete routing
• At inference: regime = argmax_m l_m(t, x)

Router Regularization

Penalty	Purpose	Formula
L_div	Prevent expert collapse	1 - H(w_bar)/log(K)
L_sparse	Per-sample commitment	1 -
L_con	Temporal coherence	Var[w_i : i in group]
L_lb	Load balance	K * sum(F_j * P_j)
L_contig	Discourage temporal scattering	Var[s * w_bar_{s,j}]

Full Objective

L = L_FM + lambdaL2L_L2 + lambda_vel*L_vel + L_routing

Key Innovations

1. Geometry-aware conditional paths: Euclidean interpolants pass through low-density regions; learned metrics keep trajectories on the data manifold
2. MoE velocity decomposition: Single velocity field cannot capture regime changes; sparse expert routing discovers latent transitions
3. Joint transport + regime discovery: Regime discovery emerges from velocity field decomposition, not post-hoc clustering
4. Multi-marginal longitudinal transport: Extends flow matching beyond two-marginal settings

Benchmarks and Results

Stanford Bunny (Geometry)

• Tests manifold transport on known 3D surface
• FLUX with learned geometry stays on surface; Euclidean methods cut through interior
• Surface deviation metric validates manifold adherence

Lorenz System (Regime Discovery)

• Controlled dynamical transition at known boundary
• FLUX recovers regime boundary from transport dynamics alone
• ARI and NMI for regime assignment

Widefield Calcium Imaging (Neural Data)

• 12 mice, 451-dim cortical population activity, 22 marginals (learning sessions)
• Router separates early/intermediate training from late training
• Regime transition coincides with behavioral divergence of CS+/CS- lick indices

Embryoid Body Differentiation (Single-Cell)

• RNA-seq profiles across developmental timepoints
• FLUX separates expression-evolution regimes: pluripotent vs differentiated
• Regime recovery depends on geometry-aware dynamics

Baselines Compared

Method	Type	Limitation
Linear interpolation	Adjacent pair	No shared velocity field
Static OT (Sinkhorn)	Adjacent pair	No continuous dynamics
Independent CFM	Adjacent pair	No parameter sharing across chain
IMMFM	Multi-marginal	Euclidean paths, no regime discovery
FLUX w/o Manifold	Multi-marginal + MoE	Euclidean paths, weak regime discovery
K-Means/GMM	Clustering	Static, no dynamics
rSLDS/SRNN	Switching dynamics	Require sequential inputs

Application Scenarios

1. Neural dynamics trajectory reconstruction -- Widefield calcium imaging across learning sessions
2. Cell differentiation modeling -- Single-cell RNA-seq developmental trajectories
3. Flow matching on manifolds -- Data concentrated near curved low-dimensional manifolds
4. Regime discovery in time series -- Unsupervised detection of dynamical phase transitions
5. Longitudinal fMRI/EEG analysis -- Tracking population-level changes across sessions

Implementation Requirements

# Key components
class GeometryBackend:  # Stage 0: Learn data-dependent metric
    def fit(self, X_pooled): ...
    def metric(self, x): ...

class BendNetwork:  # Stage 1: Geometry-aware interpolants
    def forward(self, x0, x1, alpha, G): ...

class MixtureVelocityNet:  # Stage 2: MoE velocity field
    def forward(self, t, x):
        logits = router(t, x)  # Expert logits
        experts = [v_m(t, x) for v_m in velocity_nets]  # Parallel
        weights = gumbel_softmax(logits, tau)  # ST-Gumbel
        return sum(w_m * f_m for w_m, f_m in zip(weights, experts))

Training Protocol

1. Construct or load longitudinal marginals
2. Train geometry backend on pooled samples
3. Train bend network with frozen geometry
4. Train single-velocity or MoE velocity model with frozen geometry + bend
5. Evaluate one-hop, two-hop, and full-chain transport (Wasserstein distance)
6. Evaluate segment-level regime discovery (ARI, NMI) where labels available

Key Hyperparameters

Parameter	Lorenz	Neural Data	Embryoid Body
Hidden dim	64	256/128	128
Experts K	2	3	3
Velocity layers	2	2	2
Time embed dim	16	16	32

Related Work

• Flow matching: Lipman et al. 2023, Tong et al. 2024 (Conditional Flow Matching)
• Riemannian flow matching: Chen & Lipman 2024 (known manifolds)
• Metric Flow Matching: Kapusniak et al. 2024 (learned data-dependent metrics)
• Optimal transport: Peyre & Cuturi 2019 (computational OT)
• Mixture-of-experts: Shazeer et al. 2017, Fedus et al. 2022 (Switch Transformers)
• TrajectoryNet: Tong et al. 2020 (continuous normalizing flows from snapshots)

References

• Ortega Caro, J. et al. (2026). "FLUX: Geometry-Aware Longitudinal Flow Matching with Mixture of Experts." arXiv:2605.08648
• Lipman, Y. et al. (2023). "Flow matching for generative modeling." ICLR 2023
• Chen, R.T.Q. & Lipman, Y. (2024). "Flow matching on general geometries." ICLR 2024
• Kapusniak, V. et al. (2024). "Metric Flow Matching."

Activation Keywords

• longitudinal flow matching, geometry-aware transport, unpaired snapshots
• biological trajectory reconstruction, mixture-of-experts velocity field
• regime discovery neural dynamics, calcium imaging trajectory
• cell differentiation modeling, FLUX framework, manifold-aware flow matching
• 纵向流匹配, 几何感知传输, 神经动力学重建