[Paper Review] Reinforcement Learning based Disease Progression Model for Alzheimer's Disease

Krishnakant Saboo, Anirudh Choudhary+4 more|arXiv (Cornell University)|Jun 30, 2021

Mental Health Research Topics4 citations

TL;DR

This paper proposes a novel hybrid model combining differential equations (DEs) and reinforcement learning (RL) to predict individualized 10-year Alzheimer’s disease (AD) progression, using domain knowledge to define a reward function that optimizes cognition while minimizing metabolic cost. The RL agent infers missing biological relationships not captured by DEs, resulting in more accurate and interpretable 10-year cognition trajectory predictions than state-of-the-art deep learning models, with the model uncovering implicit compensatory mechanisms in brain function.

ABSTRACT

We model Alzheimer's disease (AD) progression by combining differential equations (DEs) and reinforcement learning (RL) with domain knowledge. DEs provide relationships between some, but not all, factors relevant to AD. We assume that the missing relationships must satisfy general criteria about the working of the brain, for e.g., maximizing cognition while minimizing the cost of supporting cognition. This allows us to extract the missing relationships by using RL to optimize an objective (reward) function that captures the above criteria. We use our model consisting of DEs (as a simulator) and the trained RL agent to predict individualized 10-year AD progression using baseline (year 0) features on synthetic and real data. The model was comparable or better at predicting 10-year cognition trajectories than state-of-the-art learning-based models. Our interpretable model demonstrated, and provided insights into, "recovery/compensatory" processes that mitigate the effect of AD, even though those processes were not explicitly encoded in the model. Our framework combines DEs with RL for modelling AD progression and has broad applicability for understanding other neurological disorders.

Motivation & Objective

• To develop a predictive model for individualized 10-year Alzheimer’s disease progression using baseline clinical and neuroimaging data.
• To address the limitation of incomplete differential equation (DE) models in capturing all biological relationships in AD by integrating reinforcement learning (RL).
• To incorporate domain knowledge—specifically, the trade-off between maximizing cognition and minimizing metabolic cost—into a reward function for RL.
• To produce interpretable, mechanistic insights into brain network compensation during AD progression.
• To outperform existing learning-based models in predicting long-term cognition trajectories on real and synthetic data.

Proposed method

• The model uses a set of known DEs to simulate the evolution of brain structure, activity, and cognition over time.
• Missing relationships between cognition, brain region size, and activity are inferred via RL, which optimizes a reward function based on domain knowledge.
• The reward function balances cognitive performance (minimizing mismatch between target and actual cognition) and metabolic cost (minimizing brain activity required for cognition).
• The RL agent selects optimal information processing distributions across brain regions (e.g., hippocampus and prefrontal cortex) at each time step to maximize the reward.
• The framework uses a simulator based on DEs and trains the RL agent to optimize the reward function over 10-year time horizons.
• A control model analysis validates that the observed compensatory behavior emerges from the dual-term reward function, not from explicit encoding.

Experimental results

Research questions

• RQ1 Can a hybrid differential equation and reinforcement learning framework accurately predict 10-year individualized cognition trajectories in Alzheimer’s disease?
• RQ2 Do implicit compensatory mechanisms in brain function emerge naturally from a reward function that balances cognitive performance and metabolic cost?
• RQ3 How does the inclusion of RL improve prediction accuracy compared to state-of-the-art deep learning models like minimalRNN?
• RQ4 What role does the trade-off between cognitive mismatch and energetic cost play in generating biologically plausible recovery processes?
• RQ5 Can the model uncover interpretable, mechanistic insights into brain network adaptation during AD progression?

Key findings

• The proposed model reduced prediction error by approximately 10% compared to the state-of-the-art deep learning model (minimalRNN) on real ADNI data.
• The model produced more realistic 10-year cognition trajectories than benchmark models, including SVR and minimalRNN.
• The model uncovered implicit 'recovery/compensatory' processes in brain function that were not explicitly encoded in the model, indicating adaptive reorganization of neural networks.
• Control experiments confirmed that the compensatory behavior emerged from the dual-term reward function—cognitive mismatch and energetic cost—rather than from model architecture alone.
• The model demonstrated that RL can effectively infer biologically plausible relationships where explicit DEs are unavailable, enhancing both accuracy and interpretability.
• Grid search-based optimization without RL led to extreme, biologically implausible information processing distributions (e.g., I(t) = [10,0] or [0,10]), highlighting the necessity of RL for stable, realistic dynamics.