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

2026-05-08

AI Summary

This work addresses the challenge of reconstructing evolutionary trajectories of biological systems and identifying underlying state-switching mechanisms from unpaired longitudinal snapshots by proposing the FLUX framework. FLUX uniquely integrates geometry-aware flow matching with a mixture-of-experts architecture: it learns data-dependent manifold metrics to construct conditional evolutionary paths and decomposes the velocity field into sparse expert vector fields, with a Straight-Through Gumbel-Softmax router dynamically selecting the dominant dynamics. Evaluated across multiple biological datasets, FLUX successfully reconstructs longitudinal transport processes and uncovers interpretable latent state structures. Ablation studies confirm that geometric modeling plays a critical role in effective state discovery.

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.

Problem

• longitudinal data
• regime switching
• manifold learning
• transport dynamics
• unpaired snapshots

Innovation

• geometry-aware flow matching
• mixture of experts
• unpaired longitudinal data
• latent regime discovery
• Riemannian metric learning