Effective Latent Differential Equation Models via Attention and Multiple Shooting

GERMAN ABREVAYA; MAHTA RAMEZANIAN-PANAHI; JEAN-CHRISTOPHE GAGNON-AUDET; PABLO POLOSECKI; IRINA RISH; SILVINA PONCE DAWSON; GUILLERMO CECCHI; GUILLAUME DUMAS

Transactions on Machine Learning Research

Online

Año: 2024

2835-8856

Scientific Machine Learning (SciML) is a burgeoning field that synergistically combinesdomain-aware and interpretable models with agnostic machine learning techniques. In thiswork, we introduce GOKU-UI, an evolution of the SciML generative model GOKU-nets.GOKU-UI not only broadens the original model’s spectrum to incorporate other classes ofdifferential equations, such as Stochastic Differential Equations (SDEs), but also integratesattention mechanisms and a novel multiple shooting training strategy in the latent space.These modifications have led to a significant increase in its performance in both reconstruction and forecast tasks, as demonstrated by our evaluation on simulated and empirical data.Specifically, GOKU-UI outperformed all baseline models on synthetic datasets even with atraining set 16-fold smaller, underscoring its remarkable data efficiency. Furthermore, whenapplied to empirical human brain data, while incorporating stochastic Stuart-Landau oscillators into its dynamical core, our proposed enhancements markedly increased the model’seffectiveness in capturing complex brain dynamics. GOKU-UI demonstrated a reconstruction error five times lower than other baselines, and the multiple shooting method reducedthe GOKU-nets prediction error for future brain activity up to 15 seconds ahead. By training GOKU-UI on resting state fMRI data, we encoded whole-brain dynamics into a latentrepresentation, learning a low-dimensional dynamical system model that could offer insightsinto brain functionality and open avenues for practical applications such as the classificationof mental states or psychiatric conditions. Ultimately, our research provides further impetusfor the field of Scientific Machine Learning, showcasing the potential for advancements whenestablished scientific insights are interwoven with modern machine learning.