Tri-Scale Neural ODEs for Continuous Multi-Omics Disease Modeling
Shohaib Shaffiey ⋅ Massimiliano Pierobon
Abstract
The fields of AI-based disease fingerprinting, drug discovery and repurposing are currently among the emerging frontiers of machine learning applied to medicine. One major challenge is to obtain robust $\textit{in-silico}$ modeling of disease progression while accounting for the vastly different time scales of biochemical interactions, from gene expression to protein abundance and metabolic flux. Discrete sequence models inadequately represent such multi-scale interactions, and standard Neural Ordinary Differential Equations (NODEs) often fail to train stably under stiffness (different time scales). To address this, a Tri-Scale Stiff NODE, defined by hierarchically coupled latent differential equations that model the causal relationships from genes to proteins and metabolites, is introduced and optimized in this paper in terms of reconstruction error and information-theoretic mutual information. This enables continuous-time modeling of cellular responses to identify not only the disease dynamics, but also drug perturbations that act within narrow time windows, often invisible to discrete-time approaches. Lyapunov analysis provides a theoretical guarantee that the modeled trajectories remain stable and well-behaved even under extreme stiffness. The methodology is validated using the STATegra B-cell and Traxler macrophage datasets, with the former utilized for a proof-of-concept drug repurposing pipeline.
Lay Summary
Biological systems operate on different timescales, with slow genetic changes driving rapid metabolic reactions. Modeling this requires coupled differential equations, but the different timescales create "stiff systems" that cause standard machine learning models to crash. We introduce a new architecture that mitigates this stiffness by separating the processes into distinct spatial and temporal scales. Our framework successfully models the continuous biological processes underlying disease, allowing us to screen compounds and discover new, targeted drug repurposing candidates at specific time windows.
Successful Page Load