Trajectory Flow Matching (TFM) is a method that leverages the flow matching technique from generative modeling to model time series. This approach offers a simulation-free training process, allowing for efficient fitting of stochastic differential equations to time-series data. Augmented with memory, time interval prediction, and uncertainty prediction, TFM can better model irregularly sampled trajectories with stochastic nature, for example clinical time series.

The idea of TFM lies in using flow matching concept to predict both stochastic uncertainty and the next value in the time series. The prediction is conditioned on past data and conditional variables.

# clone project
git clone https://github.com/nZhangx/TrajectoryFlowMatching.git
cd TrajectoryFlowMatching

# [OPTIONAL] create conda environment
conda create -n tfm python=3.10
conda activate tfm

# install requirements
conda env create -f environment.yml

Under src, create new a DATA_NAME.yml under conf/data and a MODEL_NAME.yml under conf/model with desired configurations. Then replace data and model definitions in conf/config.yaml with your DATA_NAME and MODEL_NAME. Then run

python src/main.py

We have included an example of TFM modeling three crossing oscillations in a self-contained Jupyter notebook notebook/3Oscillation.ipynb.

• TFM and ablations (size ablated, uncertainty ablated)
• Aligned Flow Matching [Liu et al., 2023][Somnath et al., 2023]
• NeuralODE [Chen et al., 2018]
• NeuralSDE [Li et al., 2020] [Kidger et al., 2021]
• LatentODE [Rubanova et al. 2019]

We plan to share the clinical data we used that are from the eICU Collaborative Research Database v2.0 (ICU sepsis and ICU Cardiac Arrest) and the Medical Information Mart for Intensive Care III (MIMIC-III) critical care database (ICU GIB) on Physionet.

This repository contains the code to reproduce the main experiments and illustrations of the preprint Trajectory Flow Matching with Applications to Clinical Time Series Modeling. We are excited that it was marked as a spotlight presentation.

If you find this code useful in your research, please cite (expand for BibTeX):

@article{TFM,
	title        = {Trajectory Flow Matching with Applications to Clinical Time Series Modelling},
	author       = {Zhang, Xi and Pu, Yuan and Kawamura, Yuki and Loza, Andrew and Bengio, Yoshua and Shung, Dennis and Tong, Alexander},
	year         = 2024,
	journal      = {NeurIPS},
}

This repo is licensed under the MIT License.