Longitudinal Flow Matching for Trajectory Modeling

Simulated Trajectories

Abstract

Generative modeling of disease progression from longitudinal neuroimaging requires learning continuous, subject-specific trajectories from data that are high-dimensional, irregularly sampled, and highly sparse (e.g., occasional clinical visits). Standard pairwise transition models fall short in this setting because they do not ensure global trajectory consistency across multiple timepoints. We present Interpolative Multi-Marginal Flow Matching (IMMFM), a novel stochastic generative framework that leverages a piecewise-quadratic interpolation path as a smooth target for flow matching and jointly optimizes a data-driven uncertainty diffusion coefficient.

Figure 1: IMMFM pipeline with sparse observations and trajectory forecasting

Figure 1. IMMFM translates sparse multi-visit data into continuous, uncertainty-aware subject-specific progressions via multi-marginal path interpolation.

Methodology & Innovations

Our generative framework models trajectories as an Itô stochastic differential equation (SDE), simultaneously learning a transport flow and stochastic deviations to generate medically realistic progression.

1. MMOT as Spatial Diffeomorphic Registration

We uniquely formulate the Multi-Marginal optimal transport (MMOT) problem not as a naive pixel-matching task, but as a sequential spatial alignment and diffeomorphic registration problem. By finding smooth, invertible transformations that directly map one timepoint's image onto the next, we guarantee a topologically sound pathway prior for the continuous flow model right at the dataset foundation.

2. Piecewise-Quadratic Conditional Path

Unlike simple linear multi-marginal interpolation that creates discontinuous velocity fields, we construct a quadratic velocity formulation that smoothly blends velocities between consecutive trajectory segments. The mean path incorporates both current segment velocity and anticipated velocity from the next segment, with time-dependent blending that creates smooth transitions. This approach exhibits a beneficial low-pass filtering effect that ensures Lipschitz continuity required for stable flow matching training, effectively regularizing against ill-posed vector fields in sparse data regimes.

3. Learned Uncertainty as Trajectory Correction

Rather than predicting classical variances or calibration confidence intervals, IMMFM learns a high-dimensional directional corrective diffusion coefficient. This term operates in latent space and serves as part of the prediction mechanism rather than a separate uncertainty estimate. Jointly optimized with the drift via our theoretical equilibrium condition, it is explicitly regressed to the squared prediction error and dynamically guides the SDE trajectory toward better reconstruction.

Theoretical Guarantee: We prove that learning this uncertainty term does not bias the drift learning—the stationary points remain identical to the drift-only case. This allows stable joint optimization where the uncertainty adaptively corrects for local prediction errors while preserving the underlying trajectory dynamics.

Depending on the computational budget and physiological stochasticity, the framework acts as a versatile suite. For smoother, deterministic motions (like limb kinematics in the Starmen dataset), our ODE variant (O-IMMFM) naturally excels by focusing computational resources on drift learning. Conversely, for highly volatile progressions like tumor growth (GBM dataset), the full Stochastic SU-IMMFM formulation is deployed where the learned diffusion component effectively captures chaotic volatility that the smooth quadratic drift cannot model.

Results: Synthesis, Stability, and Insight

Synthetic Trajectory Validation

We first validate IMMFM on challenging synthetic datasets with known ground truth dynamics. The experiments involve rapidly evolving distributions with complex curvatures and crossing trajectories that stress-test the model's ability to maintain trajectory coherence.

Figure 2: IMMFM tracking on S-shaped and sigma-shaped synthetic trajectories

Figure 2. Synthetic trajectory validation on fast-changing distributions: We test on S-shaped (left) and sigma-shaped (right) trajectories where point clouds evolve through complex geometric transformations. The datasets challenge models with rapid directional changes and distribution crossings. IMMFM (red) significantly outperforms standard flow matching baselines (blue), demonstrating superior tracking of complex curvatures. The improvement stems directly from our quadratic velocity formulation that anticipates future trajectory segments, enabling smooth navigation through geometric complexity that defeats linear interpolation methods.

Following synthetic validation, we evaluated IMMFM across multiple challenging high-dimensional neuroimaging datasets including Alzheimer's Disease Neuroimaging Initiative (ADNI), Multiple Sclerosis (Brain MS), Glioblastoma (GBM), and the baseline Starmen dataset. IMMFM demonstrated state-of-the-art capability in structural synthesis and temporal consistency.

Fidelity Improvements: Models consistently beat strong baselines over extensive horizons, yielding gains of +1.0%–4.4% in anatomical Dice similarity, +1.5–2.2 dB in PSNR.
Ablation Drivers: Replacing standard linear paths with the piecewise-quadratic pathway was the single highest contributor to accuracy (up to +3.7% DSC), while our data-driven uncertainty diffusion curbed worst-case boundary errors on erratic GBM cases directly (yielding a >6 pixel Hausdorff gain).
3D Volume Scaling: Utilizing the full framework extended seamlessly up to volumetric analysis. 3D predictions lowered inter-subject variability substantially and elevated ADNI target Dice scores to 94.7% (a 3% boost over 2D).
Computational Efficiency: Despite sophisticated multi-marginal alignment, preprocessing remains practical with registration taking only 6.85 seconds per 128³ volume and 1.12 seconds per 64³ volume for real-time clinical inference.

Figure 4: Real neuroimaging forecasts with error maps

Figure 4. Left to right examples of ground truth and predictions for MS and GBM showing structural coherence and minimized spatial errors (2nd rows) across diverse modalities.

Long-Range Temporal Stability Analysis

A critical challenge in longitudinal modeling is maintaining prediction quality as forecast horizons extend beyond the training observation window. We systematically evaluate IMMFM's temporal generalization capability by extending predictions up to 4+ years beyond baseline measurements and comparing performance across different patient cohorts.

