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

TL;DR

FLUX is a geometry-aware framework that models biological state transitions from unpaired longitudinal data, enabling joint transport modeling and regime discovery with interpretable results.

Contribution

Introduces FLUX, a novel geometry-aware flow-matching method using mixture-of-experts for unsupervised regime discovery in biological longitudinal data.

Findings

FLUX accurately reconstructs longitudinal transport across various biological systems.

Geometry-aware learning improves regime discovery compared to non-geometric approaches.

Mixture-of-experts routing enhances interpretability and regime identification.

Abstract

Many biological systems evolve through continuous local dynamics while switching between latent regimes defined by learning, stimulus context, internal state, or developmental stage. These processes are often observed only as unpaired longitudinal snapshots: the same cells, neurons, or animals are not tracked as matched trajectories, even though population states are sampled across successive stages. This creates two coupled challenges. First, trajectories must respect curved low-dimensional manifolds embedded in high-dimensional biological measurements. Second, the model must identify when the transport mechanism itself changes. We introduce FLUX (FLow matching for Unpaired longitudinal data with miXture-of-experts), a geometry-aware longitudinal flow-matching framework for joint transport modeling and unsupervised regime discovery. FLUX learns a data-dependent metric from pooled labeled and unlabeled observations, uses that metric to construct geometry-aware conditional paths between adjacent marginals, and decomposes the resulting velocity field into sparse expert vector fields selected by a Straight-Through Gumbel-Softmax router. Across manifold controls, a regime-switching Lorenz system, widefield cortical calcium imaging during associative learning, and embryoid body single-cell differentiation, FLUX reconstructs longitudinal transport while recovering interpretable regime structure. Ablations show that mixture-of-experts routing alone is insufficient: FLUX without geometric learning can fit local transport but fails or weakens regime discovery when regimes are encoded in local dynamics. These results suggest that geometry-aware velocity decomposition provides a general strategy for discovering latent biological state transitions from unpaired longitudinal snapshots.