Trajectory Flow Matching with Applications to Clinical Time Series Modeling

2024-10-28

arXiv.org

AI Summary

This work addresses the instability and poor scalability of Neural Stochastic Differential Equations (Neural SDEs) in modeling irregularly sampled clinical time series. We propose Trajectory Flow Matching (TFM), a novel training paradigm that bypasses SDE simulation and backpropagation through time. We establish, for the first time, the theoretical necessary conditions for applying TFM to time-series modeling, design a stable reparameterization strategy tailored to Neural SDEs, and introduce the first simulation-free training framework specifically for clinical time-series data. Evaluated on three real-world clinical datasets, our method achieves significant improvements in predictive accuracy and uncertainty calibration, while exhibiting enhanced training stability, faster convergence, and reduced computational overhead. The approach offers a new paradigm for modeling high-noise, sparse, and asynchronous medical time series.

Technology Category

• Machine Learning: Time-Series/Data Streams
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Modeling stochastic and irregularly sampled time series is a challenging problem found in a wide range of applications, especially in medicine. Neural stochastic differential equations (Neural SDEs) are an attractive modeling technique for this problem, which parameterize the drift and diffusion terms of an SDE with neural networks. However, current algorithms for training Neural SDEs require backpropagation through the SDE dynamics, greatly limiting their scalability and stability. To address this, we propose Trajectory Flow Matching (TFM), which trains a Neural SDE in a simulation-free manner, bypassing backpropagation through the dynamics. TFM leverages the flow matching technique from generative modeling to model time series. In this work we first establish necessary conditions for TFM to learn time series data. Next, we present a reparameterization trick which improves training stability. Finally, we adapt TFM to the clinical time series setting, demonstrating improved performance on three clinical time series datasets both in terms of absolute performance and uncertainty prediction.

Problem

• Modeling irregular clinical time series
• Improving Neural SDE training stability
• Enhancing clinical dataset performance

Innovation

• Simulation-free Neural SDE training
• Flow matching technique utilization
• Reparameterization trick for stability