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: This work presents a method for modeling dynamic systems that combines mechanistically motivated differential equations with data driven models learned by RL.  The method is applied to AD progression, measured by MMSE (a cognitive function test) scores.  The key idea is that a mechanistic understanding of the system under study may be incomplete -- e.g., in a system of coupled differential equations, a quantity and its time evolution may be unspecified -- but that it may be possible to fill this gap by learning a reasonable model for the missing function.  This method is applied to AD progression in a real world dataset with mixed results - although adding the learned dynamics model improves accuracy (measured by MSE) a great deal relative to the DE baseline, it does not do as well as a pure ML based baseline (SVR).