Figure 6: Model robustness across forecasting horizons and cohorts

Figure 6. Temporal generalization robustness analysis across extended forecasting horizons. Left panels: Dice similarity scores vs. prediction horizon (months) for ADNI cohorts, showing IMMFM maintains >85% structural fidelity even at 48-month extrapolations while baseline methods degrade rapidly after 24 months. Right panels: Cross-cohort validation demonstrating consistent performance across AD, MCI, and CN populations. Unlike standard transition flows that suffer from compounding iterative errors, IMMFM's multi-marginal formulation with quadratic velocity smoothing provides inherent regularization that prevents error accumulation. The uncertainty-aware SDE formulation further stabilizes long-range predictions by adaptively modulating stochastic corrections based on local prediction confidence.

Downstream Clinical Utility: Early Alzheimer's Detection

A key clinical application of accurate trajectory forecasting is early disease detection. We test whether IMMFM's generated structural progressions contain clinically meaningful biomarkers for Alzheimer's disease conversion prediction—a critical capability for enabling earlier therapeutic intervention.

Experimental Design: Using only 18-month baseline MRI data, we generate 36-month forecasts and extract ventricle volume trajectories as a key AD biomarker. We then evaluate conversion prediction between Alzheimer's Disease (AD) and Cognitively Normal (CN) subjects using a simple threshold-based classifier applied to the forecasted ventricle expansion patterns.

Clinical Impact: Our method achieves +9.1% improvement (from 71.7% to 80.8%) in AD-CN classification accuracy using only structural image forecasts—no cognitive scores, biomarkers, or AD-specific training. This translates to correctly identifying 8 out of 100 additional future AD converters 18 months earlier, providing crucial lead time for treatment planning and family preparation. Given that dementia affects >55 million people globally and AD represents 60% of cases, even modest improvements in early detection have substantial population-level impact.

Figure 5: Ventricle progression and AD-vs-CN utility plots

Figure 5. Clinical utility demonstration for early AD detection. Panel (a): IMMFM-generated mean ventricle volume trajectories (solid lines) closely track ground truth pathological progressions (dashed lines) across 36-month horizons. AD subjects show characteristic accelerated ventricular expansion (red), while CN subjects maintain stable volumes (blue). The model captures subtle nonlinear progression patterns that linear forecasting methods miss. Panel (b): ROC analysis demonstrating clear distributional separation between AD and CN trajectory forecasts. The forecasted biomarker distributions show minimal overlap (AUC=0.84), enabling confident classification decisions. Importantly, this separation emerges purely from structural image synthesis without requiring cognitive assessments or specialized AD biomarkers, making the approach broadly applicable in clinical settings with standard MRI protocols.

Conclusions & Future Directions

Key Achievements

Quantified Clinical Achievements

Early Detection: 18-month lead time for AD diagnosis with +9.1% accuracy improvement
Anatomical Fidelity: +1.0%–4.4% Dice similarity gains and +1.5–2.2 dB PSNR improvements
Quadratic Path Benefit: Up to +3.7% DSC improvement over linear interpolation
3D Volume Analysis: ADNI Dice scores of 94.7% (3% boost over 2D)

Irregular Sampling: Handles non-uniform time intervals naturally—patients need not be imaged at identical schedules
Uncertainty-Driven Precision: >6 pixel Hausdorff gains on erratic GBM cases
Real-time Feasibility: 6.85s processing per 128³ volume

Current Limitations & Design Trade-offs

Every methodological choice involves trade-offs, and IMMFM is no exception. Our quadratic interpolation scheme introduces a beneficial low-pass filtering effect that ensures numerical stability but inherently limits the model's ability to capture some high-frequency oscillations occurring precisely between sparse observations. Eventhough this may seem like a strict limitation, in practice this prevents the ill-posed learning problems common in chaotic regimes.

Additionally, model performance depends on training set diversity and can be affected by systematic data artifacts. However, these constraints are common to most data-driven approaches and represent opportunities for future methodological development rather than fundamental barriers.

Future Research Directions

1. Enhanced Representation Learning

Temporally Aware Autoencoders

Moving beyond snapshot-independent processing to develop ordering-aware architectures that ensure latent space continuity and temporal coherence.

Multi-modal Integration

Incorporating heterogeneous data streams (genomics, biomarkers, clinical scores) for more comprehensive trajectory understanding.

Self-supervised Learning

Leveraging temporal structure as supervision signal to learn more robust and generalizable representations without additional labels.

2. Causal & Counterfactual Reasoning

Transform IMMFM from a prognostic tool into a decision support system by integrating causal representation learning. This would enable simulation of trajectories under hypothetical interventions, supporting personalized treatment planning and clinical trial design optimization.

3. Hybrid Methodology Integration

IMMFM's generic nature enables combination with specialized parametric models (like Leaspy for AD progression). Such hybrid approaches could leverage IMMFM's structural forecasting with domain-specific biomarker models for comprehensive clinical assessment.

4. Domain Knowledge Integration

In domains with established physical laws, Physics-Informed Neural Networks could constrain dynamics. For complex clinical contexts, approximate mechanistic priors (e.g., network diffusion models for pathology spread) could provide gentle regularization while preserving data-driven flexibility.

BibTeX

@article{islam2025longitudinal,
  title={Longitudinal Flow Matching for Trajectory Modeling},
  author={Islam, Mohammad Mohaiminul and Kuipers, Thijs P and Vadgama, Sharvaree and de Vente, Coen and Khan, Afsana and S{\'a}nchez, Clara I and Bekkers, Erik J},
  journal={arXiv preprint arXiv:2510.03569},
  year={2025}
}