Reinforcement Learning based Disease Progression Model for Alzheimer’s Disease

1 Introduction . Models that describe Alzheimer ’ s disease ( AD ) progression through time , i.e. , the evolution of factors involved in the disease such as brain size , brain activity , pathology , and cognition ( Fig . 1 ) , are crucial for mitigating this highly prevalent disease [ 1 , 2 ] . Such models can enhance our understanding of the disease processes and enable crucial applications like prediction of long-term cognition trajectories for early detection . Our goal is to develop a model to predict an individual ’ s future AD progression at 1-year intervals based on their baseline ( year 0 ) data . We address this goal by combining differential equations that capture the relationships between some factors , and leveraging reinforcement learning to extract the missing relationships by optimizing a domain knowledge-based reward function . Differential1 equation ( DE ) based models can describe disease progression by expressing domain knowledge as mathematical relationships between different factors [ 3 , 4 ] . DE-based models have several advantages which make them attractive for disease progression modelling . The interpretability 1For simplicity , we use the term `` differential equation '' to denote algebraic and differential equations . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of these models can enhance our understanding of disease processes [ 5 , 6 ] . These models require limited data because the data is only used for parameter estimation . DE-based models ’ mechanistic nature allows for intervention exploration [ 6 ] . However , these models provide an incomplete view of the disease because DEs describing relationships among some of the factors are unavailable [ 6 ] . We propose a framework for modeling AD progression that combines differential equations with reinforcement learning ( RL ) to overcome the above limitation . We assume that the missing relationships are the solution of an optimization problem that can be formulated based on domain knowledge . In our model , this optimization problem ’ s objective function serves as the reward function for reinforcement learning . Therefore , by optimizing the reward , RL extracts the relationships among factors for which explicit DEs are unavailable . Thus , RL combined with the DEs describes the evolution of factors pertinent for modeling AD progression . In our model , the available DEs define the simulator and the RL agent optimizes the domain-based reward function in the simulator . The parameters of the DEs are based on the available data . We overcome three main challenges for developing the proposed model . The first challenge is to identify the factors involved in each DE . We use domain knowledge to find existing causal relationships between different factors and differential equations relating them ( when available ) . Second , the DEs relating cognition , brain regions ’ sizes , and brain activity are unknown ( Fig . 1 ) . Since proposing novel DEs relating these factors requires significant scientific knowledge that is still in development , we address this challenge as follows . Multiple brain regions work together to produce cognition [ 7 , 8 ] . We assume that the working of multiple brain regions is governed by an optimization problem that maximizes cognition while minimizing the cost of cognition [ 9 ] . We represent this optimization problem ’ s objective function in terms of cognition , brain size , and brain activity . Finally , the above optimization problem needs to be solved for multiple time points for modeling disease progression with the solution at time t affecting the solution at time t ′ > t. Therefore , we use RL to optimize the above objective ( reward ) function . Our contributions are : 1 . We developed an Alzheimer ’ s disease progression model based on domain knowledge that combines DEs and RL . To the best of our knowledge , this is the first attempt at using RL to develop a disease progression model . Our framework is generic and can be used with other DEs , or to provide a basis for modelling other neurological disorders . 2 . We applied the model for predicting individualized long-term ( 10-year ) future cognition trajectories using baseline ( 0th year ) features on synthetic and real data . On real data , our model reduced the prediction error by ∼ 10 % than a state-of-the-art deep learning-based model [ 10 ] and produced more realistic trajectories than other benchmark models . Cognition trajectory prediction models are useful in a clinical setting to identify individuals at future risk of cognitive decline . 3 . Our interpretable model provided insight into how multiple brain regions together produce cognition during AD . Specifically , we observed `` recovery/compensatory '' processes that mitigate the effect of AD on cognition [ 11 ] , which could not be observed with state-of-the-art models . Recovery processes were not explicitly encoded into the model and were an outcome of the reward function ’ s form . Further investigation into the basis of recovery processes could guide the development of interventions . 1.1 Background . Alzheimer ’ s disease results from the interaction of pathology ( such as amyloid [ 12 ] ) and recovery ( compensation [ 13 , 11 ] ) processes . These processes jointly affect brain structure ( the size of different regions ) , brain function ( activity in different regions ) , and cognition [ 14 , 12 , 11 , 15 ] . Pathology leads to neurodegeneration i.e. , reduction in brain size [ 12 , 16 ] , and consequently leads to cognitive decline . Recovery processes compensate for the effect of neurodegeneration on cognition through modification of brain activity by employing other regions of the brain [ 13 ] for cognition . Ironically , the compensatory processes can lead to neurodegeneration and accelerated cognitive decline in the long term [ 15 , 11 ] . An individual ’ s demographics such as gender , education , genetic risk for disease , etc. , also play a role in determining disease progression [ 17 ] . 2 Model . First , we use domain knowledge to find existing causal relationships between different factors in Sec 2.1 . In Sec 2.2 , we define the notation for the factors and substitute their causal relationships with appropriate DEs . In Sec 2.3 , we formulate an optimization problem to find missing relationships and solve it using RL . Thus , the model predicts disease progression via the interaction between the DEs ( which constitute a simulator ) and the action of the RL agent . Finally , we describe the training of the model in Sec . 2.4 . 2.1 Identification of causal interactions . We simplify the interactions between different factors by using the causal relationships between them proposed in literature ( Fig . 1 ) . Amyloid beta ( Aβ ) ( measured using florbetapir-PET ) is a primary pathological factor in AD . It propagates from one brain region to another through tracts leading to its deposition in neighboring brain regions [ 3 ] . Accumulation of amyloid directly affects brain structure ( measured using MRI ) , i.e. , brain regions sizes , and leads to neurodegeneration [ 12 ] . Brain structure predicts brain activity ( measured using fMRI ) [ 18 , 19 ] , and brain activity results in cognition ( measured using cognitive tests for e.g. , MMSE ) . Since brain activity also depends on cognitive task difficulty [ 20 ] , we define a hypothetical variable , Ctask , which represents the cognitive demand on the brain and can be thought of as the maximum score on a cognitive test . Ctask directly influences brain activity . Brain activity has an energetic cost , and itself can lead to neurodegeneration [ 11 ] . Developing a model for AD progression is equivalent to describing the differential equation corresponding to each edge in Fig . 1 . 2.2 Specification of differential equations . Brain structure : We represent the brain as the graph GS = ( V , E ) , where a node v ∈ V represents a brain region , and an edge e ∈ E represents a tract . Let Xv ( t ) denote the size of a brain region v ∈ V at time t , and X ( t ) = [ X1 ( t ) , X2 ( t ) , ... , X|V | ( t ) ] . Pathology propagation : To incorporate the propagation and accumulation of pathological Aβ in our framework , we adapt the network diffusion based model proposed by Raj et al . [ 3 , 5 ] because it captures the propagation of Aβ through tracts . Let Dv ( t ) be the instantaneous amyloid accumulation in region v ∈ V at time t. Then propagation of amyloid is given by : dD ( t ) dt = −βHD ( t ) ( 1 ) where , D ( t ) = [ D1 ( t ) , D2 ( t ) , . . . , D|V | ( t ) ] , H is the Laplacian of the adjacency matrix of the graph GS , and β is a constant . The total amount of amyloid in a region , ϕv ( t ) , is : ϕv ( t ) = ∫ t 0 Dv ( s ) ds ( 2 ) Brain activity and cognition : Multiple brain regions work in synchrony to produce cognition [ 8 ] . Let Yv ( t ) denote the activity in region v ∈ V in support of cognition C ( t ) at time t. Although cognition , brain size ( Xv ) , and activity ( Yv ) are related , the exact relationship among them is unknown and can not be easily learned from limited data . On the other hand , we can intuitively relate a region ’ s size and activity to its `` contribution '' to cognition [ 21 , 13 ] . Therefore , we introduce a hypothetical variable termed as information processing , Iv ( t ) ∈ R≥0 , to represent region v ’ s contribution to cognition ( Fig . 2 ) . Therefore , the cognition C ( t ) supported by the brain at time t can be described as : C ( t ) = ∑ v∈V Iv ( t ) . ( 3 ) The activity in a given region depends on the amount of information it is processing and its size [ 21 , 13 ] . For a region v with size Xv ( t ) , the activity Yv ( t ) increases as Iv ( t ) increases [ 21 ] . A healthier region ( i.e. , greater region size ) can be considered more efficient and will require lesser activation than an inefficient one for comparable levels of information processing [ 20 ] . Therefore , we propose the following relationship between Xv ( t ) , Yv ( t ) , and Iv ( t ) with γ as a constant : Yv ( t ) = γ Iv ( t ) Xv ( t ) ∀v ∈ V ( 4 ) Iv ( t ) s are determined through reinforcement learning as described in Sec . 2.3 . Energetic cost : The brain consumes energy for supporting cognition . This energy consumption can be thought of as a cost to the brain . The energy consumption in a region is proportional to its activity [ 22 ] . Therefore , the energetic cost to the brain is the total brain activity across all the regions . We compute the energetic cost M ( t ) as follows : M ( t ) = ∑ v∈V Yv ( t ) ( 5 ) Degeneration of brain regions : Neurodegeneration is influenced by two factors : amyloid deposition and brain activity . Previous studies that proposed a linear relationship between the rate of brain region degeneration and Aβ deposition replicated the macroscopic neurodegeneration patterns seen in AD [ 5 , 4 ] . Moreover , brain activity supporting cognition can further accelerate neurodegeneration [ 11 ] . Therefore , we adopted the relationship proposed in [ 5 ] and modified it to incorporate the effect of brain activity as follows : dXv ( t ) dt = −α1Dv ( t ) − α2Yv ( t ) ∀v ∈ V ( 6 ) where α1 , α2 are constants . The first term on the RHS in the Eq . 6 is the degeneration due to Aβ pathology and the second term is the deterioration due to brain activity ( see Appendix C.2 for an alternate formulation of this equation ) . Parameter constants of equations : An individual ’ s demographics , such as gender , genetic risk of AD , education , etc. , influence disease progression [ 17 ] . In our model , the influence of demographics is mediated through the parameter constants α1 , α2 , β , γ introduced in Eqs 1 , 4 , and 6 [ 5 ] . For demographic factors Z0 at baseline , let ( α1 , α2 , β , γ ) = f ( Z0 ) . ( 7 ) 2.3 Reinforcement learning to determine I ( t ) Specifying the model requires determining the information processed by each region , i.e. , computing I ( t ) = [ I1 ( t ) , . . . , I|V | ( t ) ] . Multiple brain regions jointly produce cognition . Those joint relationships influence I ( t ) . We assume that those joint relationships satisfy general criteria about the working of the brain and obtain I ( t ) by solving an optimization problem capturing those criteria . Specifically , we assume that I ( t ) is chosen to distribute the workload related to cognition `` optimally '' across multiple brain regions . The objective function of that optimization problem is based on two competing criteria : ( i ) minimizing the deficit between the cognitive demand of a task and the cognition C ( t ) provided by the brain , and ( ii ) minimizing the energetic cost M ( t ) of supporting cognition [ 9 ] . The optimization is performed at every time point t with the solution at t affecting the optimization problem at t ′ > t. Therefore , we use RL to determine the optimal I ( t ) by balancing the two competing criteria . RL is a desirable method in our case since it can train with limited data and also makes the framework generic and flexible to changes/refinements in the reward function and the underlying DEs . Below we describe how RL is adopted in our model ( Fig . 2 ) . We implement the model for discrete time steps t ∈ { 1 , 2 , ... , K } . State and action : The state S ( t ) consists of the current size of the brain regions X ( t ) and the information processed by each region at the previous time point I ( t− 1 ) . The latter is provided in the state for ease of training the RL model . The action A ( t ) ∈ A specifies the change in information processed by each region from the previous time point , ∆Iv ( t ) ∈ R ∀v ∈ V . S ( t ) = ( X ( t ) , I ( t− 1 ) ) , A = { [ ∆I1 ( t ) , . . . , ∆I|V | ( t ) ] ∣∣∣∣Iv ( t ) = Iv ( t− 1 ) + ∆Iv ( t ) ; ∑ v∈V Iv ( t ) ≤ Ctask } Environment : The environment is a simulator consisting of the equations relating D ( t ) , ϕ ( t ) , X ( t ) , Y ( t ) , I ( t ) , C ( t ) , and M ( t ) ( Eqs . 1 , 2 , 3 , 4 , 5 , and 6 ) . Based on the action A ( t ) , the environment updates the state and provides a reward to the RL agent . Reward : The policy agent must balance the trade-off between the competing criteria of ( i ) reducing mismatch between Ctask and C ( t ) and ( ii ) reducing the cost M ( t ) of supporting cognition . Hence , we define the reward as follows : R ( t ) = − [ λ ( Ctask − C ( t ) ) +M ( t ) ] ( 8 ) where λ is a parameter controlling the trade-off between the mismatch and the cost . The goal of the policy agent is to maximize the reward .

Summary: The paper proposes a disease progression model for Alzheimer's disease (AD) that combines data-driven approaches capture trends in data over the course of the disease, which have proved successful over recent times, with more mechanistic biophysics-based models of how pathology spreads over the brain.  The authors introduce the use of reinforcement learning to learn update rules for regional brain quantities as a function of time thereby informing interactions between processes and measurements of neurodegeneration. Experiments with simulations confirm basic efficacy of the estimation procedure; experiments with ADNI data show some promising predictive power and demonstrate potential insights into disease biology.