Finding Patient Zero: Learning Contagion Source with Graph Neural Networks

1 INTRODUCTION . The ability to quickly identify the origin of an outbreak , or “ finding patient zero ” , is critically important in the effort to contain an emerging epidemic . The identification of early transmission chains and the reconstruction of the possible paths of diffusion of the virus can be the difference between stopping an outbreak in its infancy and letting an epidemic unfold and affect a large share of a population . Hence , solving this problem would be instrumental in informing and guiding contact tracing efforts carried out by public health authorities , allowing for optimal resource allocation that can maximize the probability of an early containment of the outbreak . Disease spreading is modeled as a contagion process on a network Stroock & Varadhan ( 2007 ) ; Pastor-Satorras et al . ( 2015 ) of human-to-human interactions where infected individuals are going to transmit the virus by infecting ( with a certain probability ) their direct contacts . In general , contagion processes can capture a wide range of phenomena , from rumor propagation on social media to virus spreading over cyber-physical networks Centola & Macy ( 2007 ) ; Baronchelli ( 2018 ) ; Wang et al . ( 2013 ) ; Mishra & Keshri ( 2013 ) . Therefore , learning the source of a contagion process would also have broader impact on various domains , from detecting sources of fake news to defending malware attacks . Learning the index case , or patient zero ( P0 ) , is a difficult problem . In this paper , we model disease spreading as a contagion process ( chains of transmissions ) over a graph . The evolution of an outbreak is noisy and highly dependent on the graph structure and disease dynamics . In addition , in real-world epidemics , there is often a delay from the start of the outbreak to when epidemic surveillance and contact tracing starts . Hence , we might only observe the state of the graph at some intermediate times without access to the complete chains of transmission . Furthermore , due to its stochastic nature , the same source node might lead to different epidemic spreading trajectories . Finally , learning P0 from noisy observations of graph snapshots is computationally intractable and the complexity grows exponentially with the size of the graph Shah & Zaman ( 2011 ) . Most work in learning the dynamics of a contagion process Rodriguez et al . ( 2011 ) ; Mei & Eisner ( 2017 ) ; Li et al . ( 2018a ) have focused on inferring the forward dynamics of the diffusion . In epidemiology , for example , Pastor-Satorras & Vespignani ( 2001 ) have studied learning the temporal dynamics of diseases spreading on mobility networks . The problem of learning the reverse dynamics and identifying diffusion sources has been largely overlooked due to the aforementioned challenges . Two of the most notable exceptions in the area are “ rumor centrality ” Shah & Zaman ( 2011 ) for contagion processes on trees and Dynamic Message-passing ( DMP ) on graphs Lokhov et al . ( 2014 ) but both require as input the parameters of the spreading dynamics simulations . Our goal is to provide fresh perspectives on the problem of finding patient zero using graph neural networks ( GNNs ) Gilmer et al . ( 2017 ) . First , we conduct a rigorous analysis of learning P0 based on the graph structure and the disease dynamics , allowing us to find conditions for identifying P0 accurately . We test our theoretical results on a set of epidemic simulations on synthetic graphs commonly used in the literature Erdös et al . ( 1959 ) ; Albert & Barabási ( 2002 ) . We also evaluate our method on a realistic co-location network for the greater Boston area , finding performance similar to the synthetic data . While collecting labeled data to train GNN to find P0 may not be possible , training GNN using simulations on real contact-tracing data can provide a fast method for inferring P0 and help with planning and resource allocation . To the best of our knowledge , our work is the first to tackle the patient zero problem with deep learning and to test the approach on a realistic contact network . In summary , we make the following contributions : • We find upper bounds on the accuracy of finding patient zero in graphs with cycles , independent of the inference algorithm used . • We show that beyond a certain time scale the inference becomes difficult , highlighting the importance of swift and early contact-tracing . • We demonstrate the superiority of GNNs over state-of-the-art message passing algorithms in terms of speed and accuracy . Most importantly , our method is model agnostic and does not require the epidemic parameters to be known . • We validate our theoretical findings using extensive experiments for different epidemic dynamics and graph structures , including a real-world co-location graph of the COVID-19 outbreak . 2 RELATED WORK . Learning contagion dynamics Learning forward dynamics of contagion processes on a graph is a well studied problem area . For instance , Rodriguez et al . ( 2011 ) ; Du et al . ( 2013 ) proposed scalable algorithms to estimate the parameters of the underlying diffusion network , a problem known as network inference . Deep learning has led to novel neural network models that can learn forward dynamics of various processes including neural Hawkes processes Mei & Eisner ( 2017 ) and Markov decision processes-based reinforcement learning Li et al . ( 2018a ) . Learning forward contagion dynamics have also been intensively studied in epidemiology Pastor-Satorras & Vespignani ( 2001 ) ; Vynnycky & White ( 2010 ) , social science Matsubara et al . ( 2012 ) , and cyber-security Prakash et al . ( 2012 ) . In contrast , research in learning the reverse dynamics of contagion processes is rather scarce . Influence maximization Kempe et al . ( 2003 ) , for instance , finds a small set of individuals that can effectively spread information in a graph , but only maximizes the number of affected nodes in the infinite time limit . Our problem is more difficult as we care not just about the number of infected nodes , but which nodes were infected . Finding patient zero In order to find patient zero , we aim to learn the reverse dynamics of contagion processes . Shah & Zaman ( 2011 ) were among the first to formalize the problem on trees in the context of modeling rumor spreading in a network . Prakash et al . ( 2012 ) ; Vosoughi et al . ( 2017 ) studied similar problems for detecting viruses in computer networks . More recent advances proposed a dynamic message passing algorithm Lokhov et al . ( 2014 ) and belief propagation Altarelli et al . ( 2014 ) to estimate the epidemic outbreak source . Fairly recently , Fanti & Viswanath ( 2017 ) reduced the deanonymization of Bitcoin to the source identification problem in an epidemic and analyzes the dynamics properties . On the theoretical side , Shah & Zaman ( 2011 ) ; Wang et al . ( 2014 ) analyzed the quality of the maximum likelihood estimator and rumor centrality , but only for the simple SI model on trees . Antulov-Fantulin et al . ( 2015 ) found detectability limits for patient zero in the SIR model using exact analytical methods and Monte Carlo estimators . Khim & Loh ( 2016 ) ; Bubeck et al . ( 2017 ) proved that it is possible to construct a confidence set for the predicted diffusion source nodes with a size independent of the number of infected nodes over a regular tree . Our work provides fresh perspectives on the patient zero problem on general graphs based on the recent development of graph neural networks Graph neural networks Graph neural networks have received considerable attention ( see several references in Bronstein et al . ( 2017 ) ; Zhang et al . ( 2018 ) ; Wu et al . ( 2019 ) ; Goyal & Ferrara ( 2018 ) ) . While most research is focused on static graphs , a few have explored dynamic graphs Li et al . ( 2018b ) ; You et al . ( 2018 ) ; Kipf et al . ( 2018 ) ; Pareja et al . ( 2019 ) ; Trivedi et al . ( 2019 ) . For example , Kipf et al . ( 2018 ) propose a deep graph model to learn both the graph attribute and structure dynamics . They use a recurrent decoder to forecast the node attributes for multiple time steps ahead . Trivedi et al . ( 2019 ) take a continuous-time modeling approach where they take the node embedding as the input and model the occurrence of an edge as a point process . Xu et al . ( 2020 ) propose a temporal graph attention layer to learn the representations of temporal graphs . However , most research is designed for link prediction tasks and none of these existing studies have studied the problem of learning the source of the dynamics on a graph . 3 CONTAGION PROCESS AND PATIENT ZERO . Finding patient zero means tracing the contagion dynamics back to its initial state and identifying the first nodes that started spreading . Here , we describe the disease dynamics on a network using Susceptible-Infected-Recovered ( SIR ) and Susceptible-Exposed-Infected-Recovered ( SEIR ) Kermack & McKendrick ( 1927 ) compartmental models that assume that infected individuals develop immunity once they recover from the infections . 3.1 CONTAGION PROCESSES ON NETWORKS . In the SIR model , the population is split into three compartments : susceptible ( S ) who are susceptible to infection by the disease ; infected ( I ) who have caught the disease and are infectious ; removed ( R ) who are removed from consideration after experiencing the full infectious period . Continuous time model For a contagion process on a graph G with N nodes , each vertex represents an individual who is in contact only with its neighbors . We can represent the graph using the adjacency matrix A ∈ RN×N , where A [ i , j ] = 1 if two individuals are connected , 0 otherwise . Let Si , Ii , Ri be the average probabilities of node i being in each of the states , with Si + Ii +Ri = 1 . The SIR dynamics on a graph is given by Newman ( 2018 ) : dSi dt = −β ∑ j AijIjSi , dRi dt = γIi , dSi dt + dIi dt + dRi dt = 0 . ( 1 ) where β is the infection rate per contact and γ the recovery/death rate . The basic reproductive rate of a disease R0 ≡ βλ1/γ is defined as the number of secondary infections created by an index case in a fully susceptible population Keeling & Rohani ( 2011 ) . Here λ1 is the largest eigenvalue of A . The disease will spread and result in an epidemic if R0 > 1 . Discrete time model We can also use an equivalent discrete time SIR model . Let xti ∈ { S , I , R } be the state of node i at time t. For a susceptible node i , its probability to become infected or removed at time t+ 1 is P ( xt+1i = I|x t i = S ) = 1− ∏ j ( 1− βAijIi ( t ) ) , P ( xt+1i = R|x t i = I ) = γ . ( 2 ) The SIR model doesn ’ t account for the incubation period , where an individual is infected but not infectious . This is remedied by introducing an “ exposed ” ( E ) state , leading to the SEIR model . For a susceptible node i , the probability to enter the exposed state , and becoming infectious at time t+ 1 is P ( xt+1i = E|x t i = S ) = 1− ∏ j ( 1− βAijIi ( t ) ) , P ( xt+1i = I|x t i = E ) = α , ( 3 ) An infected node eventually enters the removed state with probability γ , which is the same as SIR ( 2 ) . ( 2 ) and ( 3 ) yield ( 1 ) for very small β as ∏ j ( 1− βAijIi ) ≈ βAijIi ( proof in supp . B ) . Finding patient zero Finding P0 can be formulated as a maximum likelihood estimation problem for SIR and SEIR models . Specifically , we observe a snapshot of the state of the nodes at time step t as xt : = ( xt1 , · · · , xtN ) , with each node ’ s state xti ∈ { S , E , I , R } . The problem of finding P0 is to search for a set of nodes Z = { i|x0i = I , i ∈ { 1 , · · ·N } } which led to the observed state xt . A common approach is to find Z such that the likelihood of observing xt is maximized : Z ? = argmaxZ , |Z|≤kP ( xt|Z ) ( 4 ) where P ( xt|Z ) is the probability of observing xt with Z being the P0 set . We assume the number of P0s is no larger than k. Estimating the original state of the dynamic system given the future states requires computing the conditional likelihood P ( xt|Z ) exactly , which is intractable due to the combinatorics of possible transmission routes .

Summary: An S(E)IR epidemics propagates on a graph, and the goal is to detect its source (P0) only from the observation of the state (S,E,I,R) of every node of the graph at some time $T > 0$. This version of the source detection problem has been studied first by Shah and Zeman (2011) for SI epidemics, as listed in Section 2. The current paper claims to (i) establish new fundamental limits on this problem, showing in particular that after some time the source detection becomes difficult, and (ii) to demonstrate the ability of graph convolutional networks to solve the problem and validate the results on real data.