Physics-Informed Neural Koopman Machine for Interpretable Longitudinal Personalized Alzheimer's Disease Forecasting

TL;DR

The paper introduces the Physics-informed Neural Koopman Machine (NKM) to forecast longitudinal, multivariate cognitive decline in Alzheimer's disease using multimodal biomarkers. By coupling alpha-knowledge (Koopman dynamics with hierarchical attention) and beta-knowledge (modality-specific biological encoders), NKM learns observables and a Koopman operator end-to-end, yielding interpretable, patient-specific trajectories for multiple outcomes. It achieves superior predictive performance on ADNI data, reveals disease progression dynamics in latent Koopman space, and identifies neurobiological substrates and brain regions underpinning cognitive decline. The work also provides theoretical guarantees on stability and convergence, plus extensive ablations and interpretability analyses to support clinical relevance and potential precision-medicine applications.

Abstract

Early forecasting of individual cognitive decline in Alzheimer's disease (AD) is central to disease evaluation and management. Despite advances, it is as of yet challenging for existing methodological frameworks to integrate multimodal data for longitudinal personalized forecasting while maintaining interpretability. To address this gap, we present the Neural Koopman Machine (NKM), a new machine learning architecture inspired by dynamical systems and attention mechanisms, designed to forecast multiple cognitive scores simultaneously using multimodal genetic, neuroimaging, proteomic, and demographic data. NKM integrates analytical ($α$) and biological ($β$) knowledge to guide feature grouping and control the hierarchical attention mechanisms to extract relevant patterns. By implementing Fusion Group-Aware Hierarchical Attention within the Koopman operator framework, NKM transforms complex nonlinear trajectories into interpretable linear representations. To demonstrate NKM's efficacy, we applied it to study the Alzheimer's Disease Neuroimaging Initiative (ADNI) dataset. Our results suggest that NKM consistently outperforms both traditional machine learning methods and deep learning models in forecasting trajectories of cognitive decline. Specifically, NKM (1) forecasts changes of multiple cognitive scores simultaneously, (2) quantifies differential biomarker contributions to predicting distinctive cognitive scores, and (3) identifies brain regions most predictive of cognitive deterioration. Together, NKM advances personalized, interpretable forecasting of future cognitive decline in AD using past multimodal data through an explainable, explicit system and reveals potential multimodal biological underpinnings of AD progression.

Figures (3)

• Figure 1: The schematic architecture of the Neural Koopman Machine (NKM).Right panel (a-c):NKM uses biological knowledge ($\beta$-knowledge) to guide modality-specific encoders to process Magnetic Resonance Imaging (MRI), Positron Emission Tomography (PET), cerebrospinal fluid (CSF), genetic, and demographic data with time embeddings, and to fuse the resulting learned modality-specific latent representations into a joint latent state. In parallel, NKM uses analytical knowledge ($\alpha$-knowledge) to steer hierarchical attention to combine feature-group and temporal attention and generate a control vector, which drives a Koopman state transition. The decoder then maps the predicted latent state to multiple clinical outcome trajectories (e.g., MMSE, CDR-SB, and ADAS). Left panel (inset):The learned Koopman latent-space trajectories highlight distinct dynamical patterns for subjects who are stably cognitively normal (CN $\rightarrow$ CN), with mild cognitive impairment (MCI $\rightarrow$ MCI) and Alzheimer's disease (AD $\rightarrow$ AD), as well as converters (e.g., CN $\rightarrow$ MCI and MCI $\rightarrow$ AD).
• Figure 2: Using Neural Koopman Machine to forecast personalized future disease trajectory of Alzheimer's disease. The top panels illustrate personalized longitudinal predictions for two patients—one progressing from cognitively normal (CN) to mild cognitive impairment (MCI), and the other from MCI to Alzheimer's disease (AD). Colored trajectories depict historical (light green), recent (dark green), predicted future (red), and observed future (blue) values for three cognitive assessments: MMSE, CDRSB, and ADAS13. The middle-left panel shows feature importance and selection consistency for each of the three cognitive scores, while the middle-right panel displays their aggregate, group-level counterparts (i.e., feature weights averaged across the three scores and 50 runs). In the bar–circle visualization, bar length denotes feature importance and circle size reflects the frequency with which a feature appears among the top 10 features across runs. The lower-left panel reports predictive performance for the three cognitive scores, and the lower-right panel depicts neural feature importance across the brain.
• Figure 3: The Neural Koopman Machine (NKM) uncovers unsupervised, neurobiologically grounded dynamics underlying Alzheimer's disease (AD). This visualization maps the model's predictions onto a coordinate system that mirrors biology: the x-axis aligns with hippocampal health, and the y-axis with ventricular enlargement. It is important to note that, while we used brain volumes to orient these axes, the data points are generated solely by the NKM model's internal representations. To generate these trajectories, we used a sliding window approach. The model observes a patient for three time points and then projects their likely progression for the next five steps. Because we slide this window forward one step at a time, a single subject contributes multiple overlapping trajectories(samples) as they move through the study, as noted in the figure. In the plots, thick lines show the average path for each group, while thin lines show individual variation. The shaded regions left represent the density of stable CN and AD subjects in the training set. These same boundaries are replicated as dashed outlines in the validation plot right to serve as a fixed reference. This highlights the model's consistency, showing that validation trajectories fall primarily within the stable zones defined during training, despite minor distributional shifts.

Theorems & Definitions (6)

• Theorem 1: Error bound for $\varepsilon$-approximately Koopman-invariant observables
• proof
• Corollary 1.1: Asymptotic error bound
• Theorem 2: Convergence of alternating gradient descent
• proof
• proof