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

AI summary

Overview

• Research area: Machine learning for longitudinal Alzheimer's disease (AD) forecasting, combining Koopman operator theory (a dynamical-systems framework) with hierarchical attention over multimodal clinical, imaging, and molecular data.
• Technical level: Advanced. The paper assumes familiarity with dynamical systems, Koopman operator theory, attention mechanisms, and neural network training, though its clinical motivation and results are readable by a broader audience.
• Scope: This paper introduces the Neural Koopman Machine (NKM), a physics-informed architecture that forecasts multiple Alzheimer's-related cognitive scores simultaneously from past multimodal data while producing interpretable, biologically grounded explanations.

What This Paper Is About

Forecasting how an individual's cognition will decline in Alzheimer's disease is difficult because the disease is multimodal (genetic, proteomic, imaging, demographic), affects several cognitive functions at once, and follows nonlinear, highly heterogeneous trajectories across patients. Existing machine learning methods often ignore inter- and intra-modality relationships, predict only a single outcome, or act as black boxes that cannot be interpreted clinically. The authors' goal is to solve "longitudinal many-to-many forecasting" — using past multimodal data to predict future multivariate outcomes while keeping the model biologically interpretable.

Key Contributions

1. Simultaneous learning of observables and Koopman dynamics. Unlike classical Extended Dynamic Mode Decomposition (EDMD), which requires predefined observables, NKM jointly learns the lifting function (neural encoders) and the Koopman matrix in an end-to-end fashion, and extends autonomous Koopman dynamics with an attention-derived control vector for personalized evolution.
2. A biologically-informed feature architecture. The input space is decomposed into biologically meaningful modality groups — genetic, CSF, PET, MRI (including Schaefer's 17 networks), and demographics — each processed by a dedicated modality-specific encoder before fusion. This "β-knowledge" preserves modality structure and guides aggregation (e.g., Schaefer regions into 17 networks).
3. Hierarchical attention-based control. A "Fusion Group-Aware Hierarchical Attention" mechanism combines feature-group attention and temporal attention to produce a control vector that steers the Koopman state transition, allowing the model to weight recent visits more heavily and to identify which biomarkers matter for each forecast.
4. Theoretical stability guarantees. The authors prove an error bound for ε-approximately Koopman-invariant observables showing that multi-step prediction error grows in a controlled way with the number of steps τ, plus a corollary that the error approaches a finite limit as τ → ∞, motivating the spectral regularization they apply to the Koopman matrix.

Main Findings

• NKM leads on correlation and MAE. Averaged across the three cognitive scores and five folds, NKM achieved Pearson r = 0.5397 ± 0.0237, Spearman ρ = 0.5024 ± 0.0262, and MAE = 3.8810 ± 0.3710. It outperformed all baselines on Pearson r, Spearman ρ, and MAE.
• NKM's RMSE is slightly higher than two baselines. NKM's RMSE was 5.676 ± 0.549, slightly higher than LSTM (5.5955 ± 0.5097) and XGBoost (5.6461 ± 0.4888), while lower than the remaining baselines.
• Baseline comparison. LSTM: r = 0.5070 ± 0.0309, ρ = 0.4856 ± 0.0237, MAE = 4.0359 ± 0.3380, RMSE = 5.5955 ± 0.5097. AttnRNN: r = 0.5068 ± 0.0240, ρ = 0.4810 ± 0.0187, MAE = 4.0933 ± 0.2877, RMSE = 5.6918 ± 0.4038. TimesNetTiny: r = 0.4983 ± 0.0255, ρ = 0.4729 ± 0.0265, MAE = 5.1145 ± 1.9693, RMSE = 6.7831 ± 2.0315. XGBoost: r = 0.4926 ± 0.0247, ρ = 0.4727 ± 0.0121, MAE = 4.1345 ± 0.2760, RMSE = 5.6461 ± 0.4888. RandomForest: r = 0.4829 ± 0.0259, ρ = 0.4741 ± 0.0108, MAE = 4.1919 ± 0.2766, RMSE = 5.6748 ± 0.4862. LinearRegression: r = 0.4750 ± 0.0377, ρ = 0.4786 ± 0.0204, MAE = 4.2613 ± 0.2695, RMSE = 5.6804 ± 0.4709. Neural ODE: r = 0.3784 ± 0.0501, ρ = 0.4099 ± 0.0736, MAE = 4.667 ± 0.220, RMSE = 6.673 ± 0.399. EDMD (RBF basis): r = 0.3706 ± 0.0607, ρ = 0.4117 ± 0.0344, MAE = 4.539 ± 0.344, RMSE = 6.413 ± 0.642.
• Latent space separates diagnostic groups and transitions. In the learned Koopman latent space, NKM shows distinct clustered representations and distinct paths for stable groups (CN→CN, MCI→MCI, AD→AD) and for converters (CN→MCI and MCI→AD).
• Latent trajectories align with known neurobiology. One visualization maps model predictions onto axes aligned with hippocampal health (x-axis) and ventricular enlargement (y-axis). The authors emphasize the axes are used only for orientation and the data points come solely from NKM's internal representations. Validation trajectories fell primarily within the stable zones defined during training, despite minor distributional shifts.
• Simultaneous multivariate prediction. A single decoder head predicts MMSE, CDRSB, and ADAS13 together, rather than training three separate networks.
• Feature importance is quantified and stable. Feature importance and selection consistency are reported for each of the three cognitive scores, plus an aggregate group-level version averaged across the three scores and 50 runs; bar length denotes feature importance and circle size reflects how often a feature appears among the top 10 across runs.
• Brain-region mapping. The paper reports neural feature importance mapped back onto the brain, identifying brain regions most predictive of cognitive deterioration.

Methodology in Plain English

The authors collected longitudinal multimodal data from the Alzheimer's Disease Neuroimaging Initiative (ADNI) database. ADNI was launched in 2003 as a public-private partnership led by Principal Investigator Michael W. Weiner, MD, to test whether serial MRI, PET, other biological markers, and clinical and neuropsychological assessment can be combined to measure progression of mild cognitive impairment (MCI) and early Alzheimer's disease.

Each time step is described by 44 features: one genetic feature (APOE4), three CSF features (Amyloid-beta, t-tau, and p-tau protein), three PET imaging features (FDG, PIB, and AV45), 18 MRI features (17 Schaefer networks assembled from 202 cortical thickness regions plus Intracranial Volume), and demographic information (two numeric and 17 categorical features). The core biological feature set without demographics is 25 features (17 Schaefer networks, three PET features, three CSF features, one APOE4 feature, and one ICV feature).

The model works in five stages:

1. Multimodal feature encoding. Each modality is passed through its own encoder, and the resulting modality-specific latent vectors are fused into a joint latent state. The encoders learn data-driven approximations of the Koopman observables.
2. Temporal representation. A residual block (five blocks) refines the encoded latent vectors across a temporal window of width three, preserving the original signal while adding temporal structure.
3. Hierarchical attention control. Scaled dot-product attention computes a time-weighted summary (weighting recent visits more heavily) and a modality-weighted summary, which a learned sigmoid gating function fuses into a single control vector.
4. Koopman state transition. The next latent state is computed as the learned Koopman matrix applied to the refined latent state, plus the attention-derived control vector. The Koopman matrix captures short-term variations around the mean trajectory, while the control vector encodes long-term progression and patient-specific deviations.
5. Outcome decoder. A regression head maps the evolved latent state to the three cognitive scores at the next visit.

Training minimizes a composite loss: a prediction error term, a Koopman consistency term requiring the predicted latent evolution to match the true encoded next state, and a spectral regularizer that keeps the norm of the Koopman matrix below a target upper bound to ensure stable long-horizon predictions. Optimization used AdamW with early stopping on validation loss.

For evaluation, the authors included subjects with at least four time points and excluded those with fewer, yielding final sequences from 949 patients. They built temporal windows of width three from consecutive visits: features from the first three time points predict outcomes at the fourth, which are compared against observed outcomes. They used subject-stratified five-fold cross-validation, KNN imputation within each fold, and standardization parameters estimated only from training data to prevent leakage. Models were trained for up to 300 epochs with a batch size of 64, with convergence typically occurring earlier due to early stopping. Performance was measured with Pearson's and Spearman's correlation coefficients, MAE, and RMSE between predicted and observed scores at the next visit, averaged across the three cognitive scores.

Why This Matters

Impact on research. The paper shows that combining a dynamical-systems formalism (Koopman operator theory) with structured, biologically informed neural network design can beat both conventional machine learning and standard deep learning baselines on multivariate AD forecasting while remaining interpretable. It also provides theory (bounded multi-step and asymptotic prediction error) that connects spectral regularization to stable long-horizon forecasts, which is a reusable idea for other progressive-disease modeling problems. Unlike prior applications of Koopman theory in healthcare, NKM directly incorporates biological insights into the analytical framework.

Real-world applications:

• Early clinical forecasting. Personalized projections of MMSE, CDRSB, and ADAS13 trajectories could support earlier evaluation and management of cognitive decline.
• Clinical trial enrichment and stratification. Identifying likely converters (CN→MCI and MCI→AD) could help select or stratify participants in trials.
• Biomarker prioritization. Quantifying differential biomarker contributions per cognitive score could help focus data collection on the most informative modalities.
• Brain-region targeting. Identifying brain regions most predictive of cognitive deterioration could inform hypotheses about where pathology drives decline.

Industry relevance. The author affiliations include Roche Diagnostics International Ltd. (Rotkreuz, Switzerland) and F. Hoffmann-La Roche Ltd., Roche Information Solutions (Basel, Switzerland), alongside the Platform of Bioinformatics at CHUV/UNIL and the University of Basel. The paper explicitly frames its aims around precision medicine in neurodegeneration, enabling early intervention, patient stratification in clinical trials, and potentially the discovery of new therapeutic targets.

Future Directions

• External validation beyond ADNI. The current work is confined to the ADNI cohort; whether the NKM's accuracy and latent-space structure hold in other cohorts, clinical settings, or more diverse populations is not established in the provided content.
• Clinical utility testing. The paper positions NKM for early intervention and trial stratification, but no prospective or interventional evaluation is reported, so whether its forecasts change clinical decisions or trial outcomes remains open.
• Theoretical assumptions in practice. Theorem 1 requires the learned observables to be ε-approximately Koopman invariant. How well that assumption holds on real data, and how the error bound behaves empirically, is discussed theoretically but not quantified in the provided results.
• Design choices left to the appendix. The paper defers discussion of window sizes to Sec. VIII-D4 and additional convergence and error-bound analysis to Sec. VIII-A1; these sections are not included in the provided content, so the sensitivity of NKM to window size and latent dimension (the value of d_z is not numerically specified in the provided content) is not reported here.

Target Audience

This paper benefits machine learning researchers working on dynamical systems, sequence forecasting, and physics-informed neural networks; computational neuroscientists and biomedical data scientists working on Alzheimer's disease progression; and clinical or translational researchers and industry scientists interested in interpretable, personalized forecasting for neurodegeneration. Readers with a background in deep learning and time-series modeling will get the most out of the architecture and theory sections, while the results, latent-space visualizations, and feature-importance analyses are accessible to a clinically oriented audience.

Authors’ 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.